From 34110713e2dafc9de1a55c65cab804d0a3df1f27 Mon Sep 17 00:00:00 2001 From: stroblme Date: Sun, 23 Aug 2026 17:41:14 +0200 Subject: [PATCH] Let a pipeline swap one flow's nodes instead of being rebuilt Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01StpRc2C6au1WJ1EUU7fsfu --- backend/fluksio/flow/pipeline.py | 264 +++++++++++++++++------- backend/tests/flow/test_replace_flow.py | 221 ++++++++++++++++++++ 2 files changed, 416 insertions(+), 69 deletions(-) create mode 100644 backend/tests/flow/test_replace_flow.py diff --git a/backend/fluksio/flow/pipeline.py b/backend/fluksio/flow/pipeline.py index 1ee8497..b72d442 100644 --- a/backend/fluksio/flow/pipeline.py +++ b/backend/fluksio/flow/pipeline.py @@ -96,6 +96,38 @@ class NodeOutcome(BaseModel): artifacts: dict[str, dict[str, Any]] = {} +def _derive( + nodes: list[Node], +) -> tuple[dict[str, list[Node]], dict[Node, frozenset[Node]]]: + """Work out the wiring the node list implies: producers, then dependencies. + + Done over the whole list rather than one flow's share of it, because a + flow is not a subgraph — its nodes can read and write messages another + flow owns — so there is no deriving one flow's edges on their own. + """ + # A message may have several producers; every one of them is upstream + # of the nodes consuming it. + produces: dict[str, list[Node]] = {} + for node in nodes: + for msg in node.provides: + produces.setdefault(msg, []).append(node) + + # A node never depends on itself: reading a message it also provides is + # how state is carried between runs, not a cycle. An input marked + # non-triggering is the same idea across two nodes. + dependencies: dict[Node, frozenset[Node]] = { + node: frozenset( + producer + for msg, spec in node.requires.items() + if spec.trigger + for producer in produces.get(msg, ()) + if producer is not node + ) + for node in nodes + } + return produces, dependencies + + class Pipeline: """Directed graph of nodes with automatic dependency resolution.""" @@ -113,6 +145,7 @@ class Pipeline: "_paused", "_stepping", "_gate_lock", + "_graph_lock", "_queue", "_node_pool", "history_limits", @@ -140,6 +173,12 @@ class Pipeline: self._paused: set[str] = set() self._stepping: set[str] = set() self._gate_lock = threading.Lock() + # Held only while the graph maps are read or swapped together, never + # across an await, a node call or any I/O. Reentrant because the lazy + # `edges` build re-enters through `_graph`. It is never nested with + # `_gate_lock` or `state.lock()` in either order, which is what keeps + # three locks from needing an ordering rule. + self._graph_lock = threading.RLock() # An empty state backend is falsy, so this cannot be ``state or ...``: # that would quietly hand the pipeline a second, private state and # leave everyone reading the shared one seeing nothing. @@ -160,36 +199,13 @@ class Pipeline: # than the default puts its message in here. Swapped, never mutated. self.history_limits: dict[str, int] = {} - # A message may have several producers; every one of them is upstream - # of the nodes consuming it. - self.produces: dict[str, list[Node]] = {} - for node in self._nodes: - for msg in node.provides: - self.produces.setdefault(msg, []).append(node) - - # A node never depends on itself: reading a message it also provides is - # how state is carried between runs, not a cycle. An input marked - # non-triggering is the same idea across two nodes. - self.dependencies: dict[Node, frozenset[Node]] = { - node: frozenset( - producer - for msg, spec in node.requires.items() - if spec.trigger - for producer in self.produces.get(msg, ()) - if producer is not node - ) - for node in self._nodes - } + self.produces, self.dependencies = _derive(self._nodes) self._edges: dict[Node, set[Node]] | None = None self._execution_order: list[Node] | None = None self._downstream_cache: dict[Node, list[Node]] = {} - if initial_values: - with self._state.lock(): - for name, value in initial_values.items(): - if name not in self._state: - self._state[name] = value + self._seed(initial_values) for node in self._nodes: node.bind(self) @@ -209,12 +225,92 @@ class Pipeline: @property def edges(self) -> dict[Node, set[Node]]: """Producers mapped to their consumers, built on first access.""" - if self._edges is None: - self._edges = {node: set() for node in self._nodes} - for consumer, producers in self.dependencies.items(): - for producer in producers: - self._edges[producer].add(consumer) - return self._edges + edges = self._edges + if edges is None: + # Under the lock the node list and the dependencies are the same + # generation, so every producer has an entry to add a consumer to. + with self._graph_lock: + edges = self._edges + if edges is None: + edges = {node: set() for node in self._nodes} + for consumer, producers in self.dependencies.items(): + for producer in producers: + edges[producer].add(consumer) + self._edges = edges + return edges + + def _graph(self) -> tuple[dict[Node, frozenset[Node]], dict[Node, set[Node]]]: + """One matching view of the graph, for the length of a wave. + + A replace swaps these two together; reading them a moment apart is how + a wave ends up asking the new dependencies about a node the old edges + still know, which is a KeyError rather than a wrong answer. + """ + with self._graph_lock: + return self.dependencies, self.edges + + def replace_flow( + self, + flow: str, + nodes: list[Node], + initial_values: dict[str, Any] | None = None, + ) -> None: + """Swap one flow's nodes into the graph, leaving every other flow's be. + + On return the graph is what a build over the same nodes would have + produced, and the flows that were not named still hold the very node + objects they held before — started, connected, and never asked to + reconnect. An empty node list removes the flow. + """ + with self._graph_lock: + # Spliced where the old ones were, so the node list stays grouped + # by flow the way a full build lays it out. + spliced: list[Node] = [] + placed = False + for node in self._nodes: + if node.flow == flow: + if not placed: + spliced.extend(nodes) + placed = True + continue + spliced.append(node) + if not placed: + spliced.extend(nodes) + + produces, dependencies = _derive(spliced) + + # Assigned only once everything above has succeeded, and rebound + # rather than mutated: a cascade already walking the graph holds + # the old maps and finishes on them, which is the atomicity + # building a second pipeline used to give for free. + self._nodes = spliced + self.produces = produces + self.dependencies = dependencies + self._edges = None + self._execution_order = None + self._downstream_cache = {} + + for node in nodes: + node.bind(self) + self._seed(initial_values) + # A rebuild clears the debugging pause of what it rebuilt, and of + # nothing else. + with self._gate_lock: + self._paused.discard(flow) + self._stepping.discard(flow) + + def remove_flow(self, flow: str) -> None: + """Take a deleted flow out of the graph.""" + self.replace_flow(flow, []) + + def _seed(self, initial_values: dict[str, Any] | None) -> None: + """Give messages a starting value, without overwriting one already there.""" + if not initial_values: + return + with self._state.lock(): + for name, value in initial_values.items(): + if name not in self._state: + self._state[name] = value def get_node_by_id(self, nid: str) -> Node | None: return next((n for n in self._nodes if n.id == nid), None) @@ -224,39 +320,52 @@ class Pipeline: def _topological_sort(self) -> list[Node]: """Kahn's algorithm; nodes left over are part of a cycle.""" - if self._execution_order is not None: - return self._execution_order + order = self._execution_order + if order is not None: + return order - in_degree = {node: len(deps) for node, deps in self.dependencies.items()} - queue = deque(n for n, deg in in_degree.items() if deg == 0) - result: list[Node] = [] + with self._graph_lock: + deps, edges = self.dependencies, self.edges + in_degree = {node: len(d) for node, d in deps.items()} + queue = deque(n for n, deg in in_degree.items() if deg == 0) + result: list[Node] = [] - while queue: - node = queue.popleft() - result.append(node) - for consumer in self.edges[node]: - in_degree[consumer] -= 1 - if in_degree[consumer] == 0: - queue.append(consumer) + while queue: + node = queue.popleft() + result.append(node) + for consumer in edges[node]: + in_degree[consumer] -= 1 + if in_degree[consumer] == 0: + queue.append(consumer) - self._execution_order = result - return result + self._execution_order = result + return result def _get_downstream(self, start: Node) -> list[Node]: - if start not in self._downstream_cache: - reachable: set[Node] = set() - queue = deque([start]) - while queue: - for consumer in self.edges[queue.popleft()]: - if consumer not in reachable: - reachable.add(consumer) - queue.append(consumer) - order = self._topological_sort() - ordered = [n for n in order if n in reachable] - # Nodes inside a cycle never make it into the topological order. - ordered += [n for n in reachable if n not in order] - self._downstream_cache[start] = ordered - return self._downstream_cache[start] + # Bound once: a replace rebinds the cache, and writing the memo into + # the one this call started with loses it rather than corrupting it. + cache = self._downstream_cache + if start in cache: + return cache[start] + + _, edges = self._graph() + if start not in edges: + # Its flow was replaced while this cascade was on its way here. + return [] + + reachable: set[Node] = set() + queue = deque([start]) + while queue: + for consumer in edges[queue.popleft()]: + if consumer not in reachable: + reachable.add(consumer) + queue.append(consumer) + order = self._topological_sort() + ordered = [n for n in order if n in reachable] + # Nodes inside a cycle never make it into the topological order. + ordered += [n for n in reachable if n not in order] + cache[start] = ordered + return ordered # ------------------------------------------------------------------------- # Validation @@ -272,10 +381,14 @@ class Pipeline: """ declared = flow_inputs or {} issues: list[ValidationIssue] = [] + # Both callers hold the controller's rebuild lock, so nothing swaps + # these underneath; bound once so the report is of one graph either way. + nodes = self._nodes + produces = self.produces ordered = set(self._topological_sort()) - if len(ordered) != len(self._nodes): - cyclic = sorted(n.id for n in self._nodes if n not in ordered) + if len(ordered) != len(nodes): + cyclic = sorted(n.id for n in nodes if n not in ordered) issues.append( ValidationIssue( code="cycle", @@ -288,16 +401,16 @@ class Pipeline: ) ) - for node in self._nodes: + for node in nodes: for msg_name, spec in node.requires.items(): - if msg_name in self.produces: + if msg_name in produces: # A message a node both reads and writes carries state # between its runs. If the node is the only one writing it, # the first run has nothing to read unless the flow declares # a starting value. if ( spec.trigger - and self.produces[msg_name] == [node] + and produces[msg_name] == [node] and not declared.get(msg_name, False) ): issues.append( @@ -799,6 +912,10 @@ class Pipeline: def is_disabled(self, flow: str) -> bool: return flow in self._disabled + def set_disabled(self, flows: set[str]) -> None: + """Which flows are stopped. Swapped whole, so a reader never sees half.""" + self._disabled = frozenset(flows) + def paused_flows(self) -> list[str]: with self._gate_lock: return sorted(self._paused) @@ -856,13 +973,22 @@ class Pipeline: replay: bool = False, ) -> StateBackend: """Execute nodes concurrently, scheduling each as its inputs arrive.""" - target_nodes = nodes_subset if nodes_subset is not None else set(self._nodes) + # One view of the graph for the whole wave: a flow replaced halfway + # through must not have this wave asking the new dependencies about a + # node the old edges knew. + deps, edges = self._graph() + target_nodes = nodes_subset if nodes_subset is not None else set(deps) + if not target_nodes: + return state + # A wave that set out before a flow was replaced can be carrying nodes + # that are no longer in the graph. They have been stopped; running them + # is the one thing a per-flow rebuild must not let happen. + target_nodes = {n for n in target_nodes if n in deps} if not target_nodes: return state in_degree = { - n: sum(1 for dep in self.dependencies[n] if dep in target_nodes) - for n in target_nodes + n: sum(1 for dep in deps[n] if dep in target_nodes) for n in target_nodes } submitted: set[Node] = set() @@ -891,7 +1017,7 @@ class Pipeline: # outputs are still in state, so downstream carries on. skipped.add(n) submitted.discard(n) - for consumer in self.edges[n]: + for consumer in edges[n]: if consumer in target_nodes: in_degree[consumer] -= 1 continue @@ -914,7 +1040,7 @@ class Pipeline: # A node returning nothing (rate limiting, an error) stops # propagation along its branch. if result is not None: - for consumer in self.edges[n]: + for consumer in edges[n]: if consumer in target_nodes: in_degree[consumer] -= 1 submit_ready(executor) diff --git a/backend/tests/flow/test_replace_flow.py b/backend/tests/flow/test_replace_flow.py new file mode 100644 index 0000000..035c013 --- /dev/null +++ b/backend/tests/flow/test_replace_flow.py @@ -0,0 +1,221 @@ +"""Swapping one flow's nodes into a live graph, leaving the rest connected.""" + +import threading + +from fluksio.flow.messages import DType, MessageSpec +from fluksio.flow.nodes import Node +from fluksio.flow.pipeline import Pipeline +from tests.flow.test_pipeline import make_node, spec + + +def _shape(pipeline: Pipeline) -> tuple[dict[str, list[str]], dict[str, set[str]]]: + """The wiring by id, so two pipelines built differently can be compared.""" + return ( + {msg: [n.id for n in nodes] for msg, nodes in pipeline.produces.items()}, + { + node.id: {dep.id for dep in deps} + for node, deps in pipeline.dependencies.items() + }, + ) + + +def test_replacing_a_flow_leaves_the_other_flows_nodes_alone(): + b_producer = make_node( + "meter", "b", lambda params: {"power": 1.0}, provides=[spec("power")] + ) + b_consumer = make_node( + "load", "b", lambda power, params: None, requires=[spec("power")] + ) + a_old = make_node("calc", "a", lambda params: None) + pipeline = Pipeline(nodes=[a_old, b_producer, b_consumer]) + + pipeline.replace_flow("a", [make_node("calc", "a", lambda params: None)]) + + assert pipeline.get_node_by_id("b.meter") is b_producer + assert pipeline.get_node_by_id("b.load") is b_consumer + assert pipeline.dependencies[b_consumer] == frozenset({b_producer}) + + +def test_a_consumer_in_another_flow_picks_up_the_new_producer(): + seen: list[float] = [] + consumer = make_node( + "load", + "b", + lambda temp, params: seen.append(temp), + requires=[spec("a.temp")], + ) + pipeline = Pipeline( + nodes=[ + make_node( + "old", "a", lambda params: {"temp": 1.0}, provides=[spec("temp")] + ), + consumer, + ] + ) + + replacement = make_node( + "new", "a", lambda params: {"temp": 2.0}, provides=[spec("temp")] + ) + pipeline.replace_flow("a", [replacement]) + + assert pipeline.dependencies[consumer] == frozenset({replacement}) + replacement.inject() + assert seen == [2.0] + + +def test_a_consumer_loses_its_edge_when_the_producer_flow_stops_providing(): + """The crux: the consumer becomes a root, not an orphan nothing can reach.""" + consumer = make_node( + "load", "b", lambda temp, params: None, requires=[spec("a.temp")] + ) + pipeline = Pipeline( + nodes=[ + make_node( + "meter", "a", lambda params: {"temp": 1.0}, provides=[spec("temp")] + ), + consumer, + ] + ) + + pipeline.replace_flow("a", [make_node("quiet", "a", lambda params: None)]) + + assert "a.temp" not in pipeline.produces + assert pipeline.dependencies[consumer] == frozenset() + assert consumer in pipeline._topological_sort() + + +def test_removing_a_flow_takes_its_nodes_out(): + ran: list[str] = [] + a_node = make_node( + "meter", + "a", + lambda params: ran.append("a") or {"temp": 1.0}, + provides=[spec("temp")], + ) + b_node = make_node( + "beat", + "b", + lambda params: ran.append("b") or {"tick": 1.0}, + provides=[spec("tick")], + ) + pipeline = Pipeline(nodes=[a_node, b_node]) + + pipeline.remove_flow("a") + + assert pipeline.flow_nodes("a") == set() + assert pipeline.get_node_by_id("a.meter") is None + assert a_node not in pipeline.dependencies + b_node.inject() + assert ran == ["b"] + + # A caller still holding the removed node — a wave that picked its targets + # before the removal — runs nothing rather than raising or reviving it. + pipeline.run({}, nodes={a_node}) + assert ran == ["b"] + + +def test_a_replace_leaves_the_graph_a_rebuild_would_have_built(): + """A replace and a fresh build over the same nodes must agree, edge for edge.""" + + def flow_a(tag: str) -> list[Node]: + # Reads a message flow b owns, writes one flow b reads, and shares a + # third with a producer in flow b — so the comparison covers edges in + # both directions across the boundary, and the order of a fan-in whose + # producers sit on either side of it. + return [ + make_node( + "meter", + "a", + lambda power, params: {"temp": 1.0}, + requires=[spec("b.power")], + provides=[spec("temp"), MessageSpec(name=tag, dtype=DType.FLOAT)], + ), + make_node( + "inner", + "a", + lambda temp, params: None, + requires=[spec("temp")], + ), + make_node( + "mirror", + "a", + lambda params: {"power": 1.0}, + provides=[MessageSpec(name="b.power", dtype=DType.FLOAT)], + ), + ] + + def flow_b() -> list[Node]: + return [ + make_node( + "supply", "b", lambda params: {"power": 1.0}, provides=[spec("power")] + ), + make_node( + "load", "b", lambda temp, params: None, requires=[spec("a.temp")] + ), + ] + + # Both flows write it, so the order the two producers come out in is the + # node list's, and a replace that appends rather than splices gets it wrong. + b_nodes = flow_b() + replaced = Pipeline(nodes=flow_a("old") + b_nodes) + replaced.replace_flow("a", flow_a("new")) + + fresh = Pipeline(nodes=flow_a("new") + flow_b()) + + assert _shape(replaced) == _shape(fresh) + assert [n.id for n in replaced._topological_sort()] == [ + n.id for n in fresh._topological_sort() + ] + + +def test_a_cascade_running_during_a_replace_finishes_on_the_graph_it_started_with(): + started = threading.Event() + release = threading.Event() + ran: list[str] = [] + + def slow(params): + started.set() + release.wait(5) + return {"temp": 1.0} + + source = make_node("source", "a", slow, provides=[spec("temp")]) + consumer = make_node( + "load", "a", lambda temp, params: ran.append("load"), requires=[spec("temp")] + ) + pipeline = Pipeline(nodes=[source, consumer]) + + failure: list[BaseException] = [] + + def wave() -> None: + try: + pipeline.run({}) + except BaseException as exc: # noqa: BLE001 - reported, not swallowed + failure.append(exc) + + thread = threading.Thread(target=wave) + thread.start() + # The replace lands while the wave is inside the executor holding the old + # maps, which is the only moment the swap can be seen half-applied. + assert started.wait(5) + pipeline.replace_flow("a", [make_node("source", "a", lambda params: None)]) + release.set() + thread.join(10) + + assert not thread.is_alive() + assert failure == [] + assert ran == ["load"] + + +def test_replacing_a_flow_clears_its_pause_and_leaves_another_flows(): + pipeline = Pipeline( + nodes=[ + make_node("one", "a", lambda params: None), + make_node("two", "b", lambda params: None), + ] + ) + pipeline.pause("a") + pipeline.pause("b") + + pipeline.replace_flow("a", [make_node("one", "a", lambda params: None)]) + + assert pipeline.paused_flows() == ["b"]