From c71c8f042f6dac195782d7792b63e8a0b489e768 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Mon, 16 Mar 2026 11:51:57 -0400
Subject: [PATCH 01/24] Basic layout of new service

---
 alchemiscale/compute/service.py | 190 ++++++++++++++++++++++++++++++--
 1 file changed, 182 insertions(+), 8 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 3da98439..2cd30bfa 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -12,6 +12,7 @@
 import threading
 from pathlib import Path
 import shutil
+from typing import Any
 
 from gufe import Transformation
 from gufe.protocols.protocoldag import execute_DAG, ProtocolDAG, ProtocolDAGResult
@@ -401,6 +402,86 @@ def stop(self):
         self.int_sleep.interrupt()
         self._stop = True
 
+from multiprocessing import Process, Queue
+
+type NodeKey = str
+
+class JailedKeyError(Exception):
+    pass
+
+class Executor(Process):
+
+    key: NodeKey
+    queue: Queue
+
+    def __init__(self, key, queue):
+        super().__init__()
+        self._key = key
+        self._queue = queue
+
+    def run(self):
+        result = self.execute_unit()
+        self.queue.put(result)
+
+    @property
+    def key(self) -> NodeKey:
+        return self._key
+
+    @property
+    def queue(self) -> Queue:
+        return self._queue
+
+    @classmethod
+    def from_key(cls, key: NodeKey, queue: Queue):
+        return cls(key=key, queue=queue)
+
+class ExecutorStack:
+
+    stack_size: int
+    stack: list[Executor]
+    jail: dict[NodeKey, set[NodeKey]]
+    queue: Queue
+
+    def __init__(self, stack_size: int):
+        self._stack = []
+        self._stack_size = stack_size
+        self._jail = {}
+        self._queue = Queue()
+
+    @property
+    def stack(self) -> list[Executor]:
+        return self._stack
+
+    @property
+    def stack_size(self) -> int:
+        return self._stack_size
+
+    @property
+    def jail(self) -> dict[NodeKey, set[NodeKey]]:
+        return self._jail
+
+    @property
+    def queue(self) -> Queue:
+        return self._queue
+
+    def push(self, value: NodeKey):
+        if value in self._jail.keys():
+            raise JailedKeyError(value)
+
+        executor = Executor.from_key(value)
+        self._stack.append(executor)
+        self._stack[-1].start()
+
+    def pop(self):
+        if self._stack_size == 0:
+            raise IndexError("pop from empty stack")
+        popped_executor = self._stack.pop()
+        for key in self._jail.keys():
+            self._jail[key] -= popped_executor.key
+            if not self._jail[key]:
+                self._jail.pop(key)
+
+        return popped_executor
 
 class AsynchronousComputeService(SynchronousComputeService):
     """Asynchronous compute service.
@@ -410,21 +491,114 @@ class AsynchronousComputeService(SynchronousComputeService):
 
     """
 
-    def __init__(self, api_url):
+    _dag_tree: nx.DiGraph
+    _executor_stack: ExecutorStack
+
+    def __init__(self, settings: AsynchronousComputeServiceSettings):
+
+        self._dag_tree = nx.DiGraph()
+        self._dag_tree.add_node((None, "ROOT"))
+
+        self.settings = settings
+
+        self.api_url = self.settings.api_url
+        self.name = self.settings.name
+        self.sleep_interval = self.settings.sleep_interval
+        self.heartbeat_interval = self.settings.heartbeat_interval
+        self.claim_limit = self.settings.claim_limit
+
         self.scheduler = sched.scheduler(time.monotonic, time.sleep)
-        # self.loop = asyncio.get_event_loop()
+
+        self.client = AlchemiscaleComputeClient(
+            self.settings.api_url,
+            self.settings.identifier,
+            self.settings.key,
+            cache_directory=self.settings.client_cache_directory,
+            cache_size_limit=self.settings.client_cache_size_limit,
+            use_local_cache=self.settings.client_use_local_cache,
+            max_retries=self.settings.client_max_retries,
+            retry_base_seconds=self.settings.client_retry_base_seconds,
+            retry_max_seconds=self.settings.client_retry_max_seconds,
+            verify=self.settings.client_verify,
+        )
 
         self._stop = False
 
-    def get_new_tasks(self): ...
+        self.scopes = self.settings.scopes or [Scope()]
+        self.shared_basedir = Path(self.settings.shared_basedir).absolute()
+        self.shared_basedir.mkdir(exist_ok=True)
+        self.keep_shared = self.settings.keep_shared
+
+        self.scratch_basedir = Path(self.settings.scratch_basedir).absolute()
+        self.scratch_basedir.mkdir(exist_ok=True)
+        self.keep_scratch = self.settings.keep_scratch
 
-    def start(self):
-        """Start the service; will keep going until told to stop."""
+        self.compute_service_id = ComputeServiceID.new_from_name(self.name)
         self._stop = False
 
-        while True:
-            if self._stop:
-                return
+    async def async_cycle(self, max_tasks, max_time):
+
+        # (ProtocolDAG, dwindling_graph, results)
+        for task in tasks:
+            raise NotImplementedError
+
+    def blocked_units(self):
+        raise NotImplementedError
+
+    def running_units(self):
+        raise NotImplementedError
+
+    def available_units(self) -> set[tuple[string, Any]]:
+        return {node for node, degree in self._dag_tree.out_degree() if degree == 0}
+
+    def check_completed(self) -> list[str]:
+        completed = []
+        for node in self.available:
+            task_id, base_node = node
+            if base_node == "TERM":
+                completed.append(task_id)
+        return completed
+
+    def add_task(self, task):
+        """Add a ``Task`` to the ``AsynchronousComputeService`` internal DAG."""
+        task_key = str(task.key)
+
+        def node_transformation(node):
+            nonlocal task_key
+            return (task_key, node)
+
+        tagged_dag = nx.DiGraph()
+        terminating = node_transformation("TERM")
+        tagged_dag.add_node(terminating)
+
+        for child, parent in task.dag.edges:
+            tagged_child = node_transformation(child)
+            tagged_parent = node_transformation(parent)
+            tagged_dag.add_edge(tagged_child, tagged_parent)
+
+        for node, in_degree in task.dag.in_degree:
+            if in_degree == 0:
+                tagged_dag.add_edge(terminating, node_transformation(node))
+
+        self._dag_tree.add_edges_from(tagged_dag.edges)
+        self._dag_tree.add_edge((None, "ROOT"), terminating)
+
+        # TODO create necessary directories for contexts
+        raise NotImplementedError
+
+    def remove_task(self, task_key):
+        # TODO: check executor stack
+        # avoid deleting the root node
+        if task_key is None:
+            raise ValueError()
+
+        for node in self._dag_tree.nodes:
+            key, _ = node
+            if key == task_key:
+                self._dag_tree.remove_node(node)
+
+        # TODO: remove directories
+        raise NotImplementedError
 
     def stop(self):
         self._stop = True

From 8f07984b733111b307da734b49123b09f5d1c761 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Tue, 17 Mar 2026 16:04:38 -0400
Subject: [PATCH 02/24] Add queue, lock, and context to Executor

---
 alchemiscale/compute/service.py | 130 ++++++++++++++++++++++++++------
 1 file changed, 105 insertions(+), 25 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index cd9566a4..c3e035cc 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -402,9 +402,10 @@ def stop(self):
         self.int_sleep.interrupt()
         self._stop = True
 
-from multiprocessing import Process, Queue
+from multiprocessing import Process, Queue, Lock
 
-type NodeKey = str
+type TaskKey = GufeKey
+type NodeKey = (TaskKey, ProtocolUnit | None) # None covers terminating node condition
 
 class JailedKeyError(Exception):
     pass
@@ -413,15 +414,30 @@ class Executor(Process):
 
     key: NodeKey
     queue: Queue
+    lock: Lock
+    context: Context
+    unit: ProtocolUnit
 
-    def __init__(self, key, queue):
+    def __init__(self, key, queue, lock, context, unit):
         super().__init__()
         self._key = key
         self._queue = queue
+        self._lock = lock
+        self._context = context
+        self._unit = unit
 
     def run(self):
         result = self.execute_unit()
-        self.queue.put(result)
+        self.put_result(result)
+        self.cleanup()
+
+    @property
+    def context(self) -> Context:
+        return self._context
+
+    @property
+    def unit(self) -> ProtocolUnit:
+        return self._unit
 
     @property
     def key(self) -> NodeKey:
@@ -431,9 +447,27 @@ def key(self) -> NodeKey:
     def queue(self) -> Queue:
         return self._queue
 
+    @property
+    def lock(self) -> Lock:
+        return self._lock
+
     @classmethod
-    def from_key(cls, key: NodeKey, queue: Queue):
-        return cls(key=key, queue=queue)
+    def from_key(cls, key: NodeKey, queue: Queue, lock: Lock, context: Context):
+        return cls(key=key, queue=queue, lock=lock, context=context)
+
+    def put_result(self, result: ProtocolUnitResult):
+        """Acquire lock, push result to queue, release lock.
+        """
+        with self.lock:
+            self.queue.put(result)
+
+    def execute_unit(self) -> ProtocolUnitResult:
+        # this method assumes the context is in place and will be removed correctly
+        raise NotImplementedError
+
+    def cleanup(self):
+        # clean up depending on context and execution settings
+        raise NotImplementedError
 
 class ExecutorStack:
 
@@ -441,12 +475,14 @@ class ExecutorStack:
     stack: list[Executor]
     jail: dict[NodeKey, set[NodeKey]]
     queue: Queue
+    lock: Lock
 
     def __init__(self, stack_size: int):
         self._stack = []
         self._stack_size = stack_size
         self._jail = {}
         self._queue = Queue()
+        self._lock = Lock()
 
     @property
     def stack(self) -> list[Executor]:
@@ -464,15 +500,26 @@ def jail(self) -> dict[NodeKey, set[NodeKey]]:
     def queue(self) -> Queue:
         return self._queue
 
-    def push(self, value: NodeKey):
-        if value in self._jail.keys():
-            raise JailedKeyError(value)
+    def terminate_all(self, force=False):
+        if force:
+            for proc in self.stack:
+            proc.terminate()
+
+        with self.lock:
+            for proc in self.stack:
+                proc.terminate()
 
-        executor = Executor.from_key(value)
-        self._stack.append(executor)
-        self._stack[-1].start()
+    def push(self, node: NodeKey, context: Context):
+        with self.lock:
+            if node in self._jail.keys():
+                raise JailedKeyError(node)
+
+            executor = Executor.from_key(node, queue, lock, context)
+            self._stack.append(executor)
+            self._stack[-1].start()
 
     def pop(self):
+        """Remove last process in the stack. This also clears the node from the jail."""
         if self._stack_size == 0:
             raise IndexError("pop from empty stack")
         popped_executor = self._stack.pop()
@@ -483,6 +530,31 @@ def pop(self):
 
         return popped_executor
 
+    def _get_by_pid(self, pid) -> Executor | None:
+        """Get an executor by its PID.
+        """
+        for proc in self.stack:
+            if proc.pid == pid:
+                return proc
+        return None
+
+    def _get_statuses(self) -> tuple[set[Executor], set[Executor]]:
+        running = set()
+        terminated =set()
+        for proc in self.stack:
+            if proc.is_alive():
+                running.add(proc)
+            else:
+                terminated.add(proc)
+        return running, terminated
+
+from dataclasses import dataclass
+
+@dataclass
+class TaskData:
+    results: dict[GufeKey, ProtocolUnitResult]
+    context: Context
+
 class AsynchronousComputeService(SynchronousComputeService):
     """Asynchronous compute service.
 
@@ -493,11 +565,15 @@ class AsynchronousComputeService(SynchronousComputeService):
 
     _dag_tree: nx.DiGraph
     _executor_stack: ExecutorStack
+    _task_data: dict[TaskKey, TaskData]
 
     def __init__(self, settings: AsynchronousComputeServiceSettings):
 
         self._dag_tree = nx.DiGraph()
-        self._dag_tree.add_node((None, "ROOT"))
+        root_node: NodeKey = (None, "ROOT")
+        self._dag_tree.add_node(root_node)
+
+        self._task_data = dict()
 
         self.settings = settings
 
@@ -542,14 +618,14 @@ async def async_cycle(self, max_tasks, max_time):
         for task in tasks:
             raise NotImplementedError
 
-    def blocked_units(self):
-        raise NotImplementedError
+    def available_units(self) -> set[NodeKey]:
 
-    def running_units(self):
-        raise NotImplementedError
+        available = set()
+        for node, degree in self._dag_tree.out_degree():
+            if degree == 0:
+                available.add(node)
 
-    def available_units(self) -> set[tuple[string, Any]]:
-        return {node for node, degree in self._dag_tree.out_degree() if degree == 0}
+        return available
 
     def check_completed(self) -> list[str]:
         completed = []
@@ -559,11 +635,13 @@ def check_completed(self) -> list[str]:
                 completed.append(task_id)
         return completed
 
-    def add_task(self, task):
+    def add_task(self, task: ScopedKey):
         """Add a ``Task`` to the ``AsynchronousComputeService`` internal DAG."""
-        task_key = str(task.key)
+        task_key: TaskKey = str(task.gufe_key)
+
+        dag, _, _ = self.task_to_protocoldag(task)
 
-        def node_transformation(node):
+        def node_transformation(node: ProtocolUnit) -> (TaskKey, ProtocolUnit):
             nonlocal task_key
             return (task_key, node)
 
@@ -571,7 +649,7 @@ def node_transformation(node):
         terminating = node_transformation("TERM")
         tagged_dag.add_node(terminating)
 
-        for child, parent in task.dag.edges:
+        for child, parent in dag.graph.edges:
             tagged_child = node_transformation(child)
             tagged_parent = node_transformation(parent)
             tagged_dag.add_edge(tagged_child, tagged_parent)
@@ -583,7 +661,9 @@ def node_transformation(node):
         self._dag_tree.add_edges_from(tagged_dag.edges)
         self._dag_tree.add_edge((None, "ROOT"), terminating)
 
-        # TODO create necessary directories for contexts
+        context = Context()
+        self._task_data[task_key] = TaskData(results={}, context=context)
+        # TODO create necessary directories for contexts: shared and scratch
         raise NotImplementedError
 
     def remove_task(self, task_key):
@@ -597,7 +677,7 @@ def remove_task(self, task_key):
             if key == task_key:
                 self._dag_tree.remove_node(node)
 
-        # TODO: remove directories
+        # TODO: remove context directories: shared and scratch
         raise NotImplementedError
 
     def stop(self):

From ddf961570c8a9098aea217833db6576a1456166f Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Thu, 19 Mar 2026 15:38:56 -0400
Subject: [PATCH 03/24] Demonstration of Executor process behavior

---
 alchemiscale/compute/service.py | 68 ++++++++++++++++--------
 local_testing.py                | 91 +++++++++++++++++++++++++++++++++
 network.py                      | 56 ++++++++++++++++++++
 3 files changed, 193 insertions(+), 22 deletions(-)
 create mode 100644 local_testing.py
 create mode 100644 network.py

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index c3e035cc..594df953 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -16,6 +16,9 @@
 
 from gufe import Transformation
 from gufe.protocols.protocoldag import execute_DAG, ProtocolDAG, ProtocolDAGResult
+from gufe.protocols.protocolunit import Context, ProtocolUnitResult, ProtocolUnit
+from gufe.tokenization import GufeKey
+import networkx as nx
 
 from .client import AlchemiscaleComputeClient
 from .settings import ComputeServiceSettings
@@ -402,34 +405,36 @@ def stop(self):
         self.int_sleep.interrupt()
         self._stop = True
 
+
 from multiprocessing import Process, Queue, Lock
 
 type TaskKey = GufeKey
-type NodeKey = (TaskKey, ProtocolUnit | None) # None covers terminating node condition
+type NodeKey = (TaskKey, ProtocolUnit | None)  # None covers terminating node condition
+
 
 class JailedKeyError(Exception):
     pass
 
+
 class Executor(Process):
 
     key: NodeKey
     queue: Queue
     lock: Lock
     context: Context
-    unit: ProtocolUnit
+    inputs: dict
 
-    def __init__(self, key, queue, lock, context, unit):
+    def __init__(self, key, queue, lock, context, inputs):
         super().__init__()
         self._key = key
         self._queue = queue
         self._lock = lock
         self._context = context
-        self._unit = unit
+        self._inputs = inputs
 
     def run(self):
         result = self.execute_unit()
         self.put_result(result)
-        self.cleanup()
 
     @property
     def context(self) -> Context:
@@ -437,7 +442,7 @@ def context(self) -> Context:
 
     @property
     def unit(self) -> ProtocolUnit:
-        return self._unit
+        return self._key[1]
 
     @property
     def key(self) -> NodeKey:
@@ -451,23 +456,25 @@ def queue(self) -> Queue:
     def lock(self) -> Lock:
         return self._lock
 
+    @property
+    def inputs(self) -> dict:
+        return self._inputs
+
     @classmethod
-    def from_key(cls, key: NodeKey, queue: Queue, lock: Lock, context: Context):
-        return cls(key=key, queue=queue, lock=lock, context=context)
+    def from_key(
+        cls, key: NodeKey, queue: Queue, lock: Lock, context: Context, inputs: dict
+    ):
+        return cls(key=key, queue=queue, lock=lock, context=context, inputs=inputs)
 
     def put_result(self, result: ProtocolUnitResult):
-        """Acquire lock, push result to queue, release lock.
-        """
+        """Acquire lock, push key and result into the queue, release lock."""
         with self.lock:
-            self.queue.put(result)
+            self.queue.put((self._key, result))
 
     def execute_unit(self) -> ProtocolUnitResult:
         # this method assumes the context is in place and will be removed correctly
-        raise NotImplementedError
+        return self.unit.execute(context=self.context, **self._inputs)
 
-    def cleanup(self):
-        # clean up depending on context and execution settings
-        raise NotImplementedError
 
 class ExecutorStack:
 
@@ -496,6 +503,10 @@ def stack_size(self) -> int:
     def jail(self) -> dict[NodeKey, set[NodeKey]]:
         return self._jail
 
+    @property
+    def lock(self) -> Lock:
+        return self._lock
+
     @property
     def queue(self) -> Queue:
         return self._queue
@@ -503,18 +514,18 @@ def queue(self) -> Queue:
     def terminate_all(self, force=False):
         if force:
             for proc in self.stack:
-            proc.terminate()
+                proc.terminate()
 
         with self.lock:
             for proc in self.stack:
                 proc.terminate()
 
-    def push(self, node: NodeKey, context: Context):
+    def push(self, node: NodeKey, context: Context, inputs: dict):
         with self.lock:
             if node in self._jail.keys():
                 raise JailedKeyError(node)
 
-            executor = Executor.from_key(node, queue, lock, context)
+            executor = Executor.from_key(node, self.queue, self.lock, context, inputs)
             self._stack.append(executor)
             self._stack[-1].start()
 
@@ -531,8 +542,7 @@ def pop(self):
         return popped_executor
 
     def _get_by_pid(self, pid) -> Executor | None:
-        """Get an executor by its PID.
-        """
+        """Get an executor by its PID."""
         for proc in self.stack:
             if proc.pid == pid:
                 return proc
@@ -540,7 +550,7 @@ def _get_by_pid(self, pid) -> Executor | None:
 
     def _get_statuses(self) -> tuple[set[Executor], set[Executor]]:
         running = set()
-        terminated =set()
+        terminated = set()
         for proc in self.stack:
             if proc.is_alive():
                 running.add(proc)
@@ -548,13 +558,27 @@ def _get_statuses(self) -> tuple[set[Executor], set[Executor]]:
                 terminated.add(proc)
         return running, terminated
 
+
 from dataclasses import dataclass
 
+
 @dataclass
 class TaskData:
+    protocol_dag: ProtocolDAG
     results: dict[GufeKey, ProtocolUnitResult]
     context: Context
 
+    # TODO failures?
+    def to_ProtocolDAGResult(self) -> ProtocolDAGResult:
+        return ProtocolDAGResult(
+            name=self.protocol_dag.name,
+            protocol_units=self.protocol_dag.protocol_units,
+            protocol_unit_results=list(self.results.values()),
+            transformation_key=self.protocol_dag.transformation_key,
+            extends_key=self.protocol_dag.extends_key,
+        )
+
+
 class AsynchronousComputeService(SynchronousComputeService):
     """Asynchronous compute service.
 
@@ -567,7 +591,7 @@ class AsynchronousComputeService(SynchronousComputeService):
     _executor_stack: ExecutorStack
     _task_data: dict[TaskKey, TaskData]
 
-    def __init__(self, settings: AsynchronousComputeServiceSettings):
+    def __init__(self, settings: ComputeServiceSettings):
 
         self._dag_tree = nx.DiGraph()
         root_node: NodeKey = (None, "ROOT")
diff --git a/local_testing.py b/local_testing.py
new file mode 100644
index 00000000..d7d71351
--- /dev/null
+++ b/local_testing.py
@@ -0,0 +1,91 @@
+# Local Variables:
+# compile-command: "./env/bin/python local_testing.py"
+# python-shell-interpreter: "./env/bin/python"
+# End:
+
+
+from pathlib import Path
+import shutil
+
+
+
+from alchemiscale.compute.service import (
+    Executor,
+    ExecutorStack,
+    NodeKey,
+    TaskKey,
+    TaskData,
+)
+from gufe import AlchemicalNetwork
+from gufe.protocols.protocolunit import Context
+from gufe.protocols.protocoldag import _pu_to_pur, ProtocolDAGResult
+from gufe.tests.test_protocol import DummyProtocol
+from gufe.tokenization import GufeKey
+
+SCRATCH_DIR = Path("./acs_testing/scratch")
+SHARED_DIR = Path("./acs_testing/shared")
+
+SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
+SHARED_DIR.mkdir(parents=True, exist_ok=True)
+
+STACKSIZE = 2
+
+def tyk2():
+    try:
+        return AlchemicalNetwork.from_json(file="network.json")
+    except FileNotFoundError:
+        from network import network_tyk2
+        _tyk2 = network_tyk2()
+        _tyk2.to_json(file="network.json")
+        return _tyk2
+
+if __name__ == "__main__":
+
+    print("Creating network")
+    tyk2 = tyk2()
+    protocol_dag = list(tyk2.edges)[0].create()
+    # say that task_key represents the above protocol dag
+    task_key = GufeKey("FakeKey-123456")
+
+    print("Creating task directories")
+
+    task_scratch_dir = SCRATCH_DIR / f"{task_key}"
+    task_shared_dir = SHARED_DIR / f"{task_key}"
+
+    task_scratch_dir.mkdir()
+    task_shared_dir.mkdir()
+    task_context = Context(scratch=task_scratch_dir, shared=task_shared_dir)
+
+    task_data = {task_key: TaskData(results={},
+                                    context=task_context,
+                                    protocol_dag=protocol_dag,
+    )}
+    exec_stack = ExecutorStack(STACKSIZE)
+
+    print("Starting unit loop")
+    _task_data = task_data[task_key]
+    for unit in _task_data.protocol_dag.protocol_units:
+        inputs = _pu_to_pur(unit.inputs, _task_data.results)
+        key: NodeKey = (task_key, unit)
+
+        unit_scratch_dir = _task_data.context.scratch / f"{str(unit.key)}"
+        unit_shared_dir = _task_data.context.shared / f"{str(unit.key)}"
+
+        unit_scratch_dir.mkdir()
+        unit_shared_dir.mkdir()
+
+        context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
+        exec_stack.push(key, context, inputs)
+        (task_key, pu), res = exec_stack.queue.get()
+        if not res.ok():
+            raise RuntimeError
+        _task_data.results[pu.key] = res
+
+        # clean up scratch (later stderr, stdout)
+        shutil.rmtree(unit_scratch_dir)
+
+    shutil.rmtree(task_shared_dir)
+    shutil.rmtree(task_scratch_dir)
+
+    pdr = _task_data.to_ProtocolDAGResult()
+    print(pdr, pdr.ok())
diff --git a/network.py b/network.py
new file mode 100644
index 00000000..5d76fcee
--- /dev/null
+++ b/network.py
@@ -0,0 +1,56 @@
+from openfe_benchmarks import tyk2
+
+from gufe import ChemicalSystem, Transformation, NonTransformation, AlchemicalNetwork
+from gufe.tests.test_protocol import DummyProtocol
+
+class DummyProtocolA(DummyProtocol):
+    ...
+
+class DummyProtocolB(DummyProtocol):
+    ...
+
+def network_tyk2():
+    tyk2s = tyk2.get_system()
+
+    solvated = {
+        ligand.name: ChemicalSystem(
+            components={"ligand": ligand, "solvent": tyk2s.solvent_component},
+            name=f"{ligand.name}_water",
+        )
+        for ligand in tyk2s.ligand_components
+    }
+    complexes = {
+        ligand.name: ChemicalSystem(
+            components={
+                "ligand": ligand,
+                "solvent": tyk2s.solvent_component,
+                "protein": tyk2s.protein_component,
+            },
+            name=f"{ligand.name}_complex",
+        )
+        for ligand in tyk2s.ligand_components
+    }
+
+    complex_network = [
+        Transformation(
+            stateA=complexes[edge[0]],
+            stateB=complexes[edge[1]],
+            protocol=DummyProtocolA(settings=DummyProtocolA.default_settings()),
+            name=f"{edge[0]}_to_{edge[1]}_complex",
+        )
+        for edge in tyk2s.connections
+    ]
+    solvent_network = [
+        Transformation(
+            stateA=solvated[edge[0]],
+            stateB=solvated[edge[1]],
+            protocol=DummyProtocolB(settings=DummyProtocolB.default_settings()),
+            name=f"{edge[0]}_to_{edge[1]}_solvent",
+        )
+        for edge in tyk2s.connections
+    ]
+
+    return AlchemicalNetwork(
+        edges=(solvent_network + complex_network),
+        name="tyk2_relative_benchmark",
+    )

From d2dc7321a21668c36234e883cf28b8708917a692 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Thu, 19 Mar 2026 18:19:56 -0400
Subject: [PATCH 04/24] Move logic to service

---
 alchemiscale/compute/service.py | 58 ++++++++++++++++++++-------------
 local_testing.py                | 46 +++++++++++++-------------
 service.py                      | 11 +++++++
 3 files changed, 68 insertions(+), 47 deletions(-)
 create mode 100644 service.py

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 594df953..458a3c79 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -408,8 +408,11 @@ def stop(self):
 
 from multiprocessing import Process, Queue, Lock
 
-type TaskKey = GufeKey
-type NodeKey = (TaskKey, ProtocolUnit | None)  # None covers terminating node condition
+type TaskKey = ScopedKey
+type NodeKey = (
+    TaskKey | None,
+    ProtocolUnit | str,
+)  # None covers terminating node condition
 
 
 class JailedKeyError(Exception):
@@ -593,11 +596,8 @@ class AsynchronousComputeService(SynchronousComputeService):
 
     def __init__(self, settings: ComputeServiceSettings):
 
-        self._dag_tree = nx.DiGraph()
-        root_node: NodeKey = (None, "ROOT")
-        self._dag_tree.add_node(root_node)
-
         self._task_data = dict()
+        self._initialize_dag_tree()
 
         self.settings = settings
 
@@ -636,6 +636,11 @@ def __init__(self, settings: ComputeServiceSettings):
         self.compute_service_id = ComputeServiceID.new_from_name(self.name)
         self._stop = False
 
+    def _initialize_dag_tree(self):
+        self._dag_tree = nx.DiGraph()
+        root_node: NodeKey = (None, "ROOT")
+        self._dag_tree.add_node(root_node)
+
     async def async_cycle(self, max_tasks, max_time):
 
         # (ProtocolDAG, dwindling_graph, results)
@@ -659,15 +664,16 @@ def check_completed(self) -> list[str]:
                 completed.append(task_id)
         return completed
 
-    def add_task(self, task: ScopedKey):
-        """Add a ``Task`` to the ``AsynchronousComputeService`` internal DAG."""
-        task_key: TaskKey = str(task.gufe_key)
+    def add_task(self, task_scoped_key: ScopedKey):
+        dag, _, _ = self.task_to_protocoldag(task_scoped_key)
+        self.graft_dag(task_scoped_key, dag)
 
-        dag, _, _ = self.task_to_protocoldag(task)
+    def graft_dag(self, task_scoped_key: ScopedKey, dag):
+        """Add a ``Task`` to the ``AsynchronousComputeService`` internal DAG."""
 
         def node_transformation(node: ProtocolUnit) -> (TaskKey, ProtocolUnit):
-            nonlocal task_key
-            return (task_key, node)
+            nonlocal task_scoped_key
+            return (task_scoped_key, node)
 
         tagged_dag = nx.DiGraph()
         terminating = node_transformation("TERM")
@@ -678,31 +684,37 @@ def node_transformation(node: ProtocolUnit) -> (TaskKey, ProtocolUnit):
             tagged_parent = node_transformation(parent)
             tagged_dag.add_edge(tagged_child, tagged_parent)
 
-        for node, in_degree in task.dag.in_degree:
+        for node, in_degree in dag.graph.in_degree:
             if in_degree == 0:
                 tagged_dag.add_edge(terminating, node_transformation(node))
 
         self._dag_tree.add_edges_from(tagged_dag.edges)
         self._dag_tree.add_edge((None, "ROOT"), terminating)
 
-        context = Context()
-        self._task_data[task_key] = TaskData(results={}, context=context)
-        # TODO create necessary directories for contexts: shared and scratch
-        raise NotImplementedError
+        context = Context(
+            scratch=self.scratch_basedir / str(task_scoped_key),
+            shared=self.shared_basedir / str(task_scoped_key),
+        )
+        context.scratch.mkdir(exist_ok=True)
+        context.shared.mkdir(exist_ok=True)
+        self._task_data[task_scoped_key] = TaskData(
+            results={}, context=context, protocol_dag=dag
+        )
 
-    def remove_task(self, task_key):
+    def remove_task(self, task_scoped_key):
         # TODO: check executor stack
         # avoid deleting the root node
-        if task_key is None:
+        if task_scoped_key is None:
             raise ValueError()
 
-        for node in self._dag_tree.nodes:
+        for node in tuple(self._dag_tree.nodes):
             key, _ = node
-            if key == task_key:
+            if key == task_scoped_key:
                 self._dag_tree.remove_node(node)
 
-        # TODO: remove context directories: shared and scratch
-        raise NotImplementedError
+        context = self._task_data[task_scoped_key].context
+        shutil.rmtree(context.shared)
+        shutil.rmtree(context.scratch)
 
     def stop(self):
         self._stop = True
diff --git a/local_testing.py b/local_testing.py
index d7d71351..74918329 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -7,8 +7,6 @@
 from pathlib import Path
 import shutil
 
-
-
 from alchemiscale.compute.service import (
     Executor,
     ExecutorStack,
@@ -16,12 +14,17 @@
     TaskKey,
     TaskData,
 )
+from alchemiscale.models import ScopedKey
 from gufe import AlchemicalNetwork
 from gufe.protocols.protocolunit import Context
 from gufe.protocols.protocoldag import _pu_to_pur, ProtocolDAGResult
 from gufe.tests.test_protocol import DummyProtocol
 from gufe.tokenization import GufeKey
 
+import networkx as nx
+
+import service
+
 SCRATCH_DIR = Path("./acs_testing/scratch")
 SHARED_DIR = Path("./acs_testing/shared")
 
@@ -30,40 +33,35 @@
 
 STACKSIZE = 2
 
-def tyk2():
+
+def create_tyk2():
     try:
         return AlchemicalNetwork.from_json(file="network.json")
     except FileNotFoundError:
         from network import network_tyk2
+
         _tyk2 = network_tyk2()
         _tyk2.to_json(file="network.json")
         return _tyk2
 
-if __name__ == "__main__":
-
-    print("Creating network")
-    tyk2 = tyk2()
-    protocol_dag = list(tyk2.edges)[0].create()
-    # say that task_key represents the above protocol dag
-    task_key = GufeKey("FakeKey-123456")
 
-    print("Creating task directories")
+if __name__ == "__main__":
 
-    task_scratch_dir = SCRATCH_DIR / f"{task_key}"
-    task_shared_dir = SHARED_DIR / f"{task_key}"
+    mock = service.MockService(SCRATCH_DIR, SHARED_DIR, STACKSIZE)
 
-    task_scratch_dir.mkdir()
-    task_shared_dir.mkdir()
-    task_context = Context(scratch=task_scratch_dir, shared=task_shared_dir)
+    task_key = GufeKey("FakeKey-123456")
+    task_scoped_key = ScopedKey(
+        gufe_key=task_key, org="MockOrg", campaign="MockCampaign", project="MockProject"
+    )
 
-    task_data = {task_key: TaskData(results={},
-                                    context=task_context,
-                                    protocol_dag=protocol_dag,
-    )}
-    exec_stack = ExecutorStack(STACKSIZE)
+    print("Creating network")
+    tyk2 = create_tyk2()
+    protocol_dag = list(tyk2.edges)[0].create()
 
+    mock.graft_dag(task_scoped_key, protocol_dag)
+    exec_stack = mock._executor_stack
     print("Starting unit loop")
-    _task_data = task_data[task_key]
+    _task_data = mock._task_data[task_scoped_key]
     for unit in _task_data.protocol_dag.protocol_units:
         inputs = _pu_to_pur(unit.inputs, _task_data.results)
         key: NodeKey = (task_key, unit)
@@ -76,6 +74,7 @@ def tyk2():
 
         context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
         exec_stack.push(key, context, inputs)
+
         (task_key, pu), res = exec_stack.queue.get()
         if not res.ok():
             raise RuntimeError
@@ -84,8 +83,7 @@ def tyk2():
         # clean up scratch (later stderr, stdout)
         shutil.rmtree(unit_scratch_dir)
 
-    shutil.rmtree(task_shared_dir)
-    shutil.rmtree(task_scratch_dir)
+    mock.remove_task(task_scoped_key)
 
     pdr = _task_data.to_ProtocolDAGResult()
     print(pdr, pdr.ok())
diff --git a/service.py b/service.py
new file mode 100644
index 00000000..03334618
--- /dev/null
+++ b/service.py
@@ -0,0 +1,11 @@
+from alchemiscale.compute.service import AsynchronousComputeService, ExecutorStack
+
+class MockService(AsynchronousComputeService):
+
+    def __init__(self, scratch_basedir, shared_basedir, stack_size):
+        self._initialize_dag_tree()
+        self._task_data = dict()
+        self._executor_stack = ExecutorStack(stack_size)
+
+        self.scratch_basedir = scratch_basedir
+        self.shared_basedir = shared_basedir

From 078f9408556636834c21fc7639043e14e568d09e Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Fri, 20 Mar 2026 17:33:00 -0400
Subject: [PATCH 05/24] Push to stack based on nodes available and stack size

---
 alchemiscale/compute/service.py | 33 ++++++++----
 local_testing.py                | 92 +++++++++++++++++++--------------
 service.py                      |  4 ++
 utils.py                        | 32 ++++++++++++
 4 files changed, 114 insertions(+), 47 deletions(-)
 create mode 100644 utils.py

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 458a3c79..47719fac 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -544,12 +544,23 @@ def pop(self):
 
         return popped_executor
 
-    def _get_by_pid(self, pid) -> Executor | None:
-        """Get an executor by its PID."""
+    def get_result(self) -> tuple[NodeKey, ProtocolUnitResult] | None:
+        if self.queue.qsize():
+            with self.lock:
+                res = self.queue.get()
+                node_key, _ = res
+                self.remove_by_node_key(node_key)
+                return res
+
+    def remove_by_node_key(self, node_key: NodeKey):
+        to_remove = None
         for proc in self.stack:
-            if proc.pid == pid:
-                return proc
-        return None
+            if proc.key == node_key:
+                to_remove = proc
+                break
+
+        if to_remove:
+            self._stack.remove(proc)
 
     def _get_statuses(self) -> tuple[set[Executor], set[Executor]]:
         running = set()
@@ -642,13 +653,11 @@ def _initialize_dag_tree(self):
         self._dag_tree.add_node(root_node)
 
     async def async_cycle(self, max_tasks, max_time):
-
         # (ProtocolDAG, dwindling_graph, results)
         for task in tasks:
             raise NotImplementedError
 
     def available_units(self) -> set[NodeKey]:
-
         available = set()
         for node, degree in self._dag_tree.out_degree():
             if degree == 0:
@@ -656,6 +665,12 @@ def available_units(self) -> set[NodeKey]:
 
         return available
 
+    def next(self) -> set[NodeKey]:
+        running, terminated = self._executor_stack._get_statuses()
+        running = {r.key for r in running}
+        terminated = {t.key for t in terminated}
+        return self.available_units() - (running | terminated)
+
     def check_completed(self) -> list[str]:
         completed = []
         for node in self.available:
@@ -665,8 +680,8 @@ def check_completed(self) -> list[str]:
         return completed
 
     def add_task(self, task_scoped_key: ScopedKey):
-        dag, _, _ = self.task_to_protocoldag(task_scoped_key)
-        self.graft_dag(task_scoped_key, dag)
+        protocol_dag, _, _ = self.task_to_protocoldag(task_scoped_key)
+        self.graft_dag(task_scoped_key, protocol_dag)
 
     def graft_dag(self, task_scoped_key: ScopedKey, dag):
         """Add a ``Task`` to the ``AsynchronousComputeService`` internal DAG."""
diff --git a/local_testing.py b/local_testing.py
index 74918329..ad927ee1 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -24,6 +24,7 @@
 import networkx as nx
 
 import service
+import utils
 
 SCRATCH_DIR = Path("./acs_testing/scratch")
 SHARED_DIR = Path("./acs_testing/shared")
@@ -31,13 +32,14 @@
 SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
 SHARED_DIR.mkdir(parents=True, exist_ok=True)
 
-STACKSIZE = 2
+STACKSIZE = 5
 
 
 def create_tyk2():
     try:
         return AlchemicalNetwork.from_json(file="network.json")
     except FileNotFoundError:
+        print("\tCould not load from file, creating new network")
         from network import network_tyk2
 
         _tyk2 = network_tyk2()
@@ -47,43 +49,57 @@ def create_tyk2():
 
 if __name__ == "__main__":
 
-    mock = service.MockService(SCRATCH_DIR, SHARED_DIR, STACKSIZE)
+    mock_service = service.MockService(SCRATCH_DIR, SHARED_DIR, STACKSIZE)
 
-    task_key = GufeKey("FakeKey-123456")
-    task_scoped_key = ScopedKey(
-        gufe_key=task_key, org="MockOrg", campaign="MockCampaign", project="MockProject"
-    )
-
-    print("Creating network")
-    tyk2 = create_tyk2()
-    protocol_dag = list(tyk2.edges)[0].create()
-
-    mock.graft_dag(task_scoped_key, protocol_dag)
-    exec_stack = mock._executor_stack
-    print("Starting unit loop")
-    _task_data = mock._task_data[task_scoped_key]
-    for unit in _task_data.protocol_dag.protocol_units:
-        inputs = _pu_to_pur(unit.inputs, _task_data.results)
-        key: NodeKey = (task_key, unit)
-
-        unit_scratch_dir = _task_data.context.scratch / f"{str(unit.key)}"
-        unit_shared_dir = _task_data.context.shared / f"{str(unit.key)}"
-
-        unit_scratch_dir.mkdir()
-        unit_shared_dir.mkdir()
+    with utils.timer(wrap=True):
+        print("Creating network")
+        tyk2 = create_tyk2()
 
-        context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
-        exec_stack.push(key, context, inputs)
-
-        (task_key, pu), res = exec_stack.queue.get()
-        if not res.ok():
-            raise RuntimeError
-        _task_data.results[pu.key] = res
-
-        # clean up scratch (later stderr, stdout)
-        shutil.rmtree(unit_scratch_dir)
-
-    mock.remove_task(task_scoped_key)
+    transformations = tuple(tyk2.edges)
+    tasks = tuple(
+        (utils.new_task_scoped_key(), transformation)
+        for transformation in transformations
+    )
 
-    pdr = _task_data.to_ProtocolDAGResult()
-    print(pdr, pdr.ok())
+    for tsk, trans in tasks[:3]:
+        mock_service.add_task(tsk, trans)
+
+    # collect for final inspection
+    pdrs = []
+    while mock_service._task_data:
+        # only submit enough tasks to fill the stack
+        n = mock_service._executor_stack._stack_size - len(
+            mock_service._executor_stack._stack
+        )
+        for key in tuple(mock_service.next())[:n]:
+            tsk, unit = key
+            task_data = mock_service._task_data[tsk]
+
+            # TODO: if we hit the terminal node, clean up, this should
+            # live outside of this loop
+            if unit == "TERM":
+                pdr = task_data.to_ProtocolDAGResult()
+                print(f"Collected output: {pdr}")
+                data = mock_service._task_data.pop(tsk)
+                mock_service._dag_tree.remove_node(key)
+                shutil.rmtree(data.context.scratch)
+                shutil.rmtree(data.context.shared)
+                pdrs.append(pdr)
+                continue
+
+            inputs = _pu_to_pur(unit.inputs, mock_service._task_data[tsk].results)
+            unit_scratch_dir = task_data.context.scratch / f"{str(unit.key)}"
+            unit_shared_dir = task_data.context.shared / f"{str(unit.key)}"
+            unit_scratch_dir.mkdir()
+            unit_shared_dir.mkdir()
+            context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
+            mock_service._executor_stack.push(key, context, inputs)
+
+        # collect results
+        while result := mock_service._executor_stack.get_result():
+            node_key, res = result
+            task_scoped_key, pu = node_key
+            mock_service._task_data[task_scoped_key].results[pu.key] = res
+            mock_service._dag_tree.remove_node(node_key)
+
+    print(pdrs)
diff --git a/service.py b/service.py
index 03334618..ea229511 100644
--- a/service.py
+++ b/service.py
@@ -9,3 +9,7 @@ def __init__(self, scratch_basedir, shared_basedir, stack_size):
 
         self.scratch_basedir = scratch_basedir
         self.shared_basedir = shared_basedir
+
+    def add_task(self, task_scoped_key, transformation):
+        protocol_dag = transformation.create()
+        self.graft_dag(task_scoped_key, protocol_dag)
diff --git a/utils.py b/utils.py
new file mode 100644
index 00000000..791c618c
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,32 @@
+from contextlib import contextmanager
+from time import time
+from uuid import uuid4
+
+from alchemiscale.models import ScopedKey
+from gufe.tokenization import GufeKey
+
+@contextmanager
+def timer(*args, **kwargs):
+
+    wrap = kwargs.get("wrap")
+
+    if wrap:
+        print("="*20)
+
+    start = time()
+    yield
+    elapsed = time() - start
+    if wrap:
+        print("-"*20)
+
+    print(f"Time spent: {elapsed}")
+
+    if wrap:
+        print("="*20)
+
+def new_task_scoped_key():
+    task_key = GufeKey(f"FakeKey-{uuid4().hex}")
+    task_scoped_key = ScopedKey(
+        gufe_key=task_key, org="MockOrg", campaign="MockCampaign", project="MockProject"
+    )
+    return task_scoped_key

From f5781a76e42ba2c6bf9043e0b6020e783734aa40 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Fri, 20 Mar 2026 19:38:59 -0400
Subject: [PATCH 06/24] Collect results earlier

---
 local_testing.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/local_testing.py b/local_testing.py
index ad927ee1..be51c022 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -67,6 +67,14 @@ def create_tyk2():
     # collect for final inspection
     pdrs = []
     while mock_service._task_data:
+
+        # collect results
+        while result := mock_service._executor_stack.get_result():
+            node_key, res = result
+            task_scoped_key, pu = node_key
+            mock_service._task_data[task_scoped_key].results[pu.key] = res
+            mock_service._dag_tree.remove_node(node_key)
+
         # only submit enough tasks to fill the stack
         n = mock_service._executor_stack._stack_size - len(
             mock_service._executor_stack._stack
@@ -95,11 +103,5 @@ def create_tyk2():
             context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
             mock_service._executor_stack.push(key, context, inputs)
 
-        # collect results
-        while result := mock_service._executor_stack.get_result():
-            node_key, res = result
-            task_scoped_key, pu = node_key
-            mock_service._task_data[task_scoped_key].results[pu.key] = res
-            mock_service._dag_tree.remove_node(node_key)
 
     print(pdrs)

From a1bd34c6057f0ef64688f718285d188bb58edb95 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Fri, 20 Mar 2026 19:59:39 -0400
Subject: [PATCH 07/24] Clear terminated DAGs earlier

---
 alchemiscale/compute/service.py | 10 +++-------
 local_testing.py                | 26 ++++++++++++++------------
 2 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 47719fac..9020d326 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -412,7 +412,7 @@ def stop(self):
 type NodeKey = (
     TaskKey | None,
     ProtocolUnit | str,
-)  # None covers terminating node condition
+)  # None covers root node condition
 
 
 class JailedKeyError(Exception):
@@ -671,12 +671,8 @@ def next(self) -> set[NodeKey]:
         terminated = {t.key for t in terminated}
         return self.available_units() - (running | terminated)
 
-    def check_completed(self) -> list[str]:
-        completed = []
-        for node in self.available:
-            task_id, base_node = node
-            if base_node == "TERM":
-                completed.append(task_id)
+    def next_terminating_nodes(self) -> set[NodeKey]:
+        completed = {node for node in self.available_units() if node[1] == "TERM"}
         return completed
 
     def add_task(self, task_scoped_key: ScopedKey):
diff --git a/local_testing.py b/local_testing.py
index be51c022..db52a5cf 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -68,6 +68,16 @@ def create_tyk2():
     pdrs = []
     while mock_service._task_data:
 
+        for completed_node in mock_service.next_terminating_nodes():
+            tsk, _ = completed_node
+            data = mock_service._task_data.pop(tsk)
+            pdr = data.to_ProtocolDAGResult()
+            print(f"Collected output: {pdr}")
+            mock_service._dag_tree.remove_node(completed_node)
+            shutil.rmtree(data.context.scratch)
+            shutil.rmtree(data.context.shared)
+            pdrs.append(pdr)
+
         # collect results
         while result := mock_service._executor_stack.get_result():
             node_key, res = result
@@ -81,20 +91,13 @@ def create_tyk2():
         )
         for key in tuple(mock_service.next())[:n]:
             tsk, unit = key
-            task_data = mock_service._task_data[tsk]
 
-            # TODO: if we hit the terminal node, clean up, this should
-            # live outside of this loop
-            if unit == "TERM":
-                pdr = task_data.to_ProtocolDAGResult()
-                print(f"Collected output: {pdr}")
-                data = mock_service._task_data.pop(tsk)
-                mock_service._dag_tree.remove_node(key)
-                shutil.rmtree(data.context.scratch)
-                shutil.rmtree(data.context.shared)
-                pdrs.append(pdr)
+            # TODO `next` should not return TERM or ROOT nodes
+            if unit in ("TERM", "ROOT"):
                 continue
 
+            task_data = mock_service._task_data[tsk]
+
             inputs = _pu_to_pur(unit.inputs, mock_service._task_data[tsk].results)
             unit_scratch_dir = task_data.context.scratch / f"{str(unit.key)}"
             unit_shared_dir = task_data.context.shared / f"{str(unit.key)}"
@@ -103,5 +106,4 @@ def create_tyk2():
             context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
             mock_service._executor_stack.push(key, context, inputs)
 
-
     print(pdrs)

From 3248bd8341d45a4eab9724746cfc04b61b67304b Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Mon, 23 Mar 2026 18:41:53 -0400
Subject: [PATCH 08/24] Handle failures

---
 alchemiscale/compute/service.py | 86 +++++++++++++++++++++++----------
 local_testing.py                | 43 ++++++++++-------
 network.py                      |  4 +-
 3 files changed, 88 insertions(+), 45 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 9020d326..44528e65 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -11,12 +11,13 @@
 from uuid import uuid4
 import threading
 from pathlib import Path
+import queue
 import shutil
 from typing import Any
 
 from gufe import Transformation
 from gufe.protocols.protocoldag import execute_DAG, ProtocolDAG, ProtocolDAGResult
-from gufe.protocols.protocolunit import Context, ProtocolUnitResult, ProtocolUnit
+from gufe.protocols.protocolunit import Context, ProtocolUnitFailure, ProtocolUnitResult, ProtocolUnit
 from gufe.tokenization import GufeKey
 import networkx as nx
 
@@ -409,10 +410,10 @@ def stop(self):
 from multiprocessing import Process, Queue, Lock
 
 type TaskKey = ScopedKey
-type NodeKey = (
+type NodeKey = tuple[
     TaskKey | None,
     ProtocolUnit | str,
-)  # None covers root node condition
+]  # None covers root node condition
 
 
 class JailedKeyError(Exception):
@@ -424,25 +425,34 @@ class Executor(Process):
     key: NodeKey
     queue: Queue
     lock: Lock
-    context: Context
+    unit_context: Context
     inputs: dict
+    n_retries: int
 
-    def __init__(self, key, queue, lock, context, inputs):
+    def __init__(self, key, queue, lock, context, inputs, n_retries):
         super().__init__()
         self._key = key
         self._queue = queue
         self._lock = lock
-        self._context = context
+        self._unit_context = context
         self._inputs = inputs
+        assert n_retries >= 0
+        self._n_retries = n_retries
 
     def run(self):
-        result = self.execute_unit()
+        attempt = 0
+        while attempt <= self._n_retries:
+            shared_dir = self._unit_context.shared / str(attempt)
+            scratch_dir = self._unit_context.scratch / str(attempt)
+            attempt_context = Context(shared=shared_dir, scratch=scratch_dir)
+            attempt_context.shared.mkdir()
+            attempt_context.scratch.mkdir()
+            result = self.execute_unit(attempt_context)
+            if result.ok():
+                break
+            attempt = attempt + 1
         self.put_result(result)
 
-    @property
-    def context(self) -> Context:
-        return self._context
-
     @property
     def unit(self) -> ProtocolUnit:
         return self._key[1]
@@ -465,18 +475,18 @@ def inputs(self) -> dict:
 
     @classmethod
     def from_key(
-        cls, key: NodeKey, queue: Queue, lock: Lock, context: Context, inputs: dict
+            cls, key: NodeKey, queue: Queue, lock: Lock, context: Context, inputs: dict, n_retries: int
     ):
-        return cls(key=key, queue=queue, lock=lock, context=context, inputs=inputs)
+        return cls(key=key, queue=queue, lock=lock, context=context, inputs=inputs, n_retries=n_retries)
 
     def put_result(self, result: ProtocolUnitResult):
         """Acquire lock, push key and result into the queue, release lock."""
         with self.lock:
             self.queue.put((self._key, result))
 
-    def execute_unit(self) -> ProtocolUnitResult:
+    def execute_unit(self, context) -> ProtocolUnitResult | ProtocolUnitFailure:
         # this method assumes the context is in place and will be removed correctly
-        return self.unit.execute(context=self.context, **self._inputs)
+        return self.unit.execute(context=context, **self._inputs)
 
 
 class ExecutorStack:
@@ -514,23 +524,38 @@ def lock(self) -> Lock:
     def queue(self) -> Queue:
         return self._queue
 
-    def terminate_all(self, force=False):
-        if force:
+    def terminate_all(self):
+        with self.lock:
             for proc in self.stack:
                 proc.terminate()
 
+    def terminate_task(self, task_key: TaskKey):
         with self.lock:
+            to_remove = set()
             for proc in self.stack:
-                proc.terminate()
-
-    def push(self, node: NodeKey, context: Context, inputs: dict):
+                _task_key, _ = proc.key
+                if task_key == _task_key:
+                    to_remove.add(proc)
+            for proc in to_remove:
+                try:
+                    proc.close()
+                except ValueError:
+                    proc.terminate()
+                self._stack.remove(proc)
+
+    def push(self, node: NodeKey, context: Context, inputs: dict, n_retries, in_process=False):
         with self.lock:
             if node in self._jail.keys():
                 raise JailedKeyError(node)
 
-            executor = Executor.from_key(node, self.queue, self.lock, context, inputs)
+            executor = Executor.from_key(node, self.queue, self.lock, context, inputs, n_retries)
             self._stack.append(executor)
-            self._stack[-1].start()
+            if not in_process:
+                self._stack[-1].start()
+                return
+
+        self._stack[-1].run()
+        self._stack[-1].close()
 
     def pop(self):
         """Remove last process in the stack. This also clears the node from the jail."""
@@ -544,10 +569,15 @@ def pop(self):
 
         return popped_executor
 
-    def get_result(self) -> tuple[NodeKey, ProtocolUnitResult] | None:
+    def get_result(self) -> tuple[NodeKey, ProtocolUnitResult | ProtocolUnitFailure] | None:
         if self.queue.qsize():
             with self.lock:
-                res = self.queue.get()
+                # since qsize is not always reliable, we tentatively
+                # accept there might be results
+                try:
+                    res = self.queue.get_nowait()
+                except queue.Empty:
+                    return None
                 node_key, _ = res
                 self.remove_by_node_key(node_key)
                 return res
@@ -582,7 +612,6 @@ class TaskData:
     results: dict[GufeKey, ProtocolUnitResult]
     context: Context
 
-    # TODO failures?
     def to_ProtocolDAGResult(self) -> ProtocolDAGResult:
         return ProtocolDAGResult(
             name=self.protocol_dag.name,
@@ -727,5 +756,12 @@ def remove_task(self, task_scoped_key):
         shutil.rmtree(context.shared)
         shutil.rmtree(context.scratch)
 
+    def _consume_results(self, task_scoped_key) -> ProtocolDAGResult:
+        self.remove_task(task_scoped_key)
+        data = self._task_data.pop(task_scoped_key)
+        pdr = data.to_ProtocolDAGResult()
+        return pdr
+
+
     def stop(self):
         self._stop = True
diff --git a/local_testing.py b/local_testing.py
index db52a5cf..7bafb0f9 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -16,10 +16,10 @@
 )
 from alchemiscale.models import ScopedKey
 from gufe import AlchemicalNetwork
-from gufe.protocols.protocolunit import Context
-from gufe.protocols.protocoldag import _pu_to_pur, ProtocolDAGResult
-from gufe.tests.test_protocol import DummyProtocol
+from gufe.protocols.protocolunit import Context, ProtocolUnitResult, ProtocolUnitFailure
+from gufe.protocols.protocoldag import _pu_to_pur
 from gufe.tokenization import GufeKey
+from gufe.tests.test_protocol import BrokenProtocol
 
 import networkx as nx
 
@@ -28,12 +28,13 @@
 
 SCRATCH_DIR = Path("./acs_testing/scratch")
 SHARED_DIR = Path("./acs_testing/shared")
+STACKSIZE = 10
+N_RETRIES = 2
+IN_PROCESS = False
 
 SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
 SHARED_DIR.mkdir(parents=True, exist_ok=True)
 
-STACKSIZE = 5
-
 
 def create_tyk2():
     try:
@@ -61,29 +62,35 @@ def create_tyk2():
         for transformation in transformations
     )
 
-    for tsk, trans in tasks[:3]:
+    for tsk, trans in tasks:
         mock_service.add_task(tsk, trans)
 
     # collect for final inspection
     pdrs = []
     while mock_service._task_data:
-
+        # ask for terminating nodes
         for completed_node in mock_service.next_terminating_nodes():
-            tsk, _ = completed_node
-            data = mock_service._task_data.pop(tsk)
-            pdr = data.to_ProtocolDAGResult()
-            print(f"Collected output: {pdr}")
-            mock_service._dag_tree.remove_node(completed_node)
-            shutil.rmtree(data.context.scratch)
-            shutil.rmtree(data.context.shared)
+            task_scoped_key, _ = completed_node
+            pdr = mock_service._consume_results(task_scoped_key)
             pdrs.append(pdr)
 
-        # collect results
+        # collect unit results
+        failed_tasks = set()
         while result := mock_service._executor_stack.get_result():
             node_key, res = result
             task_scoped_key, pu = node_key
             mock_service._task_data[task_scoped_key].results[pu.key] = res
-            mock_service._dag_tree.remove_node(node_key)
+
+            match res:
+                case ProtocolUnitFailure():
+                    mock_service._executor_stack.terminate_task(task_scoped_key)
+                    failed_tasks.add(task_scoped_key)
+                case ProtocolUnitResult():
+                    mock_service._dag_tree.remove_node(node_key)
+
+        for failed_task in failed_tasks:
+            pdr = mock_service._consume_results(failed_task)
+            pdrs.append(pdr)
 
         # only submit enough tasks to fill the stack
         n = mock_service._executor_stack._stack_size - len(
@@ -98,12 +105,12 @@ def create_tyk2():
 
             task_data = mock_service._task_data[tsk]
 
-            inputs = _pu_to_pur(unit.inputs, mock_service._task_data[tsk].results)
+            inputs = _pu_to_pur(unit.inputs, task_data.results)
             unit_scratch_dir = task_data.context.scratch / f"{str(unit.key)}"
             unit_shared_dir = task_data.context.shared / f"{str(unit.key)}"
             unit_scratch_dir.mkdir()
             unit_shared_dir.mkdir()
             context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
-            mock_service._executor_stack.push(key, context, inputs)
+            mock_service._executor_stack.push(key, context, inputs, N_RETRIES, in_process=IN_PROCESS)
 
     print(pdrs)
diff --git a/network.py b/network.py
index 5d76fcee..0768e32a 100644
--- a/network.py
+++ b/network.py
@@ -1,7 +1,7 @@
 from openfe_benchmarks import tyk2
 
 from gufe import ChemicalSystem, Transformation, NonTransformation, AlchemicalNetwork
-from gufe.tests.test_protocol import DummyProtocol
+from gufe.tests.test_protocol import DummyProtocol, BrokenProtocol
 
 class DummyProtocolA(DummyProtocol):
     ...
@@ -44,7 +44,7 @@ def network_tyk2():
         Transformation(
             stateA=solvated[edge[0]],
             stateB=solvated[edge[1]],
-            protocol=DummyProtocolB(settings=DummyProtocolB.default_settings()),
+            protocol=BrokenProtocol(settings=BrokenProtocol.default_settings()),
             name=f"{edge[0]}_to_{edge[1]}_solvent",
         )
         for edge in tyk2s.connections

From 3191bbae6fc4a7c4d0d3e761122b97cfcddb3767 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Tue, 24 Mar 2026 14:15:06 -0400
Subject: [PATCH 09/24] Optionally keep scratch and shared after dag execution

---
 alchemiscale/compute/service.py | 20 ++++++++------------
 local_testing.py                | 16 ++++++++++++----
 service.py                      |  4 +++-
 3 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 44528e65..e8b801f4 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -537,25 +537,18 @@ def terminate_task(self, task_key: TaskKey):
                 if task_key == _task_key:
                     to_remove.add(proc)
             for proc in to_remove:
-                try:
-                    proc.close()
-                except ValueError:
-                    proc.terminate()
+                proc.terminate()
                 self._stack.remove(proc)
 
-    def push(self, node: NodeKey, context: Context, inputs: dict, n_retries, in_process=False):
+    def push(self, node: NodeKey, context: Context, inputs: dict, n_retries):
         with self.lock:
             if node in self._jail.keys():
                 raise JailedKeyError(node)
 
             executor = Executor.from_key(node, self.queue, self.lock, context, inputs, n_retries)
             self._stack.append(executor)
-            if not in_process:
-                self._stack[-1].start()
-                return
+            self._stack[-1].start()
 
-        self._stack[-1].run()
-        self._stack[-1].close()
 
     def pop(self):
         """Remove last process in the stack. This also clears the node from the jail."""
@@ -753,8 +746,11 @@ def remove_task(self, task_scoped_key):
                 self._dag_tree.remove_node(node)
 
         context = self._task_data[task_scoped_key].context
-        shutil.rmtree(context.shared)
-        shutil.rmtree(context.scratch)
+
+        if not self.keep_shared:
+            shutil.rmtree(context.shared)
+        if not self.keep_scratch:
+            shutil.rmtree(context.scratch)
 
     def _consume_results(self, task_scoped_key) -> ProtocolDAGResult:
         self.remove_task(task_scoped_key)
diff --git a/local_testing.py b/local_testing.py
index 7bafb0f9..2fd58b4c 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -30,7 +30,8 @@
 SHARED_DIR = Path("./acs_testing/shared")
 STACKSIZE = 10
 N_RETRIES = 2
-IN_PROCESS = False
+KEEP_SHARED = False
+KEEP_SCRATCH = False
 
 SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
 SHARED_DIR.mkdir(parents=True, exist_ok=True)
@@ -50,7 +51,7 @@ def create_tyk2():
 
 if __name__ == "__main__":
 
-    mock_service = service.MockService(SCRATCH_DIR, SHARED_DIR, STACKSIZE)
+    mock_service = service.MockService(SCRATCH_DIR, SHARED_DIR, STACKSIZE, KEEP_SCRATCH, KEEP_SHARED)
 
     with utils.timer(wrap=True):
         print("Creating network")
@@ -80,18 +81,25 @@ def create_tyk2():
             node_key, res = result
             task_scoped_key, pu = node_key
             mock_service._task_data[task_scoped_key].results[pu.key] = res
+            unit_context = Context(
+                shared=mock_service._task_data[task_scoped_key].context.shared / f"{pu.key}",
+                scratch=mock_service._task_data[task_scoped_key].context.scratch / f"{pu.key}",
+            )
 
             match res:
                 case ProtocolUnitFailure():
                     mock_service._executor_stack.terminate_task(task_scoped_key)
                     failed_tasks.add(task_scoped_key)
+
                 case ProtocolUnitResult():
                     mock_service._dag_tree.remove_node(node_key)
-
+            if not mock_service.keep_scratch:
+                shutil.rmtree(unit_context.scratch)
         for failed_task in failed_tasks:
             pdr = mock_service._consume_results(failed_task)
             pdrs.append(pdr)
 
+
         # only submit enough tasks to fill the stack
         n = mock_service._executor_stack._stack_size - len(
             mock_service._executor_stack._stack
@@ -111,6 +119,6 @@ def create_tyk2():
             unit_scratch_dir.mkdir()
             unit_shared_dir.mkdir()
             context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
-            mock_service._executor_stack.push(key, context, inputs, N_RETRIES, in_process=IN_PROCESS)
+            mock_service._executor_stack.push(key, context, inputs, N_RETRIES)
 
     print(pdrs)
diff --git a/service.py b/service.py
index ea229511..4cca1b9d 100644
--- a/service.py
+++ b/service.py
@@ -2,13 +2,15 @@
 
 class MockService(AsynchronousComputeService):
 
-    def __init__(self, scratch_basedir, shared_basedir, stack_size):
+    def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, keep_shared):
         self._initialize_dag_tree()
         self._task_data = dict()
         self._executor_stack = ExecutorStack(stack_size)
 
         self.scratch_basedir = scratch_basedir
         self.shared_basedir = shared_basedir
+        self.keep_scratch = keep_scratch
+        self.keep_shared = keep_shared
 
     def add_task(self, task_scoped_key, transformation):
         protocol_dag = transformation.create()

From c3d29671080055c235484254943c8f00b58a55fd Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Wed, 25 Mar 2026 13:36:18 -0400
Subject: [PATCH 10/24] Move logic into AsynchronousComputeService

---
 alchemiscale/compute/service.py | 99 +++++++++++++++++++++++++++++++--
 local_testing.py                | 82 ++++++---------------------
 service.py                      | 17 +++++-
 3 files changed, 128 insertions(+), 70 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index e8b801f4..24c3bd5d 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -16,7 +16,7 @@
 from typing import Any
 
 from gufe import Transformation
-from gufe.protocols.protocoldag import execute_DAG, ProtocolDAG, ProtocolDAGResult
+from gufe.protocols.protocoldag import _pu_to_pur, execute_DAG, ProtocolDAG, ProtocolDAGResult
 from gufe.protocols.protocolunit import Context, ProtocolUnitFailure, ProtocolUnitResult, ProtocolUnit
 from gufe.tokenization import GufeKey
 import networkx as nx
@@ -674,10 +674,99 @@ def _initialize_dag_tree(self):
         root_node: NodeKey = (None, "ROOT")
         self._dag_tree.add_node(root_node)
 
-    async def async_cycle(self, max_tasks, max_time):
-        # (ProtocolDAG, dwindling_graph, results)
-        for task in tasks:
-            raise NotImplementedError
+    def has_tasks(self) -> bool:
+        return bool(self._task_data)
+
+    def cycle(self, max_tasks, max_time) -> bool:
+        # TODO check max_tasks and max_time
+
+        match (self._stop, self.has_tasks()):
+            # should stop, but has remaining tasks, that COULD be finished
+            # 1. terminate all processes
+            # 2. Process remaining results in queue
+            # 3. Push anything that is completed or failed
+            # 4. Remove all tasks
+            # 5. break from main loop
+            case (True, True):
+                raise NotImplementedError
+                # should terminate all tasks
+                mock_service.terminate_all()
+                mock_service.process_results()
+                mock_service.remove_all()
+                return False
+            case (True, False):
+                return False
+            case (False, True):
+                pass
+            case (False, False):
+                tasks: list[ScopedKey] | None = self.claim_tasks(count=self.claim_limit)
+                should_stop = True
+
+                for task in tasks:
+                    if task is None:
+                        pass
+                    else:
+                        should_stop = False
+                        self.add_task(*task)
+                if should_stop:
+                    self.stop()
+            case _:
+                raise RuntimeError("Should never hit this")
+
+
+
+        # check for terminating nodes
+        for terminating_nodes in self.next_terminating_nodes():
+            task_scoped_key, _ = terminating_nodes
+            pdr = self._consume_results(task_scoped_key)
+            self.push_result(task_scoped_key, pdr)
+
+        # collect unit results
+        failed_tasks = set()
+        while result := self._executor_stack.get_result():
+            node_key, res = result
+            task_scoped_key, pu = node_key
+            self._task_data[task_scoped_key].results[pu.key] = res
+            unit_context = Context(
+                shared=self._task_data[task_scoped_key].context.shared / f"{pu.key}",
+                scratch=self._task_data[task_scoped_key].context.scratch / f"{pu.key}",
+            )
+
+            match res:
+                case ProtocolUnitFailure():
+                    self._executor_stack.terminate_task(task_scoped_key)
+                    failed_tasks.add(task_scoped_key)
+
+                case ProtocolUnitResult():
+                    self._dag_tree.remove_node(node_key)
+            if not self.keep_scratch:
+                shutil.rmtree(unit_context.scratch)
+
+        for failed_task in failed_tasks:
+            pdr = self._consume_results(failed_task)
+            self.push_result(task_scoped_key, pdr)
+
+        # only submit enough tasks to fill the stack
+        n = self._executor_stack._stack_size - len(
+            self._executor_stack._stack
+        )
+        for key in tuple(self.next())[:n]:
+            tsk, unit = key
+
+            # TODO `next` should not return TERM or ROOT nodes
+            if unit in ("TERM", "ROOT"):
+                continue
+
+            task_data = self._task_data[tsk]
+
+            inputs = _pu_to_pur(unit.inputs, task_data.results)
+            unit_scratch_dir = task_data.context.scratch / f"{str(unit.key)}"
+            unit_shared_dir = task_data.context.shared / f"{str(unit.key)}"
+            unit_scratch_dir.mkdir()
+            unit_shared_dir.mkdir()
+            context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
+            self._executor_stack.push(key, context, inputs, self.n_retries)
+        return True
 
     def available_units(self) -> set[NodeKey]:
         available = set()
diff --git a/local_testing.py b/local_testing.py
index 2fd58b4c..1c02106d 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -17,7 +17,7 @@
 from alchemiscale.models import ScopedKey
 from gufe import AlchemicalNetwork
 from gufe.protocols.protocolunit import Context, ProtocolUnitResult, ProtocolUnitFailure
-from gufe.protocols.protocoldag import _pu_to_pur
+from gufe.protocols.protocoldag import _pu_to_pur, ProtocolDAGResult
 from gufe.tokenization import GufeKey
 from gufe.tests.test_protocol import BrokenProtocol
 
@@ -30,8 +30,11 @@
 SHARED_DIR = Path("./acs_testing/shared")
 STACKSIZE = 10
 N_RETRIES = 2
+MAX_TASKS = 3
+MAX_TIME = None
 KEEP_SHARED = False
 KEEP_SCRATCH = False
+CLAIM_LIMIT = 3
 
 SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
 SHARED_DIR.mkdir(parents=True, exist_ok=True)
@@ -51,74 +54,25 @@ def create_tyk2():
 
 if __name__ == "__main__":
 
-    mock_service = service.MockService(SCRATCH_DIR, SHARED_DIR, STACKSIZE, KEEP_SCRATCH, KEEP_SHARED)
+    print("Creating network")
+    tyk2 = create_tyk2()
 
-    with utils.timer(wrap=True):
-        print("Creating network")
-        tyk2 = create_tyk2()
-
-    transformations = tuple(tyk2.edges)
-    tasks = tuple(
+    task_generator = (
         (utils.new_task_scoped_key(), transformation)
-        for transformation in transformations
+        for transformation in tyk2.edges
     )
 
-    for tsk, trans in tasks:
-        mock_service.add_task(tsk, trans)
+    mock_service = service.MockService(SCRATCH_DIR, SHARED_DIR, STACKSIZE, KEEP_SCRATCH, KEEP_SHARED, N_RETRIES, CLAIM_LIMIT, task_generator)
+
 
-    # collect for final inspection
+    # for local testing, override how the service pushes results
     pdrs = []
-    while mock_service._task_data:
-        # ask for terminating nodes
-        for completed_node in mock_service.next_terminating_nodes():
-            task_scoped_key, _ = completed_node
-            pdr = mock_service._consume_results(task_scoped_key)
-            pdrs.append(pdr)
-
-        # collect unit results
-        failed_tasks = set()
-        while result := mock_service._executor_stack.get_result():
-            node_key, res = result
-            task_scoped_key, pu = node_key
-            mock_service._task_data[task_scoped_key].results[pu.key] = res
-            unit_context = Context(
-                shared=mock_service._task_data[task_scoped_key].context.shared / f"{pu.key}",
-                scratch=mock_service._task_data[task_scoped_key].context.scratch / f"{pu.key}",
-            )
-
-            match res:
-                case ProtocolUnitFailure():
-                    mock_service._executor_stack.terminate_task(task_scoped_key)
-                    failed_tasks.add(task_scoped_key)
-
-                case ProtocolUnitResult():
-                    mock_service._dag_tree.remove_node(node_key)
-            if not mock_service.keep_scratch:
-                shutil.rmtree(unit_context.scratch)
-        for failed_task in failed_tasks:
-            pdr = mock_service._consume_results(failed_task)
-            pdrs.append(pdr)
-
-
-        # only submit enough tasks to fill the stack
-        n = mock_service._executor_stack._stack_size - len(
-            mock_service._executor_stack._stack
-        )
-        for key in tuple(mock_service.next())[:n]:
-            tsk, unit = key
-
-            # TODO `next` should not return TERM or ROOT nodes
-            if unit in ("TERM", "ROOT"):
-                continue
-
-            task_data = mock_service._task_data[tsk]
-
-            inputs = _pu_to_pur(unit.inputs, task_data.results)
-            unit_scratch_dir = task_data.context.scratch / f"{str(unit.key)}"
-            unit_shared_dir = task_data.context.shared / f"{str(unit.key)}"
-            unit_scratch_dir.mkdir()
-            unit_shared_dir.mkdir()
-            context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
-            mock_service._executor_stack.push(key, context, inputs, N_RETRIES)
+    def push_result(task_scoped_key: NodeKey, pdr: ProtocolDAGResult):
+        _ = task_scoped_key
+        pdrs.append(pdr)
+    mock_service.push_result = push_result
+
+    while mock_service.cycle(MAX_TASKS, MAX_TIME):
+        pass
 
     print(pdrs)
diff --git a/service.py b/service.py
index 4cca1b9d..45eafc01 100644
--- a/service.py
+++ b/service.py
@@ -2,7 +2,7 @@
 
 class MockService(AsynchronousComputeService):
 
-    def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, keep_shared):
+    def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, keep_shared, n_retries, claim_limit, task_generator):
         self._initialize_dag_tree()
         self._task_data = dict()
         self._executor_stack = ExecutorStack(stack_size)
@@ -11,7 +11,22 @@ def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, ke
         self.shared_basedir = shared_basedir
         self.keep_scratch = keep_scratch
         self.keep_shared = keep_shared
+        self.n_retries = n_retries
+        self.claim_limit = claim_limit
+        self.task_generator = task_generator
+
+        self._stop = False
 
     def add_task(self, task_scoped_key, transformation):
         protocol_dag = transformation.create()
         self.graft_dag(task_scoped_key, protocol_dag)
+
+    def claim_tasks(self, count=1):
+        claimed_tasks = []
+        remaining = count
+        for task in self.task_generator:
+            claimed_tasks.append(task)
+            remaining = remaining - 1
+            if remaining == 0:
+                return claimed_tasks
+        return claimed_tasks + [None] * remaining

From 87e8b30a26bba7bc869ae1c08d2c3521fc977f8d Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Wed, 25 Mar 2026 13:47:07 -0400
Subject: [PATCH 11/24] Process results before checking for terminal nodes

---
 alchemiscale/compute/service.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 24c3bd5d..73c2dfb2 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -715,11 +715,6 @@ def cycle(self, max_tasks, max_time) -> bool:
 
 
 
-        # check for terminating nodes
-        for terminating_nodes in self.next_terminating_nodes():
-            task_scoped_key, _ = terminating_nodes
-            pdr = self._consume_results(task_scoped_key)
-            self.push_result(task_scoped_key, pdr)
 
         # collect unit results
         failed_tasks = set()
@@ -746,6 +741,12 @@ def cycle(self, max_tasks, max_time) -> bool:
             pdr = self._consume_results(failed_task)
             self.push_result(task_scoped_key, pdr)
 
+        # check for terminating nodes
+        for terminating_nodes in self.next_terminating_nodes():
+            task_scoped_key, _ = terminating_nodes
+            pdr = self._consume_results(task_scoped_key)
+            self.push_result(task_scoped_key, pdr)
+
         # only submit enough tasks to fill the stack
         n = self._executor_stack._stack_size - len(
             self._executor_stack._stack

From 075004c65160281e9eaade484efadf0c824cb4de Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Wed, 25 Mar 2026 15:16:57 -0400
Subject: [PATCH 12/24] Respect max tasks

---
 alchemiscale/compute/service.py | 86 ++++++++++++++++-----------------
 local_testing.py                | 18 ++-----
 service.py                      | 10 ++++
 3 files changed, 57 insertions(+), 57 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 73c2dfb2..e8b18362 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -677,46 +677,13 @@ def _initialize_dag_tree(self):
     def has_tasks(self) -> bool:
         return bool(self._task_data)
 
-    def cycle(self, max_tasks, max_time) -> bool:
-        # TODO check max_tasks and max_time
-
-        match (self._stop, self.has_tasks()):
-            # should stop, but has remaining tasks, that COULD be finished
-            # 1. terminate all processes
-            # 2. Process remaining results in queue
-            # 3. Push anything that is completed or failed
-            # 4. Remove all tasks
-            # 5. break from main loop
-            case (True, True):
-                raise NotImplementedError
-                # should terminate all tasks
-                mock_service.terminate_all()
-                mock_service.process_results()
-                mock_service.remove_all()
-                return False
-            case (True, False):
-                return False
-            case (False, True):
-                pass
-            case (False, False):
-                tasks: list[ScopedKey] | None = self.claim_tasks(count=self.claim_limit)
-                should_stop = True
-
-                for task in tasks:
-                    if task is None:
-                        pass
-                    else:
-                        should_stop = False
-                        self.add_task(*task)
-                if should_stop:
-                    self.stop()
-            case _:
-                raise RuntimeError("Should never hit this")
-
-
-
+    def consume_terminated_tasks(self):
+        for terminating_nodes in self.next_terminating_nodes():
+            task_scoped_key, _ = terminating_nodes
+            pdr = self._consume_results(task_scoped_key)
+            self.push_result(task_scoped_key, pdr)
 
-        # collect unit results
+    def process_results(self):
         failed_tasks = set()
         while result := self._executor_stack.get_result():
             node_key, res = result
@@ -741,11 +708,42 @@ def cycle(self, max_tasks, max_time) -> bool:
             pdr = self._consume_results(failed_task)
             self.push_result(task_scoped_key, pdr)
 
-        # check for terminating nodes
-        for terminating_nodes in self.next_terminating_nodes():
-            task_scoped_key, _ = terminating_nodes
-            pdr = self._consume_results(task_scoped_key)
-            self.push_result(task_scoped_key, pdr)
+    def cycle(self, max_tasks, max_time) -> bool:
+        # TODO check max_tasks and max_time
+
+        if self._stop:
+            # should stop, but has remaining tasks, that COULD be finished
+            # 1. terminate all processes
+            # 2. Process remaining results in queue
+            # 3. Push anything that is completed or failed
+            # 4. Remove all tasks
+            # 5. break from main loop
+
+            if self.has_tasks():
+                mock_service.terminate_all()
+                mock_service.process_results()
+                mock_service.remove_all()
+            return False
+
+        max_less_waiting = max_tasks - len(self._task_data) # why does removing this break things?
+        max_less_claimed = max_tasks - self.tasks_claimed
+        empty_slots = self.claim_limit - len(self._task_data)
+        n_claim = min(empty_slots,
+                      #max_less_waiting,
+                      max_less_claimed,
+                      )
+        tasks = self.claim_tasks(count=n_claim)
+        for task in tasks:
+            if task is not None:
+                self.add_task(*task)
+                self.tasks_claimed = 1 + self.tasks_claimed
+
+        # collect unit results
+        self.process_results()
+        self.consume_terminated_tasks()
+        if self.tasks_finished >= max_tasks:
+            self.stop()
+            return False
 
         # only submit enough tasks to fill the stack
         n = self._executor_stack._stack_size - len(
diff --git a/local_testing.py b/local_testing.py
index 1c02106d..4394278d 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -29,12 +29,12 @@
 SCRATCH_DIR = Path("./acs_testing/scratch")
 SHARED_DIR = Path("./acs_testing/shared")
 STACKSIZE = 10
-N_RETRIES = 2
-MAX_TASKS = 3
+N_RETRIES = 0
+MAX_TASKS = 100
 MAX_TIME = None
 KEEP_SHARED = False
 KEEP_SCRATCH = False
-CLAIM_LIMIT = 3
+CLAIM_LIMIT = 10
 
 SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
 SHARED_DIR.mkdir(parents=True, exist_ok=True)
@@ -59,20 +59,12 @@ def create_tyk2():
 
     task_generator = (
         (utils.new_task_scoped_key(), transformation)
-        for transformation in tyk2.edges
+        for transformation in tuple(tyk2.edges)
     )
 
     mock_service = service.MockService(SCRATCH_DIR, SHARED_DIR, STACKSIZE, KEEP_SCRATCH, KEEP_SHARED, N_RETRIES, CLAIM_LIMIT, task_generator)
 
-
-    # for local testing, override how the service pushes results
-    pdrs = []
-    def push_result(task_scoped_key: NodeKey, pdr: ProtocolDAGResult):
-        _ = task_scoped_key
-        pdrs.append(pdr)
-    mock_service.push_result = push_result
-
     while mock_service.cycle(MAX_TASKS, MAX_TIME):
         pass
 
-    print(pdrs)
+    print(mock_service.pdrs)
diff --git a/service.py b/service.py
index 45eafc01..5f279ac3 100644
--- a/service.py
+++ b/service.py
@@ -14,6 +14,9 @@ def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, ke
         self.n_retries = n_retries
         self.claim_limit = claim_limit
         self.task_generator = task_generator
+        self.pdrs = []
+        self.tasks_claimed = 0
+        self.tasks_finished = 0
 
         self._stop = False
 
@@ -24,9 +27,16 @@ def add_task(self, task_scoped_key, transformation):
     def claim_tasks(self, count=1):
         claimed_tasks = []
         remaining = count
+        if remaining == 0:
+            return [None] * count
         for task in self.task_generator:
             claimed_tasks.append(task)
             remaining = remaining - 1
             if remaining == 0:
                 return claimed_tasks
         return claimed_tasks + [None] * remaining
+
+    def push_result(self, task_scoped_key, pdr):
+        _ = task_scoped_key
+        self.pdrs.append(pdr)
+        self.tasks_finished = self.tasks_finished + 1

From 10d768df774f076d9ca32dc272d281f493366705 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Wed, 25 Mar 2026 17:57:36 -0400
Subject: [PATCH 13/24] Respect max_time

---
 alchemiscale/compute/service.py | 69 ++++++++++++++++++++++++---------
 local_testing.py                | 13 ++++++-
 service.py                      |  3 ++
 3 files changed, 64 insertions(+), 21 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index e8b18362..ac797805 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -16,8 +16,18 @@
 from typing import Any
 
 from gufe import Transformation
-from gufe.protocols.protocoldag import _pu_to_pur, execute_DAG, ProtocolDAG, ProtocolDAGResult
-from gufe.protocols.protocolunit import Context, ProtocolUnitFailure, ProtocolUnitResult, ProtocolUnit
+from gufe.protocols.protocoldag import (
+    _pu_to_pur,
+    execute_DAG,
+    ProtocolDAG,
+    ProtocolDAGResult,
+)
+from gufe.protocols.protocolunit import (
+    Context,
+    ProtocolUnitFailure,
+    ProtocolUnitResult,
+    ProtocolUnit,
+)
 from gufe.tokenization import GufeKey
 import networkx as nx
 
@@ -475,9 +485,22 @@ def inputs(self) -> dict:
 
     @classmethod
     def from_key(
-            cls, key: NodeKey, queue: Queue, lock: Lock, context: Context, inputs: dict, n_retries: int
+        cls,
+        key: NodeKey,
+        queue: Queue,
+        lock: Lock,
+        context: Context,
+        inputs: dict,
+        n_retries: int,
     ):
-        return cls(key=key, queue=queue, lock=lock, context=context, inputs=inputs, n_retries=n_retries)
+        return cls(
+            key=key,
+            queue=queue,
+            lock=lock,
+            context=context,
+            inputs=inputs,
+            n_retries=n_retries,
+        )
 
     def put_result(self, result: ProtocolUnitResult):
         """Acquire lock, push key and result into the queue, release lock."""
@@ -528,6 +551,7 @@ def terminate_all(self):
         with self.lock:
             for proc in self.stack:
                 proc.terminate()
+            self.stack.clear()
 
     def terminate_task(self, task_key: TaskKey):
         with self.lock:
@@ -545,11 +569,12 @@ def push(self, node: NodeKey, context: Context, inputs: dict, n_retries):
             if node in self._jail.keys():
                 raise JailedKeyError(node)
 
-            executor = Executor.from_key(node, self.queue, self.lock, context, inputs, n_retries)
+            executor = Executor.from_key(
+                node, self.queue, self.lock, context, inputs, n_retries
+            )
             self._stack.append(executor)
             self._stack[-1].start()
 
-
     def pop(self):
         """Remove last process in the stack. This also clears the node from the jail."""
         if self._stack_size == 0:
@@ -562,7 +587,9 @@ def pop(self):
 
         return popped_executor
 
-    def get_result(self) -> tuple[NodeKey, ProtocolUnitResult | ProtocolUnitFailure] | None:
+    def get_result(
+        self,
+    ) -> tuple[NodeKey, ProtocolUnitResult | ProtocolUnitFailure] | None:
         if self.queue.qsize():
             with self.lock:
                 # since qsize is not always reliable, we tentatively
@@ -710,6 +737,8 @@ def process_results(self):
 
     def cycle(self, max_tasks, max_time) -> bool:
         # TODO check max_tasks and max_time
+        if (time.time() - self._start_time) >= max_time:
+            self.stop()
 
         if self._stop:
             # should stop, but has remaining tasks, that COULD be finished
@@ -720,18 +749,18 @@ def cycle(self, max_tasks, max_time) -> bool:
             # 5. break from main loop
 
             if self.has_tasks():
-                mock_service.terminate_all()
-                mock_service.process_results()
-                mock_service.remove_all()
+                self.process_results()
+                self._executor_stack.terminate_all()  # may corrupt queue, pull results out first
+                self.consume_terminated_tasks()
+                self.remove_all()
             return False
 
-        max_less_waiting = max_tasks - len(self._task_data) # why does removing this break things?
         max_less_claimed = max_tasks - self.tasks_claimed
         empty_slots = self.claim_limit - len(self._task_data)
-        n_claim = min(empty_slots,
-                      #max_less_waiting,
-                      max_less_claimed,
-                      )
+        n_claim = min(
+            empty_slots,
+            max_less_claimed,
+        )
         tasks = self.claim_tasks(count=n_claim)
         for task in tasks:
             if task is not None:
@@ -746,9 +775,7 @@ def cycle(self, max_tasks, max_time) -> bool:
             return False
 
         # only submit enough tasks to fill the stack
-        n = self._executor_stack._stack_size - len(
-            self._executor_stack._stack
-        )
+        n = self._executor_stack._stack_size - len(self._executor_stack._stack)
         for key in tuple(self.next())[:n]:
             tsk, unit = key
 
@@ -840,12 +867,16 @@ def remove_task(self, task_scoped_key):
         if not self.keep_scratch:
             shutil.rmtree(context.scratch)
 
+    def remove_all(self):
+        for task_scoped_key in self._task_data.keys():
+            self.remove_task(task_scoped_key)
+        self._task_data.clear()
+
     def _consume_results(self, task_scoped_key) -> ProtocolDAGResult:
         self.remove_task(task_scoped_key)
         data = self._task_data.pop(task_scoped_key)
         pdr = data.to_ProtocolDAGResult()
         return pdr
 
-
     def stop(self):
         self._stop = True
diff --git a/local_testing.py b/local_testing.py
index 4394278d..d669f886 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -31,7 +31,7 @@
 STACKSIZE = 10
 N_RETRIES = 0
 MAX_TASKS = 100
-MAX_TIME = None
+MAX_TIME = 25
 KEEP_SHARED = False
 KEEP_SCRATCH = False
 CLAIM_LIMIT = 10
@@ -62,7 +62,16 @@ def create_tyk2():
         for transformation in tuple(tyk2.edges)
     )
 
-    mock_service = service.MockService(SCRATCH_DIR, SHARED_DIR, STACKSIZE, KEEP_SCRATCH, KEEP_SHARED, N_RETRIES, CLAIM_LIMIT, task_generator)
+    mock_service = service.MockService(
+        SCRATCH_DIR,
+        SHARED_DIR,
+        STACKSIZE,
+        KEEP_SCRATCH,
+        KEEP_SHARED,
+        N_RETRIES,
+        CLAIM_LIMIT,
+        task_generator,
+    )
 
     while mock_service.cycle(MAX_TASKS, MAX_TIME):
         pass
diff --git a/service.py b/service.py
index 5f279ac3..3140110b 100644
--- a/service.py
+++ b/service.py
@@ -1,3 +1,5 @@
+import time
+
 from alchemiscale.compute.service import AsynchronousComputeService, ExecutorStack
 
 class MockService(AsynchronousComputeService):
@@ -19,6 +21,7 @@ def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, ke
         self.tasks_finished = 0
 
         self._stop = False
+        self._start_time = time.time()
 
     def add_task(self, task_scoped_key, transformation):
         protocol_dag = transformation.create()

From 4bc2fb0f0c686a7b4063e5921eecfd4cb19f3378 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Thu, 26 Mar 2026 14:16:42 -0400
Subject: [PATCH 14/24] Use inherited start machinery and allow None max_time
 and max_tasks

| MAX_TIME | MAX_TASKS | Limit | NOTES              |
|----------+-----------+-------+--------------------|
|       20 |      NONE | time  |                    |
|     NONE |        10 | tasks |                    |
|     NONE |      NONE | NONE  | Ctrl-C, clean stop |
|       20 |         2 | tasks |                    |
|       20 |        20 | time  |                    |
---
 alchemiscale/compute/service.py | 60 ++++++++++++++++-----------------
 local_testing.py                |  7 ++--
 service.py                      | 35 +++++++++++++++++--
 3 files changed, 64 insertions(+), 38 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index ac797805..fdbf7b05 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -509,6 +509,8 @@ def put_result(self, result: ProtocolUnitResult):
 
     def execute_unit(self, context) -> ProtocolUnitResult | ProtocolUnitFailure:
         # this method assumes the context is in place and will be removed correctly
+        import warnings
+        warnings.filterwarnings("ignore", message=r".*RDKit does not preserve.*")
         return self.unit.execute(context=context, **self._inputs)
 
 
@@ -569,6 +571,8 @@ def push(self, node: NodeKey, context: Context, inputs: dict, n_retries):
             if node in self._jail.keys():
                 raise JailedKeyError(node)
 
+            import warnings
+            warnings.filterwarnings("ignore", message=r".*This process.*is multi-threaded,.*")
             executor = Executor.from_key(
                 node, self.queue, self.lock, context, inputs, n_retries
             )
@@ -735,45 +739,40 @@ def process_results(self):
             pdr = self._consume_results(failed_task)
             self.push_result(task_scoped_key, pdr)
 
+    def stop(self):
+        if self.has_tasks():
+            self.logger.info("Cleaning up")
+            self.process_results()
+            self._executor_stack.terminate_all()  # may corrupt queue, pull results out first
+            self.consume_terminated_tasks()
+            self.remove_all()
+        super().stop()
+
     def cycle(self, max_tasks, max_time) -> bool:
-        # TODO check max_tasks and max_time
-        if (time.time() - self._start_time) >= max_time:
+
+        if max_time is not None and (time.time() - self._start_time) >= max_time:
+            self.logger.info("Exceeded maximum time")
             self.stop()
+            return False
 
-        if self._stop:
-            # should stop, but has remaining tasks, that COULD be finished
-            # 1. terminate all processes
-            # 2. Process remaining results in queue
-            # 3. Push anything that is completed or failed
-            # 4. Remove all tasks
-            # 5. break from main loop
-
-            if self.has_tasks():
-                self.process_results()
-                self._executor_stack.terminate_all()  # may corrupt queue, pull results out first
-                self.consume_terminated_tasks()
-                self.remove_all()
+        # collect unit results
+        self.process_results()
+        self.consume_terminated_tasks()
+        if max_tasks is not None and self.tasks_finished >= max_tasks:
+            self.logger.info("Exceeded maximum tasks")
+            self.stop()
             return False
 
-        max_less_claimed = max_tasks - self.tasks_claimed
-        empty_slots = self.claim_limit - len(self._task_data)
-        n_claim = min(
-            empty_slots,
-            max_less_claimed,
-        )
+        n_claim = self.claim_limit - len(self._task_data)
+        if max_tasks is not None:
+            max_less_claimed = max_tasks - self.tasks_claimed
+            n_claim = min(n_claim, max_less_claimed)
         tasks = self.claim_tasks(count=n_claim)
         for task in tasks:
             if task is not None:
                 self.add_task(*task)
                 self.tasks_claimed = 1 + self.tasks_claimed
 
-        # collect unit results
-        self.process_results()
-        self.consume_terminated_tasks()
-        if self.tasks_finished >= max_tasks:
-            self.stop()
-            return False
-
         # only submit enough tasks to fill the stack
         n = self._executor_stack._stack_size - len(self._executor_stack._stack)
         for key in tuple(self.next())[:n]:
@@ -791,7 +790,9 @@ def cycle(self, max_tasks, max_time) -> bool:
             unit_scratch_dir.mkdir()
             unit_shared_dir.mkdir()
             context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
+            self.logger.info(f"Pushing {key[1]} to the execution stack")
             self._executor_stack.push(key, context, inputs, self.n_retries)
+
         return True
 
     def available_units(self) -> set[NodeKey]:
@@ -877,6 +878,3 @@ def _consume_results(self, task_scoped_key) -> ProtocolDAGResult:
         data = self._task_data.pop(task_scoped_key)
         pdr = data.to_ProtocolDAGResult()
         return pdr
-
-    def stop(self):
-        self._stop = True
diff --git a/local_testing.py b/local_testing.py
index d669f886..c3f28795 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -30,8 +30,8 @@
 SHARED_DIR = Path("./acs_testing/shared")
 STACKSIZE = 10
 N_RETRIES = 0
-MAX_TASKS = 100
-MAX_TIME = 25
+MAX_TASKS = None
+MAX_TIME = 20
 KEEP_SHARED = False
 KEEP_SCRATCH = False
 CLAIM_LIMIT = 10
@@ -73,7 +73,6 @@ def create_tyk2():
         task_generator,
     )
 
-    while mock_service.cycle(MAX_TASKS, MAX_TIME):
-        pass
+    mock_service.start(MAX_TASKS, MAX_TIME)
 
     print(mock_service.pdrs)
diff --git a/service.py b/service.py
index 3140110b..c6722286 100644
--- a/service.py
+++ b/service.py
@@ -1,6 +1,8 @@
+import logging
 import time
 
-from alchemiscale.compute.service import AsynchronousComputeService, ExecutorStack
+from alchemiscale.compute.service import AsynchronousComputeService, ExecutorStack, InterruptableSleep
+from alchemiscale.storage.models import ComputeServiceID
 
 class MockService(AsynchronousComputeService):
 
@@ -20,13 +22,39 @@ def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, ke
         self.tasks_claimed = 0
         self.tasks_finished = 0
 
-        self._stop = False
-        self._start_time = time.time()
+        self.int_sleep = InterruptableSleep()
+
+        self.name = "MockService"
+        self.compute_service_id = ComputeServiceID.new_from_name(self.name)
+
+        # logging shim
+        extra = {"compute_service_id": "fakeid"}
+        logger = logging.getLogger("AlchemiscaleSynchronousComputeService")
+        logger.setLevel(logging.DEBUG)
+
+        formatter = logging.Formatter(
+            "[%(asctime)s] [%(compute_service_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)
+        self.logger = logging.LoggerAdapter(logger, extra)
 
     def add_task(self, task_scoped_key, transformation):
         protocol_dag = transformation.create()
         self.graft_dag(task_scoped_key, protocol_dag)
 
+    def _register(self):
+        self.logger.info("Fake register")
+
+    def _deregister(self):
+        self.logger.info("Fake deregister")
+
+    def heartbeat(self):
+        pass
+
     def claim_tasks(self, count=1):
         claimed_tasks = []
         remaining = count
@@ -41,5 +69,6 @@ def claim_tasks(self, count=1):
 
     def push_result(self, task_scoped_key, pdr):
         _ = task_scoped_key
+        self.logger.info(f"Pushing {pdr}")
         self.pdrs.append(pdr)
         self.tasks_finished = self.tasks_finished + 1

From 609412b4b59806e8edea1723cfa593ebef988460 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Thu, 26 Mar 2026 20:54:45 -0400
Subject: [PATCH 15/24] Implement resource signals

---
 alchemiscale/compute/service.py | 119 +++++++++++++++++++++++++++-----
 local_testing.py                |  11 +--
 network.py                      |  74 ++++++++++++++++++--
 3 files changed, 176 insertions(+), 28 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index fdbf7b05..72e15a5e 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -14,6 +14,7 @@
 import queue
 import shutil
 from typing import Any
+from enum import StrEnum
 
 from gufe import Transformation
 from gufe.protocols.protocoldag import (
@@ -425,6 +426,11 @@ def stop(self):
     ProtocolUnit | str,
 ]  # None covers root node condition
 
+class ResourceSignal(StrEnum):
+    MAINTAIN = "maintain"
+    SHRINK = "shrink"
+    GROW = "grow"
+    TERMINATE = "terminate"
 
 class JailedKeyError(Exception):
     pass
@@ -566,30 +572,46 @@ def terminate_task(self, task_key: TaskKey):
                 proc.terminate()
                 self._stack.remove(proc)
 
-    def push(self, node: NodeKey, context: Context, inputs: dict, n_retries):
+    def push(self, node: NodeKey, unit_context: Context, inputs: dict, n_retries):
         with self.lock:
             if node in self._jail.keys():
                 raise JailedKeyError(node)
 
+            unit_context.scratch.mkdir()
+            unit_context.shared.mkdir()
+
             import warnings
             warnings.filterwarnings("ignore", message=r".*This process.*is multi-threaded,.*")
             executor = Executor.from_key(
-                node, self.queue, self.lock, context, inputs, n_retries
+                node, self.queue, self.lock, unit_context, inputs, n_retries
             )
             self._stack.append(executor)
             self._stack[-1].start()
 
     def pop(self):
         """Remove last process in the stack. This also clears the node from the jail."""
-        if self._stack_size == 0:
-            raise IndexError("pop from empty stack")
-        popped_executor = self._stack.pop()
-        for key in self._jail.keys():
-            self._jail[key] -= popped_executor.key
-            if not self._jail[key]:
+        with self.lock:
+            if self._stack_size == 0:
+                raise IndexError("pop from empty stack")
+
+            popped_executor = self._stack.pop()
+            popped_executor.terminate()
+
+            unblocked = set()
+            for blocked_key in self._jail.keys():
+                if popped_executor.key in self._jail[blocked_key]:
+                    self._jail[blocked_key].remove(popped_executor.key)
+                if not self._jail[blocked_key]:
+                    unblocked.add(key)
+
+            for key in unblocked:
                 self._jail.pop(key)
 
-        return popped_executor
+            self._jail[popped_executor.key] = {proc.key for proc in self._stack}
+            shutil.rmtree(popped_executor._unit_context.shared)
+            shutil.rmtree(popped_executor._unit_context.scratch)
+
+            return popped_executor
 
     def get_result(
         self,
@@ -613,6 +635,14 @@ def remove_by_node_key(self, node_key: NodeKey):
                 to_remove = proc
                 break
 
+        unblock = set()
+        for key in self._jail.keys():
+            if node_key in self._jail[key]:
+                self._jail[key].remove(node_key)
+                if not self._jail[key]:
+                    unblock.add(key)
+        for key in unblock:
+            self._jail.pop(key)
         if to_remove:
             self._stack.remove(proc)
 
@@ -748,7 +778,42 @@ def stop(self):
             self.remove_all()
         super().stop()
 
+    def resource_monitor(self) -> ResourceSignal:
+        capacity = 8
+        sim_value = 3
+        fin_value = 1
+        init_value = 2
+
+        total = 0
+        for proc in self._executor_stack.stack:
+            if not proc.is_alive():
+                continue
+            pu = proc.key[1]
+            match pu.__class__.__name__:
+                case "WeightedFinishingUnit":
+                    total = total + fin_value
+                case "WeightedSimulationUnit":
+                    total = total + sim_value
+                case _:
+                    total = total + init_value
+
+        self.logger.info(f"Num jailed: {len(self._executor_stack._jail)} -- {total}/{capacity}")
+        if total == capacity:
+            self.logger.info("MAINTAINING")
+            return ResourceSignal.MAINTAIN
+
+        if total > capacity:
+            self.logger.info("SHRINKING")
+            return ResourceSignal.SHRINK
+
+        if total < capacity:
+            self.logger.info("GROWING")
+            return ResourceSignal.GROW
+
+        return ResourceSignal.TERMINATE
+
     def cycle(self, max_tasks, max_time) -> bool:
+        self.logger.info(f"Cycling: {len(self._task_data)}")
 
         if max_time is not None and (time.time() - self._start_time) >= max_time:
             self.logger.info("Exceeded maximum time")
@@ -756,13 +821,31 @@ def cycle(self, max_tasks, max_time) -> bool:
             return False
 
         # collect unit results
-        self.process_results()
-        self.consume_terminated_tasks()
+        self.process_results() # removes unit scratch
+        self.consume_terminated_tasks() # only removes if task is done
         if max_tasks is not None and self.tasks_finished >= max_tasks:
             self.logger.info("Exceeded maximum tasks")
             self.stop()
             return False
 
+        signal = self.resource_monitor()
+        match signal:
+            case ResourceSignal.MAINTAIN:
+                #self.process_results()
+                #self.consume_terminated_tasks()
+                return True
+            case ResourceSignal.GROW:
+                pass
+            case ResourceSignal.SHRINK:
+                proc = self._executor_stack.pop()
+                self.logger.info(f"Popping: {proc} -- {proc.key[1].key}")
+                return True
+            case ResourceSignal.TERMINATE:
+                self.stop()
+                return False
+            case _:
+                raise RuntimeError("Received unknown ResourceSignal")
+
         n_claim = self.claim_limit - len(self._task_data)
         if max_tasks is not None:
             max_less_claimed = max_tasks - self.tasks_claimed
@@ -775,7 +858,7 @@ def cycle(self, max_tasks, max_time) -> bool:
 
         # only submit enough tasks to fill the stack
         n = self._executor_stack._stack_size - len(self._executor_stack._stack)
-        for key in tuple(self.next())[:n]:
+        for key in tuple(self.next()):
             tsk, unit = key
 
             # TODO `next` should not return TERM or ROOT nodes
@@ -787,11 +870,15 @@ def cycle(self, max_tasks, max_time) -> bool:
             inputs = _pu_to_pur(unit.inputs, task_data.results)
             unit_scratch_dir = task_data.context.scratch / f"{str(unit.key)}"
             unit_shared_dir = task_data.context.shared / f"{str(unit.key)}"
-            unit_scratch_dir.mkdir()
-            unit_shared_dir.mkdir()
             context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
-            self.logger.info(f"Pushing {key[1]} to the execution stack")
-            self._executor_stack.push(key, context, inputs, self.n_retries)
+
+            try:
+                self._executor_stack.push(key, context, inputs, self.n_retries)
+                self.logger.info(f"Pushing {key[1]} to the execution stack")
+            except JailedKeyError:
+                continue
+            if not (n := n - 1):
+                break
 
         return True
 
diff --git a/local_testing.py b/local_testing.py
index c3f28795..a4a9e35a 100644
--- a/local_testing.py
+++ b/local_testing.py
@@ -28,13 +28,13 @@
 
 SCRATCH_DIR = Path("./acs_testing/scratch")
 SHARED_DIR = Path("./acs_testing/shared")
-STACKSIZE = 10
+STACKSIZE = 4
 N_RETRIES = 0
-MAX_TASKS = None
-MAX_TIME = 20
+MAX_TASKS = 1
+MAX_TIME = None
 KEEP_SHARED = False
 KEEP_SCRATCH = False
-CLAIM_LIMIT = 10
+CLAIM_LIMIT = 2
 
 SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
 SHARED_DIR.mkdir(parents=True, exist_ok=True)
@@ -59,7 +59,7 @@ def create_tyk2():
 
     task_generator = (
         (utils.new_task_scoped_key(), transformation)
-        for transformation in tuple(tyk2.edges)
+        for transformation in tuple(tyk2.edges)[:10]
     )
 
     mock_service = service.MockService(
@@ -76,3 +76,4 @@ def create_tyk2():
     mock_service.start(MAX_TASKS, MAX_TIME)
 
     print(mock_service.pdrs)
+    assert(all(pdr.ok() for pdr in mock_service.pdrs))
diff --git a/network.py b/network.py
index 0768e32a..b4a1e923 100644
--- a/network.py
+++ b/network.py
@@ -1,13 +1,70 @@
+import time
+
 from openfe_benchmarks import tyk2
 
 from gufe import ChemicalSystem, Transformation, NonTransformation, AlchemicalNetwork
-from gufe.tests.test_protocol import DummyProtocol, BrokenProtocol
+from gufe.tests.test_protocol import DummyProtocol, BrokenProtocol, FinishUnit, SimulationUnit, InitializeUnit, ProtocolUnit
+from gufe.protocols import ProtocolUnit
+
+class WeightedInitializeUnit(InitializeUnit):
+    value = 2
+
+class WeightedSimulationUnit(SimulationUnit):
+    value = 2
+
+    @staticmethod
+    def _execute(ctx, *, initialization, **inputs):
+        time.sleep(WeightedSimulationUnit.value * 2)
+        return SimulationUnit._execute(ctx, initialization=initialization, **inputs)
+
+class WeightedFinishUnit(FinishUnit):
+    value = 1
+
+    @staticmethod
+    def _execute(ctx, *, simulations, **inputs):
+        time.sleep(WeightedFinishUnit.value * 2)
+        return FinishUnit._execute(ctx, simulations=simulations, **inputs)
+
+
+class WeightedDummyProtocol(DummyProtocol):
+
+    def _create(
+            self,
+            stateA,
+            stateB,
+            mapping = None,
+            extends = None,
+    ):
+        if extends is not None:
+            # this is an example; wouldn't want to pass in whole ProtocolDAGResult into
+            # any ProtocolUnits below, since this could create dependency hell;
+            # instead, extract what's needed from it for starting point here
+            starting_point = extends.protocol_unit_results[-1].outputs["key_results"]
+        else:
+            starting_point = None
+
+        # convert protocol inputs into starting points for independent simulations
+        alpha = WeightedInitializeUnit(
+            name="the beginning",
+            settings=self.settings,
+            stateA=stateA,
+            stateB=stateB,
+            mapping=mapping,
+            start=starting_point,
+            some_dict={"a": 2, "b": 12},
+        )
+
+        # create several units that would each run an independent simulation
+        simulations: list[ProtocolUnit] = [
+            WeightedSimulationUnit(settings=self.settings, name=f"sim {i}", window=i, initialization=alpha)
+            for i in range(self.settings.n_repeats)  # type: ignore
+        ]
 
-class DummyProtocolA(DummyProtocol):
-    ...
+        # gather results from simulations, finalize outputs
+        omega = WeightedFinishUnit(settings=self.settings, name="the end", simulations=simulations)
 
-class DummyProtocolB(DummyProtocol):
-    ...
+        # return all `ProtocolUnit`s we created
+        return [alpha, *simulations, omega]
 
 def network_tyk2():
     tyk2s = tyk2.get_system()
@@ -35,7 +92,7 @@ def network_tyk2():
         Transformation(
             stateA=complexes[edge[0]],
             stateB=complexes[edge[1]],
-            protocol=DummyProtocolA(settings=DummyProtocolA.default_settings()),
+            protocol=WeightedDummyProtocol(settings=WeightedDummyProtocol.default_settings()),
             name=f"{edge[0]}_to_{edge[1]}_complex",
         )
         for edge in tyk2s.connections
@@ -44,11 +101,14 @@ def network_tyk2():
         Transformation(
             stateA=solvated[edge[0]],
             stateB=solvated[edge[1]],
-            protocol=BrokenProtocol(settings=BrokenProtocol.default_settings()),
+            protocol=WeightedDummyProtocol(settings=WeightedDummyProtocol.default_settings()),
             name=f"{edge[0]}_to_{edge[1]}_solvent",
         )
         for edge in tyk2s.connections
     ]
+    # breakpoint()
+    # pu = solvent_network[0].create().protocol_units[0]
+    # pu.execute(context=None, initialization=None)
 
     return AlchemicalNetwork(
         edges=(solvent_network + complex_network),

From af3be42622d21e0e592b66739c08e4f938c6b675 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Fri, 27 Mar 2026 13:27:02 -0400
Subject: [PATCH 16/24] Clean up unused code and reorganize definitions

---
 alchemiscale/compute/service.py | 117 ++++++++++++--------------------
 1 file changed, 43 insertions(+), 74 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 72e15a5e..9b833aeb 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -15,6 +15,8 @@
 import shutil
 from typing import Any
 from enum import StrEnum
+from multiprocessing import Process, Queue, Lock
+from dataclasses import dataclass
 
 from gufe import Transformation
 from gufe.protocols.protocoldag import (
@@ -418,25 +420,21 @@ def stop(self):
         self._stop = True
 
 
-from multiprocessing import Process, Queue, Lock
+## Asynchronous compute
 
 type TaskKey = ScopedKey
+
+# ProtocolUnits are paired with their parent task for book keeping
+# purposes
 type NodeKey = tuple[
     TaskKey | None,
     ProtocolUnit | str,
 ]  # None covers root node condition
 
-class ResourceSignal(StrEnum):
-    MAINTAIN = "maintain"
-    SHRINK = "shrink"
-    GROW = "grow"
-    TERMINATE = "terminate"
-
-class JailedKeyError(Exception):
-    pass
-
-
 class Executor(Process):
+    """Custom Process subclass for execution of protocol units.
+
+    """
 
     key: NodeKey
     queue: Queue
@@ -452,7 +450,6 @@ def __init__(self, key, queue, lock, context, inputs, n_retries):
         self._lock = lock
         self._unit_context = context
         self._inputs = inputs
-        assert n_retries >= 0
         self._n_retries = n_retries
 
     def run(self):
@@ -477,41 +474,10 @@ def unit(self) -> ProtocolUnit:
     def key(self) -> NodeKey:
         return self._key
 
-    @property
-    def queue(self) -> Queue:
-        return self._queue
-
-    @property
-    def lock(self) -> Lock:
-        return self._lock
-
-    @property
-    def inputs(self) -> dict:
-        return self._inputs
-
-    @classmethod
-    def from_key(
-        cls,
-        key: NodeKey,
-        queue: Queue,
-        lock: Lock,
-        context: Context,
-        inputs: dict,
-        n_retries: int,
-    ):
-        return cls(
-            key=key,
-            queue=queue,
-            lock=lock,
-            context=context,
-            inputs=inputs,
-            n_retries=n_retries,
-        )
-
     def put_result(self, result: ProtocolUnitResult):
         """Acquire lock, push key and result into the queue, release lock."""
-        with self.lock:
-            self.queue.put((self._key, result))
+        with self._lock:
+            self._queue.put((self._key, result))
 
     def execute_unit(self, context) -> ProtocolUnitResult | ProtocolUnitFailure:
         # this method assumes the context is in place and will be removed correctly
@@ -519,12 +485,18 @@ def execute_unit(self, context) -> ProtocolUnitResult | ProtocolUnitFailure:
         warnings.filterwarnings("ignore", message=r".*RDKit does not preserve.*")
         return self.unit.execute(context=context, **self._inputs)
 
+class JailedKeyError(Exception):
+    pass
 
 class ExecutorStack:
+    """Structure for coordinating the creation and management of
+    Executor processes.
+    """
 
     stack_size: int
     stack: list[Executor]
-    jail: dict[NodeKey, set[NodeKey]]
+    jail: dict[NodeKey, set[NodeKey]] # blocked node specified by a
+                                      # set of blocking nodes
     queue: Queue
     lock: Lock
 
@@ -574,6 +546,7 @@ def terminate_task(self, task_key: TaskKey):
 
     def push(self, node: NodeKey, unit_context: Context, inputs: dict, n_retries):
         with self.lock:
+            # node may be blocked from execution
             if node in self._jail.keys():
                 raise JailedKeyError(node)
 
@@ -582,7 +555,7 @@ def push(self, node: NodeKey, unit_context: Context, inputs: dict, n_retries):
 
             import warnings
             warnings.filterwarnings("ignore", message=r".*This process.*is multi-threaded,.*")
-            executor = Executor.from_key(
+            executor = Executor(
                 node, self.queue, self.lock, unit_context, inputs, n_retries
             )
             self._stack.append(executor)
@@ -657,8 +630,6 @@ def _get_statuses(self) -> tuple[set[Executor], set[Executor]]:
         return running, terminated
 
 
-from dataclasses import dataclass
-
 
 @dataclass
 class TaskData:
@@ -675,6 +646,11 @@ def to_ProtocolDAGResult(self) -> ProtocolDAGResult:
             extends_key=self.protocol_dag.extends_key,
         )
 
+class ResourceSignal(StrEnum):
+    MAINTAIN = "maintain"
+    SHRINK = "shrink"
+    GROW = "grow"
+    TERMINATE = "terminate"
 
 class AsynchronousComputeService(SynchronousComputeService):
     """Asynchronous compute service.
@@ -778,7 +754,7 @@ def stop(self):
             self.remove_all()
         super().stop()
 
-    def resource_monitor(self) -> ResourceSignal:
+    def _resource_monitor(self) -> ResourceSignal:
         capacity = 8
         sim_value = 3
         fin_value = 1
@@ -813,29 +789,27 @@ def resource_monitor(self) -> ResourceSignal:
         return ResourceSignal.TERMINATE
 
     def cycle(self, max_tasks, max_time) -> bool:
-        self.logger.info(f"Cycling: {len(self._task_data)}")
-
-        if max_time is not None and (time.time() - self._start_time) >= max_time:
-            self.logger.info("Exceeded maximum time")
-            self.stop()
-            return False
-
         # collect unit results
         self.process_results() # removes unit scratch
         self.consume_terminated_tasks() # only removes if task is done
+
+        # check if max tasks have been exceeded
         if max_tasks is not None and self.tasks_finished >= max_tasks:
             self.logger.info("Exceeded maximum tasks")
             self.stop()
             return False
 
-        signal = self.resource_monitor()
+        # Check if max time has been exceeded
+        if max_time is not None and (time.time() - self._start_time) >= max_time:
+            self.logger.info("Exceeded maximum time")
+            self.stop()
+            return False
+
+        # detemine next actions based on resource usage
+        signal = self._resource_monitor()
         match signal:
             case ResourceSignal.MAINTAIN:
-                #self.process_results()
-                #self.consume_terminated_tasks()
                 return True
-            case ResourceSignal.GROW:
-                pass
             case ResourceSignal.SHRINK:
                 proc = self._executor_stack.pop()
                 self.logger.info(f"Popping: {proc} -- {proc.key[1].key}")
@@ -843,6 +817,8 @@ def cycle(self, max_tasks, max_time) -> bool:
             case ResourceSignal.TERMINATE:
                 self.stop()
                 return False
+            case ResourceSignal.GROW:
+                pass
             case _:
                 raise RuntimeError("Received unknown ResourceSignal")
 
@@ -856,15 +832,8 @@ def cycle(self, max_tasks, max_time) -> bool:
                 self.add_task(*task)
                 self.tasks_claimed = 1 + self.tasks_claimed
 
-        # only submit enough tasks to fill the stack
-        n = self._executor_stack._stack_size - len(self._executor_stack._stack)
-        for key in tuple(self.next()):
-            tsk, unit = key
-
-            # TODO `next` should not return TERM or ROOT nodes
-            if unit in ("TERM", "ROOT"):
-                continue
-
+        for key in filter(lambda k: k[1] not in ("TERM", "ROOT"), self.next()):
+            (tsk, unit) = key
             task_data = self._task_data[tsk]
 
             inputs = _pu_to_pur(unit.inputs, task_data.results)
@@ -875,10 +844,9 @@ def cycle(self, max_tasks, max_time) -> bool:
             try:
                 self._executor_stack.push(key, context, inputs, self.n_retries)
                 self.logger.info(f"Pushing {key[1]} to the execution stack")
+                break
             except JailedKeyError:
                 continue
-            if not (n := n - 1):
-                break
 
         return True
 
@@ -894,7 +862,8 @@ def next(self) -> set[NodeKey]:
         running, terminated = self._executor_stack._get_statuses()
         running = {r.key for r in running}
         terminated = {t.key for t in terminated}
-        return self.available_units() - (running | terminated)
+        next_units = self.available_units() - (running | terminated)
+        return next_units
 
     def next_terminating_nodes(self) -> set[NodeKey]:
         completed = {node for node in self.available_units() if node[1] == "TERM"}

From 498a6d93205e561fcf41153204a0540f4dc322f3 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Mon, 30 Mar 2026 11:41:09 -0400
Subject: [PATCH 17/24] Add clarifying comments

---
 alchemiscale/compute/service.py | 48 ++++++++++++++++++++-------------
 1 file changed, 29 insertions(+), 19 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index 9b833aeb..a887876a 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -434,6 +434,8 @@ def stop(self):
 class Executor(Process):
     """Custom Process subclass for execution of protocol units.
 
+    Executors are responsible for creation of unit attempt directories
+    and data given a unit Context.
     """
 
     key: NodeKey
@@ -451,10 +453,16 @@ def __init__(self, key, queue, lock, context, inputs, n_retries):
         self._unit_context = context
         self._inputs = inputs
         self._n_retries = n_retries
+        self._validate()
+
+    def _validate(self):
+        if not self._n_retries >= 0:
+            raise ValueError("n_retries must be greater than or equal to 0")
 
     def run(self):
         attempt = 0
         while attempt <= self._n_retries:
+            # create attempt specific directories
             shared_dir = self._unit_context.shared / str(attempt)
             scratch_dir = self._unit_context.scratch / str(attempt)
             attempt_context = Context(shared=shared_dir, scratch=scratch_dir)
@@ -464,6 +472,7 @@ def run(self):
             if result.ok():
                 break
             attempt = attempt + 1
+        # put the result in the queue with lock
         self.put_result(result)
 
     @property
@@ -506,6 +515,11 @@ def __init__(self, stack_size: int):
         self._jail = {}
         self._queue = Queue()
         self._lock = Lock()
+        self._validate()
+
+    def _validate(self):
+        if not self._stack_size >= 1:
+            raise ValueError("stack_size must be greater than or equal to 1")
 
     @property
     def stack(self) -> list[Executor]:
@@ -528,15 +542,15 @@ def queue(self) -> Queue:
         return self._queue
 
     def terminate_all(self):
-        with self.lock:
-            for proc in self.stack:
+        with self._lock:
+            for proc in self._stack:
                 proc.terminate()
-            self.stack.clear()
+            self._stack.clear()
 
     def terminate_task(self, task_key: TaskKey):
-        with self.lock:
+        with self._lock:
             to_remove = set()
-            for proc in self.stack:
+            for proc in self._stack:
                 _task_key, _ = proc.key
                 if task_key == _task_key:
                     to_remove.add(proc)
@@ -545,7 +559,7 @@ def terminate_task(self, task_key: TaskKey):
                 self._stack.remove(proc)
 
     def push(self, node: NodeKey, unit_context: Context, inputs: dict, n_retries):
-        with self.lock:
+        with self._lock:
             # node may be blocked from execution
             if node in self._jail.keys():
                 raise JailedKeyError(node)
@@ -556,14 +570,14 @@ def push(self, node: NodeKey, unit_context: Context, inputs: dict, n_retries):
             import warnings
             warnings.filterwarnings("ignore", message=r".*This process.*is multi-threaded,.*")
             executor = Executor(
-                node, self.queue, self.lock, unit_context, inputs, n_retries
+                node, self._queue, self._lock, unit_context, inputs, n_retries
             )
             self._stack.append(executor)
             self._stack[-1].start()
 
     def pop(self):
         """Remove last process in the stack. This also clears the node from the jail."""
-        with self.lock:
+        with self._lock:
             if self._stack_size == 0:
                 raise IndexError("pop from empty stack")
 
@@ -589,21 +603,21 @@ def pop(self):
     def get_result(
         self,
     ) -> tuple[NodeKey, ProtocolUnitResult | ProtocolUnitFailure] | None:
-        if self.queue.qsize():
-            with self.lock:
+        if self._queue.qsize():
+            with self._lock:
                 # since qsize is not always reliable, we tentatively
                 # accept there might be results
                 try:
-                    res = self.queue.get_nowait()
+                    res = self._queue.get_nowait()
                 except queue.Empty:
                     return None
                 node_key, _ = res
-                self.remove_by_node_key(node_key)
+                self._remove_by_node_key(node_key)
                 return res
 
-    def remove_by_node_key(self, node_key: NodeKey):
+    def _remove_by_node_key(self, node_key: NodeKey):
         to_remove = None
-        for proc in self.stack:
+        for proc in self._stack:
             if proc.key == node_key:
                 to_remove = proc
                 break
@@ -622,7 +636,7 @@ def remove_by_node_key(self, node_key: NodeKey):
     def _get_statuses(self) -> tuple[set[Executor], set[Executor]]:
         running = set()
         terminated = set()
-        for proc in self.stack:
+        for proc in self._stack:
             if proc.is_alive():
                 running.add(proc)
             else:
@@ -691,9 +705,6 @@ def __init__(self, settings: ComputeServiceSettings):
             retry_max_seconds=self.settings.client_retry_max_seconds,
             verify=self.settings.client_verify,
         )
-
-        self._stop = False
-
         self.scopes = self.settings.scopes or [Scope()]
         self.shared_basedir = Path(self.settings.shared_basedir).absolute()
         self.shared_basedir.mkdir(exist_ok=True)
@@ -704,7 +715,6 @@ def __init__(self, settings: ComputeServiceSettings):
         self.keep_scratch = self.settings.keep_scratch
 
         self.compute_service_id = ComputeServiceID.new_from_name(self.name)
-        self._stop = False
 
     def _initialize_dag_tree(self):
         self._dag_tree = nx.DiGraph()

From 8cc4095737dcdbc6c04cf6875cc767ba10b55194 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Mon, 30 Mar 2026 11:55:16 -0400
Subject: [PATCH 18/24] Add monitor compute module

---
 alchemiscale/compute/monitor.py | 170 ++++++++++++++++++++++++++++++++
 alchemiscale/compute/service.py |  83 ++++++++++++----
 network.py                      |   3 -
 service.py                      |  12 +++
 4 files changed, 245 insertions(+), 23 deletions(-)
 create mode 100644 alchemiscale/compute/monitor.py

diff --git a/alchemiscale/compute/monitor.py b/alchemiscale/compute/monitor.py
new file mode 100644
index 00000000..d5cdcbc3
--- /dev/null
+++ b/alchemiscale/compute/monitor.py
@@ -0,0 +1,170 @@
+from abc import abstractmethod
+from enum import auto, IntEnum
+import os
+import subprocess
+import time
+from threading import Lock
+
+
+class ResourceSignal(IntEnum):
+    # order matters, higher priority signals should appear at top
+    TERMINATE = auto()
+    SHRINK = auto()
+    MAINTAIN = auto()
+    GROW = auto()
+
+
+class Monitor:
+
+    def __init__(self, settings):
+        self._setup(settings)
+        self._lock = Lock()
+        self._terminate = False
+
+    def signal(self) -> ResourceSignal:
+        with self._lock:
+            return self._signal()
+
+    def monitor_cycle(self):
+        """Method to be run in a thread, changing state of the monitor
+        such that the signal method can process the results.
+        """
+        while not self._terminate:
+            with self._lock:
+                self._monitor_cycle()
+            # TODO: make configurable
+            time.sleep(1)
+
+    @abstractmethod
+    def _setup(self, settings):
+        raise NotImplementedError
+
+    @abstractmethod
+    def _monitor_cycle(self):
+        """Mutating method"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _signal(self) -> ResourceSignal:
+        raise NotImplementedError
+
+
+class GPUMonitor(Monitor):
+
+    def _setup(self, settings):
+        self.history = []
+        self.gpu_index = settings.gpu_monitor_gpu_id
+
+    @staticmethod
+    def _nvidia_smi() -> int:
+        fields = ["index", "utilization.gpu"]
+        cmd = [
+            "nvidia-smi",
+            f"--query-gpu={''.join(fields)}",
+            "--format=csv",
+        ]
+        completed_process = subprocess.run(cmd, capture_output=True)
+        output = completed_process.stdout.decode()
+        num_fields = len(fields)
+        # discard header
+        gpu_entries = output.split("\n")[1:]
+        for gpu in gpu_entries:
+            idx, util = gpu.split(", ")
+            if idx == self.gpu_index:
+                util, _ = util.split(" ")
+                util = int(util)
+                return util
+
+    def _monitor_cycle(self):
+        util = self._nvidia_smi()
+        self.history.append(util)
+        self.history = self.history[-60:]
+
+    def _signal(self) -> ResourceSignal:
+        utilization = sum(self.history) / len(self.history)
+        if utilization < self.grow_limit:
+            return ResourceSignal.GROW
+        elif utilization < self.maintain_limit:
+            return ResourceSignal.MAINTAIN
+        return ResourceSignal.SHRINK
+
+
+class CPUMonitor(Monitor):
+
+    # TODO: make configurable
+    grow_limit = 0.50
+    maintain_limit = 0.75
+
+    def _setup(self, settings):
+        self.history = []
+
+    def _signal(self) -> ResourceSignal:
+        total_load = sum(self.history) / len(self.history)
+        if total_load < self.grow_limit:
+            return ResourceSignal.GROW
+        elif total_load < self.maintain_limit:
+            return ResourceSignal.MAINTAIN
+        return ResourceSignal.SHRINK
+
+    def _monitor_cycle(self):
+        load, _, _ = os.getloadavg()
+        cpu_count = os.cpu_count()
+        total_load = load / cpu_count
+        self.history.append(total_load)
+        self.history = self.history[-60:]
+
+
+class MemInfoParseError(Exception):
+    pass
+
+
+class MemoryMonitor(Monitor):
+
+    # TODO: make configurable
+    grow_limit = 0.65
+    maintain_limit = 0.85
+
+    def _setup(self, settings):
+        self.history = []
+
+    @staticmethod
+    def _get_memory() -> tuple[int, int]:
+        """Parse /proc/meminfo for memory info."""
+        total = None
+        available = None
+        with open("/proc/meminfo", "r") as f:
+            for line in f:
+                if line.startswith("MemTotal") or line.startswith("MemAvailable"):
+                    field, size_kb, _ = line.split()
+                    field = field.rstrip(":")
+                    size_kb = int(size_kb)
+                    match field:
+                        case "MemTotal":
+                            total = size_kb
+                        case "MemAvailable":
+                            available = size_kb
+                # break from for loop and avoid error-raising else block
+                if total is not None and available is not None:
+                    break
+            else:
+                # unable to determine the MemTotal or MemAvailable
+                raise MemInfoParseError
+        return total, available
+
+    def _monitor_cycle(self):
+        try:
+            total, avail = self._get_memory()
+            fraction_used = (total - avail) / total
+            self.history.append(fraction_used)
+            # roughly the last minute of entries
+            self.history = self.history[-60:]
+        except Exception:
+            self._terminate = True
+
+    def _signal(self) -> ResourceSignal:
+        fraction_used = sum(self.history) / len(self.history)
+        if fraction_used < self.grow_limit:
+            return ResourceSignal.GROW
+        elif fraction_used < self.maintain_limit:
+            return ResourceSignal.MAINTAIN
+        return ResourceSignal.SHRINK
diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index a887876a..e24ee87d 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -10,11 +10,11 @@
 import logging
 from uuid import uuid4
 import threading
+import os
 from pathlib import Path
 import queue
 import shutil
 from typing import Any
-from enum import StrEnum
 from multiprocessing import Process, Queue, Lock
 from dataclasses import dataclass
 
@@ -35,6 +35,7 @@
 import networkx as nx
 
 from .client import AlchemiscaleComputeClient
+from .monitor import ResourceSignal, MemoryMonitor, GPUMonitor, CPUMonitor
 from .settings import ComputeServiceSettings
 from ..storage.models import ComputeServiceID
 from ..models import Scope, ScopedKey
@@ -431,6 +432,7 @@ def stop(self):
     ProtocolUnit | str,
 ]  # None covers root node condition
 
+
 class Executor(Process):
     """Custom Process subclass for execution of protocol units.
 
@@ -445,7 +447,7 @@ class Executor(Process):
     inputs: dict
     n_retries: int
 
-    def __init__(self, key, queue, lock, context, inputs, n_retries):
+    def __init__(self, key, queue, lock, context, inputs, n_retries, env=None):
         super().__init__()
         self._key = key
         self._queue = queue
@@ -453,6 +455,7 @@ def __init__(self, key, queue, lock, context, inputs, n_retries):
         self._unit_context = context
         self._inputs = inputs
         self._n_retries = n_retries
+        self._env = env or {}
         self._validate()
 
     def _validate(self):
@@ -460,6 +463,8 @@ def _validate(self):
             raise ValueError("n_retries must be greater than or equal to 0")
 
     def run(self):
+        # update environment before running unit
+        os.environ |= self._env
         attempt = 0
         while attempt <= self._n_retries:
             # create attempt specific directories
@@ -491,12 +496,15 @@ def put_result(self, result: ProtocolUnitResult):
     def execute_unit(self, context) -> ProtocolUnitResult | ProtocolUnitFailure:
         # this method assumes the context is in place and will be removed correctly
         import warnings
+
         warnings.filterwarnings("ignore", message=r".*RDKit does not preserve.*")
         return self.unit.execute(context=context, **self._inputs)
 
+
 class JailedKeyError(Exception):
     pass
 
+
 class ExecutorStack:
     """Structure for coordinating the creation and management of
     Executor processes.
@@ -504,8 +512,8 @@ class ExecutorStack:
 
     stack_size: int
     stack: list[Executor]
-    jail: dict[NodeKey, set[NodeKey]] # blocked node specified by a
-                                      # set of blocking nodes
+    # blocked node specified by a set of blocking nodes
+    jail: dict[NodeKey, set[NodeKey]]
     queue: Queue
     lock: Lock
 
@@ -558,7 +566,14 @@ def terminate_task(self, task_key: TaskKey):
                 proc.terminate()
                 self._stack.remove(proc)
 
-    def push(self, node: NodeKey, unit_context: Context, inputs: dict, n_retries):
+    def push(
+        self,
+        node: NodeKey,
+        unit_context: Context,
+        inputs: dict,
+        n_retries: int,
+        env: dict[str, str],
+    ):
         with self._lock:
             # node may be blocked from execution
             if node in self._jail.keys():
@@ -568,9 +583,12 @@ def push(self, node: NodeKey, unit_context: Context, inputs: dict, n_retries):
             unit_context.shared.mkdir()
 
             import warnings
-            warnings.filterwarnings("ignore", message=r".*This process.*is multi-threaded,.*")
+
+            warnings.filterwarnings(
+                "ignore", message=r".*This process.*is multi-threaded,.*"
+            )
             executor = Executor(
-                node, self._queue, self._lock, unit_context, inputs, n_retries
+                node, self._queue, self._lock, unit_context, inputs, n_retries, env=env
             )
             self._stack.append(executor)
             self._stack[-1].start()
@@ -644,7 +662,6 @@ def _get_statuses(self) -> tuple[set[Executor], set[Executor]]:
         return running, terminated
 
 
-
 @dataclass
 class TaskData:
     protocol_dag: ProtocolDAG
@@ -660,11 +677,6 @@ def to_ProtocolDAGResult(self) -> ProtocolDAGResult:
             extends_key=self.protocol_dag.extends_key,
         )
 
-class ResourceSignal(StrEnum):
-    MAINTAIN = "maintain"
-    SHRINK = "shrink"
-    GROW = "grow"
-    TERMINATE = "terminate"
 
 class AsynchronousComputeService(SynchronousComputeService):
     """Asynchronous compute service.
@@ -677,11 +689,14 @@ class AsynchronousComputeService(SynchronousComputeService):
     _dag_tree: nx.DiGraph
     _executor_stack: ExecutorStack
     _task_data: dict[TaskKey, TaskData]
+    _child_env: dict[str, str]
 
     def __init__(self, settings: ComputeServiceSettings):
-
+        self._child_env = dict()
         self._task_data = dict()
         self._initialize_dag_tree()
+        self._resource_monitors = None
+        self._initialize_resource_monitors(settings)
 
         self.settings = settings
 
@@ -716,6 +731,23 @@ def __init__(self, settings: ComputeServiceSettings):
 
         self.compute_service_id = ComputeServiceID.new_from_name(self.name)
 
+    def _initialize_resource_monitors(self, settings):
+        self._resource_monitors = []
+
+        if settings.memory_monitor_enabled:
+            self._resource_monitors.append(MemoryMonitor(settings))
+
+        if settings.cpu_monitor_enabled:
+            self._resource_monitors.append(CPUMonitor(settings))
+
+        if settings.gpu_monitor_enabled:
+            self._resource_monitors.append(GPUMonitor(settings))
+            # reliable monitoring of the GPU requires pinning the GPU index
+            self._child_env |= {"CUDA_VISIBLE_DEVICES": settings.gpu_monitor_gpu_id}
+
+        for monitor in self._resource_monitors:
+            threading.Thread(target=monitor.monitor_cycle, daemon=True).start()
+
     def _initialize_dag_tree(self):
         self._dag_tree = nx.DiGraph()
         root_node: NodeKey = (None, "ROOT")
@@ -764,6 +796,13 @@ def stop(self):
             self.remove_all()
         super().stop()
 
+    def _get_resource_signal(self) -> ResourceSignal:
+        # without monitors, signal that the stack can always grow if there is room
+        if len(self._resource_monitors) == 0:
+            return ResourceSignal.GROW
+        # otherwise respect the highest priority signal from all monitors
+        return min(monitor.signal() for monitor in self._resource_monitors)
+
     def _resource_monitor(self) -> ResourceSignal:
         capacity = 8
         sim_value = 3
@@ -783,7 +822,9 @@ def _resource_monitor(self) -> ResourceSignal:
                 case _:
                     total = total + init_value
 
-        self.logger.info(f"Num jailed: {len(self._executor_stack._jail)} -- {total}/{capacity}")
+        self.logger.info(
+            f"Num jailed: {len(self._executor_stack._jail)} -- {total}/{capacity}"
+        )
         if total == capacity:
             self.logger.info("MAINTAINING")
             return ResourceSignal.MAINTAIN
@@ -800,8 +841,8 @@ def _resource_monitor(self) -> ResourceSignal:
 
     def cycle(self, max_tasks, max_time) -> bool:
         # collect unit results
-        self.process_results() # removes unit scratch
-        self.consume_terminated_tasks() # only removes if task is done
+        self.process_results()  # removes unit scratch
+        self.consume_terminated_tasks()  # only removes if task is done
 
         # check if max tasks have been exceeded
         if max_tasks is not None and self.tasks_finished >= max_tasks:
@@ -816,7 +857,7 @@ def cycle(self, max_tasks, max_time) -> bool:
             return False
 
         # detemine next actions based on resource usage
-        signal = self._resource_monitor()
+        signal = self._get_resource_signal()
         match signal:
             case ResourceSignal.MAINTAIN:
                 return True
@@ -843,7 +884,7 @@ def cycle(self, max_tasks, max_time) -> bool:
                 self.tasks_claimed = 1 + self.tasks_claimed
 
         for key in filter(lambda k: k[1] not in ("TERM", "ROOT"), self.next()):
-            (tsk, unit) = key
+            tsk, unit = key
             task_data = self._task_data[tsk]
 
             inputs = _pu_to_pur(unit.inputs, task_data.results)
@@ -852,7 +893,9 @@ def cycle(self, max_tasks, max_time) -> bool:
             context = Context(scratch=unit_scratch_dir, shared=unit_shared_dir)
 
             try:
-                self._executor_stack.push(key, context, inputs, self.n_retries)
+                self._executor_stack.push(
+                    key, context, inputs, self.n_retries, env=self._child_env
+                )
                 self.logger.info(f"Pushing {key[1]} to the execution stack")
                 break
             except JailedKeyError:
diff --git a/network.py b/network.py
index b4a1e923..0ed1c59d 100644
--- a/network.py
+++ b/network.py
@@ -106,9 +106,6 @@ def network_tyk2():
         )
         for edge in tyk2s.connections
     ]
-    # breakpoint()
-    # pu = solvent_network[0].create().protocol_units[0]
-    # pu.execute(context=None, initialization=None)
 
     return AlchemicalNetwork(
         edges=(solvent_network + complex_network),
diff --git a/service.py b/service.py
index c6722286..9ca1aba0 100644
--- a/service.py
+++ b/service.py
@@ -1,13 +1,25 @@
+from dataclasses import dataclass
 import logging
 import time
 
 from alchemiscale.compute.service import AsynchronousComputeService, ExecutorStack, InterruptableSleep
 from alchemiscale.storage.models import ComputeServiceID
 
+@dataclass
+class MockSettings:
+    memory_monitor_enabled: bool
+    cpu_monitor_enabled: bool
+    gpu_monitor_enabled: bool
+    gpu_monitor_gpu_id: int
+
 class MockService(AsynchronousComputeService):
 
     def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, keep_shared, n_retries, claim_limit, task_generator):
+        self._child_env = dict()
         self._initialize_dag_tree()
+        settings = MockSettings(memory_monitor_enabled=True, cpu_monitor_enabled=True, gpu_monitor_enabled=False, gpu_monitor_gpu_id="0")
+        #settings = MockSettings(memory_monitor_enabled=False, cpu_monitor_enabled=False, gpu_monitor_enabled=False, gpu_monitor_gpu_id="0")
+        self._initialize_resource_monitors(settings)
         self._task_data = dict()
         self._executor_stack = ExecutorStack(stack_size)
 

From c6b6cb1d67a95ff982cd758243fcd93a7bf44c33 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Tue, 31 Mar 2026 13:48:24 -0400
Subject: [PATCH 19/24] Expand settings

---
 alchemiscale/compute/monitor.py  | 36 ++++++++------
 alchemiscale/compute/service.py  | 85 ++++++++++++++++++++++----------
 alchemiscale/compute/settings.py | 68 +++++++++++++++++++++++++
 service.py                       | 33 +++++++++++--
 4 files changed, 178 insertions(+), 44 deletions(-)

diff --git a/alchemiscale/compute/monitor.py b/alchemiscale/compute/monitor.py
index d5cdcbc3..6a45712d 100644
--- a/alchemiscale/compute/monitor.py
+++ b/alchemiscale/compute/monitor.py
@@ -21,6 +21,11 @@ def __init__(self, settings):
         self._lock = Lock()
         self._terminate = False
 
+        if not hasattr(self, "sample_time"):
+            raise AttributeError(
+                f"{self.__class__.__name__} implementation requires definition of the `sample_time` attribute"
+            )
+
     def signal(self) -> ResourceSignal:
         with self._lock:
             return self._signal()
@@ -32,8 +37,7 @@ def monitor_cycle(self):
         while not self._terminate:
             with self._lock:
                 self._monitor_cycle()
-            # TODO: make configurable
-            time.sleep(1)
+            time.sleep(self.sample_time)
 
     @abstractmethod
     def _setup(self, settings):
@@ -53,7 +57,11 @@ class GPUMonitor(Monitor):
 
     def _setup(self, settings):
         self.history = []
-        self.gpu_index = settings.gpu_monitor_gpu_id
+        self.history_size = settings.gpu_monitor_sample_history_size
+        self.gpu_index = settings.gpu_monitor_gpu_index
+        self.grow_limit = settings.gpu_monitor_grow_limit
+        self.maintain_limit = settings.gpu_monitor_maintain_limit
+        self.sample_time = settings.gpu_monitor_sample_time
 
     @staticmethod
     def _nvidia_smi() -> int:
@@ -78,7 +86,7 @@ def _nvidia_smi() -> int:
     def _monitor_cycle(self):
         util = self._nvidia_smi()
         self.history.append(util)
-        self.history = self.history[-60:]
+        self.history = self.history[-self.history_size :]
 
     def _signal(self) -> ResourceSignal:
         utilization = sum(self.history) / len(self.history)
@@ -91,12 +99,12 @@ def _signal(self) -> ResourceSignal:
 
 class CPUMonitor(Monitor):
 
-    # TODO: make configurable
-    grow_limit = 0.50
-    maintain_limit = 0.75
-
     def _setup(self, settings):
         self.history = []
+        self.history_size = settings.cpu_monitor_sample_history_size
+        self.sample_time = settings.cpu_monitor_sample_time
+        self.grow_limit = settings.cpu_monitor_grow_limit
+        self.maintain_limit = settings.cpu_monitor_maintain_limit
 
     def _signal(self) -> ResourceSignal:
         total_load = sum(self.history) / len(self.history)
@@ -111,7 +119,7 @@ def _monitor_cycle(self):
         cpu_count = os.cpu_count()
         total_load = load / cpu_count
         self.history.append(total_load)
-        self.history = self.history[-60:]
+        self.history = self.history[-self.history_size :]
 
 
 class MemInfoParseError(Exception):
@@ -120,12 +128,12 @@ class MemInfoParseError(Exception):
 
 class MemoryMonitor(Monitor):
 
-    # TODO: make configurable
-    grow_limit = 0.65
-    maintain_limit = 0.85
-
     def _setup(self, settings):
         self.history = []
+        self.history_size = settings.memory_monitor_sample_history_size
+        self.sample_time = settings.memory_monitor_sample_time
+        self.grow_limit = settings.memory_monitor_grow_limit
+        self.maintain_limit = settings.memory_monitor_maintain_limit
 
     @staticmethod
     def _get_memory() -> tuple[int, int]:
@@ -157,7 +165,7 @@ def _monitor_cycle(self):
             fraction_used = (total - avail) / total
             self.history.append(fraction_used)
             # roughly the last minute of entries
-            self.history = self.history[-60:]
+            self.history = self.history[-self.history_size :]
         except Exception:
             self._terminate = True
 
diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index e24ee87d..d1e8032c 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -692,35 +692,26 @@ class AsynchronousComputeService(SynchronousComputeService):
     _child_env: dict[str, str]
 
     def __init__(self, settings: ComputeServiceSettings):
-        self._child_env = dict()
+        self.settings = settings
+
+        # asynccomputeservice specific data structures and resource
+        # monitors.
+        self._child_env = dict()  # mods to child process env
         self._task_data = dict()
         self._initialize_dag_tree()
-        self._resource_monitors = None
-        self._initialize_resource_monitors(settings)
-
-        self.settings = settings
+        self._initialize_resource_monitors()
 
         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
         self.claim_limit = self.settings.claim_limit
 
-        self.scheduler = sched.scheduler(time.monotonic, time.sleep)
-
-        self.client = AlchemiscaleComputeClient(
-            self.settings.api_url,
-            self.settings.identifier,
-            self.settings.key,
-            cache_directory=self.settings.client_cache_directory,
-            cache_size_limit=self.settings.client_cache_size_limit,
-            use_local_cache=self.settings.client_use_local_cache,
-            max_retries=self.settings.client_max_retries,
-            retry_base_seconds=self.settings.client_retry_base_seconds,
-            retry_max_seconds=self.settings.client_retry_max_seconds,
-            verify=self.settings.client_verify,
-        )
+        self.client = self._initialize_client()
         self.scopes = self.settings.scopes or [Scope()]
+
         self.shared_basedir = Path(self.settings.shared_basedir).absolute()
         self.shared_basedir.mkdir(exist_ok=True)
         self.keep_shared = self.settings.keep_shared
@@ -731,19 +722,59 @@ def __init__(self, settings: ComputeServiceSettings):
 
         self.compute_service_id = ComputeServiceID.new_from_name(self.name)
 
-    def _initialize_resource_monitors(self, settings):
+        self.int_sleep = InterruptableSleep()
+        self._initialize_logger()
+
+    def _initialize_logger(self):
+        extra = {"compute_service_id": str(self.compute_service_id)}
+        logger = logging.getLogger("AlchemiscaleAsynchronousComputeService")
+        logger.setLevel(self.settings.loglevel)
+
+        formatter = logging.Formatter(
+            "[%(asctime)s] [%(compute_service_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 _initialize_client(self):
+        return AlchemiscaleComputeClient(
+            api_url=self.settings.api_url,
+            identifier=self.settings.identifier,
+            key=self.settings.key,
+            cache_directory=self.settings.client_cache_directory,
+            cache_size_limit=self.settings.client_cache_size_limit,
+            use_local_cache=self.settings.client_use_local_cache,
+            max_retries=self.settings.client_max_retries,
+            retry_base_seconds=self.settings.client_retry_base_seconds,
+            retry_max_seconds=self.settings.client_retry_max_seconds,
+            verify=self.settings.client_verify,
+        )
+
+    def _initialize_resource_monitors(self):
         self._resource_monitors = []
 
-        if settings.memory_monitor_enabled:
-            self._resource_monitors.append(MemoryMonitor(settings))
+        if self.settings.memory_monitor_enabled:
+            self._resource_monitors.append(MemoryMonitor(self.settings))
 
-        if settings.cpu_monitor_enabled:
-            self._resource_monitors.append(CPUMonitor(settings))
+        if self.settings.cpu_monitor_enabled:
+            self._resource_monitors.append(CPUMonitor(self.settings))
 
-        if settings.gpu_monitor_enabled:
-            self._resource_monitors.append(GPUMonitor(settings))
+        if self.settings.gpu_monitor_enabled:
+            self._resource_monitors.append(GPUMonitor(self.settings))
             # reliable monitoring of the GPU requires pinning the GPU index
-            self._child_env |= {"CUDA_VISIBLE_DEVICES": settings.gpu_monitor_gpu_id}
+            self._child_env |= {
+                "CUDA_VISIBLE_DEVICES": self.settings.gpu_monitor_gpu_id
+            }
 
         for monitor in self._resource_monitors:
             threading.Thread(target=monitor.monitor_cycle, daemon=True).start()
diff --git a/alchemiscale/compute/settings.py b/alchemiscale/compute/settings.py
index 2791df68..e7121aa9 100644
--- a/alchemiscale/compute/settings.py
+++ b/alchemiscale/compute/settings.py
@@ -136,6 +136,74 @@ def validate_scopes(cls, values) -> list[Scope]:
         return _values
 
 
+class AsynchronousComputeServiceSettings(ComputeServiceSettings):
+
+    stack_size: int = Field(
+        2,
+        description="The number of concurrent protocol units that are able to run at once.",
+    )
+
+    gpu_monitor_enabled: bool = Field(
+        True, description="If the GPU monitor is enabled."
+    )
+    gpu_monitor_gpu_index: str = Field(
+        "0",
+        description="The GPU index to perform calculations on. This sets the CUDA_VISIBLE_DEVICES environment variable for spawned compute tasks.",
+    )
+    gpu_monitor_grow_limit: float = Field(
+        0.7,
+        description="GPU utilization percentage below this value will allow greater concurrency. See utilization.gpu in nvidia-smi for more information.",
+    )
+    gpu_monitor_maintain_limit: float = Field(
+        1.1,
+        description="GPU utilization percentage above this value will scale back concurrency. See utilization.gpu in nvidia-smi for more information.",
+    )
+    gpu_monitor_sample_time: int = Field(
+        1,
+        description="Number of seconds between collecting GPU utilization measurements.",
+    )
+    gpu_monitor_sample_history_size: int = Field(
+        60,
+        description="Maximum number of samples to use when considering reactive concurrency behavior.",
+    )
+    memory_monitor_enabled: bool = Field(
+        True, description="If the memory monitor is enabled."
+    )
+    memory_monitor_grow_limit: float = Field(
+        0.7,
+        description="Memory usage percentage below this value will allow greater concurrency.",
+    )
+    memory_monitor_maintain_limit: float = Field(
+        0.9,
+        description="Memory usage percentage above this value will scale back concurrency.",
+    )
+    memory_monitor_sample_time: int = Field(
+        1, description="Number of seconds between collecting memory usage measurements."
+    )
+    memory_monitor_sample_history_size: int = Field(
+        60,
+        description="Maximum number of samples to use when considering reactive concurrency behavior.",
+    )
+    cpu_monitor_enabled: bool = Field(
+        True, description="If the CPU monitor is enabled."
+    )
+    cpu_monitor_grow_limit: float = Field(
+        0.8,
+        description="CPU usage percentage below this value will allow greater concurrency.",
+    )
+    cpu_monitor_maintain_limit: float = Field(
+        1.2,
+        description="CPU usage percentage above this value will scale back concurrency.",
+    )
+    cpu_monitor_sample_time: int = Field(
+        1, description="Number of seconds between collecting CPU usage measurements."
+    )
+    cpu_monitor_sample_history_size: int = Field(
+        60,
+        description="Maximum number of samples to use when considering reactive concurrency behavior.",
+    )
+
+
 class ComputeManagerSettings(BaseModel):
     name: str = Field(
         ...,
diff --git a/service.py b/service.py
index 9ca1aba0..00c23c9e 100644
--- a/service.py
+++ b/service.py
@@ -8,18 +8,45 @@
 @dataclass
 class MockSettings:
     memory_monitor_enabled: bool
+    memory_monitor_sample_time: int
+    memory_monitor_sample_history_size: int
+    memory_monitor_grow_limit: float
+    memory_monitor_maintain_limit: float
     cpu_monitor_enabled: bool
+    cpu_monitor_sample_time: int
+    cpu_monitor_sample_history_size: int
+    cpu_monitor_grow_limit: float
+    cpu_monitor_maintain_limit: float
     gpu_monitor_enabled: bool
     gpu_monitor_gpu_id: int
+    gpu_monitor_sample_time: int
+    gpu_monitor_sample_history_size: int
+    gpu_monitor_grow_limit: float
+    gpu_monitor_maintain_limit: float
 
 class MockService(AsynchronousComputeService):
 
     def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, keep_shared, n_retries, claim_limit, task_generator):
         self._child_env = dict()
         self._initialize_dag_tree()
-        settings = MockSettings(memory_monitor_enabled=True, cpu_monitor_enabled=True, gpu_monitor_enabled=False, gpu_monitor_gpu_id="0")
-        #settings = MockSettings(memory_monitor_enabled=False, cpu_monitor_enabled=False, gpu_monitor_enabled=False, gpu_monitor_gpu_id="0")
-        self._initialize_resource_monitors(settings)
+        self.settings = MockSettings(memory_monitor_enabled=True,
+                                     cpu_monitor_enabled=True,
+                                     gpu_monitor_enabled=False,
+                                     gpu_monitor_gpu_id="0",
+                                     gpu_monitor_sample_time=1,
+                                     gpu_monitor_sample_history_size=60,
+                                     gpu_monitor_grow_limit=0.7,
+                                     gpu_monitor_maintain_limit=0.9,
+                                     memory_monitor_sample_time=1,
+                                     memory_monitor_sample_history_size=60,
+                                     memory_monitor_grow_limit=0.7,
+                                     memory_monitor_maintain_limit=0.9,
+                                     cpu_monitor_sample_time=1,
+                                     cpu_monitor_sample_history_size=60,
+                                     cpu_monitor_grow_limit=0.9,
+                                     cpu_monitor_maintain_limit=1.2,
+                                     )
+        self._initialize_resource_monitors()
         self._task_data = dict()
         self._executor_stack = ExecutorStack(stack_size)
 

From 6dd0dccd949b0e749d6dff3d7568ebe2fdcb436f Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Tue, 31 Mar 2026 15:31:56 -0400
Subject: [PATCH 20/24] Test task completion and result push

---
 alchemiscale/compute/service.py               | 19 +++-
 .../compute/client/test_compute_service.py    | 89 ++++++++++++++++++-
 2 files changed, 104 insertions(+), 4 deletions(-)

diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index d1e8032c..d098b165 100644
--- a/alchemiscale/compute/service.py
+++ b/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]:
diff --git a/alchemiscale/tests/integration/compute/client/test_compute_service.py b/alchemiscale/tests/integration/compute/client/test_compute_service.py
index fd0032fc..03d975cf 100644
--- a/alchemiscale/tests/integration/compute/client/test_compute_service.py
+++ b/alchemiscale/tests/integration/compute/client/test_compute_service.py
@@ -11,8 +11,14 @@
 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 +261,82 @@ 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

From 7c879419aa3d6c737dd9957cfa4ff2b325282355 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Tue, 31 Mar 2026 16:11:57 -0400
Subject: [PATCH 21/24] Add docstrings for compute.monitor

---
 alchemiscale/compute/monitor.py | 53 ++++++++++++++++++++++++++++++---
 1 file changed, 49 insertions(+), 4 deletions(-)

diff --git a/alchemiscale/compute/monitor.py b/alchemiscale/compute/monitor.py
index 6a45712d..80467253 100644
--- a/alchemiscale/compute/monitor.py
+++ b/alchemiscale/compute/monitor.py
@@ -1,3 +1,9 @@
+"""
+:mod:`alchemiscale.compute.monitor` --- resource monitoring for compute services
+================================================================================
+
+"""
+
 from abc import abstractmethod
 from enum import auto, IntEnum
 import os
@@ -6,6 +12,14 @@
 from threading import Lock
 
 
+# Signal to be issued by a resource manager to a compute
+# service. These signals are suggestions and are meant to inform the
+# service but will not force a particular behavior.
+#
+#    1. TERMINATE: tool/resource failure, shut stop all calculations
+#    2. SHRINK: resource is oversubscribed, scale down
+#    3. MAINTAIN: keep at current subscription
+#    4. GROW: resource is undersubscribed
 class ResourceSignal(IntEnum):
     # order matters, higher priority signals should appear at top
     TERMINATE = auto()
@@ -15,8 +29,25 @@ class ResourceSignal(IntEnum):
 
 
 class Monitor:
+    """Base class for a resource monitor.
+
+    This class handles provides three abstract methods for developers to define.
+
+    1. _setup: uses ComputeServiceSettings to set up datastructures
+               needed for monitoring. Note that this method must
+               define the ``sample_time`` attribute.
+    2. _monitor_cycle: method that updates the above data structures
+               to later be analyzed.
+    3. _signal: from the data collected by _monitor_cycle, return a
+               ResourceSignal.
+
+    Thread locking is handled automatically by the wrapping ``signal``
+    and ``monitor_cycle`` methods to settle race conditions and data
+    corruption.
+
+    """
 
-    def __init__(self, settings):
+    def __init__(self, settings: ComputeServiceSettings):
         self._setup(settings)
         self._lock = Lock()
         self._terminate = False
@@ -27,12 +58,13 @@ def __init__(self, settings):
             )
 
     def signal(self) -> ResourceSignal:
+        """Issue signal based on collected measurements."""
         with self._lock:
             return self._signal()
 
     def monitor_cycle(self):
-        """Method to be run in a thread, changing state of the monitor
-        such that the signal method can process the results.
+        """Continuously cycle, collecting results as defined by the
+        derived class _monitor_cycle implementation.
         """
         while not self._terminate:
             with self._lock:
@@ -41,15 +73,27 @@ def monitor_cycle(self):
 
     @abstractmethod
     def _setup(self, settings):
+        """Abstract method for setting up data structures needed for
+        storing resource measurements. These structures should be
+        mutated by ``_monitor_cycle`` and read ``_signal`` for issuing
+        resource signals.
+        """
         raise NotImplementedError
 
     @abstractmethod
     def _monitor_cycle(self):
-        """Mutating method"""
+        """Abstract method for collecting and updating internal data
+        to later be analyzed by ``_signal``.
+
+        """
         raise NotImplementedError
 
     @abstractmethod
     def _signal(self) -> ResourceSignal:
+        """Abstract method for analyzing data collected by
+        ``_monitor_cycle``.
+
+        """
         raise NotImplementedError
 
 
@@ -65,6 +109,7 @@ def _setup(self, settings):
 
     @staticmethod
     def _nvidia_smi() -> int:
+        """Collect utilization of the GPU."""
         fields = ["index", "utilization.gpu"]
         cmd = [
             "nvidia-smi",

From 89edebca0501888d7ed89cdda4c9bef6e91d4b31 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Tue, 31 Mar 2026 16:13:56 -0400
Subject: [PATCH 22/24] Remove local testing files

---
 local_testing.py |  79 ---------------------------------
 network.py       | 113 -----------------------------------------------
 service.py       | 113 -----------------------------------------------
 utils.py         |  32 --------------
 4 files changed, 337 deletions(-)
 delete mode 100644 local_testing.py
 delete mode 100644 network.py
 delete mode 100644 service.py
 delete mode 100644 utils.py

diff --git a/local_testing.py b/local_testing.py
deleted file mode 100644
index a4a9e35a..00000000
--- a/local_testing.py
+++ /dev/null
@@ -1,79 +0,0 @@
-# Local Variables:
-# compile-command: "./env/bin/python local_testing.py"
-# python-shell-interpreter: "./env/bin/python"
-# End:
-
-
-from pathlib import Path
-import shutil
-
-from alchemiscale.compute.service import (
-    Executor,
-    ExecutorStack,
-    NodeKey,
-    TaskKey,
-    TaskData,
-)
-from alchemiscale.models import ScopedKey
-from gufe import AlchemicalNetwork
-from gufe.protocols.protocolunit import Context, ProtocolUnitResult, ProtocolUnitFailure
-from gufe.protocols.protocoldag import _pu_to_pur, ProtocolDAGResult
-from gufe.tokenization import GufeKey
-from gufe.tests.test_protocol import BrokenProtocol
-
-import networkx as nx
-
-import service
-import utils
-
-SCRATCH_DIR = Path("./acs_testing/scratch")
-SHARED_DIR = Path("./acs_testing/shared")
-STACKSIZE = 4
-N_RETRIES = 0
-MAX_TASKS = 1
-MAX_TIME = None
-KEEP_SHARED = False
-KEEP_SCRATCH = False
-CLAIM_LIMIT = 2
-
-SCRATCH_DIR.mkdir(parents=True, exist_ok=True)
-SHARED_DIR.mkdir(parents=True, exist_ok=True)
-
-
-def create_tyk2():
-    try:
-        return AlchemicalNetwork.from_json(file="network.json")
-    except FileNotFoundError:
-        print("\tCould not load from file, creating new network")
-        from network import network_tyk2
-
-        _tyk2 = network_tyk2()
-        _tyk2.to_json(file="network.json")
-        return _tyk2
-
-
-if __name__ == "__main__":
-
-    print("Creating network")
-    tyk2 = create_tyk2()
-
-    task_generator = (
-        (utils.new_task_scoped_key(), transformation)
-        for transformation in tuple(tyk2.edges)[:10]
-    )
-
-    mock_service = service.MockService(
-        SCRATCH_DIR,
-        SHARED_DIR,
-        STACKSIZE,
-        KEEP_SCRATCH,
-        KEEP_SHARED,
-        N_RETRIES,
-        CLAIM_LIMIT,
-        task_generator,
-    )
-
-    mock_service.start(MAX_TASKS, MAX_TIME)
-
-    print(mock_service.pdrs)
-    assert(all(pdr.ok() for pdr in mock_service.pdrs))
diff --git a/network.py b/network.py
deleted file mode 100644
index 0ed1c59d..00000000
--- a/network.py
+++ /dev/null
@@ -1,113 +0,0 @@
-import time
-
-from openfe_benchmarks import tyk2
-
-from gufe import ChemicalSystem, Transformation, NonTransformation, AlchemicalNetwork
-from gufe.tests.test_protocol import DummyProtocol, BrokenProtocol, FinishUnit, SimulationUnit, InitializeUnit, ProtocolUnit
-from gufe.protocols import ProtocolUnit
-
-class WeightedInitializeUnit(InitializeUnit):
-    value = 2
-
-class WeightedSimulationUnit(SimulationUnit):
-    value = 2
-
-    @staticmethod
-    def _execute(ctx, *, initialization, **inputs):
-        time.sleep(WeightedSimulationUnit.value * 2)
-        return SimulationUnit._execute(ctx, initialization=initialization, **inputs)
-
-class WeightedFinishUnit(FinishUnit):
-    value = 1
-
-    @staticmethod
-    def _execute(ctx, *, simulations, **inputs):
-        time.sleep(WeightedFinishUnit.value * 2)
-        return FinishUnit._execute(ctx, simulations=simulations, **inputs)
-
-
-class WeightedDummyProtocol(DummyProtocol):
-
-    def _create(
-            self,
-            stateA,
-            stateB,
-            mapping = None,
-            extends = None,
-    ):
-        if extends is not None:
-            # this is an example; wouldn't want to pass in whole ProtocolDAGResult into
-            # any ProtocolUnits below, since this could create dependency hell;
-            # instead, extract what's needed from it for starting point here
-            starting_point = extends.protocol_unit_results[-1].outputs["key_results"]
-        else:
-            starting_point = None
-
-        # convert protocol inputs into starting points for independent simulations
-        alpha = WeightedInitializeUnit(
-            name="the beginning",
-            settings=self.settings,
-            stateA=stateA,
-            stateB=stateB,
-            mapping=mapping,
-            start=starting_point,
-            some_dict={"a": 2, "b": 12},
-        )
-
-        # create several units that would each run an independent simulation
-        simulations: list[ProtocolUnit] = [
-            WeightedSimulationUnit(settings=self.settings, name=f"sim {i}", window=i, initialization=alpha)
-            for i in range(self.settings.n_repeats)  # type: ignore
-        ]
-
-        # gather results from simulations, finalize outputs
-        omega = WeightedFinishUnit(settings=self.settings, name="the end", simulations=simulations)
-
-        # return all `ProtocolUnit`s we created
-        return [alpha, *simulations, omega]
-
-def network_tyk2():
-    tyk2s = tyk2.get_system()
-
-    solvated = {
-        ligand.name: ChemicalSystem(
-            components={"ligand": ligand, "solvent": tyk2s.solvent_component},
-            name=f"{ligand.name}_water",
-        )
-        for ligand in tyk2s.ligand_components
-    }
-    complexes = {
-        ligand.name: ChemicalSystem(
-            components={
-                "ligand": ligand,
-                "solvent": tyk2s.solvent_component,
-                "protein": tyk2s.protein_component,
-            },
-            name=f"{ligand.name}_complex",
-        )
-        for ligand in tyk2s.ligand_components
-    }
-
-    complex_network = [
-        Transformation(
-            stateA=complexes[edge[0]],
-            stateB=complexes[edge[1]],
-            protocol=WeightedDummyProtocol(settings=WeightedDummyProtocol.default_settings()),
-            name=f"{edge[0]}_to_{edge[1]}_complex",
-        )
-        for edge in tyk2s.connections
-    ]
-    solvent_network = [
-        Transformation(
-            stateA=solvated[edge[0]],
-            stateB=solvated[edge[1]],
-            protocol=WeightedDummyProtocol(settings=WeightedDummyProtocol.default_settings()),
-            name=f"{edge[0]}_to_{edge[1]}_solvent",
-        )
-        for edge in tyk2s.connections
-    ]
-
-    return AlchemicalNetwork(
-        edges=(solvent_network + complex_network),
-        name="tyk2_relative_benchmark",
-    )
diff --git a/service.py b/service.py
deleted file mode 100644
index 00c23c9e..00000000
--- a/service.py
+++ /dev/null
@@ -1,113 +0,0 @@
-from dataclasses import dataclass
-import logging
-import time
-
-from alchemiscale.compute.service import AsynchronousComputeService, ExecutorStack, InterruptableSleep
-from alchemiscale.storage.models import ComputeServiceID
-
-@dataclass
-class MockSettings:
-    memory_monitor_enabled: bool
-    memory_monitor_sample_time: int
-    memory_monitor_sample_history_size: int
-    memory_monitor_grow_limit: float
-    memory_monitor_maintain_limit: float
-    cpu_monitor_enabled: bool
-    cpu_monitor_sample_time: int
-    cpu_monitor_sample_history_size: int
-    cpu_monitor_grow_limit: float
-    cpu_monitor_maintain_limit: float
-    gpu_monitor_enabled: bool
-    gpu_monitor_gpu_id: int
-    gpu_monitor_sample_time: int
-    gpu_monitor_sample_history_size: int
-    gpu_monitor_grow_limit: float
-    gpu_monitor_maintain_limit: float
-
-class MockService(AsynchronousComputeService):
-
-    def __init__(self, scratch_basedir, shared_basedir, stack_size, keep_scratch, keep_shared, n_retries, claim_limit, task_generator):
-        self._child_env = dict()
-        self._initialize_dag_tree()
-        self.settings = MockSettings(memory_monitor_enabled=True,
-                                     cpu_monitor_enabled=True,
-                                     gpu_monitor_enabled=False,
-                                     gpu_monitor_gpu_id="0",
-                                     gpu_monitor_sample_time=1,
-                                     gpu_monitor_sample_history_size=60,
-                                     gpu_monitor_grow_limit=0.7,
-                                     gpu_monitor_maintain_limit=0.9,
-                                     memory_monitor_sample_time=1,
-                                     memory_monitor_sample_history_size=60,
-                                     memory_monitor_grow_limit=0.7,
-                                     memory_monitor_maintain_limit=0.9,
-                                     cpu_monitor_sample_time=1,
-                                     cpu_monitor_sample_history_size=60,
-                                     cpu_monitor_grow_limit=0.9,
-                                     cpu_monitor_maintain_limit=1.2,
-                                     )
-        self._initialize_resource_monitors()
-        self._task_data = dict()
-        self._executor_stack = ExecutorStack(stack_size)
-
-        self.scratch_basedir = scratch_basedir
-        self.shared_basedir = shared_basedir
-        self.keep_scratch = keep_scratch
-        self.keep_shared = keep_shared
-        self.n_retries = n_retries
-        self.claim_limit = claim_limit
-        self.task_generator = task_generator
-        self.pdrs = []
-        self.tasks_claimed = 0
-        self.tasks_finished = 0
-
-        self.int_sleep = InterruptableSleep()
-
-        self.name = "MockService"
-        self.compute_service_id = ComputeServiceID.new_from_name(self.name)
-
-        # logging shim
-        extra = {"compute_service_id": "fakeid"}
-        logger = logging.getLogger("AlchemiscaleSynchronousComputeService")
-        logger.setLevel(logging.DEBUG)
-
-        formatter = logging.Formatter(
-            "[%(asctime)s] [%(compute_service_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)
-        self.logger = logging.LoggerAdapter(logger, extra)
-
-    def add_task(self, task_scoped_key, transformation):
-        protocol_dag = transformation.create()
-        self.graft_dag(task_scoped_key, protocol_dag)
-
-    def _register(self):
-        self.logger.info("Fake register")
-
-    def _deregister(self):
-        self.logger.info("Fake deregister")
-
-    def heartbeat(self):
-        pass
-
-    def claim_tasks(self, count=1):
-        claimed_tasks = []
-        remaining = count
-        if remaining == 0:
-            return [None] * count
-        for task in self.task_generator:
-            claimed_tasks.append(task)
-            remaining = remaining - 1
-            if remaining == 0:
-                return claimed_tasks
-        return claimed_tasks + [None] * remaining
-
-    def push_result(self, task_scoped_key, pdr):
-        _ = task_scoped_key
-        self.logger.info(f"Pushing {pdr}")
-        self.pdrs.append(pdr)
-        self.tasks_finished = self.tasks_finished + 1
diff --git a/utils.py b/utils.py
deleted file mode 100644
index 791c618c..00000000
--- a/utils.py
+++ /dev/null
@@ -1,32 +0,0 @@
-from contextlib import contextmanager
-from time import time
-from uuid import uuid4
-
-from alchemiscale.models import ScopedKey
-from gufe.tokenization import GufeKey
-
-@contextmanager
-def timer(*args, **kwargs):
-
-    wrap = kwargs.get("wrap")
-
-    if wrap:
-        print("="*20)
-
-    start = time()
-    yield
-    elapsed = time() - start
-    if wrap:
-        print("-"*20)
-
-    print(f"Time spent: {elapsed}")
-
-    if wrap:
-        print("="*20)
-
-def new_task_scoped_key():
-    task_key = GufeKey(f"FakeKey-{uuid4().hex}")
-    task_scoped_key = ScopedKey(
-        gufe_key=task_key, org="MockOrg", campaign="MockCampaign", project="MockProject"
-    )
-    return task_scoped_key

From 3a2f5ae478c8f284f47c525b6870066b6b466732 Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Tue, 31 Mar 2026 17:29:34 -0400
Subject: [PATCH 23/24] Add docstrings and remove unused code

---
 alchemiscale/compute/monitor.py |   6 +-
 alchemiscale/compute/service.py | 189 +++++++++++++++++---------------
 2 files changed, 107 insertions(+), 88 deletions(-)

diff --git a/alchemiscale/compute/monitor.py b/alchemiscale/compute/monitor.py
index 80467253..84ea9cf3 100644
--- a/alchemiscale/compute/monitor.py
+++ b/alchemiscale/compute/monitor.py
@@ -11,6 +11,8 @@
 import time
 from threading import Lock
 
+from alchemiscale.compute.settings import AsynchronousComputeServiceSettings
+
 
 # Signal to be issued by a resource manager to a compute
 # service. These signals are suggestions and are meant to inform the
@@ -19,7 +21,7 @@
 #    1. TERMINATE: tool/resource failure, shut stop all calculations
 #    2. SHRINK: resource is oversubscribed, scale down
 #    3. MAINTAIN: keep at current subscription
-#    4. GROW: resource is undersubscribed
+#    4. GROW: resource is under-subscribed
 class ResourceSignal(IntEnum):
     # order matters, higher priority signals should appear at top
     TERMINATE = auto()
@@ -47,7 +49,7 @@ class Monitor:
 
     """
 
-    def __init__(self, settings: ComputeServiceSettings):
+    def __init__(self, settings: AsynchronousComputeServiceSettings):
         self._setup(settings)
         self._lock = Lock()
         self._terminate = False
diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py
index d098b165..dff42b8c 100644
--- a/alchemiscale/compute/service.py
+++ b/alchemiscale/compute/service.py
@@ -440,13 +440,6 @@ class Executor(Process):
     and data given a unit Context.
     """
 
-    key: NodeKey
-    queue: Queue
-    lock: Lock
-    unit_context: Context
-    inputs: dict
-    n_retries: int
-
     def __init__(self, key, queue, lock, context, inputs, n_retries, env=None):
         super().__init__()
         self._key = key
@@ -455,6 +448,9 @@ def __init__(self, key, queue, lock, context, inputs, n_retries, env=None):
         self._unit_context = context
         self._inputs = inputs
         self._n_retries = n_retries
+
+        # mechanism to modify child process environment variables
+        # through an update during `run`
         self._env = env or {}
         self._validate()
 
@@ -463,6 +459,13 @@ def _validate(self):
             raise ValueError("n_retries must be greater than or equal to 0")
 
     def run(self):
+        """Attempt to run a complete ProtocolUnit.
+
+        The resulting ``ProtocolUnitResult`` or
+        ``ProtocolUnitFailure`` is put into the result queue at the
+        end of execution.
+
+        """
         # update environment before running unit
         os.environ |= self._env
         attempt = 0
@@ -477,15 +480,17 @@ def run(self):
             if result.ok():
                 break
             attempt = attempt + 1
-        # put the result in the queue with lock
+        # put the result in the queue using a lock
         self.put_result(result)
 
     @property
     def unit(self) -> ProtocolUnit:
+        """The unit this Executor runs."""
         return self._key[1]
 
     @property
     def key(self) -> NodeKey:
+        """The key assigned to the executor."""
         return self._key
 
     def put_result(self, result: ProtocolUnitResult):
@@ -494,7 +499,10 @@ def put_result(self, result: ProtocolUnitResult):
             self._queue.put((self._key, result))
 
     def execute_unit(self, context) -> ProtocolUnitResult | ProtocolUnitFailure:
-        # this method assumes the context is in place and will be removed correctly
+        """Unit execution method.
+
+        This method assumes the context directories are already in place.
+        """
         import warnings
 
         warnings.filterwarnings("ignore", message=r".*RDKit does not preserve.*")
@@ -507,55 +515,37 @@ class JailedKeyError(Exception):
 
 class ExecutorStack:
     """Structure for coordinating the creation and management of
-    Executor processes.
+    ``Executor`` processes.
     """
 
-    stack_size: int
-    stack: list[Executor]
-    # blocked node specified by a set of blocking nodes
-    jail: dict[NodeKey, set[NodeKey]]
-    queue: Queue
-    lock: Lock
-
     def __init__(self, stack_size: int):
-        self._stack = []
-        self._stack_size = stack_size
-        self._jail = {}
-        self._queue = Queue()
-        self._lock = Lock()
+        self._stack: int = []
+        self._stack_size: list[Executor] = stack_size
+        self._jail: dict[NodeKey, set[NodeKey]] = {}
+        self._queue: Queue = Queue()
+        self._lock: Lock = Lock()
         self._validate()
 
     def _validate(self):
         if not self._stack_size >= 1:
             raise ValueError("stack_size must be greater than or equal to 1")
 
-    @property
-    def stack(self) -> list[Executor]:
-        return self._stack
-
-    @property
-    def stack_size(self) -> int:
-        return self._stack_size
-
-    @property
-    def jail(self) -> dict[NodeKey, set[NodeKey]]:
-        return self._jail
-
-    @property
-    def lock(self) -> Lock:
-        return self._lock
+    def terminate_all(self):
+        """Terminate all processes in the stack.
 
-    @property
-    def queue(self) -> Queue:
-        return self._queue
+        This method waits to acquire the lock before terminating
+        tasks, meaning results being written to the queue during the
+        time of the call will still be available for processing after
+        the process is terminated.
 
-    def terminate_all(self):
+        """
         with self._lock:
             for proc in self._stack:
                 proc.terminate()
             self._stack.clear()
 
     def terminate_task(self, task_key: TaskKey):
+        """Terminate any processes from a ``Task``."""
         with self._lock:
             to_remove = set()
             for proc in self._stack:
@@ -574,6 +564,23 @@ def push(
         n_retries: int,
         env: dict[str, str],
     ):
+        """Push a node to the stack.
+
+        Parameters
+        ----------
+        node
+            The ``Task`` ``ScopedKey`` and the ``ProtocolUnit`` to
+            push to the stack.
+        unit_context
+            The ``Context`` for running the ``ProtocolUnit``.
+        inputs
+            Inputs ``dict`` for running the ``ProtocolUnit``.
+        n_retries
+            The number of times to attempt to rerun a ``ProtocolUnit``
+            that raises an exception.
+        env
+            Updates to the environment of the child process.
+        """
         with self._lock:
             # node may be blocked from execution
             if node in self._jail.keys():
@@ -593,10 +600,14 @@ def push(
             self._stack.append(executor)
             self._stack[-1].start()
 
+    def full(self):
+        return len(self._stack) >= self._stack_size
+
     def pop(self):
-        """Remove last process in the stack. This also clears the node from the jail."""
+        """Remove last process in the stack. This also clears the node
+        from the jail."""
         with self._lock:
-            if self._stack_size == 0:
+            if len(self._stack) == 0:
                 raise IndexError("pop from empty stack")
 
             popped_executor = self._stack.pop()
@@ -840,42 +851,6 @@ def _get_resource_signal(self) -> ResourceSignal:
         # otherwise respect the highest priority signal from all monitors
         return min(monitor.signal() for monitor in self._resource_monitors)
 
-    def _resource_monitor(self) -> ResourceSignal:
-        capacity = 8
-        sim_value = 3
-        fin_value = 1
-        init_value = 2
-
-        total = 0
-        for proc in self._executor_stack.stack:
-            if not proc.is_alive():
-                continue
-            pu = proc.key[1]
-            match pu.__class__.__name__:
-                case "WeightedFinishingUnit":
-                    total = total + fin_value
-                case "WeightedSimulationUnit":
-                    total = total + sim_value
-                case _:
-                    total = total + init_value
-
-        self.logger.info(
-            f"Num jailed: {len(self._executor_stack._jail)} -- {total}/{capacity}"
-        )
-        if total == capacity:
-            self.logger.info("MAINTAINING")
-            return ResourceSignal.MAINTAIN
-
-        if total > capacity:
-            self.logger.info("SHRINKING")
-            return ResourceSignal.SHRINK
-
-        if total < capacity:
-            self.logger.info("GROWING")
-            return ResourceSignal.GROW
-
-        return ResourceSignal.TERMINATE
-
     def cycle(self, max_tasks, max_time) -> bool:
         # collect unit results
         self.process_results()  # removes unit scratch
@@ -893,14 +868,16 @@ def cycle(self, max_tasks, max_time) -> bool:
             self.stop()
             return False
 
-        # detemine next actions based on resource usage
+        # determine next actions based on resource usage
         signal = self._get_resource_signal()
         match signal:
             case ResourceSignal.MAINTAIN:
                 return True
             case ResourceSignal.SHRINK:
-                proc = self._executor_stack.pop()
-                self.logger.info(f"Popping: {proc} -- {proc.key[1].key}")
+                try:
+                    self._executor_stack.pop()
+                except IndexError:
+                    logger.info("Attempted to pop from an empty stack")
                 return True
             case ResourceSignal.TERMINATE:
                 self.stop()
@@ -910,6 +887,7 @@ def cycle(self, max_tasks, max_time) -> bool:
             case _:
                 raise RuntimeError("Received unknown ResourceSignal")
 
+        # determine how many tasks can be claimed and claim that many
         n_claim = self.claim_limit - len(self._task_data)
         if max_tasks is not None:
             max_less_claimed = max_tasks - self.tasks_claimed
@@ -923,11 +901,19 @@ def cycle(self, max_tasks, max_time) -> bool:
 
         self.logger.info("Claimed %d tasks", len([t for t in tasks if t is not None]))
 
+        # add claimed tasks to tree
         for task in tasks:
             if task is not None:
                 self.add_task(task)
                 self.tasks_claimed = 1 + self.tasks_claimed
 
+        # return early if no room in stack
+        if self._executor_stack.full():
+            time.sleep(self.sleep_interval)
+            return True
+
+        # iterate over all nodes that are available, less those that
+        # are already running
         for key in filter(lambda k: k[1] not in ("TERM", "ROOT"), self.next()):
             tsk, unit = key
             task_data = self._task_data[tsk]
@@ -950,6 +936,7 @@ def cycle(self, max_tasks, max_time) -> bool:
         return True
 
     def available_units(self) -> set[NodeKey]:
+        """All units with no parents."""
         available = set()
         for node, degree in self._dag_tree.out_degree():
             if degree == 0:
@@ -958,6 +945,7 @@ def available_units(self) -> set[NodeKey]:
         return available
 
     def next(self) -> set[NodeKey]:
+        """Available units, less those already running or terminated."""
         running, terminated = self._executor_stack._get_statuses()
         running = {r.key for r in running}
         terminated = {t.key for t in terminated}
@@ -965,51 +953,73 @@ def next(self) -> set[NodeKey]:
         return next_units
 
     def next_terminating_nodes(self) -> set[NodeKey]:
+        """All terminating nodes whose parents are complete."""
         completed = {node for node in self.available_units() if node[1] == "TERM"}
         return completed
 
     def add_task(self, task_scoped_key: ScopedKey):
+        """Get a ``ProtocolDAG`` given a ``ScopedKey`` and add it to the DAG tree."""
         protocol_dag, _, _ = self.task_to_protocoldag(task_scoped_key)
         self.graft_dag(task_scoped_key, protocol_dag)
 
     def graft_dag(self, task_scoped_key: ScopedKey, dag):
-        """Add a ``Task`` to the ``AsynchronousComputeService`` internal DAG."""
+        """Add a ``Task`` to the ``AsynchronousComputeService``
+        internal DAG. Additionally, create the ``Context`` directories
+        for the ``Task`` and create a ``TaskData`` record.
 
-        def node_transformation(node: ProtocolUnit) -> (TaskKey, ProtocolUnit):
+        """
+
+        # tag a node with the Task it belongs to
+        def node_transformation(
+            node: ProtocolUnit | str | None,
+        ) -> (TaskKey, ProtocolUnit):
             nonlocal task_scoped_key
             return (task_scoped_key, node)
 
+        # create the terminating node for this DAG
         tagged_dag = nx.DiGraph()
         terminating = node_transformation("TERM")
         tagged_dag.add_node(terminating)
 
+        # go over all previous nodes and add their tagged variants to
+        # the new graph
         for child, parent in dag.graph.edges:
             tagged_child = node_transformation(child)
             tagged_parent = node_transformation(parent)
             tagged_dag.add_edge(tagged_child, tagged_parent)
 
+        # find all "end" nodes and attach them to the terminating node
         for node, in_degree in dag.graph.in_degree:
             if in_degree == 0:
                 tagged_dag.add_edge(terminating, node_transformation(node))
 
         self._dag_tree.add_edges_from(tagged_dag.edges)
+        # connect the new graph to the dag tree
         self._dag_tree.add_edge((None, "ROOT"), terminating)
 
+        # establish the scratch and shared directories
         context = Context(
             scratch=self.scratch_basedir / str(task_scoped_key),
             shared=self.shared_basedir / str(task_scoped_key),
         )
         context.scratch.mkdir(exist_ok=True)
         context.shared.mkdir(exist_ok=True)
+
+        # add TaskData record for later result collection and input
+        # generation
         self._task_data[task_scoped_key] = TaskData(
             results={}, context=context, protocol_dag=dag
         )
 
     def remove_task(self, task_scoped_key):
+        """Remove nodes in the DAG tree that belong to the given task
+        and remove their context directories.
+
+        """
         # TODO: check executor stack
         # avoid deleting the root node
-        if task_scoped_key is None:
-            raise ValueError()
+        # if task_scoped_key is None:
+        # raise ValueError
 
         for node in tuple(self._dag_tree.nodes):
             key, _ = node
@@ -1024,11 +1034,18 @@ def remove_task(self, task_scoped_key):
             shutil.rmtree(context.scratch)
 
     def remove_all(self):
+        """Remove all nodes from the DAG tree (except to root) and any
+        task data.
+
+        """
         for task_scoped_key in self._task_data.keys():
             self.remove_task(task_scoped_key)
         self._task_data.clear()
 
     def _consume_results(self, task_scoped_key) -> ProtocolDAGResult:
+        """Return a ``ProtocolDAGResult`` from the collected data up
+        until this point and delete its TaskData.
+        """
         self.remove_task(task_scoped_key)
         data = self._task_data.pop(task_scoped_key)
         pdr = data.to_ProtocolDAGResult()

From 8d7991cbc4a7933600dbea5eb1ab0c5b7b2ac4eb Mon Sep 17 00:00:00 2001
From: Ian Kenney <ianmichaelkenney@gmail.com>
Date: Tue, 31 Mar 2026 17:54:11 -0400
Subject: [PATCH 24/24] Test max time termination

---
 .../compute/client/test_compute_service.py    | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/alchemiscale/tests/integration/compute/client/test_compute_service.py b/alchemiscale/tests/integration/compute/client/test_compute_service.py
index 03d975cf..c4c72997 100644
--- a/alchemiscale/tests/integration/compute/client/test_compute_service.py
+++ b/alchemiscale/tests/integration/compute/client/test_compute_service.py
@@ -340,3 +340,22 @@ def test_cycle(self, n4js_preloaded, s3os_server_fresh, service):
         results = n4js_preloaded.execute_query(q)
 
         assert results.records
+
+    def test_max_time_termination(self, n4js_preloaded, s3os_server_fresh, service):
+        allowed_time = 8  # cycle sleeping and shutdown allowance
+        max_time = 5
+
+        service.start(max_time=max_time)
+        time.sleep(2)  # give time to claim tasks
+
+        start = time.time()
+        while not service._stop:
+            assert service.has_tasks()
+            if (time.time() - start) > allowed_time:
+                raise ValueError
+            time.sleep(1)
+
+        time.sleep(2)  # give time for full termination
+
+        assert not service.has_tasks()
+        assert len(service._dag_tree.nodes) == 1  # only root is left