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 084194f77b
commit 294b72d989
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)