FEAT(net,.TC): ABORT flag

Author: ankostis

graphtik/__init__.py

@@ -10,5 +10,6 @@
 __summary__ = __doc__.splitlines()[0]
 __uri__ = "https://github.com/pygraphkit/graphtik"
 
-from .nodes import operation, compose
 from .modifiers import *  # noqa, on purpose to include any new modifiers
+from .network import abort_run, AbortedException
+from .nodes import compose, operation

graphtik/network.py

@@ -80,7 +80,27 @@
 
 log = logging.getLogger(__name__)
 
-thread_pool: ContextVar[Pool] = ContextVar("thread_pool", default=Pool(7))
+
+execution_configs: ContextVar[dict] = ContextVar(
+    "execution_configs",
+    default={"thread_pool": Pool(7), "abort": False},
+)
+
+
+class AbortedException(Exception):
+    pass
+
+
+def abort_run():
+    execution_configs.get()["abort"] = True
+
+
+def _reset_abort():
+    execution_configs.get()["abort"] = False
+
+
+def is_abort():
+    return execution_configs.get()["abort"]
 
 
 class _DataNode(str):
@@ -195,6 +215,12 @@ def _pin_data_in_solution(self, value_name, solution, inputs, overwrites):
             overwrites[value_name] = solution[value_name]
         solution[value_name] = inputs[value_name]
 
+    def _check_if_aborted(self, executed):
+        if is_abort():
+            # Restore `abort` flag for next run.
+            _reset_abort()
+            raise AbortedException({s: s in executed for s in self.steps})
+
     def _call_operation(self, op, solution):
         # Although `plan` have added to jetsam in `compute()``,
         # add it again, in case compile()/execute is called separately.
@@ -217,12 +243,13 @@ def _execute_thread_pool_barrier_method(self, solution, overwrites, executed):
             n: solution[n] for n in self.steps if isinstance(n, _PinInstruction)
         }
 
-        pool = thread_pool.get()
+        pool = execution_configs.get()["thread_pool"]
 
         # with each loop iteration, we determine a set of operations that can be
         # scheduled, then schedule them onto a thread pool, then collect their
         # results onto a memory solution for use upon the next iteration.
         while True:
+            self._check_if_aborted(executed)
 
             # the upnext list contains a list of operations for scheduling
             # in the current round of scheduling
@@ -293,6 +320,7 @@ def _execute_sequential_method(self, solution, overwrites, executed):
 
         self.times = {}
         for step in self.steps:
+            self._check_if_aborted(executed)
 
             if isinstance(step, Operation):
 

test/test_graphtik.py

@@ -10,9 +10,20 @@
 import pytest
 
 import graphtik.network as network
-from graphtik import compose, operation, optional, sideffect
+from graphtik import (
+    abort_run,
+    AbortedException,
+    compose,
+    operation,
+    optional,
+    sideffect,
+)
 from graphtik.base import Operation
-from graphtik.network import _EvictInstruction
+
+
+@pytest.fixture(params=["sequential", "parallel"])
+def exemethod(request):
+    return request.param
 
 
 def scream(*args, **kwargs):
@@ -21,8 +32,8 @@ def scream(*args, **kwargs):
     )
 
 
-def identity(x):
-    return x
+def identity(*x):
+    return x[0] if len(x) == 1 else x
 
 
 def filtdict(d, *keys):
@@ -856,7 +867,7 @@ def addplusplus(a, b, c=0):
 def test_evict_instructions_vary_with_inputs():
     # Check #21: _EvictInstructions positions vary when inputs change.
     def count_evictions(steps):
-        return sum(isinstance(n, _EvictInstruction) for n in steps)
+        return sum(isinstance(n, network._EvictInstruction) for n in steps)
 
     pipeline = compose(name="pipeline")(
         operation(name="a free without b", needs=["a"], provides=["aa"])(identity),
@@ -1048,3 +1059,23 @@ def test_compose_another_network(bools):
 
     sol = bigger_graph({"a": 2, "b": 5, "c": 5}, outputs=["a_minus_ab_minus_c"])
     assert sol == {"a_minus_ab_minus_c": -13}
+
+
+def test_abort(exemethod):
+    pipeline = compose(name="pipeline")(
+        operation(name="A", needs=["a"], provides=["b"])(identity),
+        operation(name="B", needs=["b"], provides=["c"])(lambda x: abort_run()),
+        operation(name="C", needs=["c"], provides=["d"])(identity),
+    )
+    pipeline.set_execution_method(exemethod)
+    with pytest.raises(AbortedException) as exinfo:
+        pipeline({"a": 1})
+    assert exinfo.value.jetsam["solution"] == {"a": 1, "b": 1, "c": None}
+    executed = {op.name: val for op, val in exinfo.value.args[0].items()}
+    assert executed == {"A": True, "B": True, "C": False}
+
+    pipeline = compose(name="pipeline")(
+        operation(name="A", needs=["a"], provides=["b"])(identity)
+    )
+    pipeline.set_execution_method(exemethod)
+    assert pipeline({"a": 1}) == {"a": 1, "b": 1}