perf(aws-lambda): render each chunk with as many capture browsers as the function can host (#6026)

packages/aws-lambda/src/handler.test.ts

@@ -19,7 +19,7 @@ import { afterEach, beforeEach, describe, expect, it, mock, spyOn } from "bun:te
 import { spawnSync } from "node:child_process";
 import { createHash } from "node:crypto";
 import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
-import { tmpdir } from "node:os";
+import { availableParallelism, tmpdir } from "node:os";
 import { dirname, join } from "node:path";
 import { LambdaClient } from "@aws-sdk/client-lambda";
 import { S3Client } from "@aws-sdk/client-s3";
@@ -683,6 +683,89 @@ describe("handler dispatch", () => {
     expect(renderChunkMock).toHaveBeenCalledTimes(1);
   });
 
+  it("gives a chunk as many capture browsers as the function's memory can host", async () => {
+    const tmpRoot = makeTmpRoot();
+    const s3 = new FakeS3Client();
+    s3.objects.set("s3://bucket/plan.tar.gz", await makeMinimalPlanTar());
+    const renderChunkMock = mock(
+      async (
+        _planDir: string,
+        _chunkIndex: number,
+        outputChunkPath: string,
+      ): Promise<ChunkResult> => {
+        writeFileSync(outputChunkPath, Buffer.from("FAKE-MP4-CHUNK"));
+        return {
+          outputPath: outputChunkPath,
+          outputKind: "file",
+          framesEncoded: 121,
+          sha256: "0".repeat(64),
+          durationMs: 1,
+          planHashMs: 1,
+          sessionBootMs: 1,
+          captureStageMs: 1,
+          encodeStageMs: 1,
+          workers: 5,
+          captureMode: "beginframe",
+          perfPath: outputChunkPath + ".perf.json",
+        };
+      },
+    );
+    const notCalled = mock(async () => {
+      throw new Error("should not be called");
+    });
+    const renderOnce = () =>
+      handler(
+        {
+          Action: "renderChunk",
+          PlanProtocol: "v1",
+          PlanS3Uri: "s3://bucket/plan.tar.gz",
+          PlanHash: "fakehash",
+          ChunkIndex: 0,
+          ChunkOutputS3Prefix: "s3://bucket/renders/abc/",
+          Format: "mp4",
+        },
+        {
+          s3: s3 as unknown as import("@aws-sdk/client-s3").S3Client,
+          primitives: {
+            plan: notCalled as unknown as typeof import("@hyperframes/producer/distributed").plan,
+            renderChunk:
+              renderChunkMock as unknown as typeof import("@hyperframes/producer/distributed").renderChunk,
+            assemble:
+              notCalled as unknown as typeof import("@hyperframes/producer/distributed").assemble,
+          },
+          tmpRoot,
+          skipChromeResolution: true,
+        },
+      );
+    const saved = process.env.AWS_LAMBDA_FUNCTION_MEMORY_SIZE;
+    const logged: string[] = [];
+    const logSpy = spyOn(console, "log").mockImplementation((line: string) => {
+      logged.push(line);
+    });
+    try {
+      delete process.env.AWS_LAMBDA_FUNCTION_MEMORY_SIZE;
+      await renderOnce();
+      process.env.AWS_LAMBDA_FUNCTION_MEMORY_SIZE = "10240";
+      await renderOnce();
+    } finally {
+      if (saved === undefined) delete process.env.AWS_LAMBDA_FUNCTION_MEMORY_SIZE;
+      else process.env.AWS_LAMBDA_FUNCTION_MEMORY_SIZE = saved;
+      logSpy.mockRestore();
+    }
+    const options = renderChunkMock.mock.calls.map((call) => (call as unknown[])[3]);
+    expect(options).toEqual([
+      { maxCaptureWorkers: undefined },
+      { maxCaptureWorkers: Math.max(1, Math.min(5, availableParallelism() - 1)) },
+    ]);
+    const captures = logged
+      .filter((line) => line.includes('"chunk_capture"'))
+      .map((line) => JSON.parse(line));
+    expect(captures).toEqual([
+      { event: "chunk_capture", chunkIndex: 0, workers: 5, captureMode: "beginframe" },
+      { event: "chunk_capture", chunkIndex: 0, workers: 5, captureMode: "beginframe" },
+    ]);
+  });
+
   it("rejects renderChunk when event.PlanHash diverges from plan.json", async () => {
     const tmpRoot = makeTmpRoot();
     const s3 = new FakeS3Client();

packages/aws-lambda/src/handler.ts

@@ -12,7 +12,7 @@
  */
 
 import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, statSync } from "node:fs";
-import { tmpdir } from "node:os";
+import { availableParallelism, tmpdir } from "node:os";
 import { basename, join } from "node:path";
 import { GetAccountSettingsCommand, LambdaClient } from "@aws-sdk/client-lambda";
 import { S3Client } from "@aws-sdk/client-s3";
@@ -594,6 +594,17 @@ async function handleAudio(event: AudioEvent, deps?: HandlerDeps): Promise<Audio
 
 // ── RenderChunk ─────────────────────────────────────────────────────────────
 
+// Twice a 4K capture browser's peak memory, so half the function stays free.
+const LAMBDA_MB_PER_CAPTURE_BROWSER = 2048;
+
+/** Capture browsers this function can host: by its memory, leaving one core to Node and ffmpeg. */
+function lambdaMaxCaptureWorkers(): number | undefined {
+  const memoryMb = Number(process.env.AWS_LAMBDA_FUNCTION_MEMORY_SIZE);
+  if (!(memoryMb > 0)) return undefined;
+  const byMemory = Math.floor(memoryMb / LAMBDA_MB_PER_CAPTURE_BROWSER);
+  return Math.max(1, Math.min(byMemory, availableParallelism() - 1));
+}
+
 async function handleRenderChunk(
   event: RenderChunkEvent,
   deps?: HandlerDeps,
@@ -629,7 +640,15 @@ async function handleRenderChunk(
         : `chunk-${pad(event.ChunkIndex)}${formatExtension(event.Format)}`,
     );
 
-    const result: ChunkResult = await primitive(planDir, event.ChunkIndex, chunkOutputBase);
+    const result: ChunkResult = await primitive(planDir, event.ChunkIndex, chunkOutputBase, {
+      maxCaptureWorkers: lambdaMaxCaptureWorkers(),
+    });
+    logEvent({
+      event: "chunk_capture",
+      chunkIndex: event.ChunkIndex,
+      workers: result.workers,
+      captureMode: result.captureMode,
+    });
 
     const chunkUri = await uploadChunkOutput(
       s3,
@@ -677,7 +696,15 @@ async function handleRenderChunkV2(
         ? `chunk-${pad(event.ChunkIndex)}`
         : `chunk-${pad(event.ChunkIndex)}${formatExtension(event.Format)}`,
     );
-    const result = await primitive(planDir, event.ChunkIndex, chunkOutputBase);
+    const result = await primitive(planDir, event.ChunkIndex, chunkOutputBase, {
+      maxCaptureWorkers: lambdaMaxCaptureWorkers(),
+    });
+    logEvent({
+      event: "chunk_capture",
+      chunkIndex: event.ChunkIndex,
+      workers: result.workers,
+      captureMode: result.captureMode,
+    });
     const chunkUri = await uploadChunkOutput(
       s3,
       result,

packages/producer/src/services/distributed/renderChunk.test.ts

@@ -55,6 +55,7 @@ import {
   PLAN_HASH_MISMATCH,
   renderChunk,
   RenderChunkValidationError,
+  resolveChunkCaptureWorkers,
   resolveLockedVp9CpuUsed,
   resolvePresetForLockedEncoder,
 } from "./renderChunk.js";
@@ -979,6 +980,26 @@ describe("resolvePresetForLockedEncoder", () => {
   });
 });
 
+describe("resolveChunkCaptureWorkers", () => {
+  it("keeps the machine-wide sizing when the runtime sets no budget", () => {
+    expect(resolveChunkCaptureWorkers(121, 3, undefined)).toBe(3);
+  });
+
+  it("raises a chunk to the runtime's budget", () => {
+    expect(resolveChunkCaptureWorkers(121, 3, 5)).toBe(5);
+  });
+
+  it("gives each browser at least 20 frames", () => {
+    expect(resolveChunkCaptureWorkers(121, 3, 8)).toBe(6);
+    expect(resolveChunkCaptureWorkers(99, 3, 5)).toBe(4);
+  });
+
+  it("never lowers the machine-wide sizing", () => {
+    expect(resolveChunkCaptureWorkers(121, 3, 1)).toBe(3);
+    expect(resolveChunkCaptureWorkers(30, 2, 5)).toBe(2);
+  });
+});
+
 describe("resolveLockedVp9CpuUsed", () => {
   it("uses the locked value for new VP9 planDirs", () => {
     expect(resolveLockedVp9CpuUsed({ encoder: "libvpx-vp9-software", vp9CpuUsed: 4 })).toBe(4);

packages/producer/src/services/distributed/renderChunk.ts

@@ -167,7 +167,7 @@ export interface ChunkResult {
   sessionBootMs: number;
   captureStageMs: number;
   encodeStageMs: number;
-  /** Capture workers used for this chunk (`calculateOptimalWorkers` result). */
+  /** Capture workers used for this chunk (`resolveChunkCaptureWorkers` result). */
   workers: number;
   /**
    * Effective engine mode used by every worker, after any browser fallback.
@@ -207,6 +207,25 @@ export interface RenderChunkOptions {
    * whole slice.
    */
   frameRange?: ChunkFrameRange;
+  /**
+   * Capture browsers the runtime can host. Raises the machine-wide sizing up to
+   * this many, one browser per 20 frames; never lowers it.
+   */
+  maxCaptureWorkers?: number;
+}
+
+/** A browser starts in about 3 s, a small share of 20 frames' capture. */
+const MIN_FRAMES_PER_CHUNK_BROWSER = 20;
+
+/** Capture browsers for a chunk: the machine-wide sizing, raised by the runtime's budget. */
+export function resolveChunkCaptureWorkers(
+  framesInChunk: number,
+  machineSized: number,
+  maxCaptureWorkers: number | undefined,
+): number {
+  if (maxCaptureWorkers === undefined) return machineSized;
+  const byFrames = Math.floor(framesInChunk / MIN_FRAMES_PER_CHUNK_BROWSER);
+  return Math.max(machineSized, Math.floor(Math.min(maxCaptureWorkers, byFrames)));
 }
 
 /** The frames a chunk renders: its slice, or a validated sub-range of it. */
@@ -238,6 +257,7 @@ export type ChunkRenderer = (
   planDir: string,
   chunkIndex: number,
   outputChunkPath: string,
+  options?: RenderChunkOptions,
 ) => Promise<ChunkResult>;
 
 interface DistributedCaptureSessionDependencies {
@@ -682,7 +702,11 @@ export async function renderChunk(
     // `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);
+    const chunkWorkerCount = resolveChunkCaptureWorkers(
+      framesInChunk,
+      calculateOptimalWorkers(framesInChunk, undefined, cfg),
+      options?.maxCaptureWorkers,
+    );
 
     // ── Browser + warmup ──
     let session: CaptureSession | null = null;