Let a pipeline swap one flow's nodes instead of being rebuilt

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01StpRc2C6au1WJ1EUU7fsfu

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
+        self.produces, self.dependencies = _derive(self._nodes)
-        # 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._edges: dict[Node, set[Node]] | None = None
         self._execution_order: list[Node] | None = None
         self._downstream_cache: dict[Node, list[Node]] = {}
 
-        if initial_values:
+        self._seed(initial_values)
-            with self._state.lock():
-                for name, value in initial_values.items():
-                    if name not in self._state:
-                        self._state[name] = value
 
         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:
+        edges = self._edges
-            self._edges = {node: set() for node in self._nodes}
+        if edges is None:
-            for consumer, producers in self.dependencies.items():
+            # Under the lock the node list and the dependencies are the same
-                for producer in producers:
+            # generation, so every producer has an entry to add a consumer to.
-                    self._edges[producer].add(consumer)
+            with self._graph_lock:
-        return self._edges
+                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:
+        order = self._execution_order
-            return self._execution_order
+        if order is not None:
+            return order
 
-        in_degree = {node: len(deps) for node, deps in self.dependencies.items()}
+        with self._graph_lock:
-        queue = deque(n for n, deg in in_degree.items() if deg == 0)
+            deps, edges = self.dependencies, self.edges
-        result: list[Node] = []
+            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:
+            while queue:
-            node = queue.popleft()
+                node = queue.popleft()
-            result.append(node)
+                result.append(node)
-            for consumer in self.edges[node]:
+                for consumer in edges[node]:
-                in_degree[consumer] -= 1
+                    in_degree[consumer] -= 1
-                if in_degree[consumer] == 0:
+                    if in_degree[consumer] == 0:
-                    queue.append(consumer)
+                        queue.append(consumer)
 
-        self._execution_order = result
+            self._execution_order = result
-        return result
+            return result
 
     def _get_downstream(self, start: Node) -> list[Node]:
-        if start not in self._downstream_cache:
+        # Bound once: a replace rebinds the cache, and writing the memo into
-            reachable: set[Node] = set()
+        # the one this call started with loses it rather than corrupting it.
-            queue = deque([start])
+        cache = self._downstream_cache
-            while queue:
+        if start in cache:
-                for consumer in self.edges[queue.popleft()]:
+            return cache[start]
-                    if consumer not in reachable:
+
-                        reachable.add(consumer)
+        _, edges = self._graph()
-                        queue.append(consumer)
+        if start not in edges:
-            order = self._topological_sort()
+            # Its flow was replaced while this cascade was on its way here.
-            ordered = [n for n in order if n in reachable]
+            return []
-            # Nodes inside a cycle never make it into the topological order.
+
-            ordered += [n for n in reachable if n not in order]
+        reachable: set[Node] = set()
-            self._downstream_cache[start] = ordered
+        queue = deque([start])
-        return self._downstream_cache[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):
+        if len(ordered) != len(nodes):
-            cyclic = sorted(n.id for n in self._nodes if n not in ordered)
+            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)
+            n: sum(1 for dep in deps[n] if dep in target_nodes) for n 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)

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"]