feat(worker): parallel per-model rating batches with resume via state file

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

simon-pierre boucher committed 6 days ago (Aug 5, 2026) parent fbb11ff

Showing 1 changed file with +32 and −8

modified apps/worker/src/scripts/rate.ts +32 −8

		`@@ -10,8 +10,13 @@ import { ratingJobId } from "../raters/job-id";`
10	10
11	11	// One-shot rating run: submit a Message Batch per model for every task that
12	12	// lacks ratings under the current prompt version, poll to completion, ingest.
	13	+// Models run in PARALLEL (batches are independent server-side queues).
13	14	// Usage: tsx src/scripts/rate.ts [--limit N] [--poll-seconds S]
14		−// Idempotent: already-rated tasks are excluded; re-running resumes cleanly.
	15	+// Idempotent & resumable: already-rated tasks are excluded, and an in-flight
	16	+// batch per model is remembered in a state file so a relaunch resumes polling
	17	+// instead of resubmitting (a resubmit would double the spend).
	18	+
	19	+import { readFileSync, writeFileSync } from "node:fs";
15	20
16	21	function argValue(flag: string): string \| undefined {
17	22	const index = process.argv.indexOf(flag);
		`@@ -20,6 +25,19 @@ function argValue(flag: string): string \| undefined {`
20	25
21	26	const limit = Number(argValue("--limit") ?? "0") \|\| undefined;
22	27	const pollSeconds = Number(argValue("--poll-seconds") ?? "60") \|\| 60;
	28	+const STATE_PATH = process.env.RATE_STATE ?? "/tmp/aix-batch-state.json";
	29	+
	30	+function readState(): Record<string, string> {
	31	+ try {
	32	+ return JSON.parse(readFileSync(STATE_PATH, "utf8")) as Record<string, string>;
	33	+ } catch {
	34	+ return {};
	35	+ }
	36	+}
	37	+
	38	+function writeState(state: Record<string, string>): void {
	39	+ writeFileSync(STATE_PATH, JSON.stringify(state, null, 2));
	40	+}
23	41
24	42	const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
25	43
		`@@ -48,9 +66,15 @@ async function runForModel(model: string): Promise<void> {`
48	66	ratingTasks.map((task) => [ratingJobId(task.taskId, model, RATER_PROMPT_VERSION), task.taskId]),
49	67	);
50	68
51		− console.log(`[${model}] submitting batch of ${ratingTasks.length} tasks…`);
52		− const batchId = await submitRatingBatch(ratingTasks, model);
53		− console.log(`[${model}] batch ${batchId} submitted; polling every ${pollSeconds}s`);
	69	+ let batchId = readState()[model];
	70	+ if (batchId) {
	71	+ console.log(`[${model}] resuming in-flight batch ${batchId}; polling every ${pollSeconds}s`);
	72	+ } else {
	73	+ console.log(`[${model}] submitting batch of ${ratingTasks.length} tasks…`);
	74	+ batchId = await submitRatingBatch(ratingTasks, model);
	75	+ writeState({ ...readState(), [model]: batchId });
	76	+ console.log(`[${model}] batch ${batchId} submitted; polling every ${pollSeconds}s`);
	77	+ }
54	78
55	79	for (;;) {
56	80	await sleep(pollSeconds * 1000);
		`@@ -59,6 +83,9 @@ async function runForModel(model: string): Promise<void> {`
59	83	}
60	84
61	85	const { ingested, failed } = await ingestBatchResults(batchId, model, jobIdToTaskId);
	86	+ const state = readState();
	87	+ delete state[model];
	88	+ writeState(state);
62	89	console.log(`[${model}] batch ${batchId} done: ${ingested} dimension ratings ingested, ${failed} requests failed`);
63	90	}
64	91
		`@@ -69,10 +96,7 @@ async function main(): Promise<void> {`
69	96	console.log(
70	97	`rating run — prompt ${RATER_PROMPT_VERSION}, models: ${RATER_MODELS.join(", ")}${limit ? `, limit ${limit}` : " (all unrated tasks)"}`,
71	98	);
72		− // Sequential per model keeps memory/log output simple; batches themselves are parallel server-side.
73		− for (const model of RATER_MODELS) {
74		− await runForModel(model);
75		− }
	99	+ await Promise.all(RATER_MODELS.map((model) => runForModel(model)));
76	100	console.log("rating run complete — next: pnpm score:recompute");
77	101	}
78	102
79	103