Skip to content

Commit 3468db6

Browse files
authored
feat(producer): add services/distributed/assemble.ts (#813)
## What Phase 3 of the distributed rendering plan: the third public primitive. Adds `assemble(planDir, chunkPaths, audioPath, outputPath)` and its supporting types at `packages/producer/src/services/distributed/assemble.ts`. See `DISTRIBUTED-RENDERING-PLAN.md` §11 Phase 3. ## Why `plan()` (#808) and `renderChunk()` (#809) produce the planDir and the per-chunk outputs respectively. `assemble()` is what every distributed fan-out workflow runs last: it stitches the chunks into the final deliverable. Without it, the planDir → chunks chain stops at a list of files; nothing produces the user-facing mp4/mov/png-sequence. ## How `assemble()` branches on the planDir's encoder format: **mp4 / mov**: ffmpeg `-f concat -c copy` over the ordered chunk paths. Each chunk's first frame is an IDR keyframe (PR 3.2 set `lockGopForChunkConcat: true`), so concat-copy round-trips losslessly. The concatenated output is then: 1. Passed through `padOrTrimAudioToVideoFrameCount` (PR 2.7 surface) when `audioPath` is non-null, so audio length is exactly `frameCount / fps` rather than the audio mixer's original-duration output. 2. Muxed with the normalized audio via the engine's `muxVideoWithAudio` (same helper the in-process renderer's `assembleStage` uses). 3. Passed through `applyFaststart` so the `moov` atom moves to the file's start. When no audio is present, the concat output skips mux and goes straight to `applyFaststart`. **png-sequence**: chunks are directories of `frame_NNNNNN.png` files numbered locally per chunk. `assemble()` merges them with a continuous global index so chunk 0's `frame_000000.png` lands at `frame_000001.png` in the output, chunk 1's first frame becomes `frame_(N+1)`, etc. When `audioPath` is non-null we copy it alongside as `audio.aac` so callers who need to re-mux later have it. ### Validation Both branches assert `chunkPaths.length === chunks.length` (the value read from `meta/chunks.json`) and that each chunk path exists. A missing or mismatched manifest trips a typed error before any ffmpeg invocation. ### What did NOT change No engine helpers, no Phase 1 stages, no in-process orchestrator. `assemble()` reuses `muxVideoWithAudio` / `applyFaststart` / `runFfmpeg` / `padOrTrimAudioToVideoFrameCount` exactly as they exist — the in-process `runAssembleStage` is intentionally not called because it operates on a `RenderJob` and emits `updateJobStatus` payloads, neither of which the distributed activity has. ## Test plan - [x] Unit tests added — `packages/producer/src/services/distributed/assemble.test.ts`. 5 cases: - Concat-copies two mp4 chunks and applies faststart (ffprobe asserts codec, frame count, atom order). - Muxes audio with `frame-count-derived` duration when `audio.aac` is present (ffprobe asserts audio duration within 50ms of `totalFrames / fps`). - Merges png-sequence chunk directories with continuous global numbering (asserts filenames `frame_000001..frame_000007`). - Rejects mismatched `chunkPaths.length` vs `chunks.json.length`. - Rejects a planDir missing `plan.json`. - [x] `bun test packages/producer/src/services/distributed/` — 18 pass (PRs 3.1 + 3.2 + 3.3). - [x] `bun run --filter @hyperframes/producer typecheck` — clean. - [x] `bunx oxlint` + `bunx oxfmt --check` — clean on changed files. - [ ] Producer Docker regression harness — pending CI. `executeRenderJob` is unchanged; PSNR baselines should hold. The mp4 fixture pre-renders test inputs via raw ffmpeg (`testsrc` filter + closed-GOP libx264) rather than going through the Chrome capture pipeline. This isolates concat-copy + mux + faststart from the renderChunk path that PRs 3.1/3.2 already cover, and avoids the chrome-headless-shell smoke-test gating that PR 3.2 needed. This is PR 3 of a 6-PR Phase 3 stack: - 3.1 — `services/distributed/plan.ts` (#808) - 3.2 — `services/distributed/renderChunk.ts` (#809) - **3.3 (this PR)** — `services/distributed/assemble.ts` - 3.4 — `planDir` size cap (`PLAN_TOO_LARGE`) - 3.5 — distributed format banlist (webm + HDR mp4) - 3.6 — public exports + `@hyperframes/producer/distributed` subpath 🤖 Generated with [Claude Code](https://claude.com/claude-code)
2 parents 602ef87 + b606106 commit 3468db6

2 files changed

Lines changed: 652 additions & 0 deletions

File tree

Lines changed: 356 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,356 @@
1+
/**
2+
* Unit tests for `services/distributed/assemble.ts`.
3+
*
4+
* Contracts:
5+
* - mp4/mov: pre-rendered chunks → assembled output passes ffprobe
6+
* (correct frame count, audio present and exactly `frames / fps`
7+
* long, faststart applied).
8+
* - png-sequence: chunk frame directories merge into one continuous
9+
* numbered sequence (chunk N's `frame_NNNNNN.png` files renumber
10+
* into `outputPath/frame_NNNNNN.png` with a global index).
11+
*
12+
* The mp4 fixture pre-renders chunk inputs via raw ffmpeg (test color
13+
* bars + AAC silence) so we don't need a working Chrome to exercise
14+
* assemble. The capture pipeline is covered by `renderChunk.test.ts`.
15+
*/
16+
17+
import { spawnSync } from "node:child_process";
18+
import { afterAll, beforeAll, describe, expect, it } from "bun:test";
19+
import { existsSync, mkdirSync, mkdtempSync, readdirSync, rmSync, writeFileSync } from "node:fs";
20+
import { tmpdir } from "node:os";
21+
import { join } from "node:path";
22+
import type { ChunkSliceJson } from "../render/stages/freezePlan.js";
23+
import { assemble } from "./assemble.js";
24+
25+
let runRoot: string;
26+
let hasFfmpeg = false;
27+
28+
beforeAll(() => {
29+
runRoot = mkdtempSync(join(tmpdir(), "hf-assemble-test-"));
30+
hasFfmpeg = spawnSync("ffmpeg", ["-version"]).status === 0;
31+
});
32+
33+
afterAll(() => {
34+
rmSync(runRoot, { recursive: true, force: true });
35+
});
36+
37+
/**
38+
* Build a synthetic planDir whose `meta/chunks.json` declares N chunks of
39+
* `framesPerChunk` frames each. Does NOT materialize compiled/, video-frames/,
40+
* audio.aac — assemble only reads `plan.json` + `meta/chunks.json`, and we
41+
* pass chunk paths explicitly. Keeping the dir lean speeds up the test
42+
* loop.
43+
*/
44+
function buildPlanDir(
45+
format: "mp4" | "png-sequence",
46+
chunks: ChunkSliceJson[],
47+
totalFrames: number,
48+
hasAudio: boolean,
49+
): string {
50+
const planDir = mkdtempSync(join(runRoot, `plan-${format}-`));
51+
mkdirSync(join(planDir, "meta"), { recursive: true });
52+
writeFileSync(
53+
join(planDir, "plan.json"),
54+
JSON.stringify({
55+
planHash: "fake",
56+
totalFrames,
57+
hasAudio,
58+
dimensions: { fpsNum: 30, fpsDen: 1, width: 160, height: 120, format },
59+
}),
60+
"utf-8",
61+
);
62+
writeFileSync(join(planDir, "meta", "chunks.json"), JSON.stringify(chunks), "utf-8");
63+
return planDir;
64+
}
65+
66+
/**
67+
* Encode a tiny mp4 chunk via raw ffmpeg with closed-GOP libx264 args
68+
* matching what `renderChunk` produces. Uses ffmpeg's `testsrc` filter
69+
* so the test doesn't depend on any image assets. Each chunk is
70+
* independently concatenable because GOP === frame count and the first
71+
* frame is forced as a keyframe.
72+
*/
73+
function makeMp4Chunk(outputPath: string, frameCount: number): void {
74+
const args = [
75+
"-v",
76+
"error",
77+
"-f",
78+
"lavfi",
79+
"-i",
80+
`testsrc=size=160x120:rate=30:duration=${frameCount / 30}`,
81+
"-c:v",
82+
"libx264",
83+
"-preset",
84+
"ultrafast",
85+
"-g",
86+
String(frameCount),
87+
"-keyint_min",
88+
String(frameCount),
89+
"-sc_threshold",
90+
"0",
91+
"-force_key_frames",
92+
`expr:eq(mod(n,${frameCount}),0)`,
93+
"-bf",
94+
"0",
95+
"-pix_fmt",
96+
"yuv420p",
97+
"-vframes",
98+
String(frameCount),
99+
"-y",
100+
outputPath,
101+
];
102+
const result = spawnSync("ffmpeg", args, { stdio: "pipe" });
103+
if (result.status !== 0) {
104+
throw new Error(`ffmpeg testsrc chunk failed: ${result.stderr.toString().slice(-400)}`);
105+
}
106+
}
107+
108+
/** Generate an AAC audio file of `durationSeconds` of silence. */
109+
function makeAacAudio(outputPath: string, durationSeconds: number): void {
110+
const result = spawnSync("ffmpeg", [
111+
"-v",
112+
"error",
113+
"-f",
114+
"lavfi",
115+
"-i",
116+
`anullsrc=channel_layout=stereo:sample_rate=48000`,
117+
"-t",
118+
String(durationSeconds),
119+
"-c:a",
120+
"aac",
121+
"-b:a",
122+
"128k",
123+
"-y",
124+
outputPath,
125+
]);
126+
if (result.status !== 0) {
127+
throw new Error(`ffmpeg anullsrc failed: ${result.stderr.toString().slice(-400)}`);
128+
}
129+
}
130+
131+
/** Read ffprobe JSON for one stream of `outputPath`. */
132+
function probeStream(
133+
outputPath: string,
134+
streamSelector: "v:0" | "a:0",
135+
): Record<string, unknown> | null {
136+
const result = spawnSync(
137+
"ffprobe",
138+
[
139+
"-v",
140+
"error",
141+
"-select_streams",
142+
streamSelector,
143+
"-show_entries",
144+
"stream=duration,nb_frames,nb_read_packets,codec_name,r_frame_rate",
145+
"-count_packets",
146+
"-of",
147+
"json",
148+
outputPath,
149+
],
150+
{ stdio: "pipe" },
151+
);
152+
if (result.status !== 0) return null;
153+
const parsed = JSON.parse(result.stdout.toString()) as {
154+
streams?: Array<Record<string, unknown>>;
155+
};
156+
return parsed.streams?.[0] ?? null;
157+
}
158+
159+
describe("assemble()", () => {
160+
const TIMEOUT_MS = 30_000;
161+
162+
it(
163+
"concat-copies two mp4 chunks and applies faststart",
164+
async () => {
165+
if (!hasFfmpeg) {
166+
console.warn(
167+
"[assemble.test] skipping mp4 concat test — ffmpeg not available on this host",
168+
);
169+
return;
170+
}
171+
172+
const chunks: ChunkSliceJson[] = [
173+
{ index: 0, startFrame: 0, endFrame: 5 },
174+
{ index: 1, startFrame: 5, endFrame: 10 },
175+
];
176+
const planDir = buildPlanDir("mp4", chunks, 10, false);
177+
178+
const chunkAPath = join(planDir, "chunk-0.mp4");
179+
const chunkBPath = join(planDir, "chunk-1.mp4");
180+
makeMp4Chunk(chunkAPath, 5);
181+
makeMp4Chunk(chunkBPath, 5);
182+
183+
const outputPath = join(planDir, "output.mp4");
184+
const result = await assemble(planDir, [chunkAPath, chunkBPath], null, outputPath);
185+
186+
expect(result.outputPath).toBe(outputPath);
187+
expect(existsSync(outputPath)).toBe(true);
188+
expect(result.fileSize).toBeGreaterThan(0);
189+
expect(result.framesEncoded).toBe(10);
190+
191+
// ── ffprobe: correct frame count + codec ───────────────────────────
192+
const videoStream = probeStream(outputPath, "v:0");
193+
expect(videoStream).toBeDefined();
194+
expect(videoStream?.codec_name).toBe("h264");
195+
const probedFrames = Number(videoStream?.nb_read_packets ?? videoStream?.nb_frames);
196+
expect(probedFrames).toBe(10);
197+
198+
// ── faststart applied ──────────────────────────────────────────────
199+
// Bun.file is async; resolve before asserting.
200+
const buf = await Bun.file(outputPath).arrayBuffer();
201+
const bytes = new Uint8Array(buf);
202+
let cursor = 0;
203+
let moovBeforeMdat = false;
204+
while (cursor + 8 <= bytes.length) {
205+
const size =
206+
(bytes[cursor]! << 24) |
207+
(bytes[cursor + 1]! << 16) |
208+
(bytes[cursor + 2]! << 8) |
209+
bytes[cursor + 3]!;
210+
const fourcc = String.fromCharCode(
211+
bytes[cursor + 4]!,
212+
bytes[cursor + 5]!,
213+
bytes[cursor + 6]!,
214+
bytes[cursor + 7]!,
215+
);
216+
if (fourcc === "moov") {
217+
moovBeforeMdat = true;
218+
break;
219+
}
220+
if (fourcc === "mdat") break;
221+
if (size <= 0) break;
222+
cursor += size;
223+
}
224+
expect(moovBeforeMdat).toBe(true);
225+
},
226+
TIMEOUT_MS,
227+
);
228+
229+
it(
230+
"muxes audio with frame-count-derived duration when audio.aac is present",
231+
async () => {
232+
if (!hasFfmpeg) return;
233+
234+
const chunks: ChunkSliceJson[] = [
235+
{ index: 0, startFrame: 0, endFrame: 6 },
236+
{ index: 1, startFrame: 6, endFrame: 12 },
237+
];
238+
const totalFrames = 12;
239+
const fps = 30;
240+
const planDir = buildPlanDir("mp4", chunks, totalFrames, true);
241+
242+
const chunkAPath = join(planDir, "chunk-0.mp4");
243+
const chunkBPath = join(planDir, "chunk-1.mp4");
244+
const audioPath = join(planDir, "audio.aac");
245+
makeMp4Chunk(chunkAPath, 6);
246+
makeMp4Chunk(chunkBPath, 6);
247+
// Audio is half a second longer than the video — `padOrTrimAudioToVideoFrameCount`
248+
// should trim it down to `totalFrames / fps`.
249+
makeAacAudio(audioPath, totalFrames / fps + 0.5);
250+
251+
const outputPath = join(planDir, "output-audio.mp4");
252+
const result = await assemble(planDir, [chunkAPath, chunkBPath], audioPath, outputPath);
253+
254+
expect(existsSync(outputPath)).toBe(true);
255+
expect(result.framesEncoded).toBe(totalFrames);
256+
257+
const audioStream = probeStream(outputPath, "a:0");
258+
expect(audioStream).toBeDefined();
259+
expect(audioStream?.codec_name).toBe("aac");
260+
// Audio duration should be within ~25ms of `totalFrames / fps` after
261+
// pad/trim. The 25ms tolerance absorbs AAC frame quantization (1024
262+
// samples @ 48kHz = ~21ms).
263+
const audioDuration = Number(audioStream?.duration ?? 0);
264+
const expected = totalFrames / fps;
265+
expect(Math.abs(audioDuration - expected)).toBeLessThan(0.05);
266+
},
267+
TIMEOUT_MS,
268+
);
269+
270+
it(
271+
"merges png-sequence chunk directories with continuous global numbering",
272+
() => {
273+
const chunks: ChunkSliceJson[] = [
274+
{ index: 0, startFrame: 0, endFrame: 3 },
275+
{ index: 1, startFrame: 3, endFrame: 7 },
276+
];
277+
const planDir = buildPlanDir("png-sequence", chunks, 7, false);
278+
279+
// Fabricate two chunk directories with 3 + 4 frames respectively.
280+
// Each chunk uses a 0-indexed naming scheme — `renderChunk` writes
281+
// them this way today.
282+
const chunkADir = join(planDir, "chunk-a");
283+
const chunkBDir = join(planDir, "chunk-b");
284+
mkdirSync(chunkADir, { recursive: true });
285+
mkdirSync(chunkBDir, { recursive: true });
286+
const minimalPngHeader = Buffer.from([
287+
// 8-byte PNG signature followed by an IHDR chunk for a 1×1 RGB image.
288+
0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d, 0x49, 0x48, 0x44,
289+
0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02, 0x00, 0x00, 0x00, 0x90,
290+
0x77, 0x53, 0xde,
291+
]);
292+
for (let i = 0; i < 3; i++) {
293+
// Each frame's bytes differ (suffix index) so the merged sequence's
294+
// ordering assertion has something to bite on.
295+
writeFileSync(
296+
join(chunkADir, `frame_${String(i).padStart(6, "0")}.png`),
297+
Buffer.concat([minimalPngHeader, Buffer.from([0xaa, i])]),
298+
);
299+
}
300+
for (let i = 0; i < 4; i++) {
301+
writeFileSync(
302+
join(chunkBDir, `frame_${String(i).padStart(6, "0")}.png`),
303+
Buffer.concat([minimalPngHeader, Buffer.from([0xbb, i])]),
304+
);
305+
}
306+
307+
const outputPath = join(planDir, "merged");
308+
// No await — png-sequence assemble is synchronous internally.
309+
const promise = assemble(planDir, [chunkADir, chunkBDir], null, outputPath);
310+
return promise.then((result) => {
311+
expect(result.outputPath).toBe(outputPath);
312+
expect(result.framesEncoded).toBe(7);
313+
const merged = readdirSync(outputPath).sort();
314+
expect(merged).toEqual([
315+
"frame_000001.png",
316+
"frame_000002.png",
317+
"frame_000003.png",
318+
"frame_000004.png",
319+
"frame_000005.png",
320+
"frame_000006.png",
321+
"frame_000007.png",
322+
]);
323+
});
324+
},
325+
TIMEOUT_MS,
326+
);
327+
328+
it("rejects when chunkPaths.length does not match chunks.json length", async () => {
329+
const chunks: ChunkSliceJson[] = [
330+
{ index: 0, startFrame: 0, endFrame: 5 },
331+
{ index: 1, startFrame: 5, endFrame: 10 },
332+
];
333+
const planDir = buildPlanDir("mp4", chunks, 10, false);
334+
let caught: unknown;
335+
try {
336+
await assemble(planDir, ["/tmp/nonexistent.mp4"], null, join(planDir, "out.mp4"));
337+
} catch (err) {
338+
caught = err;
339+
}
340+
expect(caught).toBeDefined();
341+
expect((caught as Error).message).toContain("does not match");
342+
});
343+
344+
it("rejects a planDir missing plan.json", async () => {
345+
const emptyDir = join(runRoot, "empty");
346+
mkdirSync(emptyDir, { recursive: true });
347+
let caught: unknown;
348+
try {
349+
await assemble(emptyDir, [], null, join(emptyDir, "out.mp4"));
350+
} catch (err) {
351+
caught = err;
352+
}
353+
expect(caught).toBeDefined();
354+
expect((caught as Error).message).toContain("plan.json");
355+
});
356+
});

0 commit comments

Comments
 (0)