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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user