[essreduce] StreamProcessor: reuse cached context nodes in finalize, fix context-derived targets read by accumulators

packages/essreduce/src/ess/reduce/streaming.py

@@ -3,7 +3,7 @@
 """This module provides tools for running workflows in a streaming fashion."""
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Mapping
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar
 
@@ -277,13 +277,8 @@ class _FedWorkflow:
     repeatedly, for the largest objects in the process. Here each input is instead
     provided by a provider reading a mapping owned by this class, so feeding a value
     is a plain dict write and the graph structure, which does not depend on the
-    values, stays as it was built.
-
-    Inputs are wired up once, at construction. Wiring assigns to the pipeline before
-    inserting the provider because only assignment prunes the branch the input
-    replaces: ``Pipeline.__setitem__`` removes the ancestors of the key, whereas
-    ``Pipeline.insert`` only cuts the key's incoming edges and leaves the orphaned
-    branch in the graph, for every later ``Pipeline.get`` to walk again.
+    values, stays as it was built. Inputs are wired up once, at construction, see
+    :func:`_wire_input`.
 
     A value stays readable until it is released. Reading a key that holds no value
     raises, so a workflow computed before its inputs arrived fails, naming the key,
@@ -297,10 +292,12 @@ def __init__(
     ) -> None:
         # Wiring mutates the pipeline, so we take a copy and own it from here on.
         self._workflow = workflow.copy()
+        # The providers hold the mapping, not this instance: capturing ``self`` would
+        # close the cycle pipeline -> provider -> closure -> self -> pipeline, which
+        # only the cyclic garbage collector can break.
         self._values: dict[sciline.typing.Key, Any] = {}
         for key in inputs:
-            self._workflow[key] = None  # prunes the branch this input replaces
-            self._insert_provider(key)
+            _wire_input(self._workflow, key, self._values)
 
     def __setitem__(self, key: sciline.typing.Key, value: Any) -> None:
         self._values[key] = value
@@ -315,23 +312,35 @@ def compute(
         """Compute ``keys`` from the values fed so far."""
         return self._workflow.compute(keys)
 
-    def _insert_provider(self, key: sciline.typing.Key) -> None:
-        # Binding the mapping to a local keeps the closure, and thus the pipeline
-        # holding it, from referencing this instance: capturing ``self`` would close
-        # the cycle pipeline -> provider -> closure -> self -> pipeline, which only
-        # the cyclic garbage collector can break.
-        values = self._values
 
-        def provide_input_value() -> Any:
-            try:
-                return values[key]
-            except KeyError:
-                raise ValueError(f"No value was fed for '{key}'") from None
+def _wire_input(
+    workflow: sciline.Pipeline,
+    key: sciline.typing.Key,
+    values: Mapping[sciline.typing.Key, Any],
+) -> None:
+    """
+    Make ``key`` an input of ``workflow`` that reads its value from ``values``.
+
+    Reading the key raises while ``values`` holds no value for it.
+
+    Wiring assigns to the pipeline before inserting the provider because only
+    assignment prunes the branch the input replaces: ``Pipeline.__setitem__`` removes
+    the ancestors of the key, whereas ``Pipeline.insert`` only cuts the key's incoming
+    edges and leaves the orphaned branch in the graph, for every later
+    ``Pipeline.get`` to walk again.
+    """
+    workflow[key] = None  # prunes the branch this input replaces
 
-        # Sciline deduces the provided key from the return annotation. The key is known
-        # only at runtime, so it has to be patched in.
-        provide_input_value.__annotations__ = {'return': key}
-        self._workflow.insert(provide_input_value)
+    def provide_input_value() -> Any:
+        try:
+            return values[key]
+        except KeyError:
+            raise ValueError(f"No value was fed for '{key}'") from None
+
+    # Sciline deduces the provided key from the return annotation. The key is known
+    # only at runtime, so it has to be patched in.
+    provide_input_value.__annotations__ = {'return': key}
+    workflow.insert(provide_input_value)
 
 
 class StreamProcessor:
@@ -425,35 +434,45 @@ def __init__(
             acc_key: nx.ancestors(graph, acc_key) & self._dynamic_keys
             for acc_key in self._accumulators
         }
-        self._keys_read_at_finalize = (
-            _find_keys_read_at_finalize(
-                graph,
-                target_keys=target_keys,
-                accumulator_keys=set(self._accumulators),
-                dynamic_keys=self._dynamic_keys,
-            )
-            if allow_bypass
-            else set()
-        )
-
-        # Chunks, accumulator values, and context are fed to the workflows rather than
-        # assigned to them, see :class:`_FedWorkflow`. Each workflow is wired up for
-        # exactly the keys the methods below feed it.
         cached_context_nodes = {
             node
             for nodes in self._context_key_to_cached_context_nodes_map.values()
             for node in nodes
         }
-        targets = set(target_keys)
+        accumulator_keys = set(self._accumulators)
+        # Each cached context node goes to the workflow(s) that read it: the chunk
+        # workflow if computing an accumulator reads it, the finalize workflow if it is
+        # a target or computing a target from the accumulators reads it, or both.
+        read_by_chunk = _find_keys_read(
+            graph, outputs=accumulator_keys, inputs=cached_context_nodes
+        )
+        read_at_finalize = _find_keys_read(
+            graph,
+            outputs=set(target_keys),
+            inputs=accumulator_keys | cached_context_nodes,
+        )
+        self._context_nodes_read_by_chunk = cached_context_nodes & read_by_chunk
+        self._context_nodes_read_at_finalize = cached_context_nodes & read_at_finalize
+        # Dynamic keys that reach a target without passing through an accumulator.
+        # Only these are fed to the finalize workflow, and only they must outlive the
+        # chunk they arrived in.
+        self._keys_read_at_finalize = (
+            self._dynamic_keys & read_at_finalize if allow_bypass else set()
+        )
+
+        # Chunks, accumulator values, and context are fed to the workflows rather than
+        # assigned to them, see :class:`_FedWorkflow`. Each workflow is wired up for
+        # exactly the keys the methods below feed it. Every path from a context key to
+        # a target passes through a cached context node, so the chunk and finalize
+        # workflows read context only through those.
         self._context_workflow = _FedWorkflow(workflow, self._context_keys)
         self._chunk_workflow = _FedWorkflow(
-            workflow, self._dynamic_keys | (cached_context_nodes - targets)
+            workflow, self._dynamic_keys | self._context_nodes_read_by_chunk
         )
         self._finalize_workflow = _FedWorkflow(
             workflow,
-            set(self._accumulators)
-            | self._context_keys
-            | (cached_context_nodes & targets)
+            accumulator_keys
+            | self._context_nodes_read_at_finalize
             | self._keys_read_at_finalize,
         )
 
@@ -473,18 +492,14 @@ def set_context(self, context: dict[sciline.typing.Key, Any]) -> None:
             needs_recompute |= self._context_key_to_cached_context_nodes_map[key]
         for key, value in context.items():
             self._context_workflow[key] = value
-            # Propagate context values to finalize workflow so providers that depend
-            # on context keys receive the updated values during finalize().
-            self._finalize_workflow[key] = value
         results = self._context_workflow.compute(needs_recompute)
         # Context values and the nodes derived from them are caches, to be reused until
         # the context changes again, so nothing is released here.
         for key, value in results.items():
-            if key in self._target_keys:
-                # Context-dependent key is direct target, independent of dynamic nodes.
-                self._finalize_workflow[key] = value
-            else:
+            if key in self._context_nodes_read_by_chunk:
                 self._chunk_workflow[key] = value
+            if key in self._context_nodes_read_at_finalize:
+                self._finalize_workflow[key] = value
 
     def add_chunk(
         self, chunks: dict[sciline.typing.Key, Any]
@@ -775,10 +790,12 @@ def _build_streaming_workflow(
     workflow = sciline.Pipeline()
     for key in target_keys:
         workflow[key] = base_workflow[key]
-    for key in dynamic_keys:
-        workflow[key] = None  # hack to prune branches
-    for key in context_keys:
-        workflow[key] = None
+    # Placeholders that raise when read. The stream processor computes on copies of
+    # this workflow, each wired up for the keys it is fed, so a computation that
+    # reaches a dynamic or context key it was not fed fails instead of reading a
+    # stand-in value.
+    for key in (*dynamic_keys, *context_keys):
+        _wire_input(workflow, key, {})
 
     nodes = _find_descendants(workflow, dynamic_keys + context_keys)
     last_static = _find_parents(workflow, nodes) - nodes
@@ -839,30 +856,26 @@ def _make_accumulators(
     }
 
 
-def _find_keys_read_at_finalize(
+def _find_keys_read(
     graph: nx.DiGraph,
     *,
-    target_keys: tuple[sciline.typing.Key, ...],
-    accumulator_keys: set[sciline.typing.Key],
-    dynamic_keys: set[sciline.typing.Key],
+    outputs: set[sciline.typing.Key],
+    inputs: set[sciline.typing.Key],
 ) -> set[sciline.typing.Key]:
     """
-    Find dynamic keys that reach a target without passing through an accumulator.
+    Find the keys that computing ``outputs`` reads, given values fed for ``inputs``.
 
-    Only these are fed to the finalize workflow, and only they must outlive the chunk
-    they arrived in. Computed on the graph with the accumulator inputs cut, which is
-    what finalize computes on. Feeding context cuts further edges, so this may name a
-    key that finalize turns out not to need, never miss one that it does.
+    Computed on the graph with the edges into ``inputs`` cut, since a fed input is not
+    computed from its ancestors. The result includes the outputs themselves.
     """
-    finalize_graph = graph.copy()
-    finalize_graph.remove_edges_from(
-        [edge for key in accumulator_keys for edge in graph.in_edges(key)]
+    cut_graph = graph.copy()
+    cut_graph.remove_edges_from(
+        [edge for key in inputs for edge in graph.in_edges(key)]
     )
-    reachable = set()
-    for key in target_keys:
-        if key in finalize_graph:
-            reachable |= nx.ancestors(finalize_graph, key)
-    return dynamic_keys & reachable
+    read = set(outputs)
+    for key in outputs:
+        read |= nx.ancestors(cut_graph, key)
+    return read
 
 
 def _find_descendants(

packages/essreduce/tests/streaming_test.py

@@ -225,8 +225,8 @@ def _make_static_a() -> StaticA:
         accumulators=(AccumA,),  # Note: No AccumB
     )
     streaming_wf.accumulate({DynamicA: 1, DynamicB: 4})
-    # Sciline passes `None` to the provider that needs AccumB.
-    with pytest.raises(TypeError, match='unsupported operand type'):
+    # Without bypass, finalize is not fed the dynamic key that AccumB depends on.
+    with pytest.raises(ValueError, match=r'No value was fed for .*DynamicB'):
         _ = streaming_wf.finalize()
 
 
@@ -1040,6 +1040,81 @@ def make_output(accumulated: Accumulated, context: Context) -> Output:
     assert sc.identical(result[Output], sc.scalar(207))
 
 
+@pytest.mark.parametrize('chunk_reads_derived', [False, True])
+@pytest.mark.parametrize(
+    ('output_reads_derived', 'derived_is_target'),
+    [(True, False), (False, True), (True, True)],
+)
+def test_StreamProcessor_computes_context_derived_node_once_per_context_change(
+    chunk_reads_derived: bool, output_reads_derived: bool, derived_is_target: bool
+) -> None:
+    Streamed = NewType('Streamed', int)
+    Context = NewType('Context', int)
+    Derived = NewType('Derived', int)
+    Accumulated = NewType('Accumulated', int)
+    Output = NewType('Output', int)
+
+    def derive(context: Context) -> Derived:
+        derive.call_count += 1
+        return Derived(context * 2)
+
+    derive.call_count = 0
+
+    def accumulate(streamed: Streamed) -> Accumulated:
+        return Accumulated(streamed)
+
+    def accumulate_scaled(streamed: Streamed, derived: Derived) -> Accumulated:
+        return Accumulated(streamed * derived)
+
+    def make_output(accumulated: Accumulated) -> Output:
+        return Output(accumulated)
+
+    def make_output_shifted(accumulated: Accumulated, derived: Derived) -> Output:
+        return Output(accumulated + derived)
+
+    def scale(derived: int) -> int:
+        return derived if chunk_reads_derived else 1
+
+    def shift(derived: int) -> int:
+        return derived if output_reads_derived else 0
+
+    wf = sciline.Pipeline(
+        (
+            derive,
+            accumulate_scaled if chunk_reads_derived else accumulate,
+            make_output_shifted if output_reads_derived else make_output,
+        )
+    )
+    streaming_wf = streaming.StreamProcessor(
+        base_workflow=wf,
+        dynamic_keys=(Streamed,),
+        context_keys=(Context,),
+        target_keys=(Output, Derived) if derived_is_target else (Output,),
+        accumulators=(Accumulated,),
+    )
+
+    streaming_wf.set_context({Context: sc.scalar(3)})
+    assert derive.call_count == 1
+    streaming_wf.accumulate({Streamed: sc.scalar(1)})
+    result = streaming_wf.finalize()
+    assert sc.identical(result[Output], sc.scalar(1 * scale(6) + shift(6)))
+    streaming_wf.accumulate({Streamed: sc.scalar(2)})
+    result = streaming_wf.finalize()
+    assert sc.identical(result[Output], sc.scalar(3 * scale(6) + shift(6)))
+    assert derive.call_count == 1
+
+    streaming_wf.set_context({Context: sc.scalar(5)})
+    assert derive.call_count == 2
+    streaming_wf.accumulate({Streamed: sc.scalar(1)})
+    result = streaming_wf.finalize()
+    expected = 3 * scale(6) + 1 * scale(10) + shift(10)
+    assert sc.identical(result[Output], sc.scalar(expected))
+    if derived_is_target:
+        assert sc.identical(result[Derived], sc.scalar(10))
+    streaming_wf.finalize()
+    assert derive.call_count == 2
+
+
 class WindowAccumulator(streaming.Accumulator[Any]):
     """Accumulator that drops its value on finalize, like a sliding window."""