diff --git a/packages/essreduce/src/ess/reduce/streaming.py b/packages/essreduce/src/ess/reduce/streaming.py index 89c16222b..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 - # 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( diff --git a/packages/essreduce/tests/streaming_test.py b/packages/essreduce/tests/streaming_test.py index 97f67d88a..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,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."""