From a93dc29f627a8c5ba7cc916f86a70e914e2f0fb9 Mon Sep 17 00:00:00 2001 From: Simon Heybrock Date: Fri, 9 Oct 2026 06:50:50 +0000 Subject: [PATCH 1/2] Feed cached context nodes to the workflows that read them StreamProcessor.set_context fed every non-target cached context node to the chunk workflow, and the finalize workflow received only the raw context keys and context-derived targets. A context-derived node read downstream of an accumulator was therefore computed in set_context for nothing, and recomputed from the raw context keys on every finalize(), even with unchanged context. Each cached context node now goes to the workflow(s) that read it: the chunk workflow if it is upstream of an accumulator, the finalize workflow if it is a target or downstream of an accumulator. Every path from a context key to a target passes through a cached context node, so the finalize workflow no longer needs the raw context keys. Co-Authored-By: Claude Opus 5.5 --- .../essreduce/src/ess/reduce/streaming.py | 89 ++++++++++--------- packages/essreduce/tests/streaming_test.py | 87 ++++++++++++++++++ 2 files changed, 133 insertions(+), 43 deletions(-) diff --git a/packages/essreduce/src/ess/reduce/streaming.py b/packages/essreduce/src/ess/reduce/streaming.py index 89c16222b..ef7b983e8 100644 --- a/packages/essreduce/src/ess/reduce/streaming.py +++ b/packages/essreduce/src/ess/reduce/streaming.py @@ -425,35 +425,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 it is upstream of an accumulator, the finalize workflow if it is + # a target or downstream of an accumulator, 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 +483,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] @@ -839,30 +845,27 @@ 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 + outputs = outputs & set(graph.nodes) + read = set(outputs) + for key in outputs: + read |= nx.ancestors(cut_graph, key) + return read def _find_descendants( diff --git a/packages/essreduce/tests/streaming_test.py b/packages/essreduce/tests/streaming_test.py index 97f67d88a..0e3b0787b 100644 --- a/packages/essreduce/tests/streaming_test.py +++ b/packages/essreduce/tests/streaming_test.py @@ -1040,6 +1040,93 @@ def make_output(accumulated: Accumulated, context: Context) -> Output: assert sc.identical(result[Output], sc.scalar(207)) +@pytest.mark.parametrize('derived_is_target', [False, True]) +def test_StreamProcessor_finalize_reuses_context_derived_node( + 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 make_output(accumulated: Accumulated, derived: Derived) -> Output: + return Output(accumulated + derived) + + wf = sciline.Pipeline((derive, accumulate, 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)}) + assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(1 + 6)) + streaming_wf.accumulate({Streamed: sc.scalar(2)}) + assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(3 + 6)) + assert derive.call_count == 1 + + streaming_wf.set_context({Context: sc.scalar(5)}) + assert derive.call_count == 2 + assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(3 + 10)) + streaming_wf.finalize() + assert derive.call_count == 2 + + +def test_StreamProcessor_reuses_context_derived_node_in_chunk_and_finalize() -> 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, derived: Derived) -> Accumulated: + return Accumulated(streamed * derived) + + def make_output(accumulated: Accumulated, derived: Derived) -> Output: + return Output(accumulated + derived) + + wf = sciline.Pipeline((derive, accumulate, make_output)) + streaming_wf = streaming.StreamProcessor( + base_workflow=wf, + dynamic_keys=(Streamed,), + context_keys=(Context,), + target_keys=(Output,), + accumulators=(Accumulated,), + ) + + streaming_wf.set_context({Context: sc.scalar(3)}) + streaming_wf.accumulate({Streamed: sc.scalar(1)}) + assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(1 * 6 + 6)) + streaming_wf.accumulate({Streamed: sc.scalar(2)}) + assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(3 * 6 + 6)) + assert derive.call_count == 1 + + streaming_wf.set_context({Context: sc.scalar(5)}) + streaming_wf.accumulate({Streamed: sc.scalar(1)}) + assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(3 * 6 + 10 + 10)) + assert derive.call_count == 2 + + class WindowAccumulator(streaming.Accumulator[Any]): """Accumulator that drops its value on finalize, like a sliding window.""" From b4e9242fd09245aa98ac0292d5f197ee799da8cc Mon Sep 17 00:00:00 2001 From: Simon Heybrock Date: Fri, 9 Oct 2026 11:50:31 +0000 Subject: [PATCH 2/2] Make unfed dynamic and context keys raise when read The workflow StreamProcessor computes on had `None` assigned to every dynamic and context key, to prune their branches. A chunk or finalize computation that reached such a key without being fed it silently computed from `None`. On main this happened for a context-derived target that an accumulator also reads: the chunk workflow was not fed the target, so it computed it from `Context = None`. The previous commit fixed that case by routing; this makes any such gap fail with "No value was fed for ''". The placeholders reuse the input wiring of _FedWorkflow, now the module function _wire_input. Without allow_bypass, finalize reading a dynamic key past the accumulators now raises that ValueError instead of a TypeError from inside a provider. The context-derived node tests become one parametrized test over who reads the node (chunk, finalize, target), which includes the case broken on main. Co-Authored-By: Claude Opus 5.5 --- .../essreduce/src/ess/reduce/streaming.py | 74 ++++++++------- packages/essreduce/tests/streaming_test.py | 90 ++++++++----------- 2 files changed, 81 insertions(+), 83 deletions(-) diff --git a/packages/essreduce/src/ess/reduce/streaming.py b/packages/essreduce/src/ess/reduce/streaming.py index ef7b983e8..f5df705f3 100644 --- a/packages/essreduce/src/ess/reduce/streaming.py +++ b/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 + + 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} - self._workflow.insert(provide_input_value) + # 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: @@ -432,8 +441,8 @@ def __init__( } accumulator_keys = set(self._accumulators) # Each cached context node goes to the workflow(s) that read it: the chunk - # workflow if it is upstream of an accumulator, the finalize workflow if it is - # a target or downstream of an accumulator, or both. + # 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 ) @@ -781,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 @@ -861,7 +872,6 @@ def _find_keys_read( cut_graph.remove_edges_from( [edge for key in inputs for edge in graph.in_edges(key)] ) - outputs = outputs & set(graph.nodes) read = set(outputs) for key in outputs: read |= nx.ancestors(cut_graph, key) diff --git a/packages/essreduce/tests/streaming_test.py b/packages/essreduce/tests/streaming_test.py index 0e3b0787b..9cda92845 100644 --- a/packages/essreduce/tests/streaming_test.py +++ b/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,9 +1040,13 @@ def make_output(accumulated: Accumulated, context: Context) -> Output: assert sc.identical(result[Output], sc.scalar(207)) -@pytest.mark.parametrize('derived_is_target', [False, True]) -def test_StreamProcessor_finalize_reuses_context_derived_node( - derived_is_target: bool, +@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) @@ -1059,71 +1063,55 @@ def derive(context: Context) -> Derived: def accumulate(streamed: Streamed) -> Accumulated: return Accumulated(streamed) - def make_output(accumulated: Accumulated, derived: Derived) -> Output: - return Output(accumulated + derived) - - wf = sciline.Pipeline((derive, accumulate, 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)}) - assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(1 + 6)) - streaming_wf.accumulate({Streamed: sc.scalar(2)}) - assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(3 + 6)) - assert derive.call_count == 1 - - streaming_wf.set_context({Context: sc.scalar(5)}) - assert derive.call_count == 2 - assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(3 + 10)) - streaming_wf.finalize() - assert derive.call_count == 2 - - -def test_StreamProcessor_reuses_context_derived_node_in_chunk_and_finalize() -> None: - Streamed = NewType('Streamed', int) - Context = NewType('Context', int) - Derived = NewType('Derived', int) - Accumulated = NewType('Accumulated', int) - Output = NewType('Output', int) + def accumulate_scaled(streamed: Streamed, derived: Derived) -> Accumulated: + return Accumulated(streamed * derived) - def derive(context: Context) -> Derived: - derive.call_count += 1 - return Derived(context * 2) + def make_output(accumulated: Accumulated) -> Output: + return Output(accumulated) - derive.call_count = 0 + def make_output_shifted(accumulated: Accumulated, derived: Derived) -> Output: + return Output(accumulated + derived) - def accumulate(streamed: Streamed, derived: Derived) -> Accumulated: - return Accumulated(streamed * derived) + def scale(derived: int) -> int: + return derived if chunk_reads_derived else 1 - def make_output(accumulated: Accumulated, derived: Derived) -> Output: - return Output(accumulated + derived) + def shift(derived: int) -> int: + return derived if output_reads_derived else 0 - wf = sciline.Pipeline((derive, accumulate, make_output)) + 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,), + 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)}) - assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(1 * 6 + 6)) + result = streaming_wf.finalize() + assert sc.identical(result[Output], sc.scalar(1 * scale(6) + shift(6))) streaming_wf.accumulate({Streamed: sc.scalar(2)}) - assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(3 * 6 + 6)) + 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)}) - assert sc.identical(streaming_wf.finalize()[Output], sc.scalar(3 * 6 + 10 + 10)) + 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