Autoscaling support in alchemiscale

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,14 +106,24 @@ 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
+        ).manager_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)
@@ -392,6 +406,124 @@ 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 ValueError as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+            details=str(e),
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_500_INTERNAL_SERVER_ERROR,
+            details=str(e),
+        )
+
+    return compute_manager_id
+
+
+@router.post("/computemanager/{compute_manager_id}/register")
+def computemanager_register(
+    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(
+        manager_name=compute_manager_id.manager_name,
+        uuid=compute_manager_id.uuid,
+        registered=now,
+        last_status_update=now,
+        status=ComputeManagerStatus.OK,
+        detail="",
+        saturation=0,
+    )
+
+    compute_manager_id_ = n4js.register_computemanager(cm_registration)
+    return compute_manager_id_
+
+
+@router.post("/computemanager/{compute_manager_id}/deregister")
+def computemanager_deregister(
+    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 computemanager_get_instruction(
+    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.from_str("*-*-*")]
+    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 computemanager_update_status(
+    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_ERRORED
+    )
+    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,
+            details=str(e),
+        )
+    except Exception as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_500_INTERNAL_SERVER_ERROR,
+            details=str(e),
+        )
+
+    return compute_manager_id
+
+
 ### 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,64 @@ 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 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 get_instruction(
+        self, compute_manager_id: ComputeManagerID
+    ) -> tuple[ComputeManagerInstruction, dict]:
+        instruction_data = self._post_resource(
+            f"/computemanager/{compute_manager_id}/instruction",
+            {"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}"
+                )

alchemiscale/compute/manager.py

@@ -0,0 +1,125 @@
+"""
+:mod:`alchemiscale.compute.manager` --- compute manager for creating compute services
+=====================================================================================
+
+"""
+
+from abc import abstractmethod
+import logging
+import time
+
+from ..storage.models import (
+    ComputeManagerID,
+    ComputeManagerInstruction,
+    ComputeManagerStatus,
+)
+from .client import AlchemiscaleComputeManagerClient
+from .settings import ComputeManagerSettings
+
+
+class ComputeManager:
+
+    compute_manager_id: ComputeManagerID
+    client: AlchemiscaleComputeManagerClient
+    service_settings_template: bytes
+    manager_settings: ComputeManagerSettings
+
+    def __init__(self, settings: ComputeManagerSettings):
+        self.settings = settings
+        self.compute_manager_id = ComputeManagerID.new_from_manager_name(
+            self.settings.name
+        )
+        self.client = AlchemiscaleComputeManagerClient(
+            api_url=self.settings.api_url,
+            identifier=self.settings.identifier,
+            key=self.settings.key,
+        )
+
+        logger = logging.getLogger("AlchemiscaleComputeManager")
+        logger.setLevel(self.settings.loglevel)
+
+        extra = {"compute_manager_id": self.compute_manager_id}
+        formatter = logging.Formatter(
+            "[%(asctime)s] [%(compute_manager_id)s] [%(levelname)s] %(message)s"
+        )
+        formatter.converter = time.gmtime  # use utc time for logging timestamps
+
+        sh = logging.StreamHandler()
+        sh.setFormatter(formatter)
+        logger.addHandler(sh)
+
+        if self.settings.logfile is not None:
+            fh = logging.FileHandler(self.settings.logfile)
+            fh.setFormatter(formatter)
+            logger.addHandler(fh)
+
+        self.logger = logging.LoggerAdapter(logger, extra)
+
+    def _register(self):
+        self.client.register(self.compute_manager_id)
+
+    def _deregister(self):
+        self.client.deregister(self.compute_manager_id)
+
+    def start(self, max_cycles: int | None = None):
+        self._register()
+        try:
+            count = 0
+            while True:
+                self.cycle()
+                count += 1
+                if max_cycles and count >= max_cycles:
+                    break
+                time.sleep(self.settings.sleep_interval)
+        except Exception as e:
+            self.client.update_status(
+                self.compute_manager_id, ComputeManagerStatus.ERRORED, detail=repr(e)
+            )
+            raise e
+        except KeyboardInterrupt:
+            self.logger.info("Caught SIGINT/Keyboard interrupt.")
+        finally:
+            self._deregister()
+
+    @abstractmethod
+    def create_compute_services(self, data) -> int:
+        """Method responsible for creating compute services based on
+        data returned with an OK ComputeManagerInstruction. This must
+        return the number of compute services started.
+        """
+        raise NotImplementedError
+
+    def cycle(self):
+        instruction, data = self.client.get_instruction(self.compute_manager_id)
+        match instruction:
+            case ComputeManagerInstruction.OK:
+                total_services = len(data["compute_service_ids"])
+                num_tasks = data["num_tasks"]
+                if (
+                    total_services < self.settings.max_compute_services
+                    and num_tasks > 0
+                ):
+                    new_services = self.create_compute_services(data)
+                    total_services += new_services
+                    if new_services:
+                        self.logger.info(
+                            f"Created {new_services} new compute service(s)"
+                        )
+                    else:
+                        self.logger.info(f"No new compute services created")
+            case ComputeManagerInstruction.SKIP:
+                total_services = len(data["compute_service_ids"])
+                self.logger.info(f"Received skip instruction")
+            case ComputeManagerInstruction.SHUTDOWN:
+                shutdown_message = data["message"]
+                self.logger.info(f'Received shutdown message: "{shutdown_message}"')
+                return
+        self.client.update_status(
+            self.compute_manager_id,
+            ComputeManagerStatus.OK,
+            saturation=total_services / self.settings.max_compute_services,
+        )
+
+    @abstractmethod
+    def create_compute_service(self):
+        raise NotImplementedError