From 22363a11c94d590dbe0163a1fe641ae166eae4d1 Mon Sep 17 00:00:00 2001 From: James Russo Date: Sat, 16 May 2026 19:29:34 -0400 Subject: [PATCH] perf(distributed): parallelize chunk capture across multiple workers (#906) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * perf(distributed): parallelize chunk capture across multiple workers The distributed `renderChunk` primitive hardcoded `workerCount: 1` and `captureStage` explicitly forbade `workerCount > 1` when `frameRange` was set, with the comment: "Distributed chunk workers fan out at the activity layer; reduce workerCount to 1 when passing frameRange." The assumption was that orchestration-layer fan-out (Temporal / Lambda / K8s Jobs / SSH) saturates the available CPU on its own. In practice adopters that deploy chunks onto multi-core hosts (8-24 vCPU is the standard producer-worker pod sizing) end up pinning only ~3-4 cores per chunk while the rest sit idle: chunk-level fan-out at the orchestration layer gives each pod one chunk at a time, and the chunk render itself was single-threaded. Validated against a real 1080p / 30fps / 22-second shader-heavy composition on a 22-vCPU Temporal pod: each chunk rendered at 165-273ms per frame (vs 94-98ms for the in-process streaming render which runs `workerCount=2` by default). The slowest chunk gates total wall-clock under parallel chunk fan-out, so the 2-3x per-frame gap compounds and `distributed` was net-slower than `in-process` on every composition smaller than ~5min of texture-class content. Lifting the restriction is a measured ~2x per-chunk speedup with no contract change at the framesDir or encoder layer. Wire-up: * `WorkerTask.outputFrameOffset` — optional offset subtracted from the absolute frame index when computing the captured file's name. Default 0 (the in-process contract; file name == absolute index). Distributed chunks set this to the chunk's startFrame so file names land 0-indexed within the chunk's range, matching the sequential chunk-capture contract and the encoder's expectation that frames are read sequentially without an `-start_number` override. * `distributeFrames(totalFrames, workerCount, workDir, rangeStart=0)` — offsets both `startFrame`/`endFrame` (used for per-frame time math on the page's virtual clock) by `rangeStart`, and threads `outputFrameOffset = rangeStart` onto each task it emits. With the default `rangeStart=0` it is a no-op for in-process renders. * `executeWorkerTask` — uses `i - (task.outputFrameOffset ?? 0)` for the captured file name, leaving the per-frame TIME computation `(i * fps.den) / fps.num` untouched so the page's virtual clock is unchanged. * `executeDiskCaptureWithAdaptiveRetry({ frameRangeStart? })` — accepts the chunk's absolute startFrame and forwards it to `distributeFrames` and `buildMissingFrameRetryBatches`. Default `undefined` preserves the in-process contract. * `buildMissingFrameRetryBatches(ranges, ..., rangeStart=0)` — `findMissingFrameRanges` walks LOCAL 0-indexed file names; the retry batch translates the local missing-range pair back to ABSOLUTE composition indices for `WorkerTask.startFrame/endFrame` and sets `outputFrameOffset = rangeStart` so the retried capture writes back to the same local file name. * `captureStage` — drops the assert; passes `frameRangeStart: frameRange?.startFrame` to the parallel branch so workers land on absolute composition frame indices for time math while file names stay 0-indexed within the chunk range. Docstring updated to reflect that the parallel branch is now supported. * `renderChunk` — `workerCount: 1` → `workerCount: 2`. The pre-warmed `probeSession` is consumed only by the sequential branch; the parallel branch closes it during stage entry and creates its own worker sessions. Documented as a follow-up: skip probeSession creation when `workerCount > 1` to recover the ~3-5s warmup cost. Backwards compatibility: every change is gated on a parameter that defaults to the prior behavior. In-process callers (`executeRenderJob`) pass no `frameRangeStart`, so `rangeStart === 0`, `outputFrameOffset` defaults to 0, and the file-name math collapses to the prior `i` value. The framesDir contract (`frame_0..frame_(totalFrames-1)`) and the WorkerTask interface are extended, not replaced. Tests: 24 pass / 0 fail across the distributed test suite (renderChunk, plan, assemble, planFormatBanlist, planSizeCap, publicExports). 7 pass / 0 fail in `parallelCoordinator.test.ts`. The renderOrchestrator suite has one pre-existing Windows-only failure (`writeCompiledArtifacts — external assets on Windows drive-letter paths`) unrelated to this change; the other 56 tests pass. Refs: distributed-vs-inprocess benchmark thread at heygen-com/experiment-framework#36950 * perf(distributed): auto-size chunk workerCount via calculateOptimalWorkers Match the in-process renderer's worker selection instead of hardcoding 2. `calculateOptimalWorkers(framesInChunk, undefined, cfg)` is the same call `resolveRenderWorkerCount` makes under the hood, minus the capture-cost calibration reduction (which would require plumbing the chunk's compiled metadata through — left as a follow-up). For a typical 22-vCPU producer-worker pod with `cfg.concurrency: "auto"` this resolves to ~6 workers for a 240-frame chunk (capped by `defaultSafeMaxWorkers() = max(6, min(16, floor(cpuCount/8)))`), matching what `executeRenderJob` (the in-process path) already does. The prior hardcoded `workerCount: 2` was a safe-minimum starting point that undersized chunks vs prod's auto behavior. Tests: 12/12 pass in `renderChunk.test.ts` (unchanged — the test suite mocks the inner runCaptureStage call so workerCount selection is opaque to it). * refactor(distributed): /simplify pass on PR #906 Review pass on the parallel-capture frame-range change. Four targeted cleanups identified by code-quality and efficiency review agents: 1. Add the missing `frameRange.endFrame - frameRange.startFrame === totalFrames` assert. The parallel branch forwards `totalFrames` separately from `frameRangeStart`; a caller passing mismatched values would have got a silently wrong distribution. The sequential branch already implicitly relied on this via its `rangeFrames = rangeEnd - rangeStart` arithmetic. 2. Collapse three near-duplicate docstrings (on `WorkerTask.outputFrameOffset`, `executeDiskCaptureWithAdaptiveRetry.frameRangeStart`, and `runCaptureStage`'s `frameRange`) so only the WorkerTask field carries the full contract. The other two cross-reference it. 3. Drop the WHAT-narrating comments inside `executeWorkerTask`'s per-frame loop. The variable names (`fileFrameIdx = i - outputOffset`) already say what the line does; the only remaining comment flags the non-obvious contract that the streaming callback gets the absolute index. 4. Trim the 30-line `chunkWorkerCount` block in `renderChunk` to one paragraph explaining the one non-obvious thing (why we use `calculateOptimalWorkers` directly instead of `resolveRenderWorkerCount`). The probeSession-wasted-on- parallel acknowledgement stays as a 3-line follow-up flag — investigated skipping it in this pass, but the SwiftShader probe is safety-critical and has no per-worker equivalent, so deferred to a separate change with proper per-worker assertion plumbing. Tests + format + lint clean: * `bun test parallelCoordinator.test.ts` — 7/7 * `bun test distributed/{renderChunk,plan}.test.ts` — 24/24 * `bunx oxfmt` + `bunx oxlint` — clean --- .../src/services/parallelCoordinator.ts | 28 ++++++++--- .../src/services/distributed/renderChunk.ts | 18 ++++--- .../services/render/stages/captureStage.ts | 47 ++++++++++--------- .../src/services/renderOrchestrator.ts | 28 +++++++++-- 4 files changed, 81 insertions(+), 40 deletions(-) diff --git a/packages/engine/src/services/parallelCoordinator.ts b/packages/engine/src/services/parallelCoordinator.ts index 4b21019bd..cf405476c 100644 --- a/packages/engine/src/services/parallelCoordinator.ts +++ b/packages/engine/src/services/parallelCoordinator.ts @@ -29,6 +29,15 @@ export interface WorkerTask { startFrame: number; endFrame: number; outputDir: string; + /** + * Offset subtracted from the absolute frame index when naming the captured + * file (`frame_.{ext}`). Default 0. Distributed + * chunks set this to the chunk's absolute startFrame so file names land + * 0-indexed within the chunk's range — the encoder reads frames + * sequentially without an `-start_number` override. The per-frame TIME + * calculation still uses the absolute frame index. + */ + outputFrameOffset?: number; } export interface WorkerResult { @@ -148,20 +157,22 @@ export function distributeFrames( totalFrames: number, workerCount: number, workDir: string, + rangeStart: number = 0, ): WorkerTask[] { const tasks: WorkerTask[] = []; const framesPerWorker = Math.ceil(totalFrames / workerCount); for (let i = 0; i < workerCount; i++) { - const startFrame = i * framesPerWorker; - const endFrame = Math.min((i + 1) * framesPerWorker, totalFrames); - if (startFrame >= totalFrames) break; + const startFrame = rangeStart + i * framesPerWorker; + const endFrame = Math.min(rangeStart + (i + 1) * framesPerWorker, rangeStart + totalFrames); + if (startFrame >= rangeStart + totalFrames) break; tasks.push({ workerId: i, startFrame, endFrame, outputDir: join(workDir, `worker-${i}`), + outputFrameOffset: rangeStart, }); } @@ -196,6 +207,7 @@ async function executeWorkerTask( ); await initializeSession(session); + const outputOffset = task.outputFrameOffset ?? 0; for (let i = task.startFrame; i < task.endFrame; i++) { if (signal?.aborted) { throw new Error("Parallel worker cancelled"); @@ -204,14 +216,16 @@ async function executeWorkerTask( // frame-index → time math. The 1-in-1001 ULP loss for NTSC is invisible // at our scales (frame count tops out at single-digit thousands). const time = (i * captureOptions.fps.den) / captureOptions.fps.num; + const fileFrameIdx = i - outputOffset; if (onFrameBuffer) { - // Streaming mode: capture to buffer and invoke callback - const { buffer } = await captureFrameToBuffer(session, i, time); + // The streaming-encode callback receives the absolute index `i` + // (not `fileFrameIdx`) so the encoder sequences frames against the + // composition's timeline. + const { buffer } = await captureFrameToBuffer(session, fileFrameIdx, time); await onFrameBuffer(i, buffer); } else { - // Disk mode: capture to file - await captureFrame(session, i, time); + await captureFrame(session, fileFrameIdx, time); } framesCaptured++; diff --git a/packages/producer/src/services/distributed/renderChunk.ts b/packages/producer/src/services/distributed/renderChunk.ts index 0cc494d8b..d0f9013e6 100644 --- a/packages/producer/src/services/distributed/renderChunk.ts +++ b/packages/producer/src/services/distributed/renderChunk.ts @@ -42,6 +42,7 @@ import { assertSwiftShader, type BeforeCaptureHook, BROWSER_GPU_NOT_SOFTWARE, + calculateOptimalWorkers, type CaptureOptions, type CaptureSession, closeCaptureSession, @@ -531,7 +532,13 @@ export async function renderChunk( // would deadlock Chrome's compositor by issuing a second beginFrame // at a `frameTimeTicks` it had just advanced to. - // ── Capture the chunk's range via runCaptureStage ── + // Capture-cost calibration based on shader transitions / + // renderModeHints is not threaded through to chunks yet; the in-process + // renderer's `resolveRenderWorkerCount` wraps this with that reduction, + // but `PlanJson` doesn't carry the compiled hints needed to call it + // directly. The existing adaptive-retry path reduces workers if + // compositor contention surfaces as CDP timeouts. + const chunkWorkerCount = calculateOptimalWorkers(framesInChunk, undefined, cfg); await runCaptureStage({ fileServer, workDir, @@ -541,11 +548,10 @@ export async function renderChunk( cfg, forceScreenshot: encoder.forceScreenshot, log, - workerCount: 1, - // Pass the pre-warmed session through as `probeSession` so captureStage - // reuses it via `prepareCaptureSessionForReuse` instead of spinning up - // a fresh browser. The stage closes the session in its `finally`, - // so we MUST clear our own reference here to avoid a double-close. + workerCount: chunkWorkerCount, + // The parallel branch closes this session and spins up its own + // worker sessions, wasting the ~3-5s of pre-warmed setup. Worth a + // follow-up to skip pre-warmup when the resolved workerCount > 1. probeSession: session, needsAlpha: plan.dimensions.format !== "mp4", captureAttempts: [], diff --git a/packages/producer/src/services/render/stages/captureStage.ts b/packages/producer/src/services/render/stages/captureStage.ts index c52ad2d7a..176498b70 100644 --- a/packages/producer/src/services/render/stages/captureStage.ts +++ b/packages/producer/src/services/render/stages/captureStage.ts @@ -96,23 +96,15 @@ export interface CaptureStageInput { onProgress?: ProgressCallback; /** * Capture a sub-range `[startFrame, endFrame)` of the composition's - * timeline. Used by distributed `renderChunk` workers to render only - * their assigned chunk. Captured frames are written with file names - * normalized to start at zero (`frame_000000.{ext}`) so the encoder - * doesn't need an `-start_number` override; per-frame TIMES still - * reflect the absolute frame index via `(absIdx * fps.den) / fps.num`, - * keeping the page's virtual clock identical to what an in-process - * render at that frame would see. + * timeline. Used by distributed `renderChunk` to render only its chunk. + * Captured file names are 0-indexed within the range; per-frame TIMES use + * the absolute frame index so the page's virtual clock matches an + * in-process render at that frame. Supported on both the sequential and + * parallel branches; the parallel branch threads `frameRange.startFrame` + * through as `frameRangeStart`. See `WorkerTask.outputFrameOffset`. * - * Only honored on the sequential capture branch (workerCount === 1). - * The parallel branch in this stage targets in-process renders where - * adaptive retry across the whole timeline is the contract, and chunk - * workers fan out at the activity layer instead. Passing `frameRange` - * with `workerCount > 1` throws — the caller should reduce - * `workerCount` to 1. - * - * Default `undefined`: the stage captures `[0, totalFrames)` (the - * in-process contract). + * Default `undefined`: capture `[0, totalFrames)` (in-process contract). + * When set, `endFrame - startFrame` MUST equal `totalFrames`. */ frameRange?: { startFrame: number; endFrame: number }; } @@ -155,12 +147,6 @@ export async function runCaptureStage(input: CaptureStageInput): Promise 1) { - throw new Error( - `[captureStage] frameRange capture requires workerCount === 1 (received workerCount=${workerCount}). ` + - `Distributed chunk workers fan out at the activity layer; reduce workerCount to 1 when passing frameRange.`, - ); - } if (frameRange !== undefined) { if ( !Number.isFinite(frameRange.startFrame) || @@ -173,10 +159,24 @@ export async function runCaptureStage(input: CaptureStageInput): Promise 1) { - // Parallel capture + // Parallel capture. When `frameRange` is set (distributed chunk), pass + // `frameRangeStart` so workers land on absolute composition frame indices + // for time math while file names stay 0-indexed within the chunk range. const attempts = await executeDiskCaptureWithAdaptiveRetry({ serverUrl: fileServer.url, workDir, @@ -188,6 +188,7 @@ export async function runCaptureStage(input: CaptureStageInput): Promise { job.framesRendered = progress.capturedFrames; const frameProgress = progress.capturedFrames / progress.totalFrames; diff --git a/packages/producer/src/services/renderOrchestrator.ts b/packages/producer/src/services/renderOrchestrator.ts index 4b598b7e2..10164939e 100644 --- a/packages/producer/src/services/renderOrchestrator.ts +++ b/packages/producer/src/services/renderOrchestrator.ts @@ -536,17 +536,24 @@ export function buildMissingFrameRetryBatches( maxWorkers: number, workDir: string, attempt: number, + rangeStart: number = 0, ): WorkerTask[][] { const workersPerBatch = Math.max(1, Math.floor(maxWorkers)); const batches: WorkerTask[][] = []; + // `ranges` are 0-indexed within the chunk's frame range (or full timeline + // when `rangeStart === 0`); translate to absolute composition indices so + // `WorkerTask`'s per-frame time math lands on the page's actual virtual + // clock, and propagate `outputFrameOffset` so the retry captures back at + // the same local file name `findMissingFrameRanges` was looking for. for (let i = 0; i < ranges.length; i += workersPerBatch) { const batchIndex = batches.length; const batch = ranges.slice(i, i + workersPerBatch).map((range, workerId) => ({ workerId, - startFrame: range.startFrame, - endFrame: range.endFrame, + startFrame: rangeStart + range.startFrame, + endFrame: rangeStart + range.endFrame, outputDir: join(workDir, `retry-${attempt}-batch-${batchIndex}-worker-${workerId}`), + outputFrameOffset: rangeStart, })); batches.push(batch); } @@ -605,11 +612,18 @@ export async function executeDiskCaptureWithAdaptiveRetry(options: { onProgress?: (progress: ParallelProgress) => void; cfg: EngineConfig; log: ProducerLogger; + /** + * Forwarded to each `WorkerTask`'s `outputFrameOffset` and to the + * `buildMissingFrameRetryBatches` translation. Default 0 (in-process + * contract: `[0, totalFrames)`). See `WorkerTask.outputFrameOffset`. + */ + frameRangeStart?: number; }): Promise { const attempts: CaptureAttemptSummary[] = []; let currentWorkers = options.initialWorkerCount; let missingRanges: FrameRange[] | null = null; let attempt = 0; + const rangeStart = options.frameRangeStart ?? 0; while (true) { const frameCount = missingRanges ? countFrameRanges(missingRanges) : options.totalFrames; @@ -622,8 +636,14 @@ export async function executeDiskCaptureWithAdaptiveRetry(options: { const attemptWorkDir = join(options.workDir, `capture-attempt-${attempt}`); const batches = missingRanges - ? buildMissingFrameRetryBatches(missingRanges, currentWorkers, attemptWorkDir, attempt) - : [distributeFrames(options.totalFrames, currentWorkers, attemptWorkDir)]; + ? buildMissingFrameRetryBatches( + missingRanges, + currentWorkers, + attemptWorkDir, + attempt, + rangeStart, + ) + : [distributeFrames(options.totalFrames, currentWorkers, attemptWorkDir, rangeStart)]; try { for (const tasks of batches) {