REFACT(SOL) Callback from Solution (not marshaled)

Author: ankostis

docs/source/arch.rst

@@ -708,22 +708,24 @@ Architecture
         value.  All of them default to ``None`` (false).
 
     callbacks
-        x2 user-defined optional callables called before & after each function,
-        installed on :meth:`.Pipeline.compute()`, that they must have this signature::
+        x2 optional callables called before/after each `operation` :meth:`.Pipeline.compute()`.
+        Attention, any errors will abort the pipeline execution.
 
-            callbacks(op_cb) -> None
+        pre-op-callback
+            Called from solution code before :term:`marshalling`.
+            A use case would be to validate `solution`, or
+            :ref:`trigger a breakpoint by some condition <break_with_pre_callback>`.
+
+        post-op-callback:
+            Called after solution have been populated with `operation` results.
+            A use case would be to validate operation `outputs` and/or
+            solution after results have been populated.
 
-        ... where ``op_cb`` is an instance of the :class:`.OpTask` :func:`~collections.namedtuple`.
+        Callbacks must have this signature::
 
-        pre-callback
-            Called after :term:`marshalling` and before `matching inputs`.
-            A use case would be to :ref:`a breakpoint triggered by some condition
-            <break_with_pre_callback>`.
+            callbacks(op_cb) -> None
 
-        post-callback:
-            Called after :term:`zipping outputs` and before :term:`marshalling` results
-            back into `solution`.
-            A use case would be to validate the solution and scream immediately on errors.
+        ... where ``op_cb`` is an instance of the :class:`.OpTask`.
 
     jetsam
         When a pipeline or an operation fails, the original exception gets annotated

graphtik/execution.py

@@ -118,7 +118,19 @@ def __init__(
         is_layered=None,
     ):
         super().__init__(input_values)
+        ## Make callbacks a 2-tuple with possible None callables.
+        #
+        if callable(callbacks):
+            callbacks = (callbacks,)
+        elif not callbacks:
+            callbacks = ()
+        else:
+            callbacks = tuple(callbacks)
+        n_cbs = len(callbacks)
+        if n_cbs < 2:
+            callbacks = callbacks + ((None,) * (2 - n_cbs))
         self.callbacks = callbacks
+
         is_layered = first_solid(is_layered_solution(), is_layered)
 
         ##: By default, disable layers if network contains :term:`jsonp` dependencies.
@@ -466,37 +478,24 @@ class OpTask:
     Mimic :class:`concurrent.futures.Future` for :term:`sequential` execution.
 
     This intermediate class is needed to solve pickling issue with process executor.
-
-    Use also as argument passed in a :term:`callbacks` callable
-    before executing each operation, and contains fields to identify
-    the operation call and results:
     """
 
-    __slots__ = ("op", "sol", "solid", "callbacks", "result")
+    __slots__ = ("op", "sol", "solid", "result")
     logname = __name__
 
-    def __init__(self, op, sol, solid, callbacks=None, result=UNSET):
+    def __init__(self, op, sol, solid, result=UNSET):
         #: the operation about to be computed.
         self.op = op
         #: the solution (might be just a plain dict if it has been marshalled).
         self.sol = sol
         #: the operation identity, needed if `sol` is a plain dict.
         self.solid = solid
-
-        ## Make callbacks a 2-tuple with possible None callables.
-        #
-        if callbacks is None:
-            callbacks = (None, None)
-        elif callable(callbacks):
-            callbacks = (callbacks, None)
-        else:
-            callbacks = tuple(callbacks)
-            if len(callbacks) < 2:
-                callbacks = (*callbacks, None)
-        self.callbacks = callbacks
-        #: Initially would :data:`.UNSET`, will be set after execution.
+        #: Initially would :data:`.UNSET`, will be set after execution
+        #: with operation's outputs or exception.
         self.result = result
 
+        # if
+
     def marshalled(self):
         import dill
 
@@ -508,13 +507,8 @@ def __call__(self):
             log = logging.getLogger(self.logname)
             log.debug("+++ (%s) Executing %s...", self.solid, self)
             token = task_context.set(self)
-            callbacks = self.callbacks
             try:
-                if callbacks[0]:
-                    callbacks[0](self)
                 self.result = self.op.compute(self.sol)
-                if callbacks[1]:
-                    callbacks[1](self)
             finally:
                 task_context.reset(token)
 
@@ -716,7 +710,7 @@ def prep_task(op):
                 # Mark start time here, to include also marshalling overhead.
                 solution.elapsed_ms[op] = time.time()
 
-                task = OpTask(op, input_values, solution.solid, solution.callbacks)
+                task = OpTask(op, input_values, solution.solid)
                 if first_solid(global_marshal, getattr(op, "marshalled", None)):
                     task = task.marshalled()
 
@@ -743,7 +737,7 @@ def prep_task(op):
 
         return [prep_task(op) for op in operations]
 
-    def _handle_task(self, future, op, solution) -> None:
+    def _handle_task(self, future: Union[OpTask, "AsyncResult"], op, solution) -> None:
         """Un-dill parallel task results (if marshalled), and update solution / handle failure."""
 
         def elapsed_ms(op):
@@ -752,14 +746,18 @@ def elapsed_ms(op):
 
             return elapsed
 
+        result = UNSET
         try:
             ## Reset start time for Sequential tasks
             #  (bummer, they will miss marshalling overhead).
             #
             if isinstance(future, OpTask):
                 solution.elapsed_ms[op] = time.time()
 
-            outputs = future.get()
+                if solution.callbacks[0]:
+                    solution.callbacks[0](future)
+
+            outputs = result = future.get()
             if isinstance(outputs, bytes):
                 import dill
 
@@ -772,6 +770,7 @@ def elapsed_ms(op):
                 "... (%s) op(%s) completed in %sms.", solution.solid, op.name, elapsed
             )
         except Exception as ex:
+            result = ex
             is_endured = first_solid(
                 solution.is_endurance, getattr(op, "endured", None)
             )
@@ -801,6 +800,9 @@ def elapsed_ms(op):
                 # add it again, in case compile()/execute() is called separately.
                 save_jetsam(ex, locals(), "solution", task="future", plan="self")
                 raise
+        finally:
+            if isinstance(future, OpTask) and solution.callbacks[1]:
+                solution.callbacks[1](future)
 
     def _execute_thread_pool_barrier_method(self, solution: Solution):
         """
@@ -904,7 +906,7 @@ def _execute_sequential_method(self, solution: Solution):
                 if step in solution.canceled:
                     continue
 
-                task = OpTask(step, solution, solution.solid, solution.callbacks)
+                task = OpTask(step, solution, solution.solid)
                 self._handle_task(task, step, solution)
 
             elif isinstance(step, str):
@@ -927,7 +929,7 @@ def execute(
         outputs=None,
         *,
         name="",
-        callbacks: Callable[[OpTask], None] = None,
+        callbacks: Tuple[Callable[[OpTask], None], ...] = None,
         solution_class=None,
         layered_solution=None,
     ) -> Solution:
@@ -943,9 +945,10 @@ def execute(
         :param name:
             name of the pipeline used for logging
         :param callbacks:
-            If given, a 2-tuple with (optional) x2 :term:`callbacks` to call before & after
-            each operation, with :class:`.OpTask` as argument containing the op & solution.
-            Less or no elements accepted.
+            If given, a 2-tuple with (optional) :term:`callbacks` to call
+            before/after computing operation, with :class:`.OpTask` as argument
+            containing the op & solution.
+            Can be one (scalar), less than 2, or nothing/no elements accepted.
         :param solution_class:
             a custom solution factory to use
         :param layered_solution:

graphtik/pipeline.py

@@ -412,10 +412,10 @@ def compute(
             filter-out nodes before compiling
             If not given, those set by a previous call to :meth:`withset()` or cstor are used.
         :param callbacks:
-            If given, a 4-tuple with (optional) x2 :term:`callbacks`,
-            2 to call before & after each operation, and another 2 before/after batch,
-            with :class:`.OpTask` as argument containing the op & solution.
-            One (scalar), less than 4, or no elements accepted.
+            If given, a 2-tuple with (optional) :term:`callbacks` to call
+            before/after computing operation, with :class:`.OpTask` as argument
+            containing the op & solution.
+            Can be one (scalar), less than 2, or nothing/no elements accepted.
         :param solution_class:
             a custom solution factory to use
         :param layered_solution:

test/test_graphtik.py

@@ -945,26 +945,40 @@ def test_rescheduling_NO_RESULT(exemethod):
         assert sol.scream_if_incomplete()
 
 
+@pytest.mark.xfail(
+    sys.version_info < (3, 8), reason="unittest.nock.call.args is trange in PY37-"
+)
 def test_pre_callback(quarantine_pipeline, exemethod):
-    pipeline = compose("covid19", quarantine_pipeline, parallel=exemethod)
-    called_ops = []
-
-    def op_called(op_cb):
-        assert op_cb.sol["quarantine"]
-        called_ops.append(op_cb.op.name)
+    # Cannot import top-level `unittest.mock.call`, due to
+    # https://bugs.python.org/issue35753
+    from unittest.mock import MagicMock, call
+
+    pipeline = compose("covid19", quarantine_pipeline)
+
+    callbacks = [MagicMock() for _ in range(2)]
+    sol = pipeline.compute({"quarantine": True}, callbacks=callbacks)
+
+    cbs_count = [cb.call_count for cb in callbacks]
+    assert cbs_count == [2, 2]
+    ops_called = [[call.args[0].op for call in cb.call_args_list] for cb in callbacks]
+    assert ops_called == [
+        ["get_out_or_stay_home", "read_book"],
+        ["get_out_or_stay_home", "read_book"],
+    ]
+    results_called = [
+        [call.args[0].result for call in cb.call_args_list] for cb in callbacks
+    ]
+    assert results_called == [
+        [{"time": "1h"}, {"fun": "relaxed", "brain": "popular physics"}],
+        [{"time": "1h"}, {"fun": "relaxed", "brain": "popular physics"}],
+    ]
 
-    sol = pipeline.compute({"quarantine": True}, callbacks=op_called)
     assert sol == {
         "quarantine": True,
         "time": "1h",
         "fun": "relaxed",
         "brain": "popular physics",
     }
-    if exe_params.marshal == 1:
-        # Marshaled `called_ops` is not itself :-)
-        assert called_ops == []
-    else:
-        assert called_ops == ["get_out_or_stay_home", "read_book"]
 
 
 ##########