perf(distributed): parallelize chunk capture across multiple workers (#906)

* 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

Author: jrusso1020

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_<i - outputFrameOffset>.{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++;
 

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: [],

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<Capture
   const captureCfg: EngineConfig =
     cfg.forceScreenshot === forceScreenshot ? cfg : { ...cfg, forceScreenshot };
 
-  if (frameRange !== undefined && workerCount > 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<Capture
           `Expected non-negative startFrame strictly less than endFrame.`,
       );
     }
+    // The parallel branch passes `totalFrames` to executeDiskCaptureWithAdaptiveRetry
+    // (which drives `distributeFrames` partitioning and `findMissingFrameRanges`
+    // completion checks) AND `frameRangeStart` separately. They must describe the
+    // same window: callers passing `totalFrames=100, frameRange={50, 200}` would
+    // get a silently wrong distribution.
+    const rangeFrames = frameRange.endFrame - frameRange.startFrame;
+    if (rangeFrames !== totalFrames) {
+      throw new Error(
+        `[captureStage] frameRange size (${rangeFrames}) must equal totalFrames (${totalFrames}). ` +
+          `Received frameRange=${JSON.stringify(frameRange)}.`,
+      );
+    }
   }
 
   if (workerCount > 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<Capture
       captureOptions: buildCaptureOptions(),
       createBeforeCaptureHook: createRenderVideoFrameInjector,
       abortSignal,
+      frameRangeStart: frameRange?.startFrame,
       onProgress: (progress) => {
         job.framesRendered = progress.capturedFrames;
         const frameProgress = progress.capturedFrames / progress.totalFrames;

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<CaptureAttemptSummary[]> {
   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) {