Autoscaling support in alchemiscale (#424)

Co-authored-by: David L. Dotson <dotsdl@gmail.com>

Author: ianmkenney

alchemiscale/compute/api.py

@@ -9,7 +9,8 @@
 from datetime import timedelta
 import random
 
-from fastapi import FastAPI, APIRouter, Body, Depends, Request
+from fastapi import FastAPI, APIRouter, Body, Depends, HTTPException, Request
+from fastapi import status as http_status
 from fastapi.middleware.gzip import GZipMiddleware
 from gufe.tokenization import JSON_HANDLER
 from gufe.protocols import ProtocolDAGResult
@@ -40,7 +41,10 @@
 from ..storage.models import (
     ProtocolDAGResultRef,
     ComputeServiceID,
+    ComputeManagerID,
+    ComputeManagerRegistration,
     ComputeServiceRegistration,
+    ComputeManagerStatus,
 )
 from ..models import Scope, ScopedKey
 from ..security.models import (
@@ -102,17 +106,31 @@ def list_scopes(
 @router.post("/computeservice/{compute_service_id}/register")
 def register_computeservice(
     compute_service_id,
+    *,
+    compute_manager_id: str | None = Body(None, embed=True),
     n4js: Neo4jStore = Depends(get_n4js_depends),
 ):
     now = datetime.datetime.now(tz=datetime.UTC)
+    if compute_manager_id:
+        manager_name = process_compute_manager_id_string(compute_manager_id).name
+    else:
+        manager_name = None
+
     csreg = ComputeServiceRegistration(
         identifier=ComputeServiceID(compute_service_id),
         registered=now,
         heartbeat=now,
         failure_times=[],
+        manager_name=manager_name,
     )
 
-    compute_service_id_ = n4js.register_computeservice(csreg)
+    try:
+        compute_service_id_ = n4js.register_computeservice(csreg)
+    except ValueError as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=str(e),
+        )
 
     return compute_service_id_
 
@@ -392,6 +410,136 @@ async def set_task_result(
     return result_sk
 
 
+def process_compute_manager_id_string(
+    compute_manager_id_string: str,
+) -> ComputeManagerID:
+    """Try creating a ComputeManagerID from a string representation. Raise HTTPException."""
+    try:
+        compute_manager_id = ComputeManagerID(compute_manager_id_string)
+    except Exception as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=str(e),
+        )
+
+    return compute_manager_id
+
+
+@router.post("/computemanager/{compute_manager_id}/register")
+def register_computemanager(
+    compute_manager_id,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+):
+
+    compute_manager_id = process_compute_manager_id_string(compute_manager_id)
+
+    now = datetime.datetime.now(tz=datetime.UTC)
+    cm_registration = ComputeManagerRegistration(
+        name=compute_manager_id.name,
+        uuid=compute_manager_id.uuid,
+        registered=now,
+        last_status_update=now,
+        status=ComputeManagerStatus.OK,
+        detail="",
+        saturation=0,
+    )
+
+    try:
+        compute_manager_id_ = n4js.register_computemanager(cm_registration)
+    except ValueError as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=str(e),
+        )
+
+    return compute_manager_id_
+
+
+@router.post("/computemanager/{compute_manager_id}/deregister")
+def deregister_computemanager(
+    compute_manager_id,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+):
+    compute_manager_id = process_compute_manager_id_string(compute_manager_id)
+    n4js.deregister_computemanager(compute_manager_id)
+    return compute_manager_id
+
+
+@router.post("/computemanager/{compute_manager_id}/instruction")
+def get_instruction_computemanager(
+    compute_manager_id,
+    *,
+    scopes: list[Scope] = Body([], embed=True),
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    settings: ComputeAPISettings = Depends(get_base_api_settings),
+    token: TokenData = Depends(get_token_data_depends),
+):
+    scopes = scopes or [Scope()]
+    scopes_reduced = minimize_scope_space(scopes)
+    query_scopes = []
+    for scope in scopes_reduced:
+        query_scopes.extend(validate_scopes_query(scope, token))
+
+    compute_manager_id = process_compute_manager_id_string(compute_manager_id)
+    now = datetime.datetime.now(tz=datetime.UTC)
+    instruction, payload = n4js.get_computemanager_instruction(
+        compute_manager_id,
+        now - timedelta(seconds=settings.ALCHEMISCALE_COMPUTE_API_FORGIVE_TIME_SECONDS),
+        settings.ALCHEMISCALE_COMPUTE_API_MAX_FAILURES,
+        query_scopes,
+    )
+    payload["instruction"] = str(instruction)
+    return payload
+
+
+@router.post("/computemanager/{compute_manager_id}/status")
+def update_status_computemanager(
+    compute_manager_id,
+    *,
+    status: str = Body(),
+    detail: str | None = Body(None),
+    saturation: float | None = Body(None),
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    settings: ComputeAPISettings = Depends(get_base_api_settings),
+):
+    expire_seconds = settings.ALCHEMISCALE_COMPUTE_API_MANAGER_EXPIRE_SECONDS
+    expire_seconds_errored = (
+        settings.ALCHEMISCALE_COMPUTE_API_MANAGER_EXPIRE_SECONDS_ERROR
+    )
+    compute_manager_id = process_compute_manager_id_string(compute_manager_id)
+    try:
+        n4js.update_compute_manager_status(
+            compute_manager_id, status, detail, saturation
+        )
+        now = datetime.datetime.now(tz=datetime.UTC)
+        n4js.expire_computemanager_registrations(
+            now - timedelta(seconds=expire_seconds),
+            now - timedelta(seconds=expire_seconds_errored),
+        )
+    except ValueError as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_400_BAD_REQUEST,
+            detail=str(e),
+        )
+
+    return compute_manager_id
+
+
+@router.post("/computemanager/{compute_manager_name}/clear_error")
+def clear_error_computemanager(
+    compute_manager_name: str,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+):
+    if not compute_manager_name.isalnum():
+        raise ValueError("Provided manager name is not alphanumeric")
+
+    with n4js.transaction() as tx:
+        compute_manager_id = n4js.get_compute_manager_id(
+            name=compute_manager_name, tx=tx
+        )
+        n4js.clear_errored_computemanager(compute_manager_id, tx=tx)
+
+
 ### add router
 
 app.include_router(router)

alchemiscale/compute/client.py

@@ -17,7 +17,14 @@
 )
 from ..compression import compress_gufe_zstd, decompress_gufe_zstd
 from ..models import Scope, ScopedKey
-from ..storage.models import TaskHub, Task, ComputeServiceID
+from ..storage.models import (
+    TaskHub,
+    Task,
+    ComputeServiceID,
+    ComputeManagerID,
+    ComputeManagerInstruction,
+    ComputeManagerStatus,
+)
 
 
 class AlchemiscaleComputeClientError(AlchemiscaleBaseClientError): ...
@@ -28,8 +35,15 @@ class AlchemiscaleComputeClient(AlchemiscaleBaseClient):
 
     _exception = AlchemiscaleComputeClientError
 
-    def register(self, compute_service_id: ComputeServiceID):
-        res = self._post_resource(f"/computeservice/{compute_service_id}/register", {})
+    def register(
+        self,
+        compute_service_id: ComputeServiceID,
+        compute_manager_id: ComputeManagerID | None = None,
+    ):
+        res = self._post_resource(
+            f"/computeservice/{compute_service_id}/register",
+            {"compute_manager_id": compute_manager_id},
+        )
         return ComputeServiceID(res)
 
     def deregister(self, compute_service_id: ComputeServiceID):
@@ -144,3 +158,72 @@ def set_task_result(
         pdr_sk = self._post_resource(f"/tasks/{task}/results", data)
 
         return ScopedKey.from_dict(pdr_sk)
+
+
+class AlchemiscaleComputeManagerClientError(AlchemiscaleBaseClientError): ...
+
+
+class AlchemiscaleComputeManagerClient(AlchemiscaleBaseClient):
+
+    _exception = AlchemiscaleComputeManagerClientError
+
+    def register(self, compute_manager_id: ComputeManagerID) -> ComputeManagerID:
+        res = self._post_resource(f"/computemanager/{compute_manager_id}/register", {})
+        return ComputeManagerID(res)
+
+    def deregister(self, compute_manager_id: ComputeManagerID) -> ComputeManagerID:
+        res = self._post_resource(
+            f"/computemanager/{compute_manager_id}/deregister", {}
+        )
+        return ComputeManagerID(res)
+
+    def get_instruction(
+        self,
+        scopes: list[Scope],
+        compute_manager_id: ComputeManagerID,
+    ) -> tuple[ComputeManagerInstruction, dict]:
+        instruction_data = self._post_resource(
+            f"/computemanager/{compute_manager_id}/instruction",
+            {"scopes": [scope.to_dict() for scope in scopes]},
+        )
+
+        match instruction_data:
+            case {
+                "instruction": "OK",
+                "compute_service_ids": ids,
+                "num_tasks": num_tasks,
+            }:
+                return ComputeManagerInstruction.OK, {
+                    "compute_service_ids": ids,
+                    "num_tasks": num_tasks,
+                }
+            case {"instruction": "SKIP", "compute_service_ids": ids}:
+                return ComputeManagerInstruction.SKIP, {"compute_service_ids": ids}
+            case {"instruction": "SHUTDOWN", "message": message}:
+                return ComputeManagerInstruction.SHUTDOWN, {"message": message}
+            case _:
+                raise self._exception(
+                    f"Received unknown instruction pattern: {instruction_data}"
+                )
+
+    def update_status(
+        self,
+        compute_manager_id: ComputeManagerID,
+        status: ComputeManagerStatus,
+        *,
+        detail: str | None = None,
+        saturation: float | None = None,
+    ) -> ComputeManagerID:
+        payload = {"detail": detail, "saturation": saturation, "status": str(status)}
+        res = self._post_resource(
+            f"/computemanager/{compute_manager_id}/status",
+            payload,
+        )
+
+        return ComputeManagerID(res)
+
+    def clear_error(self, compute_manager_name: str):
+        res = self._post_resource(
+            f"/computemanager/{compute_manager_name}/clear_error",
+            {},
+        )