diff --git a/src/work_graph_orchestrate.py b/src/work_graph_orchestrate.py index d2562f5..ebc9622 100644 --- a/src/work_graph_orchestrate.py +++ b/src/work_graph_orchestrate.py @@ -64,7 +64,9 @@ exception, re-raised unchanged after logging), `quota_limit` (the actually happened (N attempts), not a bug. """ +import concurrent.futures import json +import threading import time import uuid @@ -83,9 +85,61 @@ class ConsistencyCheckFailedError(Exception): 'blocked' by the time this is raised.""" -def topological_order(graph_manifest): - """Kahn's algorithm. Dependencies are read from BOTH `node['depends_on']` - and `graph_manifest['edges']` and merged -- `plan_work_graph()`'s own +class WorkGraphLayerError(Exception): + """One or more nodes in a concurrently-executed layer failed -- + `errors` is the full [(node_id, exception), ...] list (not just the + first failure), so a caller can tell whether every failure was the + same recoverable kind or something more serious.""" + + def __init__(self, errors): + self.errors = errors + summary = "; ".join("{}: {}".format(node_id, e) for node_id, e in errors) + super().__init__("{} node(s) failed in one layer: {}".format(len(errors), summary)) + + +class _LockedConn: + """sqlite3.Connection-compatible thread-safe proxy: execute/executemany/ + commit each take a shared lock for the duration of the call. SQLite only + ever allows one writer at a time regardless -- this doesn't create new + concurrency, it just makes the existing many-scattered-call-sites code + (`work_node_run.py`/`work_graph_orchestrate.py`/`execution_log.py`) safe + to call from multiple threads without rewriting every call site, since + the lock is held only per-statement, not around the (slow) `claude -p` + call itself. Mirrors `backend_server.py`'s existing shared-connection + + `threading.Lock()` pattern, generalized to per-statement granularity. + Attribute access not explicitly wrapped here (cursor()/rollback()/ + close()) passes through unlocked -- nothing on the concurrent path in + this module uses them.""" + + def __init__(self, conn, lock=None): + self._conn = conn + self._lock = lock or threading.Lock() + + def execute(self, *args, **kwargs): + with self._lock: + return self._conn.execute(*args, **kwargs) + + def executemany(self, *args, **kwargs): + with self._lock: + return self._conn.executemany(*args, **kwargs) + + def commit(self): + with self._lock: + return self._conn.commit() + + def __getattr__(self, name): + return getattr(self._conn, name) + + +def topological_layers(graph_manifest): + """Kahn's algorithm, layered: returns List[List[node]], where each inner + list is one round's batch of in-degree-zero nodes -- nodes within a + layer are, by construction, mutually independent (neither depends on + the other), so they're safe to run concurrently. Layers are ordered + (layer N+1 only becomes ready once layer N is fully accounted for); + each layer's nodes are sorted by node_id for deterministic test + assertions. Dependencies are read from BOTH `node['depends_on']` and + `graph_manifest['edges']` and merged -- `plan_work_graph()`'s own output isn't schema-validated (it can't be, `--json-schema` doesn't work with `--agent`, see work_graph_plan.py), so a real LLM response might populate one and not the other; trusting only one would silently @@ -106,21 +160,30 @@ def topological_order(graph_manifest): if dep in dependents: dependents[dep].append(node_id) - ready = sorted(node_id for node_id, degree in in_degree.items() if degree == 0) - ordered = [] - while ready: - node_id = ready.pop(0) - ordered.append(nodes[node_id]) - for dependent in sorted(dependents[node_id]): - in_degree[dependent] -= 1 - if in_degree[dependent] == 0: - ready.append(dependent) - ready.sort() - if len(ordered) != len(nodes): + layers = [] + current = sorted(node_id for node_id, degree in in_degree.items() if degree == 0) + seen_count = 0 + while current: + layers.append([nodes[node_id] for node_id in current]) + seen_count += len(current) + next_layer = set() + for node_id in current: + for dependent in dependents[node_id]: + in_degree[dependent] -= 1 + if in_degree[dependent] == 0: + next_layer.add(dependent) + current = sorted(next_layer) + if seen_count != len(nodes): raise ValueError( "graph_manifest has a dependency cycle or a depends_on/edge references an unknown node_id" ) - return ordered + return layers + + +def topological_order(graph_manifest): + """Flat view of topological_layers() -- kept for callers (e.g. + e2e_pipeline.py) that only need a total order, not the layer grouping.""" + return [node for layer in topological_layers(graph_manifest) for node in layer] def ensure_placeholder_skill(conn, skill_id, owner_org_node): @@ -193,25 +256,37 @@ def _make_run_fn(conn, target_folder_path, node_exec_id, org_node, prompt, outpu return run_fn +_DEFAULT_MAX_WORKERS = 4 # see this module's own docstring section on G for rationale + + def run_graph( conn, run_id, org_node, skill_id, target_folder_path, graph_manifest, build_prompt, input_schema, output_schema, timeout=600, needs_worktree=False, - verify_fn=None, sleep=time.sleep, now=time.time, + verify_fn=None, sleep=time.sleep, now=time.time, max_workers=_DEFAULT_MAX_WORKERS, ): - """Runs every node in `graph_manifest` (topologically ordered) through - `run_work_node()`, each backed by a real one-node dynamic-workflow - `agent()` call using `org_node`'s migrated persona. + """Runs every node in `graph_manifest` through `run_work_node()`, each + backed by a real one-node dynamic-workflow `agent()` call using + `org_node`'s migrated persona. Nodes are grouped into topological + layers (`topological_layers()`) and run layer-by-layer with a barrier + in between; within a layer -- nodes that are mutually independent by + construction -- a layer with 2+ nodes runs them concurrently via a + `ThreadPoolExecutor` (bounded by `max_workers`), a layer with exactly 1 + node runs it synchronously (no thread-pool overhead for the common + case: every real graph_manifest observed so far is a single-node-per- + layer chain). `build_prompt(node, outputs)` is called right before running that - node -- `outputs` contains every already-completed node's - output_payload so far (topological order guarantees a node's - dependencies are always already present), so prompts needing upstream - content (e.g. "here is the approved BRD...") can be built just-in-time - instead of requiring the whole graph's prompts precomputed before - anything runs. This also makes resume-after-restart work correctly: - `run_work_node` short-circuits already-`done` nodes and returns their - persisted output, which lands in `outputs` the same as a freshly-run - one would. + node -- `outputs` contains every already-completed EARLIER layer's + output_payload (a layer never starts until every node in the previous + layer has finished, so this is always safe to read without locking), + so prompts needing upstream content (e.g. "here is the approved + BRD...") can be built just-in-time instead of requiring the whole + graph's prompts precomputed before anything runs. `build_prompt` never + sees a same-layer sibling's output -- by definition, same-layer nodes + don't depend on each other. This also makes resume-after-restart work + correctly: `run_work_node` short-circuits already-`done` nodes and + returns their persisted output, which lands in `outputs` the same as a + freshly-run one would. `needs_worktree=True` (03-C §3.6: "并行写代码的节点必须先申请 git_worktree_ref") isolates each node's actual file changes in its own @@ -239,37 +314,89 @@ def run_graph( passes: `work_graph_run.status` -> 'verifying' and the regenerated output replaces the original in `outputs`. + **G (2026-08-06, Work Graph parallel execution)**: a multi-node layer's + failures are collected from every node in that layer (a sibling's + failure never cancels an already-running sibling -- they're + independent work, no reason to throw away completed work), then raised + together as `WorkGraphLayerError` once the whole layer has settled. + One real limitation, accepted deliberately (see + `docs/todo/work-graph-parallel-execution-design.md` §6): a node + waiting on a quota-backoff sleep inside `run_work_node` blocks the + thread pool worker running it, which in turn blocks the layer barrier + even after every sibling has finished -- Python threads have no clean + preemptible cancel, and building one is out of scope for this pass. + Returns {node_id: output_payload}. """ ensure_placeholder_skill(conn, skill_id, org_node) effective_needs_worktree = needs_worktree or paths.is_self_target(target_folder_path) - ordered = topological_order(graph_manifest) + layers = topological_layers(graph_manifest) outputs = {} - for node in ordered: - node_id = node["node_id"] - node_exec_id = _ensure_node_row( - conn, run_id, node_id, org_node, skill_id, - "work_graph_node_input.schema.json", "work_graph_node_output.schema.json", - {"node_id": node_id}, - ) - prompt = build_prompt(node, outputs) - run_fn = _make_run_fn(conn, target_folder_path, node_exec_id, org_node, prompt, output_schema, timeout) - output = work_node_run.run_work_node( - conn, node_exec_id, {"node_id": node_id}, input_schema, output_schema, run_fn, - needs_worktree=effective_needs_worktree, repo_path=target_folder_path, - sleep=sleep, now=now, - ) - - if verify_fn is not None: - output = _verify_and_maybe_regenerate( - conn, run_id, org_node, skill_id, target_folder_path, node, node_id, prompt, - output, input_schema, output_schema, timeout, effective_needs_worktree, verify_fn, sleep, now, + locked_conn = _LockedConn(conn) + + for layer in layers: + if len(layer) == 1: + node = layer[0] + outputs[node["node_id"]] = _run_one_node( + conn, run_id, org_node, skill_id, target_folder_path, node, outputs, + build_prompt, input_schema, output_schema, timeout, + effective_needs_worktree, verify_fn, sleep, now, ) + continue + + worker_count = min(max_workers, len(layer)) + errors = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=worker_count) as pool: + futures = { + pool.submit( + _run_one_node, locked_conn, run_id, org_node, skill_id, target_folder_path, + node, outputs, build_prompt, input_schema, output_schema, timeout, + effective_needs_worktree, verify_fn, sleep, now, + ): node + for node in layer + } + done, _ = concurrent.futures.wait(futures, return_when=concurrent.futures.ALL_COMPLETED) + for future in done: + node = futures[future] + try: + outputs[node["node_id"]] = future.result() + except Exception as e: + errors.append((node["node_id"], e)) + + if errors: + raise WorkGraphLayerError(errors) - outputs[node_id] = output return outputs +def _run_one_node( + conn, run_id, org_node, skill_id, target_folder_path, node, outputs, + build_prompt, input_schema, output_schema, timeout, needs_worktree, verify_fn, sleep, now, +): + """One node's full lifecycle (row creation, prompt, run, optional + verify+regenerate) -- shared by both the synchronous single-node-layer + path and the ThreadPoolExecutor multi-node-layer path in run_graph(), + so there's exactly one implementation of "run a node" to maintain.""" + node_id = node["node_id"] + node_exec_id = _ensure_node_row( + conn, run_id, node_id, org_node, skill_id, + "work_graph_node_input.schema.json", "work_graph_node_output.schema.json", + {"node_id": node_id}, + ) + prompt = build_prompt(node, outputs) + run_fn = _make_run_fn(conn, target_folder_path, node_exec_id, org_node, prompt, output_schema, timeout) + output = work_node_run.run_work_node( + conn, node_exec_id, {"node_id": node_id}, input_schema, output_schema, run_fn, + needs_worktree=needs_worktree, repo_path=target_folder_path, sleep=sleep, now=now, + ) + if verify_fn is not None: + output = _verify_and_maybe_regenerate( + conn, run_id, org_node, skill_id, target_folder_path, node, node_id, prompt, + output, input_schema, output_schema, timeout, needs_worktree, verify_fn, sleep, now, + ) + return output + + def _verify_and_maybe_regenerate( conn, run_id, org_node, skill_id, target_folder_path, node, node_id, prompt, output, input_schema, output_schema, timeout, needs_worktree, verify_fn, sleep, now, diff --git a/src/work_node_run.py b/src/work_node_run.py index 5854973..44c23e8 100644 --- a/src/work_node_run.py +++ b/src/work_node_run.py @@ -32,6 +32,7 @@ import json import os import re import subprocess +import threading import time import uuid from datetime import datetime, timezone @@ -105,15 +106,24 @@ def _quota_limit_type(text): return match.group(1).lower() if match else None +_worktree_creation_lock = threading.Lock() + + def create_worktree(repo_path, node_exec_id): """`git worktree add` for an isolated code-writing node (01-M - work_node_execution.git_worktree_ref).""" + work_node_execution.git_worktree_ref). Serialized across threads via a + module-level lock -- `git worktree add` mutates shared `.git/worktrees/` + metadata, and concurrent Work Graph layers (G) may call this from + multiple threads for the same repo_path at once. The lock only wraps + this one git invocation (milliseconds), not the node's actual work, so + it doesn't meaningfully reduce the parallelism gained elsewhere.""" worktree_path = os.path.join(repo_path, ".graphlemon-worktrees", node_exec_id) branch = "graphlemon/" + node_exec_id - subprocess.run( - ["git", "worktree", "add", "-b", branch, worktree_path], - cwd=repo_path, check=True, capture_output=True, text=True, - ) + with _worktree_creation_lock: + subprocess.run( + ["git", "worktree", "add", "-b", branch, worktree_path], + cwd=repo_path, check=True, capture_output=True, text=True, + ) return worktree_path