Implement tests for compute manager client

Author: ianmkenney

alchemiscale/compute/api.py

@@ -425,7 +425,8 @@ def computemanager_register(
 
     now = datetime.utcnow()
     cm_registration = ComputeManagerRegistration(
-        identifier=ComputeManagerID(compute_manager_id),
+        manager_id=compute_manager_id.manager_id,
+        uuid=compute_manager_id.uuid,
         registered=now,
         last_status_update=now,
         status=ComputeManagerStatus.OK,
@@ -434,7 +435,6 @@ def computemanager_register(
     )
 
     compute_manager_id_ = n4js.register_computemanager(cm_registration)
-
     return compute_manager_id_
 
 
@@ -457,12 +457,13 @@ def computemanager_get_instruction(
 ):
     compute_manager_id = process_compute_manager_id_string(compute_manager_id)
     now = datetime.utcnow()
-    instruction = n4js.get_computemanager_instruction(
+    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,
     )
-    return instruction
+    payload["instruction"] = str(instruction)
+    return payload
 
 
 @router.post("/computemanager/{compute_manager_id}/update_status")

alchemiscale/compute/client.py

@@ -153,8 +153,13 @@ def set_task_result(
         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)
@@ -179,9 +184,29 @@ def update_status(
 
         return ComputeManagerID(res)
 
-    def get_instruction(self, compute_manager_id: ComputeManagerID) -> ComputeManagerID:
-        res = self._post_resource(
+    def get_instruction(
+        self, compute_manager_id: ComputeManagerID
+    ) -> tuple[ComputeManagerInstruction, dict]:
+        instruction_data = self._post_resource(
             f"/computemanager/{compute_manager_id}/instruction",
             {},
         )
-        return ComputeManagerInstruction(res)
+
+        match instruction_data:
+            case {
+                "instruction": "OK",
+                "num_registered": num_registered,
+                "compute_service_ids": ids,
+            }:
+                return ComputeManagerInstruction.OK, {
+                    "num_registered": num_registered,
+                    "compute_service_ids": ids,
+                }
+            case {"instruction": "SKIP"}:
+                return ComputeManagerInstruction.SKIP, {}
+            case {"instruction": "SHUTDOWN", "message": message}:
+                return ComputeManagerInstruction.SHUTDOWN, {message: message}
+            case _:
+                raise self._exception(
+                    f"Received unknown instruction pattern: {instruction_data}"
+                )

alchemiscale/compute/settings.py

@@ -118,8 +118,31 @@ class Config:
 
 class ComputeManagerSettings(BaseModel):
 
-    api_url: str
-    name: str
-    status_update_interval: int
-    logfile: Path | None
-    max_compute_services: int
+    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."
+        ),
+    )
+
+    status_update_interval: int = Field(
+        ..., description="Time in seconds to send a status update to the compute API."
+    )
+    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/storage/models.py

@@ -73,7 +73,7 @@ class ComputeManagerStatus(StrEnum):
 class ComputeManagerID(str):
 
     def __init__(self, value):
-        super().__init__(value)
+        super().__init__()
 
         parts = self.split("-")
 
@@ -116,6 +116,35 @@ class ComputeManagerRegistration(BaseModel):
     detail: str
     saturation: float
 
+    def __repr__(self):  # pragma: no cover
+        return f"<ComputeManagerRegistration('{str(self)}')>"
+
+    def __str__(self):
+        return "-".join([self.manager_id, self.uuid])
+
+    @classmethod
+    def from_now(cls, identifier: ComputeManagerID):
+        now = datetime.utcnow()
+        return cls(
+            manager_id=identifier.manager_id,
+            uuid=identifier.uuid,
+            last_status_update=now,
+            status=ComputeManagerStatus.OK,
+            detail="",
+            saturation=0,
+        )
+
+    def to_dict(self):
+        dct = self.dict()
+        dct["manager_id"] = str(self.manager_id)
+        dct["uuid"] = str(self.uuid)
+
+        return dct
+
+    @classmethod
+    def from_dict(cls, dct):
+        return cls(**dct_)
+
 
 class TaskProvenance(BaseModel):
     computeserviceid: ComputeServiceID

alchemiscale/storage/statestore.py

@@ -1447,12 +1447,12 @@ def get_computemanager_instruction(
         compute_manager_id: ComputeManagerID,
         forgive_time: datetime,
         max_failures: int,
-    ) -> ComputeManagerInstruction:
+    ) -> tuple[ComputeManagerInstruction, dict]:
 
         manager_id, uuid = compute_manager_id.manager_id, compute_manager_id.uuid
 
         query = """
-        MATCH (csm: ComputeServiceManager {manager_id: $manager_id, uuid: $uuid})
+        MATCH (csm: ComputeManagerRegistration {manager_id: $manager_id, uuid: $uuid})
         OPTIONAL MATCH (csm)-[rel:MANAGES]->(csr: ComputeServiceRegistration)
         RETURN csm, csr.identifier as csr_id
         """
@@ -1461,21 +1461,29 @@ def get_computemanager_instruction(
 
         # no compute manager was found the given name and UUID
         if len(results.records) == 0:
-            return ComputeManagerInstruction.SHUTDOWN
+            msg = "no compute manager was found the given name and UUID"
+            return ComputeManagerInstruction.SHUTDOWN, {"message": msg}
 
         # TODO: very chatty, try and do this in a single query
         # this would require a new state store method that requests
         # failure times in bulk
+        csr_ids = []
         for record in results.records:
-            csr_id = record["csr_id"]
-            if not self.compute_service_can_claim(
-                csr_id,
-                forgive_time,
-                max_failures,
-            ):
-                return ComputeManagerInstruction.SKIP
+            if csr_id := record["csr_id"]:
+                if not self.compute_service_can_claim(
+                    csr_id,
+                    forgive_time,
+                    max_failures,
+                ):
+                    return ComputeManagerInstruction.SKIP, {}
+            else:
+                break
+            csr_ids.append(ComputeServiceID(csr_id))
 
-        return ComputeManagerInstruction.OK
+        return ComputeManagerInstruction.OK, {
+            "num_registered": len(csr_ids),
+            "compute_service_ids": csr_ids,
+        }
 
     def update_compute_manager_status(
         self,

alchemiscale/tests/integration/compute/client/conftest.py

@@ -70,3 +70,27 @@ def compute_client_wrong_credential(uvicorn_server, compute_identity):
         identifier=compute_identity["identifier"],
         key="wrong credential",
     )
+
+
+@pytest.fixture(scope="module")
+def compute_manager_client(
+    uvicorn_server,
+    compute_identity,
+    single_scoped_credentialed_compute,
+):
+    return client.AlchemiscaleComputeManagerClient(
+        api_url="http://127.0.0.1:8000/",
+        # use the identifier for the single-scoped user who should have access to some things
+        identifier=single_scoped_credentialed_compute.identifier,
+        # all the test users are based on compute_identity who use the same password
+        key=compute_identity["key"],
+    )
+
+
+@pytest.fixture(scope="module")
+def manager_client_wrong_credential(uvicorn_server, compute_identity):
+    return client.AlchemiscaleComputeManagerClient(
+        api_url="http://127.0.0.1:8000/",
+        identifier=compute_identity["identifier"],
+        key="wrong credential",
+    )

alchemiscale/tests/integration/compute/client/test_manager_client.py

@@ -0,0 +1,91 @@
+from time import sleep
+
+import pytest
+
+from alchemiscale.compute import client
+from alchemiscale.storage.models import ComputeManagerID, ComputeManagerInstruction
+from alchemiscale.tests.integration.compute.utils import get_compute_settings_override
+
+
+class TestComputeManager:
+
+    def test_wrong_credential(
+        self,
+        scope_test,
+        n4js_preloaded,
+        manager_client_wrong_credential: client.AlchemiscaleComputeManagerClient,
+        uvicorn_server,
+    ):
+        with pytest.raises(client.AlchemiscaleComputeManagerClientError):
+            manager_client_wrong_credential.get_info()
+
+    def test_refresh_credential(
+        self,
+        n4js_preloaded,
+        compute_client: client.AlchemiscaleComputeManagerClient,
+        uvicorn_server,
+    ):
+        settings = get_compute_settings_override()
+        assert compute_client._jwtoken is None
+        compute_client._get_token()
+
+        token = compute_client._jwtoken
+        assert token is not None
+
+        # token shouldn't change this fast
+        compute_client.get_info()
+        assert token == compute_client._jwtoken
+
+        # should change if we wait a bit
+        sleep(settings.JWT_EXPIRE_SECONDS + 2)
+        compute_client.get_info()
+        assert token != compute_client._jwtoken
+
+    def test_api_check(
+        self,
+        n4js_preloaded,
+        compute_manager_client: client.AlchemiscaleComputeManagerClient,
+        uvicorn_server,
+    ):
+        compute_manager_client._api_check()
+
+    def test_registration(
+        self,
+        n4js_preloaded,
+        compute_manager_client: client.AlchemiscaleComputeManagerClient,
+    ):
+        compute_manager_id = ComputeManagerID.from_manager_id("testmanager")
+        returned_id = compute_manager_client.register(compute_manager_id)
+
+        assert compute_manager_id == returned_id
+
+    def test_deregistration(
+        self,
+        n4js_preloaded,
+        compute_manager_client: client.AlchemiscaleComputeManagerClient,
+    ):
+        compute_manager_id = ComputeManagerID.from_manager_id("testmanager")
+        compute_manager_client.register(compute_manager_id)
+        returned_id = compute_manager_client.deregister(compute_manager_id)
+        assert compute_manager_id == returned_id
+
+    def test_get_instruction(
+        self,
+        n4js_preloaded,
+        compute_manager_client: client.AlchemiscaleComputeManagerClient,
+    ):
+        compute_manager_id = ComputeManagerID.from_manager_id("testmanager")
+        compute_manager_client.register(compute_manager_id)
+        instruction, payload = compute_manager_client.get_instruction(
+            compute_manager_id
+        )
+
+        assert instruction == ComputeManagerInstruction.OK, (instruction, payload)
+        assert payload == {"compute_service_ids": [], "num_registered": 0}
+
+    def test_update_status(
+        self,
+        n4js_preloaded,
+        computer_manager_client: client.AlchemiscaleComputeManagerClient,
+    ):
+        raise NotImplementedError