Add missing features, polish to SynchronousComputeService

alchemiscale/cli.py

@@ -5,11 +5,12 @@
 """
 
 import click
+import yaml
+import signal
 import gunicorn.app.base
 from typing import Type
 
-from .models import Scope
-from .security.auth import hash_key, authenticate, AuthenticationError
+from .security.auth import hash_key
 from .security.models import (
     CredentialedEntity,
     CredentialedUserIdentity,
@@ -228,7 +229,7 @@ def cli():
     name="api",
     help="Start the user-facing API service",
 )
-@api_starting_params("FA_API_HOST", "FA_API_PORT", "FA_API_LOGLEVEL")
+@api_starting_params("ALCHEMISCALE_API_HOST", "ALCHEMISCALE_API_PORT", "ALCHEMISCALE_API_LOGLEVEL")
 @db_params
 @s3os_params
 @jwt_params
@@ -268,7 +269,7 @@ def get_settings_override():
 
     app.dependency_overrides[get_base_api_settings] = get_settings_override
 
-    start_api(app, workers, host["FA_API_HOST"], port["FA_API_PORT"])
+    start_api(app, workers, host["ALCHEMISCALE_API_HOST"], port["ALCHEMISCALE_API_PORT"])
 
 
 @cli.group(help="Subcommands for the compute service")
@@ -278,7 +279,7 @@ def compute():
 
 @compute.command(help="Start the compute API service.")
 @api_starting_params(
-    "FA_COMPUTE_API_HOST", "FA_COMPUTE_API_PORT", "FA_COMPUTE_API_LOGLEVEL"
+    "ALCHEMISCALE_COMPUTE_API_HOST", "ALCHEMISCALE_COMPUTE_API_PORT", "ALCHEMISCALE_COMPUTE_API_LOGLEVEL"
 )
 @db_params
 @s3os_params
@@ -316,12 +317,41 @@ def get_settings_override():
 
     app.dependency_overrides[get_base_api_settings] = get_settings_override
 
-    start_api(app, workers, host["FA_COMPUTE_API_HOST"], port["FA_COMPUTE_API_PORT"])
+    start_api(app, workers, host["ALCHEMISCALE_COMPUTE_API_HOST"], port["ALCHEMISCALE_COMPUTE_API_PORT"])
 
 
 @compute.command(help="Start the synchronous compute service.")
-def synchronous():
-    ...
+@click.option(
+        "--config-file", 
+        "-c",
+        type=click.File(),
+        help="YAML-based configuration file giving the settings for this service",
+        required=True
+        )
+def synchronous(config_file):
+    from alchemiscale.models import Scope
+    from alchemiscale.compute.service import SynchronousComputeService
+
+    params = yaml.load(config_file, Loader=yaml.Loader)
+
+    if 'scopes' in params:
+        params['scopes'] = [Scope.from_str(scope) for scope in params['scopes']]
+
+    service = SynchronousComputeService(**params)
+
+    # add signal handling
+    for signame in {"SIGHUP", "SIGINT", "SIGTERM"}:
+
+        def stop(*args, **kwargs):
+            service.stop()
+            raise KeyboardInterrupt()
+
+        signal.signal(getattr(signal, signame), stop)
+
+    try:
+        service.start()
+    except KeyboardInterrupt:
+        pass
 
 
 @cli.group(help="Subcommands for the database")
@@ -491,6 +521,7 @@ def remove(url, user, password, dbname, identity_type, identifier):
 @scope
 def add_scope(url, user, password, dbname, identity_type, identifier, scope):
     """Add a scope for the given identity."""
+    from .models import Scope
     from .storage.statestore import get_n4js
     from .settings import Neo4jStoreSettings
 
@@ -532,6 +563,7 @@ def list_scope(url, user, password, dbname, identity_type, identifier):
 @scope
 def remove_scope(url, user, password, dbname, identity_type, identifier, scope):
     """Remove a scope for the given identity(s)."""
+    from .models import Scope
     from .storage.statestore import get_n4js
     from .settings import Neo4jStoreSettings
 

alchemiscale/compute/api.py

@@ -8,6 +8,7 @@
 from typing import Any, Dict, List
 import os
 import json
+from datetime import datetime
 
 from fastapi import FastAPI, APIRouter, Body, Depends, HTTPException, status
 from gufe.tokenization import GufeTokenizable, JSON_HANDLER
@@ -29,7 +30,12 @@
 from ..settings import get_base_api_settings, get_compute_api_settings
 from ..storage.statestore import Neo4jStore
 from ..storage.objectstore import S3ObjectStore
-from ..storage.models import ProtocolDAGResultRef, TaskStatusEnum
+from ..storage.models import (
+    ProtocolDAGResultRef,
+    ComputeServiceID,
+    ComputeServiceRegistration,
+    TaskStatusEnum,
+)
 from ..models import Scope, ScopedKey
 from ..security.auth import get_token_data, oauth2_scheme
 from ..security.models import (
@@ -92,6 +98,42 @@ async def list_scopes(
     return [str(scope) for scope in scopes]
 
 
+@router.post("/computeservice/{compute_service_id}/register")
+async def register_computeservice(
+    compute_service_id,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+):
+    now = datetime.utcnow()
+    csreg = ComputeServiceRegistration(
+        identifier=compute_service_id, registered=now, heartbeat=now
+    )
+
+    n4js.register_computeservice(csreg)
+
+    return compute_service_id
+
+
+@router.post("/computeservice/{compute_service_id}/deregister")
+async def deregister_computeservice(
+    compute_service_id,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+):
+    n4js.deregister_computeservice(ComputeServiceID(compute_service_id))
+
+    return compute_service_id
+
+
+@router.post("/computeservice/{compute_service_id}/heartbeat")
+async def heartbeat_computeservice(
+    compute_service_id,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+):
+    now = datetime.utcnow()
+    n4js.heartbeat_computeservice(compute_service_id, now)
+
+    return compute_service_id
+
+
 @router.get("/taskhubs")
 async def query_taskhubs(
     *,
@@ -117,18 +159,11 @@ async def query_taskhubs(
     return taskhubs_handler.format_return()
 
 
-# @app.get("/taskhubs/{scoped_key}")
-# async def get_taskhub(scoped_key: str,
-#                        *,
-#                        n4js: Neo4jStore = Depends(get_n4js_depends)):
-#    return
-
-
 @router.post("/taskhubs/{taskhub_scoped_key}/claim")
 async def claim_taskhub_tasks(
     taskhub_scoped_key,
     *,
-    claimant: str = Body(),
+    compute_service_id: str = Body(),
     count: int = Body(),
     n4js: Neo4jStore = Depends(get_n4js_depends),
     token: TokenData = Depends(get_token_data_depends),
@@ -137,7 +172,9 @@ async def claim_taskhub_tasks(
     validate_scopes(sk.scope, token)
 
     tasks = n4js.claim_taskhub_tasks(
-        taskhub=taskhub_scoped_key, claimant=claimant, count=count
+        taskhub=taskhub_scoped_key,
+        compute_service_id=ComputeServiceID(compute_service_id),
+        count=count,
     )
 
     return [str(t) if t is not None else None for t in tasks]
@@ -197,8 +234,12 @@ def set_task_result(
         task=task_sk, protocoldagresultref=protocoldagresultref
     )
 
-    # TODO: if success, set task complete, remove from all hubs
+    # if success, set task complete, remove from all hubs
     # otherwise, set as errored, leave in hubs
+    if protocoldagresultref.ok:
+        n4js.set_task_complete(tasks=[task_sk])
+    else:
+        n4js.set_task_error(tasks=[task_sk])
 
     return result_sk
 

alchemiscale/compute/client.py

@@ -23,7 +23,7 @@
     json_to_gufe,
 )
 from ..models import Scope, ScopedKey
-from ..storage.models import TaskHub, Task, TaskStatusEnum
+from ..storage.models import TaskHub, Task, ComputeServiceID, TaskStatusEnum
 
 
 class AlchemiscaleComputeClientError(AlchemiscaleBaseClientError):
@@ -35,6 +35,18 @@ class AlchemiscaleComputeClient(AlchemiscaleBaseClient):
 
     _exception = AlchemiscaleComputeClientError
 
+    def register(self, compute_service_id: ComputeServiceID):
+        res = self._post_resource(f"computeservice/{compute_service_id}/register", {})
+        return ComputeServiceID(res)
+
+    def deregister(self, compute_service_id: ComputeServiceID):
+        res = self._post_resource(f"computeservice/{compute_service_id}/deregister", {})
+        return ComputeServiceID(res)
+
+    def heartbeat(self, compute_service_id: ComputeServiceID):
+        res = self._post_resource(f"computeservice/{compute_service_id}/heartbeat", {})
+        return ComputeServiceID(res)
+
     def list_scopes(self) -> List[Scope]:
         scopes = self._get_resource(
             f"/identities/{self.identifier}/scopes",
@@ -62,10 +74,10 @@ def query_taskhubs(
         return taskhubs
 
     def claim_taskhub_tasks(
-        self, taskhub: ScopedKey, claimant: str, count: int = 1
+        self, taskhub: ScopedKey, compute_service_id: ComputeServiceID, count: int = 1
     ) -> Task:
         """Claim a `Task` from the specified `TaskHub`"""
-        data = dict(claimant=claimant, count=count)
+        data = dict(compute_service_id=str(compute_service_id), count=count)
         tasks = self._post_resource(f"taskhubs/{taskhub}/claim", data)
 
         return [ScopedKey.from_str(t) if t is not None else None for t in tasks]