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
This commit is contained in:
2026-08-23 17:41:14 +02:00
co-authored by Claude Opus 5
parent 715e9a9e31
commit 34110713e2
2 changed files with 416 additions and 69 deletions
+195 -69
View File
@@ -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)
+221
View File
@@ -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"]