Test task completion and result push

Author: ianmkenney

alchemiscale/compute/service.py

@@ -698,6 +698,9 @@ def __init__(self, settings: ComputeServiceSettings):
         # monitors.
         self._child_env = dict()  # mods to child process env
         self._task_data = dict()
+        self._executor_stack = ExecutorStack(self.settings.stack_size)
+        self.tasks_claimed = 0
+        self.tasks_finished = 0
         self._initialize_dag_tree()
         self._initialize_resource_monitors()
 
@@ -723,6 +726,7 @@ def __init__(self, settings: ComputeServiceSettings):
         self.compute_service_id = ComputeServiceID.new_from_name(self.name)
 
         self.int_sleep = InterruptableSleep()
+        self._stop = False
         self._initialize_logger()
 
     def _initialize_logger(self):
@@ -792,6 +796,7 @@ def consume_terminated_tasks(self):
             task_scoped_key, _ = terminating_nodes
             pdr = self._consume_results(task_scoped_key)
             self.push_result(task_scoped_key, pdr)
+            self.tasks_finished = 1 + self.tasks_finished
 
     def process_results(self):
         failed_tasks = set()
@@ -817,6 +822,7 @@ def process_results(self):
         for failed_task in failed_tasks:
             pdr = self._consume_results(failed_task)
             self.push_result(task_scoped_key, pdr)
+            self.tasks_finished = 1 + tasks_finished
 
     def stop(self):
         if self.has_tasks():
@@ -909,9 +915,17 @@ def cycle(self, max_tasks, max_time) -> bool:
             max_less_claimed = max_tasks - self.tasks_claimed
             n_claim = min(n_claim, max_less_claimed)
         tasks = self.claim_tasks(count=n_claim)
+
+        if tasks is None:
+            self.logger.info("No tasks claimed. Compute API denied request.")
+            time.sleep(self.deep_sleep_interval)
+            return
+
+        self.logger.info("Claimed %d tasks", len([t for t in tasks if t is not None]))
+
         for task in tasks:
             if task is not None:
-                self.add_task(*task)
+                self.add_task(task)
                 self.tasks_claimed = 1 + self.tasks_claimed
 
         for key in filter(lambda k: k[1] not in ("TERM", "ROOT"), self.next()):
@@ -925,13 +939,14 @@ def cycle(self, max_tasks, max_time) -> bool:
 
             try:
                 self._executor_stack.push(
-                    key, context, inputs, self.n_retries, env=self._child_env
+                    key, context, inputs, self.settings.n_retries, env=self._child_env
                 )
                 self.logger.info(f"Pushing {key[1]} to the execution stack")
                 break
             except JailedKeyError:
                 continue
 
+        time.sleep(self.sleep_interval)
         return True
 
     def available_units(self) -> set[NodeKey]:

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

@@ -11,8 +11,8 @@
 from alchemiscale.storage.statestore import Neo4jStore
 from alchemiscale.storage.objectstore import S3ObjectStore
 from alchemiscale.compute.client import AlchemiscaleComputeClientError
-from alchemiscale.compute.service import SynchronousComputeService
-from alchemiscale.compute.settings import ComputeServiceSettings
+from alchemiscale.compute.service import AsynchronousComputeService, SynchronousComputeService
+from alchemiscale.compute.settings import AsynchronousComputeServiceSettings, ComputeServiceSettings
 
 
 class TestSynchronousComputeService:
@@ -255,3 +255,81 @@ def test_missing_compute_manager(self, n4js_preloaded, service):
             match="Could not find ComputeManagerRegistration",
         ):
             service._register()
+
+class TestAsynchronousComputeService:
+
+    @pytest.fixture
+    def service(self, n4js_preloaded, compute_client, tmpdir):
+        with tmpdir.as_cwd():
+            return AsynchronousComputeService(
+                AsynchronousComputeServiceSettings(
+                    gpu_monitor_enabled=False,
+                    api_url=compute_client.api_url,
+                    identifier=compute_client.identifier,
+                    key=compute_client.key,
+                    name="test_compute_service",
+                    shared_basedir=Path("shared").absolute(),
+                    scratch_basedir=Path("scratch").absolute(),
+                    heartbeat_interval=1,
+                    sleep_interval=1,
+                    deep_sleep_interval=1,
+                )
+            )
+
+    def test_heartbeat(self, n4js_preloaded, service):
+        n4js: Neo4jStore = n4js_preloaded
+
+        # register service; normally happens on service start, but needed
+        # for heartbeats
+        service._register()
+
+        # start up heartbeat thread
+        heartbeat_thread = threading.Thread(target=service.heartbeat, daemon=True)
+        heartbeat_thread.start()
+
+        # give time for a heartbeat
+        time.sleep(2)
+
+        q = f"""
+        match (csreg:ComputeServiceRegistration {{identifier: '{service.compute_service_id}'}})
+        return csreg
+        """
+        csreg = n4js.execute_query(q).records[0]["csreg"]
+
+        assert csreg["registered"] < csreg["heartbeat"]
+
+        # stop the service; should trigger heartbeat to stop
+        service.stop()
+        time.sleep(2)
+        assert not heartbeat_thread.is_alive()
+
+    def test_cycle(self, n4js_preloaded, s3os_server_fresh, service):
+        service._register()
+
+        q = """
+        match (pdr:ProtocolDAGResultRef)
+        return pdr
+        """
+
+        # preconditions
+        protocoldagresultref = n4js_preloaded.execute_query(q)
+        assert not protocoldagresultref.records
+
+        # note that non-None max_time will fail due to _start_time
+        # never being set
+        while service.cycle(max_tasks=1, max_time=None):
+            pass
+
+        # postconditions
+        protocoldagresultref = n4js_preloaded.execute_query(q)
+        assert protocoldagresultref.records
+        assert protocoldagresultref.records[0]["pdr"]["ok"] is True
+
+        q = """
+        match (t:Task {status: 'complete'})
+        return t
+        """
+
+        results = n4js_preloaded.execute_query(q)
+
+        assert results.records