perf(aws-lambda): small-chunk cloud renders stitch their output in parallel

examples/aws-lambda/template.yaml

@@ -170,9 +170,21 @@ Resources:
           TMPDIR: /tmp
           HYPERFRAMES_RENDER_BUCKET: !Ref RenderBucket
           HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY: !Ref ReservedConcurrency
+          # Assemble spreads the parts S3 cannot copy across the warm function that ran the chunks.
+          HYPERFRAMES_STITCH_FUNCTION: !If
+            - HasVideoFunction
+            - !Sub "${ProjectName}-video"
+            - !Sub "${ProjectName}-render"
       Policies:
         - S3CrudPolicy:
             BucketName: !Ref RenderBucket
+        - Statement:
+            - Effect: Allow
+              Action: lambda:InvokeFunction
+              Resource: !If
+                - HasVideoFunction
+                - !Sub "arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${ProjectName}-video"
+                - !Sub "arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${ProjectName}-render"
         - Statement:
             - Effect: Allow
               Action: s3:AbortMultipartUpload

examples/aws-lambda/video-function.yaml

@@ -62,6 +62,10 @@ Resources:
               - Effect: Allow
                 Action: lambda:GetAccountSettings
                 Resource: "*"
+              # Stitch helpers, should this function ever run an assemble.
+              - Effect: Allow
+                Action: lambda:InvokeFunction
+                Resource: !Sub "arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${ProjectName}-video"
       Tags:
         - Key: Project
           Value: !Ref ProjectName
@@ -95,6 +99,7 @@ Resources:
           TMPDIR: /tmp
           HYPERFRAMES_RENDER_BUCKET: !Ref RenderBucket
           HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY: !Ref ReservedConcurrency
+          HYPERFRAMES_STITCH_FUNCTION: !Sub "${ProjectName}-video"
       Tags:
         - Key: Project
           Value: !Ref ProjectName

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

@@ -493,6 +493,40 @@ describe("HyperframesRenderStack — snapshot", () => {
     expect(samActions).toContain("lambda:GetAccountSettings");
   }, 30000);
 
+  it("lets assemble invoke exactly the stitch helper it names, in CDK and SAM", () => {
+    const cdkFunction = Object.values(
+      SYNTHED.template.findResources("AWS::Lambda::Function"),
+    )[0] as {
+      Properties: { Environment: { Variables: Record<string, unknown> } };
+    };
+    expect(cdkFunction.Properties.Environment.Variables.HYPERFRAMES_STITCH_FUNCTION).toBe(
+      "hyperframes-render",
+    );
+    const cdkInvoke = Object.values(SYNTHED.template.findResources("AWS::IAM::Policy"))
+      .flatMap(policyStatements)
+      .filter((statement) => statement.Action === "lambda:InvokeFunction");
+    expect(JSON.stringify(cdkInvoke.map((s) => s.Resource))).toContain(
+      ":function:hyperframes-render",
+    );
+
+    const samFunction = requireRecordProperty(
+      requireRecordProperty(readSamResources(), "RenderFunction", "SAM RenderFunction"),
+      "Properties",
+      "SAM function properties",
+    ) as { Environment: { Variables: Record<string, unknown> }; Policies: unknown[] };
+    // The YAML parser drops the !If and !Sub tags and keeps [condition, then, else].
+    const [, withVideo, without] = samFunction.Environment.Variables
+      .HYPERFRAMES_STITCH_FUNCTION as string[];
+    expect([withVideo, without]).toEqual(["${ProjectName}-video", "${ProjectName}-render"]);
+    const samInvoke = samFunction.Policies.flatMap(
+      (policy) => (policy as { Statement?: Record<string, unknown>[] }).Statement ?? [],
+    ).filter((statement) => statement.Action === "lambda:InvokeFunction");
+    expect(samInvoke).toHaveLength(1);
+    const [, grantWithVideo, grantWithout] = samInvoke[0]!.Resource as string[];
+    expect(grantWithVideo!.endsWith(`:function:${withVideo}`)).toBe(true);
+    expect(grantWithout!.endsWith(`:function:${without}`)).toBe(true);
+  }, 30000);
+
   it("fails at once when a chunk cannot fit the worker timeout, in CDK and SAM", () => {
     for (const definition of [SYNTHED.definition, readSamDefinition()]) {
       for (const name of ["Plan", "PlanV2"]) {

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

@@ -21,7 +21,7 @@
 
 import { dirname, resolve } from "node:path";
 import { fileURLToPath } from "node:url";
-import { Duration, RemovalPolicy, Size } from "aws-cdk-lib";
+import { ArnFormat, Duration, RemovalPolicy, Size, Stack } from "aws-cdk-lib";
 import * as cloudwatch from "aws-cdk-lib/aws-cloudwatch";
 import * as iam from "aws-cdk-lib/aws-iam";
 import * as lambda from "aws-cdk-lib/aws-lambda";
@@ -120,6 +120,7 @@ export class HyperframesRenderStack extends Construct {
         HYPERFRAMES_LAMBDA_CHROME_SOURCE: chromeSource,
         HYPERFRAMES_RENDER_BUCKET: this.bucket.bucketName,
         HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY: String(props.reservedConcurrency ?? -1),
+        HYPERFRAMES_STITCH_FUNCTION: `${projectName}-render`,
       },
     });
 
@@ -131,6 +132,18 @@ export class HyperframesRenderStack extends Construct {
     this.renderFunction.addToRolePolicy(
       new iam.PolicyStatement({ actions: ["lambda:GetAccountSettings"], resources: ["*"] }),
     );
+    const renderFunctionArnWithoutRoleCycle = Stack.of(this).formatArn({
+      service: "lambda",
+      resource: "function",
+      resourceName: `${projectName}-render`,
+      arnFormat: ArnFormat.COLON_RESOURCE_NAME,
+    });
+    this.renderFunction.addToRolePolicy(
+      new iam.PolicyStatement({
+        actions: ["lambda:InvokeFunction"],
+        resources: [renderFunctionArnWithoutRoleCycle],
+      }),
+    );
 
     const stateMachineLogGroup = new logs.LogGroup(this, "RenderStateMachineLogGroup", {
       logGroupName: `/aws/states/${projectName}-render`,

packages/aws-lambda/src/events.ts

@@ -28,6 +28,14 @@ export type LambdaAction = "plan" | "audio" | "renderChunk" | "assemble";
 export interface WarmEvent {
   Action: "warm";
 }
+/** Assemble hands a helper part of the final mp4's multipart upload: each part's S3 ranges, in order. */
+export interface StitchPartsEvent {
+  Action: "stitchParts";
+  Bucket: string;
+  Key: string;
+  UploadId: string;
+  Parts: { partNumber: number; ranges: { uri: string; start: number; end: number }[] }[];
+}
 /** Transport protocol selected for one complete distributed render. */
 export type LambdaPlanProtocol = "v1" | "v2";
 
@@ -45,6 +53,7 @@ export type LambdaEvent =
   | RenderChunkEvent
   | AssembleEvent
   | WarmEvent
+  | StitchPartsEvent
   | { Payload: LambdaEvent }
   | { Input: LambdaEvent };
 
@@ -239,8 +248,14 @@ export interface WarmLambdaResult {
   Action: "warm";
 }
 
+export interface StitchPartsLambdaResult {
+  Action: "stitchParts";
+  ETags: string[];
+}
+
 export type LambdaResult =
   | WarmLambdaResult
+  | StitchPartsLambdaResult
   | PlanLambdaResult
   | AudioLambdaResult
   | RenderChunkLambdaResult

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

@@ -38,8 +38,14 @@ import {
 } from "@hyperframes/producer/distributed";
 import { resolveMaxParallelChunks } from "../../producer/src/services/distributed/plan.js";
 import { recomputePlanHashFromPlanDir } from "../../producer/src/services/render/stages/freezePlan.js";
-import type { AssembleEvent, LambdaEvent, PlanEvent, RenderChunkEvent } from "./events.js";
-import { handler as handleEvent, unwrapEvent } from "./handler.js";
+import type {
+  AssembleEvent,
+  LambdaEvent,
+  PlanEvent,
+  RenderChunkEvent,
+  StitchPartsEvent,
+} from "./events.js";
+import { handler as handleEvent, stitchHelper, unwrapEvent } from "./handler.js";
 import { S3PlanV2ArtifactPublisher } from "./s3PlanV2Publisher.js";
 
 interface FakeS3Op {
@@ -1626,6 +1632,98 @@ describe("handler — S3 URI allowlist (security: F-004)", () => {
     await expect(handler(event, deps)).rejects.toMatchObject({ name: "S3_URI_NOT_ALLOWED" });
     expect(s3.ops).toHaveLength(0);
   });
+
+  it("rejects a stitchParts range outside the allowed bucket", async () => {
+    process.env.HYPERFRAMES_RENDER_BUCKET = "good-bucket";
+    const s3 = new FakeS3Client();
+    const event: StitchPartsEvent = {
+      Action: "stitchParts",
+      Bucket: "good-bucket",
+      Key: "renders/abc/output.mp4",
+      UploadId: "upload-1",
+      Parts: [{ partNumber: 2, ranges: [{ uri: "s3://evil-bucket/secret", start: 0, end: 4 }] }],
+    };
+    await expect(
+      handler(event, { s3: s3 as unknown as S3Client, skipChromeResolution: true }),
+    ).rejects.toMatchObject({ name: "S3_URI_NOT_ALLOWED" });
+    expect(s3.ops).toHaveLength(0);
+  });
+});
+
+describe("handler — stitch helpers", () => {
+  const event = (uri: string): StitchPartsEvent => ({
+    Action: "stitchParts",
+    Bucket: "bucket",
+    Key: "renders/abc/output.mp4",
+    UploadId: "upload-1",
+    Parts: [
+      { partNumber: 3, ranges: [{ uri, start: 1, end: 4 }] },
+      {
+        partNumber: 4,
+        ranges: [
+          { uri, start: 0, end: 2 },
+          { uri, start: 4, end: 6 },
+        ],
+      },
+    ],
+  });
+
+  it("builds the parts a stitchParts call names and returns their ETags in order", async () => {
+    const s3 = new FakeS3Client();
+    s3.objects.set("s3://bucket/chunks/0.mp4", Buffer.from("abcdef"));
+    const result = await handler(event("s3://bucket/chunks/0.mp4"), {
+      s3: s3 as unknown as S3Client,
+      skipChromeResolution: true,
+    });
+    expect(result).toEqual({ Action: "stitchParts", ETags: ['"part-3"', '"part-4"'] });
+    expect([s3.parts.get(3)?.toString(), s3.parts.get(4)?.toString()]).toEqual(["bcd", "abef"]);
+  });
+
+  it("the stitch helper invokes the function the stack names and fails with its error", async () => {
+    const prev = process.env.HYPERFRAMES_STITCH_FUNCTION;
+    const prevReserved = process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY;
+    delete process.env.HYPERFRAMES_STITCH_FUNCTION;
+    expect(stitchHelper()).toBeUndefined();
+    process.env.HYPERFRAMES_STITCH_FUNCTION = "proj-video";
+    try {
+      // A reservation of 8 leaves 7 slots once this assemble holds one.
+      process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY = "8";
+      expect(stitchHelper()!.maxInvocations).toBe(7);
+      process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY = "-1";
+      expect(stitchHelper()!.maxInvocations).toBe(Infinity);
+      const s3 = new FakeS3Client();
+      s3.objects.set("s3://bucket/chunks/0.mp4", Buffer.from("abcdef"));
+      const invoked: string[] = [];
+      const lambda = {
+        send: async (command: { input: { FunctionName: string; Payload: string } }) => {
+          invoked.push(command.input.FunctionName);
+          const sent = JSON.parse(command.input.Payload) as StitchPartsEvent;
+          if (sent.Parts[0]!.ranges[0]!.uri.endsWith("missing")) {
+            const payload = JSON.stringify({ errorMessage: "NoSuchKey", errorType: "NoSuchKey" });
+            return { FunctionError: "Unhandled", Payload: new TextEncoder().encode(payload) };
+          }
+          const out = await handler(sent, { s3: s3 as unknown as S3Client });
+          return { Payload: new TextEncoder().encode(JSON.stringify(out)) };
+        },
+      } as unknown as LambdaClient;
+      const { send } = stitchHelper(lambda)!;
+      const target = { Bucket: "bucket", Key: "renders/abc/output.mp4", UploadId: "upload-1" };
+      expect(await send(target, event("s3://bucket/chunks/0.mp4").Parts)).toEqual([
+        '"part-3"',
+        '"part-4"',
+      ]);
+      await expect(send(target, event("s3://bucket/missing").Parts)).rejects.toMatchObject({
+        name: "NoSuchKey",
+        message: expect.stringContaining("stitchParts on proj-video failed: NoSuchKey"),
+      });
+      expect(invoked).toEqual(["proj-video", "proj-video"]);
+    } finally {
+      if (prev === undefined) delete process.env.HYPERFRAMES_STITCH_FUNCTION;
+      else process.env.HYPERFRAMES_STITCH_FUNCTION = prev;
+      if (prevReserved === undefined) delete process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY;
+      else process.env.HYPERFRAMES_LAMBDA_RESERVED_CONCURRENCY = prevReserved;
+    }
+  });
 });
 
 // ── helpers ─────────────────────────────────────────────────────────────────