feat(workflow): ParallelWorker and JoinNode (Part 4) - #591
Conversation
fc27f1f to
980ee86
Compare
Part 4/9 of the feature/workflows split, stacked on the built-in nodes.
- nodes/parallel_worker: runs a wrapped node once per item of a list input,
order-preserving, bounded by maxParallelWorkers, cancelling on first error
(a non-list input is treated as a single-element list). Registers a factory
with the engine (registerParallelWorkerFactory) so
buildNode(..., {parallelWorker: true}) works without a static import.
- nodes/join_node: a fan-in barrier that requires all predecessors and emits
the aggregated predecessor outputs as its output.
Tests (9): ParallelWorker mapping/order, single-item + empty-list handling,
concurrency bounding, first-error propagation, the registry factory
(buildNode + parallelWorker / maxParallelWorkers guard), and JoinNode
passthrough — all driven directly against a NodeContext. The graph-level
parallel/fan-in integration tests land in Part 6 with the runner. Full core
suite green (2384 tests).
980ee86 to
013670a
Compare
AmaadMartin
left a comment
There was a problem hiding this comment.
Order is deterministic and that part is right: results[i] is keyed by item index rather than completion order, and the shared nextIndex claim can't double-issue an item, so the output list always matches the input list. The things I'd want resolved before this lands are cancellation (the worker loop never checks ctx.abortSignal, so a timeout or abort stops nothing), the unlimited default concurrency, retryConfig/timeout being applied to both the wrapper and the node it wraps, and a failing item discarding every sibling's completed output. Separately, JoinNode's barrier flag has no reader anywhere in the tree at this head, and waitForOutput is a second unread flag for the same idea. Static review only — I did not run anything.
| constructor(inner: BaseNode, config: ParallelWorkerConfig = {}) { | ||
| super({ | ||
| name: inner.name, | ||
| rerunOnResume: true, | ||
| retryConfig: config.retryConfig, | ||
| timeout: config.timeout, | ||
| }); |
There was a problem hiding this comment.
Not a nit. The wrapper takes the same retryConfig and timeout as the node it wraps, so both levels apply them and the effects multiply.
super({
name: inner.name,
rerunOnResume: true,
retryConfig: config.retryConfig,
timeout: config.timeout,
});buildNode hands one options bag to both sides: buildInnerNode(nodeLike, options) at workflow_graph_utils.ts:123 (which reaches new FunctionNode(name, handler, options) at function_node.ts:218, whose ctor does super({name, ...config})), and then the same retryConfig/timeout again to the factory at workflow_graph_utils.ts:135-139. Two consequences:
- Retries compose.
{maxAttempts: 3}gives up to 3 inner attempts per item and 3 outer attempts of the whole fan-out. Worse,resultsis rebuilt on every entry intorunImpl, so an outer retry re-executes every item, including the ones that already succeeded. One flaky item in a 20-item LLM fan-out costs 20 more model calls per outer attempt. - The timeout means two different things. On the inner node it is a per-item budget; on the wrapper it is the budget for the entire fan-out.
{maxParallelWorkers: 2, timeout: 30}over 10 items that each take ~20s can only ever hit the outer deadline.
Simplest fix: don't forward them to the wrapper — drop retryConfig/timeout from this super() call and from ParallelWorkerFactory's options, and leave them on the inner node. If the wrapper genuinely needs its own, they should be separate options so a caller can set the two independently. The buildNode half is Part 2 code, so either side can carry the fix. Nothing in parallel_worker_test.ts constructs a ParallelWorker with a retryConfig or a timeout, which is why this is invisible today.
| for (;;) { | ||
| if (firstError !== undefined) { | ||
| return; | ||
| } |
There was a problem hiding this comment.
Not a nit. The loop never observes cancellation, so a timeout or an aborted invocation does not stop the fan-out.
for (;;) {
if (firstError !== undefined) {
return;
}node_context.ts:61-69 documents ctx.abortSignal as exactly the hook a node body is meant to poll, and node_runner.ts:192-193 sets it when the node declares a timeout (which this class accepts). On the deadline the runner stops consuming and fire-and-forgets iterator.return() (node_runner.ts:222) — but that request queues behind the in-flight await Promise.all(...), so the generator body runs to completion and every remaining item is still claimed and executed. Their events don't stop either: each child pushes into the shared channel through its own executeChildNode/runOnce, not through the parent's consume, so the runner's "nothing is pushed past the deadline" guarantee does not hold for this node.
for (;;) {
if (firstError !== undefined || isAborted(ctx)) {
return;
}with
const isAborted = (ctx: NodeContext): boolean =>
ctx.abortSignal?.aborted === true ||
ctx.invocationContext.abortSignal?.aborted === true;Both checks are needed: ctx.abortSignal is only populated when the ParallelWorker itself declares a timeout (node_runner.ts:176), so on a plain invocation abort it is undefined. This only stops scheduling — already-started items still run to completion, because ctx.runNode has no way to take a signal. Worth stating that limit in the class doc rather than leaving it implied.
Smaller point on the same guard: firstError !== undefined doubles as the failure flag, so an item that rejects with undefined (a bare Promise.reject()) leaves it unset — that item is skipped, results keeps a hole, and the node resolves successfully with a silently wrong list. A separate let failed = false; next to firstError removes the ambiguity.
| if (firstError !== undefined) { | ||
| throw firstError; | ||
| } | ||
| yield results; |
There was a problem hiding this comment.
Not a nit. One failing item throws away every sibling's completed output, and the caller never sees them.
if (firstError !== undefined) {
throw firstError;
}
yield results;At this point results holds the output of every item that finished, and all of it is discarded. With an LLM or remote-tool inner node those are results you have already paid for — and combined with the outer retry noted on the constructor, the next attempt pays for them again. Either emit a partial-result event before rethrowing, or state in the class doc that ParallelWorker is all-or-nothing so a graph author knows to make individual items failure-tolerant. What matters is that it is a deliberate choice; right now it reads as a side effect of the control flow.
Related, at line 22 the class doc says "cancelling on first error" and nothing is cancelled — the loop stops claiming new items while in-flight items run to completion (and, per the abort note, keep running past an abort too). "stopping on first error" is what the code does.
| /** Maximum number of items processed concurrently. `undefined` = unlimited. */ | ||
| maxParallelWorkers?: number; |
There was a problem hiding this comment.
Not a nit. The default is unlimited concurrency.
/** Maximum number of items processed concurrently. `undefined` = unlimited. */
maxParallelWorkers?: number;Unset, poolSize at line 61 becomes items.length, so a node fanned out over a 500-element list issues 500 concurrent inner runs. When the inner node is an LLM or a remote tool that is a rate-limit and cost incident, and the list length is normally data-driven, not author-chosen — the person who gets burned is not the person who wrote the graph. A named default:
const DEFAULT_MAX_PARALLEL_WORKERS = 8;
...
const poolSize = Math.min(
this.maxParallelWorkers ?? DEFAULT_MAX_PARALLEL_WORKERS,
items.length,
);Callers who really want unbounded can pass Infinity, which Math.min handles. If _parallel_worker.py is unbounded and you are matching it deliberately, please say so in the PR body — this is a case where diverging from the reference for a safe default seems worth it.
| // Key each child run by its item index (not call order) so the | ||
| // run id -> item mapping is deterministic. On resume this lets each | ||
| // item fast-forward from its own cached run rather than being matched | ||
| // to a differently-ordered run id. | ||
| const child = await ctx.runNode(this.inner, items[i], { | ||
| useSubBranch: true, | ||
| runId: String(i), | ||
| }); |
There was a problem hiding this comment.
Nit. The runId is set but the node path is not, so outside the scheduler path the children are not actually distinguishable.
const child = await ctx.runNode(this.inner, items[i], {
useSubBranch: true,
runId: String(i),
});executeChildNode computes nodePath as ${parent.nodePath}.${nodeName} (node_runner.ts:60-62) unless overrideNodePath is passed, and RunNodeOptions.overrideNodePath is documented as the field "used by the dynamic scheduler to embed the run id, e.g. wf.node@1, so distinct runs are distinguishable on resume" (node_runner.ts:36-41). With no scheduler set — the only path this PR's tests exercise — all N children emit events stamped with the same nodeInfo.path, so per-item attribution in the event stream is lost. The branch is unique here (useSubBranch + runId); the path is not.
The resume claim in the comment above also rests on ScheduleDynamicNode keying its runs map by (nodeName, runId), and only the interface is in this branch, so it can't be verified from here. Either pass overrideNodePath as well, or trim the comment to what this part actually guarantees and make the resume claim in the part that lands the scheduler.
| * A fan-in barrier node: it waits for ALL of its predecessors to complete, then | ||
| * emits the aggregated inputs (a map of predecessor name → output) as its | ||
| * output. | ||
| * | ||
| * Ported from `google/adk-python` `workflow/_join_node.py`. | ||
| */ | ||
| export class JoinNode extends BaseNode { | ||
| override get requiresAllPredecessors(): boolean { | ||
| return true; | ||
| } |
There was a problem hiding this comment.
Not a nit. Nothing reads requiresAllPredecessors, so this class is a barrier in name only — and there is a second flag for the same concept.
export class JoinNode extends BaseNode {
override get requiresAllPredecessors(): boolean {
return true;
}Grepping the tree at this head, requiresAllPredecessors appears exactly three times: the base getter (base_node.ts:110), this override, and parallel_worker_test.ts:93 asserting the getter returns true. There is no consumer anywhere. And waitForOutput — declared, assigned, never read at base_node.ts:41/71/95 — is documented as "the node only produces its output once all of its predecessors have triggered it (fan-in / join semantics)", which is the same sentence as this getter's doc. So there are two carriers for one concept with no stated precedence, and JoinNode sets one while leaving the other false: new JoinNode({name: 'j', waitForOutput: true}) is accepted today and means nothing. Please collapse them to one flag now, while there is no consumer to disagree with, and say in the PR body which part lands the reader.
Same file, lines 12-14: the doc says the node "emits the aggregated inputs (a map of predecessor name → output)", but runImpl yields input unchanged whatever its shape — the aggregation belongs to the engine and the engine side isn't here yet (the test hand-builds {a: 1, b: 2} and asserts it comes back). "Emits its input unchanged; the engine supplies the predecessor-name → output map" would describe the code that exists.
Since it came up on Part 1: I checked both carrier issues against this part and neither lands here — nothing in this diff reads Event.route/EventActions.route, and joinCompleted does not appear anywhere in the tree, so JoinNode is not writing it.
| [1, 2, 3, 4, 5], | ||
| ); | ||
| expect(output).toEqual([1, 2, 3, 4, 5]); | ||
| expect(peak).toBeLessThanOrEqual(2); |
There was a problem hiding this comment.
Nit. This assertion passes on an implementation that isn't parallel at all.
expect(peak).toBeLessThanOrEqual(2);peak would be 1 if the pool loop regressed to running items one at a time, and this test would stay green — so the only test of the class's headline feature can't fail in the direction that matters. expect(peak).toBe(2) pins both halves (never more than 2, and actually 2 at some point). It shouldn't be flaky the way timing assertions usually are: both workers are started before either item awaits, and the 5ms body means the first active-- is many microtask turns away.
Please ensure you have read the contribution guide before creating a pull request.
Link to Issue or Description of Change
1. Link to an existing issue (if applicable):
2. Or, if no issue exists, describe the change:
Problem:
Continuing the stacked split of the large
feature/workflowsbranch. With the engine core (Part 2) and the built-in Function/Tool nodes (Part 3) in place, the engine needs its concurrency primitives.Solution:
This is Part 4 of 9 — ParallelWorker and JoinNode — stacked on Part 3.
Stacked on: #part3_pr_number (Part 3 — Function/Tool nodes). Please merge Part 3 first.
Included:
nodes/parallel_worker.ts— runs a wrapped node once per item of a list input; order-preserving, bounded bymaxParallelWorkers, cancelling on the first error (a non-list input is treated as a single-element list). It registers a factory with the engine (registerParallelWorkerFactory) sobuildNode(..., {parallelWorker: true})works without a static import.nodes/join_node.ts— a fan-in barrier that requires all predecessors and emits the aggregated predecessor outputs as its output.Intentionally deferred: the graph-level parallel / fan-in integration tests (which use the runner) land in Part 6. Dynamic scheduling (Part 5), LLM-as-node (Part 7), and HITL (Part 8) follow.
Testing Plan
Unit Tests:
Bundled tests (9):
workflow/parallel_worker_test.ts— ParallelWorker mapping/order, single-item and empty-list handling, concurrency bounding (maxParallelWorkers), first-error propagation, the registry factory (buildNode+parallelWorker, and themaxParallelWorkers-without-parallelWorkerguard), and JoinNode passthrough — all driven directly against aNodeContext.Full core suite green (2384 tests). Typecheck clean:
npx tsc --noEmit -p core/tsconfig.json.Manual End-to-End (E2E) Tests:
N/A — node-level units; graph/runner E2E coverage lands in Part 6.
Checklist
Additional context
Stacked split — merge in order (…Part 3 → Part 4 → Part 5 → …). Diff: 3 files, +242.