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,114 @@ 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),
+):
+    compute_manager_id = process_compute_manager_id_string(compute_manager_id)
+    try:
+        n4js.update_compute_manager_status(
+            compute_manager_id, status, detail, saturation
+        )
+    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,48 @@
+"""
+:mod:`alchemiscale.compute.manager` --- compute manager for creating compute services
+=====================================================================================
+
+"""
+
+from abc import abstractmethod
+
+from ..storage.models import ComputeManagerID, ComputeManagerInstruction
+from .client import AlchemiscaleComputeManagerClient
+from .settings import ComputeManagerSettings
+
+
+class ComputeManager:
+
+    compute_manager_id: ComputeManagerID
+    sleep_interval: int
+    client: AlchemiscaleComputeManagerClient
+    service_settings_template: bytes
+    manager_settings: ComputeManagerSettings
+
+    def _register(self):
+        self.client.register(self.compute_manager_id)
+
+    def _deregister(self):
+        self.client.deregister(self.compute_manager_id)
+
+    def request_instruction(self) -> ComputeManagerInstruction:
+        instruction = self.client._post_resource(
+            f"/computemanager/{self.compute_manager_id}/update_status", {}
+        )
+        return instruction
+
+    @abstractmethod
+    def start(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def stop(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def cycle(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_compute_service(self):
+        raise NotImplementedError

alchemiscale/compute/service.py

@@ -73,6 +73,7 @@ def __init__(self, settings: ComputeServiceSettings):
 
         self.api_url = self.settings.api_url
         self.name = self.settings.name
+        self.compute_manager_id = self.settings.compute_manager_id
         self.sleep_interval = self.settings.sleep_interval
         self.deep_sleep_interval = self.settings.deep_sleep_interval
         self.heartbeat_interval = self.settings.heartbeat_interval
@@ -135,7 +136,7 @@ def __init__(self, settings: ComputeServiceSettings):
 
     def _register(self):
         """Register this compute service with the compute API."""
-        self.client.register(self.compute_service_id)
+        self.client.register(self.compute_service_id, self.settings.compute_manager_id)
 
     def _deregister(self):
         """Deregister this compute service with the compute API."""

alchemiscale/compute/settings.py

@@ -27,6 +27,13 @@ class Config:
             "resources, e.g. different hosts or HPC clusters."
         ),
     )
+    compute_manager_id: str | None = Field(
+        None,
+        description=(
+            "The ID of the autoscaling compute manager responsible for the creation "
+            "of this compute service. Default is None."
+        ),
+    )
     shared_basedir: Path = Field(
         ..., description="Filesystem path to use for `ProtocolDAG` `shared` space."
     )
@@ -114,3 +121,30 @@ class Config:
         True,
         description="Whether to verify SSL certificate presented by the API server.",
     )
+
+
+class ComputeManagerSettings(BaseModel):
+    api_url: str | None = Field(
+        ..., description="URL of the compute API to manager services for."
+    )
+    identifier: str | None = Field(
+        ..., description="Identifier for the compute identity used for authentication."
+    )
+    key: str | None = Field(
+        ..., description="Credential for the compute identity used for authentication."
+    )
+    name: str = Field(
+        ...,
+        description=(
+            "The name to give this compute manager. This value should be distinct from all"
+            "other compute managers."
+        ),
+    )
+    logfile: Path | None = Field(..., description="File path to write logs to.")
+    max_compute_services: int = Field(
+        ...,
+        description="Maximum number of compute services the manager is allowed to have running at a time.",
+    )
+    sleep_interval: int = Field(
+        1800, description="Time in seconds to wait for another instruction."
+    )

alchemiscale/models.py

@@ -5,6 +5,7 @@
 """
 
 from typing import Any
+from enum import StrEnum
 from pydantic import BaseModel, field_validator, model_validator, ConfigDict
 from gufe.tokenization import GufeKey
 from re import fullmatch