perf(aws-lambda): warm the audio function while the plan runs (#6148)

* perf(aws-lambda): warm the audio function while the plan runs

The audio task runs on the x86 render function, which is cold when the plan ran on the arm64 video function. A fire-and-forget invoke at the start of a render now wakes it, so the audio no longer outlasts the chunks.

* perf(aws-lambda): warm the render function only on a stack with a video function

An x86-only stack has no cold function to warm, so it takes the old path with no extra step. The CDK stack has no video function and is unchanged.

* perf(aws-lambda): wrap a template comment

* style(aws-lambda): format the warm-up template test

Author: miguel-heygen

examples/aws-lambda/template.yaml

@@ -254,6 +254,11 @@ Resources:
           SelectPlanProtocol:
             Type: Choice
             Choices:
+              # Only a stack with the arm64 video function has a cold x86 function to warm.
+              # Without one the compared value can never match.
+              - Variable: $.PlanProtocol
+                StringEquals: !If [HasVideoFunction, v2, v2-without-video-function]
+                Next: WarmRender
               - Variable: $.PlanProtocol
                 StringEquals: v2
                 Next: PlanV2
@@ -265,6 +270,23 @@ Resources:
                 Next: UnsupportedPlanProtocol
             Default: PlanV2
 
+          # The audio task runs on this function once the plan is done; starting its container now hides the cold start.
+          WarmRender:
+            Type: Task
+            Resource: arn:aws:states:::lambda:invoke
+            Parameters:
+              FunctionName: !GetAtt RenderFunction.Arn
+              InvocationType: Event
+              Payload:
+                Action: warm
+            ResultPath: null
+            Catch:
+              - ErrorEquals:
+                  - States.ALL
+                ResultPath: null
+                Next: PlanV2
+            Next: PlanV2
+
           UnsupportedPlanProtocol:
             Type: Fail
             Error: PLAN_PROTOCOL_UNSUPPORTED

packages/aws-lambda/src/cdk/HyperframesRenderStack.snapshot.test.ts

@@ -161,6 +161,24 @@ describe("HyperframesRenderStack — snapshot", () => {
     expect(actualStates.sort()).toEqual([...EXPECTED_STATE_NAMES].sort());
   });
 
+  it("warms the render function before a v2 plan only on a stack with a video function, and a failed warm-up is caught", () => {
+    const states = readSamDefinition().States;
+    const warm = requireRecord(states.WarmRender, "WarmRender state");
+    expect(sentPayload(warm)).toEqual({ Action: "warm" });
+    expect(requireRecordProperty(warm, "Parameters", "parameters").InvocationType).toBe("Event");
+    expect(warm.Next).toBe("PlanV2");
+    expect(JSON.stringify(warm.Catch)).toContain('"Next":"PlanV2"');
+    const choices = requireRecord(states.SelectPlanProtocol, "SelectPlanProtocol")
+      .Choices as Record<string, unknown>[];
+    const toWarm = choices.filter((choice) => choice.Next === "WarmRender");
+    expect(toWarm).toHaveLength(1);
+    expect(toWarm[0]).toEqual({
+      Variable: "$.PlanProtocol",
+      StringEquals: ["HasVideoFunction", "v2", "v2-without-video-function"],
+      Next: "WarmRender",
+    });
+  });
+
   it("defaults omitted plan protocol to v2 and preserves the explicit v1 branch", () => {
     for (const definition of [SYNTHED.definition, readSamDefinition()]) {
       const selection = requireRecord(
@@ -334,7 +352,7 @@ describe("HyperframesRenderStack — snapshot", () => {
     // CDK drops null-valued payload fields when it renders the definition (#5259), so compare
     // what each deployed task sends rather than the construct source.
     const cdkTasks = lambdaTasks(SYNTHED.definition.States);
-    const samTasks = lambdaTasks(readSamDefinition().States);
+    const { WarmRender: _warm, ...samTasks } = lambdaTasks(readSamDefinition().States);
     expect(Object.keys(cdkTasks).sort()).toEqual(Object.keys(samTasks).sort());
     expect(Object.keys(cdkTasks).length).toBeGreaterThanOrEqual(6);
     for (const [taskName, cdkTask] of Object.entries(cdkTasks)) {
@@ -346,14 +364,14 @@ describe("HyperframesRenderStack — snapshot", () => {
   });
 
   it("runs plan and renderChunk on the SAM template's optional arm64 video function", () => {
-    // A chunk refuses a plan made with another ffmpeg build; audio and assemble stay on x86.
+    // A chunk refuses a plan made with another ffmpeg build; warm-up, audio and assemble stay on x86.
     const targets = Object.values(lambdaTasks(readSamDefinition().States)).map((task) => ({
       action: sentPayload(task).Action,
       functionName: requireRecordProperty(requireRecord(task, "task"), "Parameters", "task params")
         .FunctionName,
     }));
     expect(new Set(targets.map((target) => target.action))).toEqual(
-      new Set(["plan", "renderChunk", "audio", "assemble"]),
+      new Set(["warm", "plan", "renderChunk", "audio", "assemble"]),
     );
     for (const { action, functionName } of targets) {
       expect({ action, functionName }).toEqual({
@@ -491,6 +509,7 @@ describe("HyperframesRenderStack — snapshot", () => {
   it("gives Lambda throttles a minutes-long backoff ahead of the catch-all, in CDK and SAM", () => {
     for (const definition of [SYNTHED.definition, readSamDefinition()]) {
       for (const [taskName, task] of Object.entries(lambdaTasks(definition.States))) {
+        if (taskName === "WarmRender") continue;
         expect({ taskName, ...throttleRetry(task) }).toEqual({
           taskName,
           beforeCatchAll: true,

packages/aws-lambda/src/events.ts

@@ -25,6 +25,9 @@ export type { SerializableDistributedRenderConfig } from "@hyperframes/producer/
 
 /** Discriminator for the three roles the one Lambda image fulfills. */
 export type LambdaAction = "plan" | "audio" | "renderChunk" | "assemble";
+export interface WarmEvent {
+  Action: "warm";
+}
 /** Transport protocol selected for one complete distributed render. */
 export type LambdaPlanProtocol = "v1" | "v2";
 
@@ -41,6 +44,7 @@ export type LambdaEvent =
   | AudioEvent
   | RenderChunkEvent
   | AssembleEvent
+  | WarmEvent
   | { Payload: LambdaEvent }
   | { Input: LambdaEvent };
 
@@ -231,7 +235,12 @@ export interface AssembleLambdaResult {
   DurationMs: number;
 }
 
+export interface WarmLambdaResult {
+  Action: "warm";
+}
+
 export type LambdaResult =
+  | WarmLambdaResult
   | PlanLambdaResult
   | AudioLambdaResult
   | RenderChunkLambdaResult

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

@@ -255,6 +255,19 @@ describe("unwrapEvent", () => {
 });
 
 describe("handler dispatch", () => {
+  it("answers a warm event without touching S3 or the project", async () => {
+    const s3 = new FakeS3Client();
+    expect(
+      await handleEvent(
+        { Action: "warm" },
+        { s3: s3 as unknown as import("@aws-sdk/client-s3").S3Client, skipChromeResolution: true },
+      ),
+    ).toEqual({
+      Action: "warm",
+    });
+    expect(s3.ops).toEqual([]);
+  });
+
   it("preserves explicit v1 plan compatibility", async () => {
     const tmpRoot = makeTmpRoot();
     const s3 = new FakeS3Client();

packages/aws-lambda/src/handler.ts

@@ -121,6 +121,11 @@ export interface HandlerDeps {
  * shape; unwrap until we hit a discriminated event.
  */
 export async function handler(event: LambdaEvent, deps?: HandlerDeps): Promise<LambdaResult> {
+  if ((event as { Action?: unknown }).Action === "warm") {
+    primeRuntimeEnv();
+    await ensureChromePath(deps);
+    return { Action: "warm" };
+  }
   const unwrapped = unwrapEvent(event);
   validatePlanProtocolShape(unwrapped);
   validateEventS3Uris(unwrapped);