diff --git a/.changeset/workflow-eval-submissions.md b/.changeset/workflow-eval-submissions.md new file mode 100644 index 000000000..9935e6629 --- /dev/null +++ b/.changeset/workflow-eval-submissions.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Turn batch evals API into API for evals deferred AI provider calls diff --git a/e2e/config/pr-comment-scenarios.json b/e2e/config/pr-comment-scenarios.json index adead5482..4267601f1 100644 --- a/e2e/config/pr-comment-scenarios.json +++ b/e2e/config/pr-comment-scenarios.json @@ -684,12 +684,12 @@ ] }, { - "scenarioDirName": "durable-eval-webhook", - "label": "Durable Eval Webhook", - "metadataScenario": "durable-eval-webhook", + "scenarioDirName": "workflow-eval-webhook", + "label": "Workflow Eval Webhook", + "metadataScenario": "workflow-eval-webhook", "evals": [ { - "experimentNameTemplate": "durable-eval-webhook-{testRunId}", + "experimentNameTemplate": "workflow-eval-webhook-{testRunId}", "label": "webhook" } ] diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts deleted file mode 100644 index 53b861efd..000000000 --- a/e2e/scenarios/durable-eval-webhook/scenario.ts +++ /dev/null @@ -1,132 +0,0 @@ -import { - BatchScorer, - BatchTask, - defineDurableEval, - DurableEvalMemoryStore, -} from "braintrust"; -import { - getTestRunId, - runMain, - scopedName, -} from "../../helpers/scenario-runtime"; - -async function main() { - const testRunId = getTestRunId(); - const scenario = "durable-eval-webhook"; - const store = new DurableEvalMemoryStore(); - const jobs = new Map(); - const webhookCompletion = { - mode: "webhook" as const, - getExternalId: (submissionData: { id: string }) => submissionData.id, - }; - const task = new BatchTask< - number, - number, - number, - { testRunId: string; kind: string }, - Record - >({ - batchSize: 2, - async submit(items) { - const id = `task-${jobs.size + 1}`; - jobs.set(id, items); - return { id }; - }, - completion: webhookCompletion, - async collect(submissionData) { - const items = (jobs.get(submissionData.id) ?? []) as Array<{ - id: string; - input: number; - }>; - return items.map((item) => ({ - id: item.id, - output: item.input * 2, - })); - }, - }); - const scorer = new BatchScorer< - number, - number, - number, - { testRunId: string; kind: string }, - { id: string } - >({ - name: "batch_exact", - batchSize: 2, - async submit(items) { - const id = `score-${jobs.size + 1}`; - jobs.set(id, items); - return { id }; - }, - completion: webhookCompletion, - async collect(submissionData) { - const items = (jobs.get(submissionData.id) ?? []) as Array<{ - id: string; - output: number; - expected: number; - }>; - return items.map((item) => ({ - id: item.id, - score: { - score: item.output === item.expected ? 1 : 0, - metadata: { method: "batch-provider" }, - }, - })); - }, - }); - const definition = defineDurableEval( - scopedName("e2e-durable-eval-webhook-project", testRunId), - { - store, - experimentName: `${scenario}-${testRunId}`, - data: [1, 2, 3].map((input) => ({ - id: `case-${input}`, - input, - expected: input * 2, - metadata: { scenario, testRunId, kind: "webhook" }, - })), - task, - scores: [ - function exact({ output, expected }) { - return { - score: output === expected ? 1 : 0, - metadata: { method: "shared-eval-runtime" }, - }; - }, - scorer, - ], - classifiers: [ - function quality({ output, expected }) { - return { - name: "quality", - id: output === expected ? "pass" : "fail", - label: output === expected ? "Pass" : "Fail", - }; - }, - ], - }, - ); - - const waiting = await definition.start(); - if (waiting.status !== "waiting" || jobs.size !== 2) { - throw new Error("Durable eval did not pause with two webhook batches"); - } - - let completed = false; - const completedJobs = new Set(); - while (completedJobs.size < jobs.size || !completed) { - const externalId = [...jobs.keys()].find((id) => !completedJobs.has(id)); - if (!externalId) throw new Error("Durable eval stopped before completion"); - completedJobs.add(externalId); - const processed = await definition.processBatchResult({ - runId: waiting.runId, - externalId, - }); - completed = processed.status === "completed"; - } - if ([...jobs.keys()].filter((id) => id.startsWith("score-")).length !== 2) { - throw new Error("Batch scorer did not split three cases into two batches"); - } -} - -runMain(main); diff --git a/e2e/scenarios/durable-eval-webhook/scenario.test.ts b/e2e/scenarios/workflow-eval-webhook/scenario.test.ts similarity index 68% rename from e2e/scenarios/durable-eval-webhook/scenario.test.ts rename to e2e/scenarios/workflow-eval-webhook/scenario.test.ts index 4a0b3eae9..a81e7785c 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.test.ts +++ b/e2e/scenarios/workflow-eval-webhook/scenario.test.ts @@ -10,7 +10,7 @@ const scenarioDir = await prepareScenarioDir({ scenarioDir: resolveScenarioDir(import.meta.url), }); -test("durable eval collects task and scorer webhook sub-batches", async () => { +test("workflow eval advances individual task and scorer submissions", async () => { await withScenarioHarness( async ({ events, runScenarioDir, testRunEvents }) => { await runScenarioDir({ scenarioDir }); @@ -30,17 +30,17 @@ test("durable eval collects task and scorer webhook sub-batches", async () => { JSON.stringify(left).localeCompare(JSON.stringify(right)), ), ).toEqual([ - { batch_exact: 1, exact: 1 }, - { batch_exact: 1, exact: 1 }, - { batch_exact: 1, exact: 1 }, + { workflow_exact: 1, exact: 1 }, + { workflow_exact: 1, exact: 1 }, + { workflow_exact: 1, exact: 1 }, + ]); + expect( + webhookSpans.map((event) => event.metadata?.workflow_eval), + ).toEqual([ + expect.objectContaining({ run_id: expect.any(String) }), + expect.objectContaining({ run_id: expect.any(String) }), + expect.objectContaining({ run_id: expect.any(String) }), ]); - expect(webhookSpans.map((event) => event.metadata?.durable_eval)).toEqual( - [ - expect.objectContaining({ run_id: expect.any(String) }), - expect.objectContaining({ run_id: expect.any(String) }), - expect.objectContaining({ run_id: expect.any(String) }), - ], - ); const taskSpans = findAllSpans(events(), "task"); expect(taskSpans).toHaveLength(3); @@ -59,18 +59,16 @@ test("durable eval collects task and scorer webhook sub-batches", async () => { "shared-eval-runtime", ]); - const batchScoreSpans = findAllSpans(events(), "batch_exact"); - expect(batchScoreSpans).toHaveLength(3); - expect(batchScoreSpans.map((event) => event.scores)).toEqual([ - { batch_exact: 1 }, - { batch_exact: 1 }, - { batch_exact: 1 }, - ]); - expect(batchScoreSpans.map((event) => event.metadata?.method)).toEqual([ - "batch-provider", - "batch-provider", - "batch-provider", + const workflowScoreSpans = findAllSpans(events(), "workflow_exact"); + expect(workflowScoreSpans).toHaveLength(3); + expect(workflowScoreSpans.map((event) => event.scores)).toEqual([ + { workflow_exact: 1 }, + { workflow_exact: 1 }, + { workflow_exact: 1 }, ]); + expect(workflowScoreSpans.map((event) => event.metadata?.method)).toEqual( + ["workflow-provider", "workflow-provider", "workflow-provider"], + ); const classifierSpans = findAllSpans(events(), "quality"); expect(classifierSpans).toHaveLength(3); diff --git a/e2e/scenarios/workflow-eval-webhook/scenario.ts b/e2e/scenarios/workflow-eval-webhook/scenario.ts new file mode 100644 index 000000000..2ef67869d --- /dev/null +++ b/e2e/scenarios/workflow-eval-webhook/scenario.ts @@ -0,0 +1,163 @@ +import { + WorkflowScorer, + WorkflowTask, + defineWorkflowEval, + WorkflowEvalMemoryStore, +} from "braintrust"; +import { + getTestRunId, + runMain, + scopedName, +} from "../../helpers/scenario-runtime"; + +async function main() { + const testRunId = getTestRunId(); + const scenario = "workflow-eval-webhook"; + const store = new WorkflowEvalMemoryStore(); + const jobs = new Map< + string, + { input: number; output?: number; expected: number } + >(); + const webhookCompletion = { + mode: "webhook" as const, + getExternalId: (submissionData: { id: string }) => submissionData.id, + }; + const task = new WorkflowTask< + number, + number, + number, + { testRunId: string; kind: string }, + Record, + { id: string } + >({ + async submit(item) { + const id = `task-${jobs.size + 1}`; + jobs.set(id, item); + return { id }; + }, + completion: webhookCompletion, + async collect(submissionData) { + return { output: jobs.get(submissionData.id)!.input * 2 }; + }, + }); + const scorer = new WorkflowScorer< + number, + number, + number, + { testRunId: string; kind: string }, + { id: string } + >({ + name: "workflow_exact", + async submit(item) { + const id = `score-${jobs.size + 1}`; + jobs.set(id, item); + return { id }; + }, + completion: webhookCompletion, + async collect(submissionData) { + const item = jobs.get(submissionData.id)!; + return { + score: { + score: item.output === item.expected ? 1 : 0, + metadata: { method: "workflow-provider" }, + }, + }; + }, + }); + const definition = defineWorkflowEval( + scopedName("e2e-workflow-eval-webhook-project", testRunId), + { + store, + maxConcurrency: 2, + experimentName: `${scenario}-${testRunId}`, + data: [1, 2, 3].map((input) => ({ + id: `case-${input}`, + input, + expected: input * 2, + metadata: { scenario, testRunId, kind: "webhook" }, + })), + task, + scores: [ + function exact({ output, expected }) { + localScoreCount++; + return { + score: output === expected ? 1 : 0, + metadata: { method: "shared-eval-runtime" }, + }; + }, + scorer, + ], + classifiers: [ + function quality({ output, expected }) { + classifierCount++; + return { + name: "quality", + id: output === expected ? "pass" : "fail", + label: output === expected ? "Pass" : "Fail", + }; + }, + ], + }, + ); + + let localScoreCount = 0; + let classifierCount = 0; + const waiting = await definition.start(); + if (waiting.status !== "waiting" || jobs.size !== 3) { + throw new Error( + "Workflow eval did not pause with three webhook submissions", + ); + } + + // Finish the second case, including scoring, while the other tasks wait. + const taskIds = [...jobs.keys()]; + const first = await definition.processSubmissionResult({ + runId: waiting.runId, + externalId: taskIds[1], + }); + if ( + first.status !== "waiting" || + first.pending.webhook !== 3 || + localScoreCount !== 1 || + classifierCount !== 1 + ) { + throw new Error( + "Completed case did not advance its scorers and classifier independently", + ); + } + const firstScoreId = [...jobs.keys()].find((id) => id.startsWith("score-")); + if (!firstScoreId) + throw new Error("Completed task did not submit its workflow scorer"); + const scored = await definition.processSubmissionResult({ + runId: waiting.runId, + externalId: firstScoreId, + }); + if (scored.status !== "waiting" || scored.pending.webhook !== 2) { + throw new Error("Eval completed before the remaining tasks"); + } + // A repeated delivery must not produce another scorer or log duplicate results. + await definition.processSubmissionResult({ + runId: waiting.runId, + externalId: taskIds[1], + }); + const completedJobs = new Set([taskIds[1], firstScoreId]); + for (const externalId of jobs.keys()) { + if (completedJobs.has(externalId)) continue; + await definition.processSubmissionResult({ + runId: waiting.runId, + externalId, + }); + completedJobs.add(externalId); + } + const completed = await definition.status({ runId: waiting.runId }); + if ( + completed.status !== "completed" || + [...jobs.keys()].filter((id) => id.startsWith("score-")).length !== 3 + ) { + throw new Error( + "Workflow eval did not complete three individual workflow scorers", + ); + } +} + +runMain(main); diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts deleted file mode 100644 index f814e638c..000000000 --- a/js/src/durable-eval.test.ts +++ /dev/null @@ -1,737 +0,0 @@ -import { describe, expect, test, vi } from "vitest"; -import { configureNode } from "./node/config"; -import { - BatchScorer, - BatchTask, - defineDurableEval, - DurableEvalMemoryStore, - DurableEvalRedisStore, - type DurableBatchScorerItem, - type DurableBatchTaskItem, - type DurableEvalStore, -} from "./durable-eval"; - -configureNode(); - -describe("durable eval stores", () => { - test("memory store copies values on read and write", async () => { - const store = new DurableEvalMemoryStore(); - const value = new Uint8Array([1, 2, 3]); - - await store.write("run", value); - value[0] = 9; - - const firstRead = await store.read("run"); - expect(firstRead).toEqual(new Uint8Array([1, 2, 3])); - firstRead![1] = 9; - expect(await store.read("run")).toEqual(new Uint8Array([1, 2, 3])); - expect(await store.read("missing")).toBeUndefined(); - - const [first, second] = await Promise.all([ - store.getOrSet("claim", new Uint8Array([1])), - store.getOrSet("claim", new Uint8Array([2])), - ]); - expect([first.created, second.created]).toEqual([true, false]); - expect(first.value).toEqual(new Uint8Array([1])); - expect(second.value).toEqual(new Uint8Array([1])); - }); - - test("redis store uses prefixed string operations", async () => { - const values = new Map(); - const client = { - get: vi.fn(async (key: string) => values.get(key) ?? null), - set: vi.fn(async (key: string, value: string, _options?: unknown) => { - values.set(key, value); - return "OK"; - }), - sendCommand: vi.fn(), - }; - const store = new DurableEvalRedisStore({ - client, - keyPrefix: "evals:", - ttlMs: 1_234, - }); - - await store.write("run", new Uint8Array([0, 255, 1])); - - expect(client.set).toHaveBeenCalledWith("evals:run", "AP8B", { - PX: 1_234, - }); - expect(await store.read("run")).toEqual(new Uint8Array([0, 255, 1])); - expect(client.get).toHaveBeenCalledWith("evals:run"); - expect(await store.read("missing")).toBeUndefined(); - - await expect( - new DurableEvalRedisStore({ - client: { - get: async () => 42, - set: async () => "OK", - }, - }).read("invalid"), - ).rejects.toThrow("expected GET to return a string"); - expect(() => new DurableEvalRedisStore({ client, ttlMs: 0 })).toThrow( - "ttlMs must be a positive integer", - ); - }); - - test("redis store uses node-redis atomic SET options", async () => { - const values = new Map(); - const client = { - get: vi.fn(async (key: string) => values.get(key) ?? null), - set: vi.fn( - async ( - key: string, - value: string, - options?: { PX?: number; NX?: boolean; GET?: boolean }, - ) => { - expect(options).toEqual({ PX: 1_234, NX: true, GET: true }); - const existing = values.get(key) ?? null; - if (existing === null) values.set(key, value); - return existing; - }, - ), - sendCommand: vi.fn(), - }; - const store = new DurableEvalRedisStore({ client, ttlMs: 1_234 }); - - await expect(store.getOrSet("claim", new Uint8Array([1]))).resolves.toEqual( - { value: new Uint8Array([1]), created: true }, - ); - await expect(store.getOrSet("claim", new Uint8Array([2]))).resolves.toEqual( - { value: new Uint8Array([1]), created: false }, - ); - }); - - test("redis store uses ioredis atomic SET arguments", async () => { - const values = new Map(); - const client = { - get: vi.fn(async (key: string) => values.get(key) ?? null), - set: vi.fn(async (key: string, value: string, ...options: unknown[]) => { - expect(options).toEqual(["PX", 1_234, "NX", "GET"]); - const existing = values.get(key) ?? null; - if (existing === null) values.set(key, value); - return existing; - }), - defineCommand: vi.fn(), - }; - const store = new DurableEvalRedisStore({ client, ttlMs: 1_234 }); - - await expect(store.getOrSet("claim", new Uint8Array([1]))).resolves.toEqual( - { value: new Uint8Array([1]), created: true }, - ); - await expect(store.getOrSet("claim", new Uint8Array([2]))).resolves.toEqual( - { value: new Uint8Array([1]), created: false }, - ); - }); - - test("redis store uses Upstash atomic SET options", async () => { - const values = new Map(); - const client = { - get: vi.fn(async (key: string) => values.get(key) ?? null), - set: vi.fn( - async ( - key: string, - value: string, - options?: { px?: number; nx?: boolean; get?: boolean }, - ) => { - expect(options).toEqual({ px: 1_234, nx: true, get: true }); - const existing = values.get(key) ?? null; - if (existing === null) values.set(key, value); - return existing; - }, - ), - createScript: vi.fn(), - }; - const store = new DurableEvalRedisStore({ client, ttlMs: 1_234 }); - - await expect(store.getOrSet("claim", new Uint8Array([1]))).resolves.toEqual( - { value: new Uint8Array([1]), created: true }, - ); - await expect(store.getOrSet("claim", new Uint8Array([2]))).resolves.toEqual( - { value: new Uint8Array([1]), created: false }, - ); - }); -}); - -describe("defineDurableEval", () => { - test("runs ordinary tasks and scorers", async () => { - const task = vi.fn((input: number) => input * 2); - const result = await defineDurableEval("local", { - store: new DurableEvalMemoryStore(), - data: [ - { id: "one", input: 1, expected: 2 }, - { id: "two", input: 2, expected: 4 }, - ], - task, - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }).start({ noSendLogs: true }); - - expect(result).toMatchObject({ - status: "completed", - summary: { scores: { exact: { score: 1 } } }, - }); - expect(task).toHaveBeenCalledTimes(2); - }); - - test("generates a new run id for every start", async () => { - const durable = defineDurableEval("generated-runs", { - store: new DurableEvalMemoryStore(), - data: [{ input: 1 }], - task: (input) => input, - scores: [() => 1], - }); - - const first = await durable.start({ noSendLogs: true }); - const second = await durable.start({ noSendLogs: true }); - - expect(first.runId).not.toBe(second.runId); - }); - - test("merges batch task metadata with case metadata", async () => { - type Metadata = { fromCase?: boolean; fromTask?: boolean }; - let taskItems: DurableBatchTaskItem< - number, - void, - Metadata, - Record - >[] = []; - let scorerMetadata: Metadata | undefined; - const task = new BatchTask< - number, - number, - void, - Metadata, - Record - >({ - async submit(items) { - taskItems = items; - return null; - }, - completion: { - mode: "poll", - async poll() { - return { status: "complete" }; - }, - }, - async collect() { - return taskItems.map((item) => ({ - id: item.id, - output: item.input, - metadata: { fromTask: true }, - })); - }, - }); - const durable = defineDurableEval("merged-batch-metadata", { - store: new DurableEvalMemoryStore(), - data: [{ id: "one", input: 1, metadata: { fromCase: true } }], - task, - scores: [ - ({ metadata }) => { - scorerMetadata = metadata; - return 1; - }, - ], - }); - - const waiting = await durable.start({ noSendLogs: true }); - await durable.poll({ runId: waiting.runId }); - - expect(scorerMetadata).toEqual({ fromCase: true, fromTask: true }); - }); - - test("stores run, case, and batch records separately", async () => { - const values = new Map(); - const store: DurableEvalStore = { - async read(key) { - return values.get(key); - }, - async write(key, value) { - values.set(key, value); - }, - async getOrSet(key, value) { - const existing = values.get(key); - if (existing) return { value: existing, created: false }; - values.set(key, value); - return { value, created: true }; - }, - }; - const task = new BatchTask< - number, - number, - void, - void, - Record, - { id: string } - >({ - async submit() { - return { id: "provider-job" }; - }, - completion: { - mode: "poll", - async poll() { - return { status: "pending" }; - }, - }, - async collect() { - return []; - }, - }); - - await defineDurableEval("normalized-store", { - store, - data: [ - { id: "one", input: 1 }, - { id: "two", input: 2 }, - ], - task, - }).start({ noSendLogs: true }); - - const decoder = new TextDecoder(); - const records = [...values] - .filter(([key]) => !key.includes("/claims/")) - .map( - ([key, value]) => - [ - key, - JSON.parse(decoder.decode(value)) as Record, - ] as const, - ); - const run = records.find( - ([key]) => !key.includes("/cases/") && !key.includes("/batches/"), - )?.[1]; - expect(run).toMatchObject({ - status: "running", - caseIds: ["one:trial:0", "two:trial:0"], - }); - expect(run).not.toHaveProperty("schemaVersion"); - expect(run).not.toHaveProperty("projectName"); - expect(run).not.toHaveProperty("evalName"); - expect(run).not.toHaveProperty("cases"); - expect(run).not.toHaveProperty("batches"); - expect(run).not.toHaveProperty("batchIds"); - expect(records.filter(([key]) => key.includes("/cases/"))).toHaveLength(2); - expect(records.filter(([key]) => key.includes("/batches/"))).toHaveLength( - 1, - ); - expect( - [...values.keys()].filter((key) => key.includes("/claims/")), - ).toHaveLength(1); - }); - - test("polls each existing task and scorer sub-batch once", async () => { - const taskJobs = new Map< - string, - DurableBatchTaskItem>[] - >(); - const scoreJobs = new Map< - string, - DurableBatchScorerItem[] - >(); - const taskPoll = vi.fn(async () => ({ status: "complete" as const })); - const scorePoll = vi.fn(async () => ({ status: "complete" as const })); - - const task = new BatchTask< - number, - number, - number, - void, - Record, - { id: string } - >({ - batchSize: 2, - async submit(items) { - const id = `task-${taskJobs.size + 1}`; - taskJobs.set(id, items); - return { id }; - }, - completion: { - mode: "poll", - poll: taskPoll, - }, - async collect(submissionData) { - return (taskJobs.get(submissionData.id) ?? []).map((item) => ({ - id: item.id, - output: item.input * 2, - })); - }, - }); - const scorer = new BatchScorer< - number, - number, - number, - void, - { id: string } - >({ - name: "exact", - batchSize: 2, - async submit(items) { - const id = `score-${scoreJobs.size + 1}`; - scoreJobs.set(id, items); - return { id }; - }, - completion: { - mode: "poll", - poll: scorePoll, - }, - async collect(submissionData) { - return (scoreJobs.get(submissionData.id) ?? []).map((item) => ({ - id: item.id, - score: item.output === item.expected ? 1 : 0, - })); - }, - }); - - const store = new DurableEvalMemoryStore(); - const durable = defineDurableEval("polling-batches", { - store, - data: [1, 2, 3].map((input) => ({ - id: `case-${input}`, - input, - expected: input * 2, - })), - task, - scores: [scorer], - }); - const waiting = await durable.start({ noSendLogs: true }); - expect(waiting).toMatchObject({ - status: "waiting", - runId: expect.any(String), - pending: { poll: 2, webhook: 0 }, - }); - const options = { runId: waiting.runId }; - expect(taskJobs.size).toBe(2); - expect(scoreJobs.size).toBe(0); - await expect(durable.status(options)).resolves.toEqual({ - status: "waiting", - runId: waiting.runId, - pending: { poll: 2, webhook: 0 }, - }); - expect(taskJobs.size).toBe(2); - expect(taskPoll).not.toHaveBeenCalled(); - - await expect(durable.poll(options)).resolves.toEqual({ - status: "waiting", - runId: waiting.runId, - pending: { poll: 2, webhook: 0 }, - }); - expect(scoreJobs.size).toBe(2); - expect(taskPoll).toHaveBeenCalledTimes(2); - expect(scorePoll).not.toHaveBeenCalled(); - - const result = await durable.poll(options); - - expect(result).toMatchObject({ - status: "completed", - pending: { poll: 0, webhook: 0 }, - summary: { scores: { exact: { score: 1 } } }, - }); - await expect(durable.status(options)).resolves.toMatchObject({ - status: "completed", - pending: { poll: 0, webhook: 0 }, - summary: { scores: { exact: { score: 1 } } }, - }); - expect(scorePoll).toHaveBeenCalledTimes(2); - expect([...taskJobs.values()].map((items) => items.length)).toEqual([2, 1]); - expect([...scoreJobs.values()].map((items) => items.length)).toEqual([ - 2, 1, - ]); - }); - - test("processes task and scorer webhook batches through one method", async () => { - const store = new DurableEvalMemoryStore(); - const taskJobs = new Map< - string, - DurableBatchTaskItem>[] - >(); - const scoreJobs = new Map< - string, - DurableBatchScorerItem[] - >(); - let taskCollectCount = 0; - let releaseTaskCollect!: () => void; - const taskBatchesCollecting = new Promise((resolve) => { - releaseTaskCollect = resolve; - }); - const task = new BatchTask< - number, - number, - number, - void, - Record, - { id: string } - >({ - batchSize: 2, - async submit(items) { - const id = `task-provider-${taskJobs.size + 1}`; - taskJobs.set(id, items); - return { id }; - }, - completion: { - mode: "webhook", - getExternalId: (submissionData) => submissionData.id, - }, - async collect(submissionData) { - taskCollectCount++; - if (taskCollectCount === 2) releaseTaskCollect(); - await taskBatchesCollecting; - return (taskJobs.get(submissionData.id) ?? []).map((item) => ({ - id: item.id, - output: item.input * 2, - })); - }, - }); - const scorer = new BatchScorer< - number, - number, - number, - void, - { id: string } - >({ - name: "exact", - batchSize: 2, - async submit(items) { - const id = `score-provider-${scoreJobs.size + 1}`; - scoreJobs.set(id, items); - return { id }; - }, - completion: { - mode: "webhook", - getExternalId: (submissionData) => submissionData.id, - }, - async collect(submissionData) { - return (scoreJobs.get(submissionData.id) ?? []).map((item) => ({ - id: item.id, - score: item.output === item.expected ? 1 : 0, - })); - }, - }); - const durable = defineDurableEval("webhook-batches", { - store, - data: [1, 2, 3].map((input) => ({ - id: `case-${input}`, - input, - expected: input * 2, - })), - task, - scores: [scorer], - }); - - const waiting = await durable.start({ noSendLogs: true }); - expect(waiting).toMatchObject({ - status: "waiting", - runId: expect.any(String), - pending: { poll: 0, webhook: 2 }, - }); - expect(taskJobs.size).toBe(2); - const runId = waiting.runId; - - const taskIds = [...taskJobs.keys()]; - await expect( - durable.processBatchResult({ - runId: "missing-run", - externalId: taskIds[0], - }), - ).rejects.toThrow("Durable eval run missing-run is missing"); - await Promise.all( - taskIds.map((externalId) => - durable.processBatchResult({ runId, externalId }), - ), - ); - expect(scoreJobs.size).toBe(2); - await expect(durable.status({ runId })).resolves.toMatchObject({ - status: "waiting", - pending: { poll: 0, webhook: 2 }, - }); - - const scoreIds = [...scoreJobs.keys()]; - let result; - for (const externalId of scoreIds) { - result = await durable.processBatchResult({ runId, externalId }); - } - expect(result).toMatchObject({ - status: "completed", - pending: { poll: 0, webhook: 0 }, - summary: { scores: { exact: { score: 1 } } }, - }); - }); - - test("claims downstream work once across concurrent webhook deliveries", async () => { - let taskItems: DurableBatchTaskItem< - number, - number, - void, - Record - >[] = []; - let collectCount = 0; - let releaseCollect!: () => void; - const bothCollecting = new Promise((resolve) => { - releaseCollect = resolve; - }); - const task = new BatchTask< - number, - number, - number, - void, - Record, - { id: string } - >({ - async submit(items) { - taskItems = items; - return { id: "task-provider" }; - }, - completion: { - mode: "webhook", - getExternalId: (submissionData) => submissionData.id, - }, - async collect() { - collectCount++; - if (collectCount === 2) releaseCollect(); - await bothCollecting; - return taskItems.map((item) => ({ - id: item.id, - output: item.input * 2, - })); - }, - }); - const scoreSubmit = vi.fn(async () => ({ id: "score-provider" })); - const scorer = new BatchScorer< - number, - number, - number, - void, - { id: string } - >({ - name: "exact", - submit: scoreSubmit, - completion: { - mode: "webhook", - getExternalId: (submissionData) => submissionData.id, - }, - async collect() { - return []; - }, - }); - const durable = defineDurableEval("concurrent-webhooks", { - store: new DurableEvalMemoryStore(), - data: [{ id: "one", input: 2, expected: 4 }], - task, - scores: [scorer], - }); - const waiting = await durable.start({ noSendLogs: true }); - - await Promise.all([ - durable.processBatchResult({ - runId: waiting.runId, - externalId: "task-provider", - }), - durable.processBatchResult({ - runId: waiting.runId, - externalId: "task-provider", - }), - ]); - - expect(scoreSubmit).toHaveBeenCalledTimes(1); - await expect( - durable.status({ runId: waiting.runId }), - ).resolves.toMatchObject({ - status: "waiting", - pending: { poll: 0, webhook: 1 }, - }); - }); - - test("supports scorer names inherited from Object.prototype", async () => { - const jobs = new Map< - string, - Array<{ id: string; input: number; expected?: number; output?: number }> - >(); - const completion = { - mode: "poll" as const, - async poll() { - return { status: "complete" as const }; - }, - }; - const task = new BatchTask< - number, - number, - number, - void, - Record - >({ - async submit(items) { - jobs.set("task", items); - return { id: "task" }; - }, - completion, - async collect() { - return (jobs.get("task") ?? []).map((item) => ({ - id: item.id, - output: item.input * 2, - })); - }, - }); - const scorer = new BatchScorer< - number, - number, - number, - void, - { id: string } - >({ - name: "__proto__", - async submit(items) { - jobs.set("score", items); - return { id: "score" }; - }, - completion, - async collect() { - return (jobs.get("score") ?? []).map((item) => ({ - id: item.id, - score: item.output === item.expected ? 1 : 0, - })); - }, - }); - const durable = defineDurableEval("prototype-names", { - store: new DurableEvalMemoryStore(), - data: [{ id: "one", input: 2, expected: 4 }], - task, - scores: [scorer], - }); - - const waiting = await durable.start({ noSendLogs: true }); - await expect(durable.poll({ runId: waiting.runId })).resolves.toMatchObject( - { - status: "waiting", - }, - ); - const result = await durable.poll({ runId: waiting.runId }); - expect(result.status).toBe("completed"); - if (result.status !== "completed") throw new Error("Eval did not complete"); - expect(Object.hasOwn(result.summary.scores, "__proto__")).toBe(true); - expect(result.summary.scores.__proto__?.score).toBe(1); - }); - - test("requires stable case ids", async () => { - await expect( - defineDurableEval("missing-ids", { - store: new DurableEvalMemoryStore(), - data: [{ input: "hello" }], - task: new BatchTask({ - async submit() { - return { id: "unused" }; - }, - completion: { - mode: "webhook", - getExternalId: (submissionData) => submissionData.id, - }, - async collect() { - return []; - }, - }), - scores: [], - }).start({ noSendLogs: true }), - ).rejects.toThrow("requires id, upsert_id, or caseId"); - }); -}); diff --git a/js/src/exports.ts b/js/src/exports.ts index 501051b4f..36a916fe9 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -273,15 +273,15 @@ export { defaultErrorScoreHandler, } from "./framework"; -export type { DurableEvalStore } from "./durable-eval"; +export type { WorkflowEvalStore } from "./workflow-eval"; export { - BatchScorer, - BatchTask, - defineDurableEval, - DurableEvalMemoryStore, - DurableEvalRedisStore, -} from "./durable-eval"; + WorkflowScorer, + WorkflowTask, + defineWorkflowEval, + WorkflowEvalMemoryStore, + WorkflowEvalRedisStore, +} from "./workflow-eval"; export { agentAssertionScorer } from "./agent-assertions"; diff --git a/js/src/workflow-eval.test.ts b/js/src/workflow-eval.test.ts new file mode 100644 index 000000000..434c7138f --- /dev/null +++ b/js/src/workflow-eval.test.ts @@ -0,0 +1,1066 @@ +import { describe, expect, expectTypeOf, test, vi } from "vitest"; +import { configureNode } from "./node/config"; +import { + WorkflowScorer, + WorkflowTask, + defineWorkflowEval, + WorkflowEvalMemoryStore, + WorkflowEvalRedisStore, + type WorkflowScorerItem, + type WorkflowTaskItem, + type WorkflowEvalStore, +} from "./workflow-eval"; + +configureNode(); + +describe("workflow eval stores", () => { + test("progress sets deduplicate concurrent additions", async () => { + const store = new WorkflowEvalMemoryStore(); + expect(await store.getSetSize("progress")).toBe(0); + await Promise.all( + ["one", "two", "one"].map((id) => store.addToSet("progress", id)), + ); + expect(await store.getSetSize("progress")).toBe(2); + }); + + test.each(["node-redis", "ioredis", "upstash"])( + "%s progress sets use atomic Redis scripts", + async (variant) => { + const evalCommand = vi.fn(async () => 2); + const client = { + get: async () => null, + set: async () => "OK", + eval: evalCommand, + ...(variant === "node-redis" + ? { sendCommand() {} } + : variant === "ioredis" + ? { defineCommand() {} } + : { createScript() {} }), + }; + const store = new WorkflowEvalRedisStore({ + client, + keyPrefix: "test:", + ttlMs: 1234, + }); + await store.addToSet("progress", "case-one"); + expect(await store.getSetSize("progress")).toBe(2); + const script = + "redis.call('SADD', KEYS[1], ARGV[1]); redis.call('PEXPIRE', KEYS[1], ARGV[2]); return 1"; + expect(evalCommand.mock.calls[0]).toEqual( + variant === "node-redis" + ? [ + script, + { keys: ["test:progress"], arguments: ["case-one", "1234"] }, + ] + : variant === "ioredis" + ? [script, 1, "test:progress", "case-one", "1234"] + : [script, ["test:progress"], ["case-one", "1234"]], + ); + expect(evalCommand.mock.calls[1]).toEqual( + variant === "node-redis" + ? [ + "return redis.call('SCARD', KEYS[1])", + { keys: ["test:progress"], arguments: [] }, + ] + : variant === "ioredis" + ? ["return redis.call('SCARD', KEYS[1])", 1, "test:progress"] + : ["return redis.call('SCARD', KEYS[1])", ["test:progress"], []], + ); + }, + ); + + test("memory store copies values on read and write", async () => { + const store = new WorkflowEvalMemoryStore(); + const value = new Uint8Array([1, 2, 3]); + + await store.write("run", value); + value[0] = 9; + + const firstRead = await store.read("run"); + expect(firstRead).toEqual(new Uint8Array([1, 2, 3])); + firstRead![1] = 9; + expect(await store.read("run")).toEqual(new Uint8Array([1, 2, 3])); + expect(await store.read("missing")).toBeUndefined(); + + const [first, second] = await Promise.all([ + store.getOrSet("claim", new Uint8Array([1])), + store.getOrSet("claim", new Uint8Array([2])), + ]); + expect([first.created, second.created]).toEqual([true, false]); + expect(first.value).toEqual(new Uint8Array([1])); + expect(second.value).toEqual(new Uint8Array([1])); + }); + + test("redis store uses prefixed string operations", async () => { + const values = new Map(); + const client = { + get: vi.fn(async (key: string) => values.get(key) ?? null), + set: vi.fn(async (key: string, value: string, _options?: unknown) => { + values.set(key, value); + return "OK"; + }), + sendCommand: vi.fn(), + }; + const store = new WorkflowEvalRedisStore({ + client, + keyPrefix: "evals:", + ttlMs: 1_234, + }); + + await store.write("run", new Uint8Array([0, 255, 1])); + + expect(client.set).toHaveBeenCalledWith("evals:run", "AP8B", { + PX: 1_234, + }); + expect(await store.read("run")).toEqual(new Uint8Array([0, 255, 1])); + expect(client.get).toHaveBeenCalledWith("evals:run"); + expect(await store.read("missing")).toBeUndefined(); + + await expect( + new WorkflowEvalRedisStore({ + client: { + get: async () => 42, + set: async () => "OK", + }, + }).read("invalid"), + ).rejects.toThrow("expected GET to return a string"); + expect(() => new WorkflowEvalRedisStore({ client, ttlMs: 0 })).toThrow( + "ttlMs must be a positive integer", + ); + }); + + test("redis store uses node-redis atomic SET options", async () => { + const values = new Map(); + const client = { + get: vi.fn(async (key: string) => values.get(key) ?? null), + set: vi.fn( + async ( + key: string, + value: string, + options?: { PX?: number; NX?: boolean; GET?: boolean }, + ) => { + expect(options).toEqual({ PX: 1_234, NX: true, GET: true }); + const existing = values.get(key) ?? null; + if (existing === null) values.set(key, value); + return existing; + }, + ), + sendCommand: vi.fn(), + }; + const store = new WorkflowEvalRedisStore({ client, ttlMs: 1_234 }); + + await expect(store.getOrSet("claim", new Uint8Array([1]))).resolves.toEqual( + { value: new Uint8Array([1]), created: true }, + ); + await expect(store.getOrSet("claim", new Uint8Array([2]))).resolves.toEqual( + { value: new Uint8Array([1]), created: false }, + ); + }); + + test("redis store uses ioredis atomic SET arguments", async () => { + const values = new Map(); + const client = { + get: vi.fn(async (key: string) => values.get(key) ?? null), + set: vi.fn(async (key: string, value: string, ...options: unknown[]) => { + expect(options).toEqual(["PX", 1_234, "NX", "GET"]); + const existing = values.get(key) ?? null; + if (existing === null) values.set(key, value); + return existing; + }), + defineCommand: vi.fn(), + }; + const store = new WorkflowEvalRedisStore({ client, ttlMs: 1_234 }); + + await expect(store.getOrSet("claim", new Uint8Array([1]))).resolves.toEqual( + { value: new Uint8Array([1]), created: true }, + ); + await expect(store.getOrSet("claim", new Uint8Array([2]))).resolves.toEqual( + { value: new Uint8Array([1]), created: false }, + ); + }); + + test("redis store uses Upstash atomic SET options", async () => { + const values = new Map(); + const client = { + get: vi.fn(async (key: string) => values.get(key) ?? null), + set: vi.fn( + async ( + key: string, + value: string, + options?: { px?: number; nx?: boolean; get?: boolean }, + ) => { + expect(options).toEqual({ px: 1_234, nx: true, get: true }); + const existing = values.get(key) ?? null; + if (existing === null) values.set(key, value); + return existing; + }, + ), + createScript: vi.fn(), + }; + const store = new WorkflowEvalRedisStore({ client, ttlMs: 1_234 }); + + await expect(store.getOrSet("claim", new Uint8Array([1]))).resolves.toEqual( + { value: new Uint8Array([1]), created: true }, + ); + await expect(store.getOrSet("claim", new Uint8Array([2]))).resolves.toEqual( + { value: new Uint8Array([1]), created: false }, + ); + }); +}); + +describe("defineWorkflowEval", () => { + test.each(["poll", "rejected poll", "collect"])( + "a scorer %s failure does not block other cases", + async (failure) => { + const f = workflowEval("poll"); + const { runId } = await f.definition.start({ noSendLogs: true }); + f.ready.add("task-one:trial:0"); + await f.definition.poll({ runId }); + f.ready.add("task-two:trial:0"); + f.ready.add("task-three:trial:0"); + f.ready.add("score-one:trial:0"); + if (failure !== "collect") { + f.poll.mockImplementation(async ({ id }) => { + if (id === "score-one:trial:0") { + if (failure === "rejected poll") throw new Error("scorer failed"); + return { status: "failed", error: "scorer failed" }; + } + return { status: f.ready.has(id) ? "complete" : "pending" }; + }); + } else { + f.scoreCollect.mockRejectedValueOnce(new Error("scorer failed")); + } + await expect(f.definition.poll({ runId })).rejects.toThrow( + "scorer failed", + ); + expect(f.taskCollect).toHaveBeenCalledTimes(3); + expect(f.scoreSubmit).toHaveBeenCalledTimes(3); + expect(f.localScore).toHaveBeenCalledTimes(3); + expect(f.classifier).toHaveBeenCalledTimes(3); + await expect(f.definition.status({ runId })).resolves.toMatchObject({ + pending: { poll: 3, webhook: 0 }, + }); + f.poll.mockImplementation(async () => ({ status: "complete" })); + await expect(f.definition.poll({ runId })).resolves.toMatchObject({ + status: "completed", + }); + }, + ); + + test("polling retries an interrupted progress update", async () => { + const store = new WorkflowEvalMemoryStore(); + const f = workflowEval("poll", store); + const { runId } = await f.definition.start({ noSendLogs: true }); + const addToSet = store.addToSet.bind(store); + let interrupted = false; + vi.spyOn(store, "addToSet").mockImplementation(async (key, member) => { + if (!interrupted && key.endsWith("/poll/complete")) { + interrupted = true; + throw new Error("store unavailable"); + } + return addToSet(key, member); + }); + f.ready.add("task-one:trial:0"); + await expect(f.definition.poll({ runId })).rejects.toThrow( + "store unavailable", + ); + expect(f.taskCollect).toHaveBeenCalledTimes(1); + await f.definition.poll({ runId }); + expect(f.taskCollect).toHaveBeenCalledTimes(2); + expect(f.scoreSubmit).toHaveBeenCalledTimes(1); + await expect(f.definition.status({ runId })).resolves.toMatchObject({ + pending: { poll: 3, webhook: 0 }, + }); + f.poll.mockImplementation(async () => ({ status: "complete" })); + await f.definition.poll({ runId }); + await expect(f.definition.poll({ runId })).resolves.toMatchObject({ + status: "completed", + }); + }); + + test.each([undefined, 1, 3])( + "bounds submission and polling concurrency at %s", + async (maxConcurrency) => { + const active = { submit: 0, poll: 0, collect: 0 }; + const peaks = { ...active }; + const task = new WorkflowTask({ + async submit() { + peaks.submit = Math.max(peaks.submit, ++active.submit); + await new Promise((resolve) => setTimeout(resolve, 1)); + active.submit--; + return null; + }, + completion: { + mode: "poll", + async poll() { + peaks.poll = Math.max(peaks.poll, ++active.poll); + await new Promise((resolve) => setTimeout(resolve, 1)); + active.poll--; + return { status: "complete" }; + }, + }, + async collect() { + peaks.collect = Math.max(peaks.collect, ++active.collect); + await new Promise((resolve) => setTimeout(resolve, 1)); + active.collect--; + return { output: 1 }; + }, + }); + const definition = defineWorkflowEval("concurrency", { + store: new WorkflowEvalMemoryStore(), + maxConcurrency, + data: Array.from({ length: 12 }, (_, input) => ({ + id: String(input), + input, + })), + task, + }); + const { runId } = await definition.start({ noSendLogs: true }); + await expect(definition.poll({ runId })).resolves.toMatchObject({ + status: "completed", + }); + expect(peaks.submit).toBe(maxConcurrency ?? 10); + expect(peaks.poll).toBe(maxConcurrency ?? 10); + expect(peaks.collect).toBeGreaterThan(0); + expect(peaks.collect).toBeLessThanOrEqual(maxConcurrency ?? 10); + }, + ); + + test.each([0, -1, 1.5, Infinity, NaN])( + "rejects invalid concurrency %s", + (maxConcurrency) => { + expect(() => + defineWorkflowEval("invalid", { + store: new WorkflowEvalMemoryStore(), + maxConcurrency, + data: [], + task: () => 1, + }), + ).toThrow("maxConcurrency must be a positive integer"); + }, + ); + + test("webhook record reads stay constant as the dataset grows", async () => { + const readCounts: number[] = []; + for (const size of [3, 100]) { + const store = new WorkflowEvalMemoryStore(); + const read = vi.spyOn(store, "read"); + const f = workflowEval("webhook", store); + const definition = defineWorkflowEval("indexed-webhooks", { + store, + data: Array.from({ length: size }, (_, input) => ({ + id: String(input), + input, + expected: input * 2, + })), + task: f.task, + scores: [f.scorer], + }); + const { runId } = await definition.start({ noSendLogs: true }); + read.mockClear(); + await expect( + definition.processSubmissionResult({ + runId, + externalId: "task-0:trial:0", + }), + ).resolves.toMatchObject({ pending: { webhook: size } }); + readCounts.push(read.mock.calls.length); + expect(read.mock.calls.some(([key]) => key.endsWith("/case-ids"))).toBe( + false, + ); + await definition.processSubmissionResult({ + runId, + externalId: "task-0:trial:0", + }); + await definition.processSubmissionResult({ + runId, + externalId: "score-0:trial:0", + }); + await definition.processSubmissionResult({ + runId, + externalId: "score-0:trial:0", + }); + read.mockClear(); + await expect(definition.status({ runId })).resolves.toMatchObject({ + pending: { webhook: size - 1 }, + }); + expect(read).toHaveBeenCalledTimes(1); + } + expect(readCounts[1]).toBe(readCounts[0]); + expect(readCounts[1]).toBeLessThan(50); + }); + + test("runs ordinary tasks and scorers", async () => { + const task = vi.fn((input: number) => input * 2); + const result = await defineWorkflowEval("local", { + store: new WorkflowEvalMemoryStore(), + data: [ + { id: "one", input: 1, expected: 2 }, + { id: "two", input: 2, expected: 4 }, + ], + task, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }).start({ noSendLogs: true }); + + expect(result).toMatchObject({ + status: "completed", + summary: { scores: { exact: { score: 1 } } }, + }); + expect(task).toHaveBeenCalledTimes(2); + }); + + test("generates a new run id for every start", async () => { + const workflow = defineWorkflowEval("generated-runs", { + store: new WorkflowEvalMemoryStore(), + data: [{ input: 1 }], + task: (input) => input, + scores: [() => 1], + }); + + const first = await workflow.start({ noSendLogs: true }); + const second = await workflow.start({ noSendLogs: true }); + + expect(first.runId).not.toBe(second.runId); + }); + + test("stores run, cases, and individual submissions separately", async () => { + const values = new Map(); + const progressStore = new WorkflowEvalMemoryStore(); + const store: WorkflowEvalStore = { + addToSet: (key, member) => progressStore.addToSet(key, member), + getSetSize: (key) => progressStore.getSetSize(key), + async read(key) { + return values.get(key); + }, + async write(key, value) { + values.set(key, value); + }, + async getOrSet(key, value) { + const existing = values.get(key); + if (existing) return { value: existing, created: false }; + values.set(key, value); + return { value, created: true }; + }, + }; + const { definition } = workflowEval("poll", store); + await definition.start({ noSendLogs: true }); + const records = [...values] + .filter(([key]) => !key.includes("/claims/")) + .map( + ([key, value]) => + [key, JSON.parse(new TextDecoder().decode(value))] as const, + ); + const run = records.find(([key]) => + /^workflow-eval\/v1\/runs\/[^/]+$/.test(key), + )!; + expect(run[0]).toMatch(/^workflow-eval\/v1\/runs\//); + expect(run[1]).toMatchObject({ + caseCount: 3, + }); + expect(run[1]).not.toHaveProperty("cases"); + expect(run[1]).not.toHaveProperty("submissions"); + const submissions = records.filter(([key]) => + key.includes("/submissions/"), + ); + expect(submissions).toHaveLength(3); + expect(submissions.map(([, record]) => record.itemId)).toEqual([ + "one:trial:0", + "two:trial:0", + "three:trial:0", + ]); + for (const [, record] of submissions) + expect(record).not.toHaveProperty("itemIds"); + expect(records.filter(([key]) => key.includes("/cases/"))).toHaveLength(3); + }); + + test("polls existing submissions once and starts scoring only ready cases", async () => { + const f = workflowEval("poll"); + const waiting = await f.definition.start({ noSendLogs: true }); + const options = { runId: waiting.runId }; + expect(waiting).toMatchObject({ + status: "waiting", + pending: { poll: 3, webhook: 0 }, + }); + expect(f.taskSubmit).toHaveBeenCalledTimes(3); + await expect(f.definition.status(options)).resolves.toEqual(waiting); + expect(f.poll).not.toHaveBeenCalled(); + + f.ready.add("task-two:trial:0"); + await expect(f.definition.poll(options)).resolves.toMatchObject({ + status: "waiting", + pending: { poll: 3, webhook: 0 }, + }); + expect(f.poll).toHaveBeenCalledTimes(3); + expect(f.scoreSubmit).toHaveBeenCalledTimes(1); + expect(f.localScore).toHaveBeenCalledTimes(1); + expect(f.classifier).toHaveBeenCalledTimes(1); + expect(f.scoreSubmit.mock.calls[0][0]).toMatchObject({ + id: "two:trial:0", + output: 4, + }); + expect(f.localScore.mock.calls[0][0]).toMatchObject({ + input: 2, + output: 4, + }); + + f.ready.add("score-two:trial:0"); + await expect(f.definition.poll(options)).resolves.toMatchObject({ + status: "waiting", + }); + expect(f.poll).toHaveBeenCalledTimes(6); + expect(f.scoreSubmit).toHaveBeenCalledTimes(1); + expect(f.localScore).toHaveBeenCalledTimes(1); + + f.ready.add("task-one:trial:0"); + f.ready.add("task-three:trial:0"); + await f.definition.poll(options); + expect(f.scoreSubmit).toHaveBeenCalledTimes(3); + f.ready.add("score-one:trial:0"); + f.ready.add("score-three:trial:0"); + const completed = await f.definition.poll(options); + expect(completed).toMatchObject({ + status: "completed", + pending: { poll: 0, webhook: 0 }, + summary: { + scores: { workflow_exact: { score: 1 }, extra: { score: 0.5 } }, + }, + }); + const callCount = f.poll.mock.calls.length; + await expect(f.definition.status(options)).resolves.toEqual(completed); + await expect(f.definition.poll(options)).resolves.toEqual(completed); + expect(f.poll).toHaveBeenCalledTimes(callCount); + expect(f.taskSubmit).toHaveBeenCalledTimes(3); + expect(f.classifier).toHaveBeenCalledTimes(3); + }); + + test("resumes webhook submissions with a fresh definition and supports both locators", async () => { + const store = new WorkflowEvalMemoryStore(); + const first = workflowEval("webhook", store); + const { runId } = await first.definition.start({ noSendLogs: true }); + const f = workflowEval("webhook", store); + // Simulate fetching provider results in a later process, without rerunning submit. + for (const [id, job] of first.taskJobs) f.taskJobs.set(id, job); + const [externalId, { context }] = [...first.taskJobs][1]; + await expect( + f.definition.processSubmissionResult({ + runId, + submissionId: context.submissionId, + }), + ).resolves.toMatchObject({ + status: "waiting", + pending: { poll: 0, webhook: 3 }, + }); + expect(f.taskSubmit).not.toHaveBeenCalled(); + expect(f.scoreSubmit).toHaveBeenCalledTimes(1); + expect(f.scoreSubmit.mock.calls[0][0].id).toBe("two:trial:0"); + expect(f.taskCollect.mock.calls[0][1]).toEqual(context); + await f.definition.processSubmissionResult({ runId, externalId }); + expect(f.taskCollect).toHaveBeenCalledTimes(1); + expect(f.scoreSubmit).toHaveBeenCalledTimes(1); + + for (const id of first.taskJobs.keys()) { + await f.definition.processSubmissionResult({ runId, externalId: id }); + } + for (const id of f.scoreJobs.keys()) { + await f.definition.processSubmissionResult({ runId, externalId: id }); + } + const completed = await f.definition.status({ runId }); + expect(completed.status).toBe("completed"); + await expect( + f.definition.processSubmissionResult({ runId, externalId }), + ).resolves.toEqual(completed); + expect(f.taskCollect).toHaveBeenCalledTimes(3); + expect(f.localScore).toHaveBeenCalledTimes(3); + }); + + test("claims downstream work once across concurrent webhook deliveries", async () => { + const f = workflowEval("webhook"); + const { runId } = await f.definition.start({ noSendLogs: true }); + let release!: () => void; + const collecting = new Promise((resolve) => { + release = resolve; + }); + let count = 0; + f.taskCollect.mockImplementation(async ({ id }) => { + if (++count === 2) release(); + await collecting; + return { output: f.taskJobs.get(id)!.item.input * 2 }; + }); + await Promise.all([ + f.definition.processSubmissionResult({ + runId, + externalId: "task-one:trial:0", + }), + f.definition.processSubmissionResult({ + runId, + externalId: "task-one:trial:0", + }), + ]); + expect(f.scoreSubmit).toHaveBeenCalledTimes(1); + expect(f.localScore).toHaveBeenCalledTimes(1); + expect(f.classifier).toHaveBeenCalledTimes(1); + await expect(f.definition.status({ runId })).resolves.toMatchObject({ + status: "waiting", + pending: { poll: 0, webhook: 3 }, + }); + }); + + test("preserves results across concurrent completion of different cases", async () => { + const f = workflowEval("webhook"); + const { runId } = await f.definition.start({ noSendLogs: true }); + await Promise.all( + [...f.taskJobs.keys()].map((externalId) => + f.definition.processSubmissionResult({ runId, externalId }), + ), + ); + expect(f.scoreSubmit).toHaveBeenCalledTimes(3); + expect(f.localScore).toHaveBeenCalledTimes(3); + await Promise.all( + [...f.scoreJobs.keys()].map((externalId) => + f.definition.processSubmissionResult({ runId, externalId }), + ), + ); + await expect(f.definition.status({ runId })).resolves.toMatchObject({ + status: "completed", + }); + }); + + test("validates completion locators", async () => { + const f = workflowEval("webhook"); + const { runId } = await f.definition.start({ noSendLogs: true }); + await expect( + f.definition.processSubmissionResult({ runId }), + ).rejects.toThrow("require submissionId or externalId"); + await expect( + f.definition.processSubmissionResult({ runId, externalId: "missing" }), + ).rejects.toThrow("No submission matches"); + await expect( + f.definition.processSubmissionResult({ + runId: "missing", + externalId: "task-one:trial:0", + }), + ).rejects.toThrow("run missing is missing"); + const submissionId = [...f.taskJobs.values()][0].context.submissionId; + await expect( + f.definition.processSubmissionResult({ + runId, + submissionId, + externalId: "task-two:trial:0", + }), + ).rejects.toThrow("identify different submissions"); + expect(f.taskCollect).not.toHaveBeenCalled(); + }); + + test("propagates callback errors and polling failures", async () => { + const f = workflowEval("poll"); + const { runId } = await f.definition.start({ noSendLogs: true }); + f.poll.mockRejectedValueOnce(new Error("provider unavailable")); + await expect(f.definition.poll({ runId })).rejects.toThrow( + "provider unavailable", + ); + f.poll.mockResolvedValueOnce({ + status: "failed", + error: "provider failed", + }); + await expect(f.definition.poll({ runId })).rejects.toThrow( + "provider failed", + ); + expect(f.taskCollect).not.toHaveBeenCalled(); + f.ready.add("task-one:trial:0"); + f.taskCollect.mockRejectedValueOnce(new Error("collection failed")); + await expect(f.definition.poll({ runId })).rejects.toThrow( + "collection failed", + ); + expect(f.scoreSubmit).not.toHaveBeenCalled(); + }); + + test("rejects array collection results", async () => { + const f = workflowEval("webhook"); + const { runId } = await f.definition.start({ noSendLogs: true }); + // Exercise the runtime boundary for JavaScript consumers. + // @ts-expect-error Collection returns one result envelope, never an array. + f.taskCollect.mockResolvedValueOnce([{ output: 2 }]); + await expect( + f.definition.processSubmissionResult({ + runId, + externalId: "task-one:trial:0", + }), + ).rejects.toThrow("must return a result object"); + await f.definition.processSubmissionResult({ + runId, + externalId: "task-one:trial:0", + }); + // @ts-expect-error Score arrays must be nested inside the score envelope. + f.scoreCollect.mockResolvedValueOnce([{ score: 1 }]); + await expect( + f.definition.processSubmissionResult({ + runId, + externalId: "score-one:trial:0", + }), + ).rejects.toThrow("must return a result object"); + }); + + test("propagates metadata, tags, parameters, and distinct trials", async () => { + type Metadata = { fromCase?: boolean; fromTask?: boolean }; + const items: WorkflowTaskItem< + number, + number, + Metadata, + Record + >[] = []; + const contexts: Array<{ runId: string; submissionId: string }> = []; + const scoreItems: WorkflowScorerItem[] = + []; + const task = new WorkflowTask< + number, + number, + number, + Metadata, + Record, + { input: number } + >({ + async submit(item, context) { + items.push(item); + contexts.push(context); + return { input: item.input }; + }, + completion: { + mode: "poll", + async poll() { + return { status: "complete" }; + }, + }, + async collect({ input }) { + return { + output: input * 2, + metadata: { fromTask: true }, + tags: ["updated"], + }; + }, + }); + const scorer = new WorkflowScorer({ + name: "__proto__", + async submit(item) { + scoreItems.push(item); + return null; + }, + completion: { + mode: "poll", + async poll() { + return { status: "complete" }; + }, + }, + async collect() { + return { score: 1 }; + }, + }); + const localScore = vi.fn(({ metadata, tags }) => { + expect(metadata).toEqual({ fromCase: true, fromTask: true }); + expect(tags).toEqual(["updated"]); + return 1; + }); + const definition = defineWorkflowEval("trials", { + store: new WorkflowEvalMemoryStore(), + data: [ + { + input: 2, + expected: 4, + metadata: { fromCase: true }, + tags: ["original"], + }, + ], + caseId: () => "one", + trialCount: 2, + task, + scores: [scorer, localScore], + }); + const { runId } = await definition.start({ noSendLogs: true }); + expect(items.map(({ id, trialIndex }) => ({ id, trialIndex }))).toEqual([ + { id: "one:trial:0", trialIndex: 0 }, + { id: "one:trial:1", trialIndex: 1 }, + ]); + expect(items[0]).toMatchObject({ + input: 2, + expected: 4, + parameters: {}, + tags: ["original"], + }); + expect(new Set(contexts.map(({ submissionId }) => submissionId)).size).toBe( + 2, + ); + await definition.poll({ runId }); + expect(scoreItems).toHaveLength(2); + expect(scoreItems[0]).toMatchObject({ + metadata: { fromCase: true, fromTask: true }, + tags: ["updated"], + output: 4, + }); + const result = await definition.poll({ runId }); + expect(result.status).toBe("completed"); + if (result.status !== "completed") throw new Error("Eval did not complete"); + expect(Object.hasOwn(result.summary.scores, "__proto__")).toBe(true); + expect(result.summary.scores.__proto__?.score).toBe(1); + }); + + test("supports ordinary tasks with workflow scorers and empty workflow datasets", async () => { + const f = workflowEval("poll"); + const definition = defineWorkflowEval("ordinary-task", { + store: new WorkflowEvalMemoryStore(), + data: [{ id: "one", input: 2, expected: 4 }], + task: (input) => input * 2, + scores: [f.scorer], + }); + const { runId } = await definition.start({ noSendLogs: true }); + expect(f.scoreSubmit.mock.calls[0][0]).toMatchObject({ + input: 2, + output: 4, + }); + f.ready.add("score-one:trial:0"); + await expect(definition.poll({ runId })).resolves.toMatchObject({ + status: "completed", + }); + await expect( + defineWorkflowEval("empty", { + store: new WorkflowEvalMemoryStore(), + data: [], + task: f.task, + scores: [f.scorer], + }).start({ noSendLogs: true }), + ).resolves.toMatchObject({ status: "completed" }); + expect(f.taskSubmit).not.toHaveBeenCalled(); + }); + + test("keeps task-only runs waiting until every task completes", async () => { + const f = workflowEval("webhook"); + const definition = defineWorkflowEval("task-only", { + store: new WorkflowEvalMemoryStore(), + data: [ + { id: "one", input: 1, expected: 2 }, + { id: "two", input: 2, expected: 4 }, + ], + task: f.task, + }); + const { runId } = await definition.start({ noSendLogs: true }); + await expect( + definition.processSubmissionResult({ + runId, + externalId: "task-one:trial:0", + }), + ).resolves.toMatchObject({ + status: "waiting", + pending: { poll: 0, webhook: 1 }, + }); + await expect( + definition.processSubmissionResult({ + runId, + externalId: "task-two:trial:0", + }), + ).resolves.toMatchObject({ status: "completed" }); + }); + + test("mixes polling tasks with webhook scorers", async () => { + const f = workflowEval("poll"); + f.scorer.processor.completion = { + mode: "webhook", + getExternalId: ({ id }) => id, + }; + const { runId } = await f.definition.start({ noSendLogs: true }); + f.ready.add("task-one:trial:0"); + await expect(f.definition.poll({ runId })).resolves.toMatchObject({ + status: "waiting", + pending: { poll: 2, webhook: 1 }, + }); + await expect( + f.definition.processSubmissionResult({ + runId, + externalId: "score-one:trial:0", + }), + ).resolves.toMatchObject({ + status: "waiting", + pending: { poll: 2, webhook: 0 }, + }); + expect(f.poll).toHaveBeenCalledTimes(3); + for (const [id, { context }] of f.taskJobs) { + expect(f.poll).toHaveBeenCalledWith({ id }, context); + } + }); + + test("requires stable case ids", async () => { + const f = workflowEval("poll"); + await expect( + defineWorkflowEval("missing-ids", { + store: new WorkflowEvalMemoryStore(), + data: [{ input: 1, expected: 2 }], + task: f.task, + }).start({ noSendLogs: true }), + ).rejects.toThrow("requires id, upsert_id, or caseId"); + }); + + test("infers submission data and callback result types", () => { + const task = new WorkflowTask({ + async submit( + item: WorkflowTaskItem>, + ) { + expectTypeOf(item.input).toEqualTypeOf(); + return { providerId: "request", attempt: 1 }; + }, + completion: { + mode: "webhook", + getExternalId(submission) { + expectTypeOf(submission).toEqualTypeOf<{ + providerId: string; + attempt: number; + }>(); + return submission.providerId; + }, + }, + async collect(submission) { + return { output: submission.attempt }; + }, + }); + expectTypeOf(task.processor.collect).returns.resolves.toMatchTypeOf<{ + output: number; + }>(); + const scorer = new WorkflowScorer({ + name: "score", + async submit() { + return { providerId: "score" }; + }, + completion: { + mode: "poll", + async poll(submission) { + expectTypeOf(submission).toEqualTypeOf<{ providerId: string }>(); + return { status: "pending" }; + }, + }, + async collect(submission) { + expectTypeOf(submission.providerId).toEqualTypeOf(); + return { score: 1 }; + }, + }); + expect(scorer.name).toBe("score"); + }); +}); + +function workflowEval( + mode: "poll" | "webhook", + store: WorkflowEvalStore = new WorkflowEvalMemoryStore(), +) { + type TaskItem = WorkflowTaskItem>; + type ScoreItem = WorkflowScorerItem; + type Context = { runId: string; submissionId: string }; + const taskJobs = new Map(); + const scoreJobs = new Map(); + const ready = new Set(); + const poll = vi.fn( + async ({ + id, + }: { + id: string; + }): Promise< + | { status: "pending" } + | { status: "complete" } + | { status: "failed"; error: unknown } + > => ({ status: ready.has(id) ? "complete" : "pending" }), + ); + const completion = + mode === "poll" + ? { mode, poll } + : { mode, getExternalId: ({ id }: { id: string }) => id }; + const taskSubmit = vi.fn(async (item: TaskItem, context: Context) => { + const id = `task-${item.id}`; + taskJobs.set(id, { item, context }); + return { id }; + }); + const taskCollect = vi.fn( + async ({ id }: { id: string }, _context: Context) => ({ + output: taskJobs.get(id)!.item.input * 2, + }), + ); + const task = new WorkflowTask< + number, + number, + number, + void, + Record, + { id: string } + >({ + submit: taskSubmit, + completion, + collect: taskCollect, + }); + const scoreSubmit = vi.fn(async (item: ScoreItem, context: Context) => { + const id = `score-${item.id}`; + scoreJobs.set(id, { item, context }); + return { id }; + }); + const scoreCollect = vi.fn(async ({ id }: { id: string }) => { + const { item } = scoreJobs.get(id)!; + return { + score: [ + { + name: "workflow_exact", + score: item.output === item.expected ? 1 : 0, + }, + { name: "extra", score: 0.5 }, + ], + }; + }); + const scorer = new WorkflowScorer< + number, + number, + number, + void, + { id: string } + >({ + name: "workflow_exact", + submit: scoreSubmit, + completion, + collect: scoreCollect, + }); + const localScore = vi.fn( + ({ + output, + expected, + }: { + input: number; + output: number; + expected: number; + }) => (output === expected ? 1 : 0), + ); + const classifier = vi.fn(() => ({ + name: "quality", + id: "pass", + label: "Pass", + })); + const definition = defineWorkflowEval("workflow", { + store, + data: ["one", "two", "three"].map((id, index) => ({ + id, + input: index + 1, + expected: (index + 1) * 2, + })), + task, + scores: [scorer, localScore], + classifiers: [classifier], + }); + return { + definition, + task, + scorer, + taskJobs, + scoreJobs, + ready, + poll, + taskSubmit, + taskCollect, + scoreSubmit, + scoreCollect, + localScore, + classifier, + }; +} diff --git a/js/src/durable-eval.ts b/js/src/workflow-eval.ts similarity index 60% rename from js/src/durable-eval.ts rename to js/src/workflow-eval.ts index 7a0c45e70..7d38bf66b 100644 --- a/js/src/durable-eval.ts +++ b/js/src/workflow-eval.ts @@ -1,3 +1,4 @@ +import { queue } from "async"; import { base64ToUint8Array, makeScorerPropagatedEvent, @@ -45,21 +46,20 @@ import { const encoder = new TextEncoder(); const decoder = new TextDecoder(); -const BATCH_TASK_KIND = "braintrust.durable.batch-task"; -const BATCH_SCORER_KIND = "braintrust.durable.batch-scorer"; -const DEFAULT_BATCH_SIZE = 1_000; +const WORKFLOW_TASK_KIND = "braintrust.workflow.task"; +const WORKFLOW_SCORER_KIND = "braintrust.workflow.scorer"; type JsonPrimitive = string | number | boolean | null; type JsonValue = JsonPrimitive | JsonValue[] | { [key: string]: JsonValue }; /** - * Minimal persistence used to reconnect provider webhooks with submitted - * batches. Each run, case, and batch is stored under its own key. Durable + * Minimal persistence used to reconnect provider webhooks with provider + * submissions. Each run, case, and submission is stored under its own key. Workflow * evaluations do not require any Braintrust backend changes. * * @experimental - The API for this interface is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. */ -export interface DurableEvalStore { +export interface WorkflowEvalStore { read(key: string): Promise; write(key: string, value: Uint8Array): Promise; /** Atomically stores `value` when `key` is absent and returns its stored value. */ @@ -67,16 +67,24 @@ export interface DurableEvalStore { key: string, value: Uint8Array, ): Promise<{ value: Uint8Array; created: boolean }>; + /** + * Atomically adds a unique member to a set. Repeated additions are harmless. + * Set keys are separate from byte-record keys. Implementations must retain sets + * for the same lifetime as run records. + */ + addToSet(key: string, member: string): Promise; + getSetSize(key: string): Promise; } /** - * Stores durable evaluation state in memory. State is lost when the current + * Stores workflow evaluation state in memory. State is lost when the current * JavaScript process exits. * * @experimental - The API for this class is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. */ -export class DurableEvalMemoryStore implements DurableEvalStore { +export class WorkflowEvalMemoryStore implements WorkflowEvalStore { private readonly values = new Map(); + private readonly sets = new Map>(); async read(key: string): Promise { return this.values.get(key)?.slice(); @@ -92,17 +100,26 @@ export class DurableEvalMemoryStore implements DurableEvalStore { this.values.set(key, value.slice()); return { value: value.slice(), created: true }; } + + async addToSet(key: string, member: string) { + let members = this.sets.get(key); + if (!members) this.sets.set(key, (members = new Set())); + members.add(member); + } + + async getSetSize(key: string) { + return this.sets.get(key)?.size ?? 0; + } } /** - * Stores durable evaluation state in Redis using an existing Redis client. - * Values are base64 encoded so only string `GET` and `SET` operations are - * required from the client. Clients from `redis` (node-redis), `ioredis`, and + * Stores workflow evaluation state in Redis using an existing Redis client. + * Records are base64 encoded. Atomic progress sets use Redis Lua scripts. Clients from `redis` (node-redis), `ioredis`, and * `@upstash/redis` can be passed directly. * * @experimental - The API for this class is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. */ -export class DurableEvalRedisStore implements DurableEvalStore { +export class WorkflowEvalRedisStore implements WorkflowEvalStore { private readonly client: { get(key: string): Promise; set(key: string, value: string): Promise; @@ -125,7 +142,9 @@ export class DurableEvalRedisStore implements DurableEvalStore { this.keyPrefix = options.keyPrefix ?? "braintrust-eval:"; this.ttlMs = options.ttlMs ?? 1000 * 60 * 60 * 24 * 7; if (!Number.isInteger(this.ttlMs) || this.ttlMs < 1) { - throw new Error("DurableEvalRedisStore ttlMs must be a positive integer"); + throw new Error( + "WorkflowEvalRedisStore ttlMs must be a positive integer", + ); } } @@ -133,7 +152,7 @@ export class DurableEvalRedisStore implements DurableEvalStore { const value = await this.client.get(`${this.keyPrefix}${key}`); if (value == null) return undefined; if (typeof value !== "string") { - throw new Error("DurableEvalRedisStore expected GET to return a string"); + throw new Error("WorkflowEvalRedisStore expected GET to return a string"); } return base64ToUint8Array(value); } @@ -157,7 +176,7 @@ export class DurableEvalRedisStore implements DurableEvalStore { await set.call(client, redisKey, encoded, { px: this.ttlMs }); } else { throw new Error( - "DurableEvalRedisStore requires a node-redis, ioredis, or @upstash/redis client", + "WorkflowEvalRedisStore requires a node-redis, ioredis, or @upstash/redis client", ); } } @@ -182,55 +201,89 @@ export class DurableEvalRedisStore implements DurableEvalStore { setOptions = [{ px: this.ttlMs, nx: true, get: true }]; } else { throw new Error( - "DurableEvalRedisStore getOrSet requires a node-redis, ioredis, or @upstash/redis client", + "WorkflowEvalRedisStore getOrSet requires a node-redis, ioredis, or @upstash/redis client", ); } const existing = await set.call(client, redisKey, encoded, ...setOptions); if (existing === null) return { value: value.slice(), created: true }; if (typeof existing !== "string") { throw new Error( - "DurableEvalRedisStore expected atomic SET to return a string or null", + "WorkflowEvalRedisStore expected atomic SET to return a string or null", ); } return { value: base64ToUint8Array(existing), created: false }; } + + async addToSet(key: string, member: string) { + await this.evalSet( + "redis.call('SADD', KEYS[1], ARGV[1]); redis.call('PEXPIRE', KEYS[1], ARGV[2]); return 1", + key, + [member, String(this.ttlMs)], + ); + } + + async getSetSize(key: string) { + return this.evalSet("return redis.call('SCARD', KEYS[1])", key, []); + } + + private async evalSet(script: string, key: string, args: string[]) { + const client = this.client as typeof this.client & { + eval: (...args: unknown[]) => Promise; + defineCommand?: unknown; + sendCommand?: unknown; + createScript?: unknown; + }; + const redisKey = `${this.keyPrefix}${key}`; + const result = + typeof client.defineCommand === "function" + ? await client.eval(script, 1, redisKey, ...args) + : typeof client.sendCommand === "function" + ? await client.eval(script, { keys: [redisKey], arguments: args }) + : await client.eval(script, [redisKey], args); + if (typeof result !== "number") { + throw new Error( + "WorkflowEvalRedisStore expected EVAL to return a number", + ); + } + return result; + } } -interface DurableBatchContext { +interface WorkflowSubmissionContext { runId: string; - batchId: string; + submissionId: string; } -type DurableBatchPoll = +type WorkflowSubmissionPoll = | { status: "pending" } | { status: "complete" } | { status: "failed"; error: unknown }; -type DurableBatchCompletion = +type WorkflowSubmissionCompletion = | { mode: "poll"; - /** Checks whether the submitted provider batch is ready to collect. */ + /** Checks whether the provider submission is ready to collect. */ poll( submissionData: SubmissionData, - context: DurableBatchContext, - ): Promise; + context: WorkflowSubmissionContext, + ): Promise; } | { mode: "webhook"; - /** Returns the provider ID used to match an incoming webhook to this batch. */ + /** Returns the provider ID used to match an incoming webhook to this submission. */ getExternalId( submissionData: SubmissionData, - context: DurableBatchContext, + context: WorkflowSubmissionContext, ): string; }; -export interface DurableBatchTaskItem< +export interface WorkflowTaskItem< Input, Expected, Metadata extends BaseMetadata, Parameters extends EvalParameters, > { - /** Stable identifier for this case and trial within the durable run. */ + /** Stable identifier for this case and trial within the workflow run. */ id: string; /** Input value from the evaluation case. */ input: Input; @@ -240,15 +293,13 @@ export interface DurableBatchTaskItem< metadata: Metadata; /** Tags associated with the evaluation case. */ tags: string[] | undefined; - /** Parameters supplied when the durable evaluation started. */ + /** Parameters supplied when the workflow evaluation started. */ parameters: InferParameters; /** Zero-based trial index for this case. */ trialIndex: number; } -type DurableBatchTaskResult = { - /** ID of the submitted item this result belongs to. */ - id: string; +type WorkflowTaskResult = { /** Task output for the item. */ output: Output; /** Metadata to merge into the evaluation case. */ @@ -257,44 +308,43 @@ type DurableBatchTaskResult = { tags?: string[]; }; -export type DurableBatchScorerItem< +export type WorkflowScorerItem< Input, Output, Expected, Metadata extends BaseMetadata, > = EvalScorerArgs & { - /** Stable identifier for this case and trial within the durable run. */ + /** Stable identifier for this case and trial within the workflow run. */ id: string; /** Zero-based trial index for this case. */ trialIndex: number; }; -type DurableBatchScorerResult = { - /** ID of the submitted item this result belongs to. */ - id: string; +type WorkflowScorerResult = { /** Score or named scores produced for the item. */ score: OneOrMoreScores; }; -interface DurableBatchProcessor< +interface WorkflowSubmissionProcessor< Item, Result, SubmissionData extends JsonValue, > { - /** Maximum items submitted in one provider batch. Defaults to 1,000. */ - batchSize?: number; - /** Submits a batch and returns the provider-specific submission data. */ - submit(items: Item[], context: DurableBatchContext): Promise; - /** Configures how the SDK learns that the submitted batch completed. */ - completion: DurableBatchCompletion; - /** Collects one result for every item in a completed provider batch. */ + /** Submits one case/trial and returns JSON-serializable provider data. */ + submit( + item: Item, + context: WorkflowSubmissionContext, + ): Promise; + /** Configures how the SDK learns that the submission completed. */ + completion: WorkflowSubmissionCompletion; + /** Collects the result for a completed submission. May be called again on replay. */ collect( submissionData: SubmissionData, - context: DurableBatchContext, - ): Promise; + context: WorkflowSubmissionContext, + ): Promise; } -interface DurableBatchTask< +interface WorkflowTaskDefinition< Input, Output, Expected, @@ -302,43 +352,43 @@ interface DurableBatchTask< Parameters extends EvalParameters, SubmissionData extends JsonValue, > { - readonly kind: typeof BATCH_TASK_KIND; - readonly processor: DurableBatchProcessor< - DurableBatchTaskItem, - DurableBatchTaskResult, + readonly kind: typeof WORKFLOW_TASK_KIND; + readonly processor: WorkflowSubmissionProcessor< + WorkflowTaskItem, + WorkflowTaskResult, SubmissionData >; } -interface DurableBatchScorer< +interface WorkflowScorerDefinition< Input, Output, Expected, Metadata extends BaseMetadata, SubmissionData extends JsonValue, > { - readonly kind: typeof BATCH_SCORER_KIND; + readonly kind: typeof WORKFLOW_SCORER_KIND; name: string; - readonly processor: DurableBatchProcessor< - DurableBatchScorerItem, - DurableBatchScorerResult, + readonly processor: WorkflowSubmissionProcessor< + WorkflowScorerItem, + WorkflowScorerResult, SubmissionData >; } /** - * Defines a task that runs through asynchronous provider batch operations. + * Defines a task that submits one asynchronous provider operation per case/trial. * * @experimental - The API for this class is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. */ -export class BatchTask< +export class WorkflowTask< Input, Output, Expected = void, Metadata extends BaseMetadata = DefaultMetadataType, Parameters extends EvalParameters = EvalParameters, SubmissionData extends JsonValue = JsonValue, -> implements DurableBatchTask< +> implements WorkflowTaskDefinition< Input, Output, Expected, @@ -346,42 +396,42 @@ export class BatchTask< Parameters, SubmissionData > { - readonly kind: typeof BATCH_TASK_KIND = BATCH_TASK_KIND; + readonly kind: typeof WORKFLOW_TASK_KIND = WORKFLOW_TASK_KIND; constructor( - readonly processor: DurableBatchProcessor< - DurableBatchTaskItem, - DurableBatchTaskResult, + readonly processor: WorkflowSubmissionProcessor< + WorkflowTaskItem, + WorkflowTaskResult, SubmissionData >, ) {} } /** - * Defines a scorer that runs through asynchronous provider batch operations. + * Defines a scorer that submits one asynchronous provider operation per case/trial. * * @experimental - The API for this class is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. */ -export class BatchScorer< +export class WorkflowScorer< Input, Output, Expected = void, Metadata extends BaseMetadata = DefaultMetadataType, SubmissionData extends JsonValue = JsonValue, -> implements DurableBatchScorer< +> implements WorkflowScorerDefinition< Input, Output, Expected, Metadata, SubmissionData > { - readonly kind: typeof BATCH_SCORER_KIND = BATCH_SCORER_KIND; + readonly kind: typeof WORKFLOW_SCORER_KIND = WORKFLOW_SCORER_KIND; readonly name: string; constructor( - readonly processor: DurableBatchProcessor< - DurableBatchScorerItem, - DurableBatchScorerResult, + readonly processor: WorkflowSubmissionProcessor< + WorkflowScorerItem, + WorkflowScorerResult, SubmissionData > & { name: string }, ) { @@ -389,7 +439,7 @@ export class BatchScorer< } } -type DurableEvaluator< +type WorkflowEvaluator< Input, Output, Expected = void, @@ -399,7 +449,9 @@ type DurableEvaluator< Evaluator, "task" | "scores" | "timeout" | "signal" | "maxConcurrency" | "update" > & { - store: DurableEvalStore; + store: WorkflowEvalStore; + /** Maximum concurrent provider callbacks per invocation. Defaults to 10. */ + maxConcurrency?: number; /** * Returns a stable case ID when a data item has neither `id` nor `upsert_id`. * The ID is shared by all trials of the same case. @@ -409,7 +461,7 @@ type DurableEvaluator< ) => string | Promise; task: | EvalTask - | DurableBatchTask< + | WorkflowTaskDefinition< Input, Output, Expected, @@ -419,24 +471,24 @@ type DurableEvaluator< >; scores?: Array< | EvalScorer - | DurableBatchScorer + | WorkflowScorerDefinition >; }; -interface DurableEvalStartOptions< +interface WorkflowEvalStartOptions< Parameters extends EvalParameters = EvalParameters, > { parameters?: InferParameters; noSendLogs?: boolean; } -type DurableBatchResult = { +type WorkflowSubmissionResult = { runId: string; - batchId?: string; + submissionId?: string; externalId?: string; }; -type DurableEvalResult = +type WorkflowEvalResult = | { status: "waiting"; runId: string; @@ -455,7 +507,7 @@ type DurableEvalResult = summary: ExperimentSummary; }; -interface DurableEvalRuntimeDefinition< +interface WorkflowEvalRuntimeDefinition< Input, Output, Expected = void, @@ -464,7 +516,7 @@ interface DurableEvalRuntimeDefinition< > { readonly projectName: string; readonly evalName: string; - readonly evaluator: DurableEvaluator< + readonly evaluator: WorkflowEvaluator< Input, Output, Expected, @@ -473,16 +525,18 @@ interface DurableEvalRuntimeDefinition< >; } -interface DurableEvalDefinition { +interface WorkflowEvalDefinition { start( - options?: DurableEvalStartOptions, - ): Promise; - status(options: { runId: string }): Promise; - poll(options: { runId: string }): Promise; - processBatchResult(result: DurableBatchResult): Promise; + options?: WorkflowEvalStartOptions, + ): Promise; + status(options: { runId: string }): Promise; + poll(options: { runId: string }): Promise; + processSubmissionResult( + result: WorkflowSubmissionResult, + ): Promise; } -type DurableCaseRecord = { +type WorkflowCaseRecord = { id: string; caseId: string; trialIndex: number; @@ -499,68 +553,67 @@ type DurableCaseRecord = { loggedClassifications: Record; }; -type DurableCaseBaseRecord = Pick< - DurableCaseRecord, +type WorkflowCaseBaseRecord = Pick< + WorkflowCaseRecord, "id" | "caseId" | "trialIndex" | "datum" | "metadata" | "tags" >; -type DurableTaskResultRecord = Pick< - DurableCaseRecord, +type WorkflowTaskResultRecord = Pick< + WorkflowCaseRecord, "output" | "metadata" | "tags" > & { taskComplete: true }; -type DurableTaskLogRecord = Pick & { +type WorkflowTaskLogRecord = Pick & { taskLogged: true; }; -type DurableBatchRecord = { +type WorkflowSubmissionRecord = { id: string; kind: "task" | "score"; scorerName?: string; - itemIds: string[]; + itemId: string; submissionData: JsonValue; externalId?: string; status: "submitted" | "complete"; + completionMode: "poll" | "webhook"; }; -type DurableRunState = { +type WorkflowRunState = { runId: string; experimentName: string; noSendLogs: boolean; parameters: JsonValue; status: "running" | "completed"; summary?: ExperimentSummary; - cases: DurableCaseRecord[]; - batches: DurableBatchRecord[]; + caseCount: number; + cases: WorkflowCaseRecord[]; + submissions: WorkflowSubmissionRecord[]; }; -type DurableRunRecord = Omit & { - caseIds: string[]; -}; +type WorkflowRunRecord = Omit; /* * Internal usage notes. Keep these out of the public README while - * defineDurableEval() is experimental. + * defineWorkflowEval() is experimental. * - * ## Durable evaluations + * ## Workflow evaluations * - * `defineDurableEval()` runs tasks and scorers through asynchronous provider - * batch APIs. - * `batchSize` splits a dataset into provider-sized sub-batches. A small external - * store connects submitted jobs with later webhook callbacks; it is required on + * `defineWorkflowEval()` runs tasks and scorers through asynchronous provider + * operations, one submission per case/trial. A small external store connects + * submitted jobs with later webhook callbacks; it is required on * the eval definition so every invocation uses the same persistence authority. * No Braintrust backend changes are required. * * Every case needs a stable `id` (or a `caseId` function). * - * For local or single-process runs, use the built-in memory store. For durable + * For local or single-process runs, use the built-in memory store. For workflow * deployments, the Redis adapter accepts any existing client with asynchronous * `get(key)` and `set(key, value)` methods. The adapter itself adds no Redis * dependency, so install and configure whichever client your application already * uses. * * The following popular clients can be passed directly to - * `DurableEvalRedisStore`: + * `WorkflowEvalRedisStore`: * * - [`redis`](https://github.com/redis/node-redis) (node-redis), including the * lower-level `@redis/client` package @@ -574,190 +627,104 @@ type DurableRunRecord = Omit & { * * ```typescript * import { createClient } from "redis"; - * import { DurableEvalRedisStore } from "braintrust"; + * import { WorkflowEvalRedisStore } from "braintrust"; * * const nodeRedis = await createClient({ url: process.env.REDIS_URL! }).connect(); - * const redisStore = new DurableEvalRedisStore({ client: nodeRedis }); + * const redisStore = new WorkflowEvalRedisStore({ client: nodeRedis }); * ``` * * ioredis: * * ```typescript * import Redis from "ioredis"; - * import { DurableEvalRedisStore } from "braintrust"; + * import { WorkflowEvalRedisStore } from "braintrust"; * * const ioRedis = new Redis(process.env.REDIS_URL!); - * const redisStore = new DurableEvalRedisStore({ client: ioRedis }); + * const redisStore = new WorkflowEvalRedisStore({ client: ioRedis }); * ``` * * Upstash: * * ```typescript * import { Redis } from "@upstash/redis"; - * import { DurableEvalRedisStore } from "braintrust"; + * import { WorkflowEvalRedisStore } from "braintrust"; * * const upstashRedis = Redis.fromEnv(); - * const redisStore = new DurableEvalRedisStore({ client: upstashRedis }); + * const redisStore = new WorkflowEvalRedisStore({ client: upstashRedis }); * ``` * * Redis entries expire after seven days by default. Set `ttlMs` in the store * options to use a different lifetime. For local testing, - * `new DurableEvalMemoryStore()` requires no external client, but is + * `new WorkflowEvalMemoryStore()` requires no external client, but is * process-local and loses its state when the process exits, so it should not be * used to reconnect webhooks across serverless invocations. * * ```typescript - * import { BatchTask, defineDurableEval } from "braintrust"; + * import { WorkflowTask, defineWorkflowEval } from "braintrust"; * - * const supportEval = defineDurableEval("Support bot", { + * const supportEval = defineWorkflowEval("Support bot", { * store: redisStore, - * data: [ - * { - * id: "password-reset", - * input: "How do I reset my password?", - * expected: "Open account settings...", - * }, - * ], - * task: new BatchTask({ - * // Each provider job contains at most 500 eval cases. - * batchSize: 500, - * - * // Submit one sub-batch and return JSON-serializable submission data. - * async submit(items, context) { - * const batch = await provider.submit({ - * idempotencyKey: context.batchId, - * metadata: { - * durableRunId: context.runId, - * durableBatchId: context.batchId, - * }, - * items, + * data: [{ id: "password-reset", input: "How do I reset my password?" }], + * task: new WorkflowTask({ + * async submit(item, { runId, submissionId }) { + * const request = await provider.submit({ + * input: item.input, + * idempotencyKey: submissionId, + * metadata: { runId, submissionId }, * }); - * return { id: batch.id }; + * return { id: request.id }; * }, - * * completion: { - * // "webhook" waits for processBatchResult(). Use "poll" with a poll() - * // callback when the provider does not send completion events. * mode: "webhook", - * getExternalId: (submissionData) => submissionData.id, + * getExternalId: (submission) => submission.id, * }, - * - * async collect(submissionData) { - * return (await provider.results(submissionData.id)).map((item) => ({ - * id: item.id, - * output: item.output, - * })); + * async collect(submission) { + * return { output: await provider.result(submission.id) }; * }, * }), - * scores: [ - * function exact({ output, expected }) { - * return output === expected ? 1 : 0; - * }, - * ], + * scores: [({ output }) => output.length > 0 ? 1 : 0], * }); * - * const result = await supportEval.start(); - * const { runId } = result; + * const { runId } = await supportEval.start(); + * await supportEval.processSubmissionResult({ runId, externalId: event.id }); * ``` * - * `start()` initializes the run, submits every ready task sub-batch, and returns. - * It never waits in a polling loop. When all task results are available, scoring - * begins. `BatchScorer` uses the same `batchSize`, `submit`, `completion`, and - * array-returning `collect` contract. + * Each submission belongs to one case/trial. `WorkflowScorer` has the same + * lifecycle, requires a `name`, and collects `{ score }` instead of `{ output }`. + * Task results may include `metadata` to merge and `tags` to replace case tags. + * Submission data and collected values must be JSON serializable. * - * ### Polling + * Use `completion: { mode: "poll", poll }` for polling providers. The callback + * receives submission data and `{ runId, submissionId }`, and returns + * `{ status: "pending" }`, `{ status: "complete" }`, or + * `{ status: "failed", error }`. Invoke `poll({ runId })` from a cron or worker; + * it checks each existing polling submission once without sleeping. Newly + * submitted work is polled on a later invocation. * - * Polling adapters report the provider's current status through `completion`: + * `processSubmissionResult()` accepts either the provider's `externalId` or the + * SDK's `submissionId`, plus `runId`. Collection callbacks must tolerate repeated + * invocation, including concurrent webhook deliveries. Provider webhook failure + * handling remains the application's responsibility. * - * ```typescript - * completion: { - * mode: "poll", - * async poll(submissionData) { - * const batch = await provider.getBatch(submissionData.id); - * if (batch.status === "completed") return { status: "complete" }; - * if (batch.status === "failed") { - * return { status: "failed", error: batch.error }; - * } - * return { status: "pending" }; - * }, - * }, - * ``` + * `start()`, `poll()`, and `processSubmissionResult()` return the current status. + * Waiting statuses include `pending: { poll, webhook }`, counting outstanding + * submissions. Completed statuses include the saved experiment summary and zero + * pending submissions. `status({ runId })` reads status without advancing work. * - * Call `poll()` from a cron, queue worker, or another short-lived invocation. It - * checks every previously submitted polling batch once, collects completed - * results, submits newly ready work, and returns without sleeping: + * Once a task result is persisted and logged, that case's scorers and classifiers + * can start even while other tasks are pending. The run completes only after all + * cases finish. Ordinary task and scorer functions are also supported. * - * ```typescript - * const result = await supportEval.poll({ - * runId, - * }); - * - * if (result.status === "waiting" && result.pending.poll > 0) { - * scheduleAnotherPoll(); - * } - * ``` - * - * `start()`, `poll()`, and `processBatchResult()` return the current eval status. - * A waiting result includes the number of submitted batches using each completion - * mode: - * - * ```typescript - * { - * status: "waiting", - * runId, - * pending: { poll: 2, webhook: 1 }, - * } - * ``` - * - * Use `status()` to read the same information without polling providers, - * collecting results, or advancing the evaluation: - * - * ```typescript - * const status = await supportEval.status({ - * runId, - * }); - * ``` - * - * Completed statuses have zero pending batches and include the saved experiment - * summary. They can be read repeatedly without logging the eval again. - * - * ### Webhook processing - * - * When the provider reports that any task or scorer batch completed, fetch and - * store its results through `processBatchResult()`: - * - * ```typescript - * app.post("/webhooks/provider", async (request, response) => { - * const event = request.body; - * const batch = await provider.getBatch(event.batchId); - * const runId = batch.metadata.durableRunId; - * - * const result = await supportEval.processBatchResult({ - * // Returned by start() and saved alongside the provider job. - * runId, - * // The provider's batch ID. The durable eval saved it from submit()'s result. - * externalId: batch.id, - * // The SDK-generated ID passed to submit(); include it in provider metadata - * // when the webhook cannot provide the external ID used by the submission data. - * batchId: batch.metadata?.durableBatchId, - * }); - * - * response.status(result.status === "waiting" ? 202 : 200).end(); - * }); - * ``` - * - * The method accepts either `externalId` or `batchId`. The stored batch locator - * identifies the task or scorer batch, whose `collect()` results are stored - * before the eval advances. Provider failure handling remains the application's - * responsibility for now. + * Provider submissions and polling use `maxConcurrency` (default 10). A failed + * provider callback is reported after independent submissions have advanced. */ /** - * Defines a durable evaluation backed by a user-provided store. + * Defines a workflow evaluation backed by a user-provided store. * * @experimental - The API for this function is not yet stabilized and may change or be removed across non-major versions. Functionality is not guaranteed. */ -export function defineDurableEval< +export function defineWorkflowEval< Input, Output, Expected = void, @@ -765,9 +732,16 @@ export function defineDurableEval< Parameters extends EvalParameters = EvalParameters, >( projectName: string, - evaluator: DurableEvaluator, -): DurableEvalDefinition { - const definition: DurableEvalRuntimeDefinition< + evaluator: WorkflowEvaluator, +): WorkflowEvalDefinition { + if ( + evaluator.maxConcurrency !== undefined && + (!Number.isInteger(evaluator.maxConcurrency) || + evaluator.maxConcurrency < 1) + ) { + throw new Error("maxConcurrency must be a positive integer"); + } + const definition: WorkflowEvalRuntimeDefinition< Input, Output, Expected, @@ -779,30 +753,30 @@ export function defineDurableEval< evaluator, }; return { - start: (options = {}) => startDurableEval(definition, options), - status: (options) => getDurableEvalStatus(definition, options), - poll: (options) => pollDurableEval(definition, options), - processBatchResult: (result) => - processDurableBatchResult(definition, result), + start: (options = {}) => startWorkflowEval(definition, options), + status: (options) => getWorkflowEvalStatus(definition, options), + poll: (options) => pollWorkflowEval(definition, options), + processSubmissionResult: (result) => + processWorkflowSubmissionResult(definition, result), }; } -async function startDurableEval< +async function startWorkflowEval< Input, Output, Expected, Metadata extends BaseMetadata, Parameters extends EvalParameters, >( - definition: DurableEvalRuntimeDefinition< + definition: WorkflowEvalRuntimeDefinition< Input, Output, Expected, Metadata, Parameters >, - options: DurableEvalStartOptions, -): Promise { + options: WorkflowEvalStartOptions, +): Promise { const store = definition.evaluator.store; const runId = newId(); const key = runKey(definition.projectName, definition.evalName, runId); @@ -830,8 +804,8 @@ async function startDurableEval< }, ); if ( - !isBatchTask(definition.evaluator.task) && - !(definition.evaluator.scores ?? []).some(isBatchScorer) + !isWorkflowTask(definition.evaluator.task) && + !(definition.evaluator.scores ?? []).some(isWorkflowScorer) ) { const result = await runEvaluator( experiment, @@ -858,42 +832,50 @@ async function startDurableEval< true, true, ); - const state: DurableRunState = { + const state: WorkflowRunState = { runId, experimentName, noSendLogs: options.noSendLogs ?? false, parameters: assertJsonValue(parameters, "eval parameters"), status: "completed", summary: result.summary, + caseCount: 0, cases: [], - batches: [], + submissions: [], }; await experiment?.flush(); await writeRunRecord(store, key, state); return currentStatus(definition, state); } - const state: DurableRunState = { + const cases = await materializeCases(definition, data, experiment); + const state: WorkflowRunState = { runId, experimentName, noSendLogs: options.noSendLogs ?? false, parameters: assertJsonValue(parameters, "eval parameters"), status: "running", - cases: await materializeCases(definition, data, experiment), - batches: [], + caseCount: cases.length, + cases, + submissions: [], }; + await writeJson( + store, + `${key}/case-ids`, + cases.map(({ id }) => id), + ); await writeCaseBaseRecords(store, key, state.cases); await writeRunRecord(store, key, state); - return advanceDurableEval(definition, state, store, key, experiment); + return advanceWorkflowEval(definition, state, store, key, experiment); } -async function getDurableEvalStatus< +async function getWorkflowEvalStatus< Input, Output, Expected, Metadata extends BaseMetadata, Parameters extends EvalParameters, >( - definition: DurableEvalRuntimeDefinition< + definition: WorkflowEvalRuntimeDefinition< Input, Output, Expected, @@ -901,75 +883,96 @@ async function getDurableEvalStatus< Parameters >, options: { runId: string }, -): Promise { +): Promise { const store = definition.evaluator.store; - const state = await readRunState( - definition, - store, - runKey(definition.projectName, definition.evalName, options.runId), + const key = runKey( + definition.projectName, + definition.evalName, + options.runId, ); - if (!state) throw new Error(`Durable eval run ${options.runId} is missing`); + const state = await readJson(store, key); + if (!state) throw new Error(`Workflow eval run ${options.runId} is missing`); return currentStatus(definition, state); } -async function processDurableBatchResult< +async function processWorkflowSubmissionResult< Input, Output, Expected, Metadata extends BaseMetadata, Parameters extends EvalParameters, >( - definition: DurableEvalRuntimeDefinition< + definition: WorkflowEvalRuntimeDefinition< Input, Output, Expected, Metadata, Parameters >, - result: DurableBatchResult, -): Promise { - if (!result.batchId && !result.externalId) { - throw new Error("Batch results require batchId or externalId"); + result: WorkflowSubmissionResult, +): Promise { + if (!result.submissionId && !result.externalId) { + throw new Error("Submission results require submissionId or externalId"); } const store = definition.evaluator.store; const key = runKey(definition.projectName, definition.evalName, result.runId); - const state = await readRunState(definition, store, key); - if (!state) throw new Error(`Durable eval run ${result.runId} is missing`); - const byBatch = result.batchId - ? state.batches.find((candidate) => candidate.id === result.batchId) + const run = await readJson(store, key); + if (!run) throw new Error(`Workflow eval run ${result.runId} is missing`); + const externalSubmissionId = result.externalId + ? await readJson( + store, + `${key}/external/${encodedKeyPart(result.externalId)}`, + ) : undefined; - const byExternal = result.externalId - ? state.batches.find( - (candidate) => candidate.externalId === result.externalId, + if ( + result.submissionId && + externalSubmissionId && + result.submissionId !== externalSubmissionId + ) { + throw new Error( + "submissionId and externalId identify different submissions", + ); + } + const submissionId = result.submissionId ?? externalSubmissionId; + const submission = submissionId + ? await readJson( + store, + submissionRecordKey(key, submissionId), ) : undefined; - if (byBatch && byExternal && byBatch.id !== byExternal.id) { - throw new Error("batchId and externalId identify different batches"); + if (!submission) throw new Error("No submission matches this result"); + if (result.externalId && submission.externalId !== result.externalId) { + throw new Error( + "submissionId and externalId identify different submissions", + ); } - const batch = byBatch ?? byExternal; - if (!batch) throw new Error("No submitted batch matches this result"); - if (batch.status !== "complete") { - const records = await collectBatch(definition, state, batch); - batch.status = "complete"; - await writeCaseRecords(store, key, records); - await writeBatchRecords(store, key, [batch]); + if (run.status === "completed") return currentStatus(definition, run); + const state = (await readRunState(definition, store, key, [ + submission.itemId, + ]))!; + if (submission.status !== "complete") { + const record = await collectSubmission(definition, state, submission); + await writeCaseRecords(store, key, [record]); + submission.status = "complete"; } - return advanceDurableEval( + // Repeat progress writes on replay to recover an interrupted persistence step. + await writeSubmissionRecords(store, key, [submission]); + return advanceWorkflowEval( definition, - (await readRunState(definition, store, key))!, + (await readRunState(definition, store, key, [submission.itemId]))!, store, key, ); } -async function pollDurableEval< +async function pollWorkflowEval< Input, Output, Expected, Metadata extends BaseMetadata, Parameters extends EvalParameters, >( - definition: DurableEvalRuntimeDefinition< + definition: WorkflowEvalRuntimeDefinition< Input, Output, Expected, @@ -977,7 +980,7 @@ async function pollDurableEval< Parameters >, options: { runId: string }, -): Promise { +): Promise { const store = definition.evaluator.store; const key = runKey( definition.projectName, @@ -985,55 +988,61 @@ async function pollDurableEval< options.runId, ); const state = await readRunState(definition, store, key); - if (!state) throw new Error(`Durable eval run ${options.runId} is missing`); + if (!state) throw new Error(`Workflow eval run ${options.runId} is missing`); - const batches = state.batches.filter((batch) => { - if (batch.status === "complete") return false; + const submissions = state.submissions.filter((submission) => { + if (submission.status === "complete") return false; return ( - processorForStage(definition, batch.kind, batch.scorerName).completion - .mode === "poll" + processorForStage(definition, submission.kind, submission.scorerName) + .completion.mode === "poll" ); }); - const results = await Promise.all( - batches.map(async (batch) => ({ - batch, - result: await ( - processorForStage(definition, batch.kind, batch.scorerName) - .completion as Extract< - DurableBatchCompletion, - { mode: "poll" } - > - ).poll(batch.submissionData, { - runId: state.runId, - batchId: batch.id, - }), - })), - ); - const changedCases = new Map(); - const changedBatches: DurableBatchRecord[] = []; - for (const { batch, result } of results) { + const workers = queue(async (submission: WorkflowSubmissionRecord) => { + const completion = processorForStage( + definition, + submission.kind, + submission.scorerName, + ).completion; + if (completion.mode !== "poll") return; + const result = await completion.poll(submission.submissionData, { + runId: state.runId, + submissionId: submission.id, + }); if (result.status === "failed") throw asError(result.error); - if (result.status !== "complete") continue; - for (const record of await collectBatch(definition, state, batch)) { - changedCases.set(record.id, record); + if (result.status === "complete") { + const record = await collectSubmission(definition, state, submission); + await writeCaseRecords(store, key, [record]); + submission.status = "complete"; + await writeSubmissionRecords(store, key, [submission]); } - batch.status = "complete"; - changedBatches.push(batch); - } - if (changedBatches.length > 0) { - await writeCaseRecords(store, key, [...changedCases.values()]); - await writeBatchRecords(store, key, changedBatches); + }, definition.evaluator.maxConcurrency ?? 10); + const results = await Promise.allSettled( + submissions.map((submission) => workers.pushAsync(submission)), + ); + // Advance persisted work even when an unrelated provider callback failed. + let status: WorkflowEvalResult | undefined; + const errors = results.flatMap((result) => + result.status === "rejected" ? [asError(result.reason)] : [], + ); + try { + status = await advanceWorkflowEval( + definition, + (await readRunState(definition, store, key))!, + store, + key, + ); + } catch (error) { + errors.push(asError(error)); } - const currentState = - changedBatches.length > 0 - ? (await readRunState(definition, store, key))! - : state; - return advanceDurableEval(definition, currentState, store, key); + if (errors.length === 1) throw errors[0]; + if (errors.length > 1) + throw new AggregateError(errors, "Workflow submission callbacks failed"); + return status!; } -async function openDurableExperiment( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, +async function openWorkflowExperiment( + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunState, ) { const data: EvalCase[] = []; return await _internalInitEvaluatorExperiment( @@ -1054,80 +1063,82 @@ async function openDurableExperiment( ); } -async function advanceDurableEval< +async function advanceWorkflowEval< Input, Output, Expected, Metadata extends BaseMetadata, Parameters extends EvalParameters, >( - definition: DurableEvalRuntimeDefinition< + definition: WorkflowEvalRuntimeDefinition< Input, Output, Expected, Metadata, Parameters >, - state: DurableRunState, - store: DurableEvalStore, + state: WorkflowRunState, + store: WorkflowEvalStore, key: string, existingExperiment?: Experiment | null, -): Promise { +): Promise { if (state.status === "completed") return currentStatus(definition, state); const experiment = existingExperiment === undefined - ? await openDurableExperiment(definition, state) + ? await openWorkflowExperiment(definition, state) : existingExperiment; + const caseIds = state.cases.map(({ id }) => id); await runTaskStage(definition, state, store, key, experiment); await logCompletedTasks(definition, state, store, key, experiment); - state = (await readRunState(definition, store, key)) ?? state; - if ( - state.cases.some((record) => !record.taskComplete || !record.taskLogged) - ) { - return currentStatus(definition, state); - } - + state = (await readRunState(definition, store, key, caseIds)) ?? state; await runScoreStages(definition, state, store, key, experiment); - state = (await readRunState(definition, store, key)) ?? state; + state = (await readRunState(definition, store, key, caseIds)) ?? state; const scorerNames = resolveScorers(definition.evaluator.scores ?? []).map( ({ name }) => name, ); const classifierNames = (definition.evaluator.classifiers ?? []).map( classifierName, ); - if ( - state.cases.some( - (record) => - scorerNames.some( + await Promise.all( + state.cases.map(async (record) => { + if ( + record.taskComplete && + record.taskLogged && + scorerNames.every( (name) => - !Object.hasOwn(record.scores, name) || - !Object.hasOwn(record.loggedScores, name), - ) || - classifierNames.some( - (name) => !Object.hasOwn(record.loggedClassifications, name), - ), - ) - ) { + Object.hasOwn(record.scores, name) && + Object.hasOwn(record.loggedScores, name), + ) && + classifierNames.every((name) => + Object.hasOwn(record.loggedClassifications, name), + ) + ) { + await store.addToSet(`${key}/progress/cases`, record.id); + } + }), + ); + if ((await store.getSetSize(`${key}/progress/cases`)) !== state.caseCount) { return currentStatus(definition, state); } if (!(await claimAction(store, key, "finish"))) { - const latest = await readRunState(definition, store, key); + const latest = await readJson(store, key); return currentStatus(definition, latest ?? state); } + state = (await readRunState(definition, store, key))!; state.summary = await finishExperiment(definition, state, experiment); state.status = "completed"; await writeRunRecord(store, key, state); return currentStatus(definition, state); } -function currentStatus( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, -): DurableEvalResult { +async function currentStatus( + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunRecord, +): Promise { if (state.status === "completed") { if (!state.summary) { - throw new Error(`Durable eval run ${state.runId} has no saved summary`); + throw new Error(`Workflow eval run ${state.runId} has no saved summary`); } return { status: "completed", @@ -1136,21 +1147,28 @@ function currentStatus( summary: state.summary, }; } + const key = runKey(definition.projectName, definition.evalName, state.runId); + const store = definition.evaluator.store; const pending = { poll: 0, webhook: 0 }; - for (const batch of state.batches) { - if (batch.status === "complete") continue; - pending[ - processorForStage(definition, batch.kind, batch.scorerName).completion - .mode - ]++; - } + await Promise.all( + (["poll", "webhook"] as const).map(async (mode) => { + // Read completions first so an in-flight submission cannot yield a negative count. + const completed = await store.getSetSize( + `${key}/progress/${mode}/complete`, + ); + const submitted = await store.getSetSize( + `${key}/progress/${mode}/submitted`, + ); + pending[mode] = submitted - completed; + }), + ); return { status: "waiting", runId: state.runId, pending }; } async function startCaseRoot( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, - record: DurableCaseRecord, + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunState, + record: WorkflowCaseRecord, experiment: Experiment | null, ): Promise { if (!experiment) return NOOP_SPAN; @@ -1174,9 +1192,9 @@ async function startCaseRoot( } async function logTaskResult( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, - record: DurableCaseRecord, + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunState, + record: WorkflowCaseRecord, experiment: Experiment | null, task?: EvalTask, ) { @@ -1220,7 +1238,7 @@ async function logTaskResult( expected: "expected" in datum ? datum.expected : undefined, metadata: { ...(record.metadata as Record), - durable_eval: { + workflow_eval: { run_id: state.runId, case_id: record.caseId, trial_index: record.trialIndex, @@ -1239,13 +1257,13 @@ async function logTaskResult( } async function logCompletedTasks( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, - store: DurableEvalStore, + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunState, + store: WorkflowEvalStore, key: string, experiment: Experiment | null, ) { - const changed: DurableCaseRecord[] = []; + const changed: WorkflowCaseRecord[] = []; for (const record of state.cases) { if (!record.taskComplete || record.taskLogged) continue; if (!(await claimAction(store, key, "task-log", record.id))) continue; @@ -1265,20 +1283,20 @@ async function runTaskStage< Metadata extends BaseMetadata, Parameters extends EvalParameters, >( - definition: DurableEvalRuntimeDefinition< + definition: WorkflowEvalRuntimeDefinition< Input, Output, Expected, Metadata, Parameters >, - state: DurableRunState, - store: DurableEvalStore, + state: WorkflowRunState, + store: WorkflowEvalStore, key: string, experiment: Experiment | null, ) { - if (isBatchTask(definition.evaluator.task)) { - await ensureBatches(definition, state, store, key, "task"); + if (isWorkflowTask(definition.evaluator.task)) { + await ensureSubmissions(definition, state, store, key, "task"); return; } @@ -1289,7 +1307,7 @@ async function runTaskStage< Metadata, Parameters >; - const changed: DurableCaseRecord[] = []; + const changed: WorkflowCaseRecord[] = []; for (const record of state.cases) { if (record.taskComplete) continue; if (!(await claimAction(store, key, "task", record.id))) continue; @@ -1309,20 +1327,20 @@ async function runScoreStages< Metadata extends BaseMetadata, Parameters extends EvalParameters, >( - definition: DurableEvalRuntimeDefinition< + definition: WorkflowEvalRuntimeDefinition< Input, Output, Expected, Metadata, Parameters >, - state: DurableRunState, - store: DurableEvalStore, + state: WorkflowRunState, + store: WorkflowEvalStore, key: string, experiment: Experiment | null, ) { const scorers = resolveScorers(definition.evaluator.scores ?? []); - const changed = new Map(); + const changed = new Map(); const persistChangedCases = async () => { if (changed.size === 0) return; await experiment?.flush(); @@ -1330,8 +1348,9 @@ async function runScoreStages< changed.clear(); }; for (const { name, scorer } of scorers) { - if (isBatchScorer(scorer)) { + if (isWorkflowScorer(scorer)) { for (const record of state.cases) { + if (!record.taskComplete || !record.taskLogged) continue; if ( Object.hasOwn(record.scores, name) && !Object.hasOwn(record.loggedScores, name) @@ -1350,10 +1369,11 @@ async function runScoreStages< } } await persistChangedCases(); - await ensureBatches(definition, state, store, key, "score", name); + await ensureSubmissions(definition, state, store, key, "score", name); continue; } for (const record of state.cases) { + if (!record.taskComplete || !record.taskLogged) continue; if (Object.hasOwn(record.loggedScores, name)) continue; if (!(await claimAction(store, key, "score", record.id, name))) continue; await evaluateAndLogScore( @@ -1373,6 +1393,7 @@ async function runScoreStages< ).entries()) { const name = classifierName(classifier, index); for (const record of state.cases) { + if (!record.taskComplete || !record.taskLogged) continue; if (Object.hasOwn(record.loggedClassifications, name)) continue; if (!(await claimAction(store, key, "classification", record.id, name))) { continue; @@ -1391,23 +1412,24 @@ async function runScoreStages< await persistChangedCases(); } -function scorerArgs(record: DurableCaseRecord) { +function scorerArgs(record: WorkflowCaseRecord) { const datum = record.datum as EvalCase; return { ...datum, metadata: record.metadata, + tags: record.tags, output: record.output, } as EvalScorerArgs; } function resumeCaseRoot( - definition: DurableEvalRuntimeDefinition, - record: DurableCaseRecord, + definition: WorkflowEvalRuntimeDefinition, + record: WorkflowCaseRecord, experiment: Experiment | null, ) { if (!experiment) return NOOP_SPAN; if (!record.rootSpan) { - throw new Error(`Durable eval case ${record.caseId} has no root span`); + throw new Error(`Workflow eval case ${record.caseId} has no root span`); } return _internalResumeSpan({ exported: record.rootSpan, @@ -1416,9 +1438,9 @@ function resumeCaseRoot( } async function evaluateAndLogScore( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, - record: DurableCaseRecord, + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunState, + record: WorkflowCaseRecord, name: string, experiment: Experiment | null, scorer?: EvalScorer, @@ -1466,9 +1488,9 @@ async function evaluateAndLogScore( } async function evaluateAndLogClassification( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, - record: DurableCaseRecord, + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunState, + record: WorkflowCaseRecord, name: string, classifier: EvalClassifier, experiment: Experiment | null, @@ -1514,136 +1536,133 @@ async function evaluateAndLogClassification( } } -async function ensureBatches( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, - store: DurableEvalStore, +async function ensureSubmissions( + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunState, + store: WorkflowEvalStore, key: string, kind: "task" | "score", scorerName?: string, ) { const processor = processorForStage(definition, kind, scorerName); - const plans = plannedBatches( + const plans = plannedSubmissions( definition, state.runId, state.cases.map(({ id }) => id), ); const casesById = new Map(state.cases.map((record) => [record.id, record])); - for (const plan of plans) { - if (plan.kind !== kind || plan.scorerName !== scorerName) continue; - if (state.batches.some(({ id }) => id === plan.id)) continue; - const records = plan.itemIds.map((id) => casesById.get(id)!); - const ready = records.every((record) => + const existingIds = new Set(state.submissions.map(({ id }) => id)); + const workers = queue(async (plan: WorkflowSubmissionPlan) => { + if (plan.kind !== kind || plan.scorerName !== scorerName) return; + if (existingIds.has(plan.id)) return; + const record = casesById.get(plan.itemId)!; + const ready = kind === "task" ? !record.taskComplete - : !Object.hasOwn(record.scores, scorerName!), - ); - if (!ready) continue; - const batchId = plan.id; + : record.taskComplete && + record.taskLogged && + !Object.hasOwn(record.scores, scorerName!); + if (!ready) return; + const submissionId = plan.id; const claim = await store.getOrSet( - claimRecordKey(key, "batch", batchId), - encoder.encode(batchId), + claimRecordKey(key, "submission", submissionId), + encoder.encode(submissionId), ); - if (!claim.created) continue; - const context = { runId: state.runId, batchId }; - const items = records.map((record) => + if (!claim.created) return; + const context = { runId: state.runId, submissionId }; + const item = kind === "task" - ? taskBatchItem(record, state.parameters) - : scorerBatchItem(record), - ); + ? taskSubmissionItem(record, state.parameters) + : scorerSubmissionItem(record); const submissionData = assertJsonValue( - await processor.submit(items, context), - `submission data for batch ${batchId}`, + await processor.submit(item, context), + `submission data for submission ${submissionId}`, ); const externalId = processor.completion.mode === "webhook" ? processor.completion.getExternalId(submissionData, context) : undefined; if (externalId !== undefined && !externalId.trim()) { - throw new Error(`Batch ${batchId} produced an empty externalId`); + throw new Error( + `Submission ${submissionId} produced an empty externalId`, + ); } - const batch: DurableBatchRecord = { - id: batchId, + const submission: WorkflowSubmissionRecord = { + id: submissionId, kind, scorerName, - itemIds: records.map((record) => record.id), + itemId: record.id, submissionData, externalId, status: "submitted", + completionMode: processor.completion.mode, }; - state.batches.push(batch); - await writeBatchRecords(store, key, [batch]); - } + state.submissions.push(submission); + await writeSubmissionRecords(store, key, [submission]); + }, definition.evaluator.maxConcurrency ?? 10); + const results = await Promise.allSettled( + plans.map((plan) => workers.pushAsync(plan)), + ); + const errors = results.flatMap((result) => + result.status === "rejected" ? [asError(result.reason)] : [], + ); + if (errors.length === 1) throw errors[0]; + if (errors.length > 1) + throw new AggregateError(errors, "Workflow submissions failed"); } -async function collectBatch( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, - batch: DurableBatchRecord, +async function collectSubmission( + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunState, + submission: WorkflowSubmissionRecord, ) { - const processor = processorForStage(definition, batch.kind, batch.scorerName); - const context = { runId: state.runId, batchId: batch.id }; - const results = await processor.collect(batch.submissionData, context); - if (!Array.isArray(results)) { - throw new Error(`collect for batch ${batch.id} must return an array`); + const processor = processorForStage( + definition, + submission.kind, + submission.scorerName, + ); + const context = { runId: state.runId, submissionId: submission.id }; + const result = await processor.collect(submission.submissionData, context); + if (typeof result !== "object" || result === null || Array.isArray(result)) { + throw new Error( + `collect for submission ${submission.id} must return a result object`, + ); } - const expectedIds = new Set(batch.itemIds); - const seen = new Set(); - const records: DurableCaseRecord[] = []; - for (const result of results) { - const id = resultItemId(result); - if (!expectedIds.has(id)) { - throw new Error(`Batch ${batch.id} returned unknown item ${id}`); - } - if (seen.has(id)) { - throw new Error(`Batch ${batch.id} returned item ${id} more than once`); - } - seen.add(id); - const record = state.cases.find((candidate) => candidate.id === id)!; - records.push(record); - if (batch.kind === "task") { - record.output = assertJsonValue( - result.output, - `task output for item ${id}`, - ); - if ("metadata" in result && result.metadata !== undefined) { - record.metadata = assertJsonValue( - { - ...(record.metadata as Record), - ...result.metadata, - }, - `metadata for ${id}`, - ); - } - if ("tags" in result && result.tags !== undefined) - record.tags = result.tags; - record.taskComplete = true; - } else { - record.scores[batch.scorerName!] = assertJsonValue( - result.score, - `score output for item ${id}`, + const record = state.cases.find( + (candidate) => candidate.id === submission.itemId, + )!; + if (submission.kind === "task") { + record.output = assertJsonValue( + result.output, + `task output for item ${record.id}`, + ); + if (result.metadata !== undefined) { + record.metadata = assertJsonValue( + { ...(record.metadata as Record), ...result.metadata }, + `metadata for ${record.id}`, ); } - } - const missing = batch.itemIds.filter((id) => !seen.has(id)); - if (missing.length > 0) { - throw new Error( - `Batch ${batch.id} did not return results for: ${missing.join(", ")}`, + if (result.tags !== undefined) record.tags = result.tags; + record.taskComplete = true; + } else { + record.scores[submission.scorerName!] = assertJsonValue( + result.score, + `score output for item ${record.id}`, ); } - return records; + return record; } function processorForStage( - definition: DurableEvalRuntimeDefinition, + definition: WorkflowEvalRuntimeDefinition, kind: "task" | "score", scorerName?: string, -): DurableBatchProcessor { +): WorkflowSubmissionProcessor { if (kind === "task") { - if (!isBatchTask(definition.evaluator.task)) { - throw new Error("Definition no longer contains the batch task"); + if (!isWorkflowTask(definition.evaluator.task)) { + throw new Error("Definition no longer contains the submission task"); } - return definition.evaluator.task.processor as DurableBatchProcessor< + return definition.evaluator.task.processor as WorkflowSubmissionProcessor< any, any, JsonValue @@ -1652,10 +1671,10 @@ function processorForStage( const scorer = resolveScorers(definition.evaluator.scores ?? []).find( ({ name }) => name === scorerName, )?.scorer; - if (!isBatchScorer(scorer)) { + if (!isWorkflowScorer(scorer)) { throw new Error(`Definition no longer contains scorer ${scorerName}`); } - return scorer.processor as DurableBatchProcessor; + return scorer.processor as WorkflowSubmissionProcessor; } async function materializeCases< @@ -1665,7 +1684,7 @@ async function materializeCases< Metadata extends BaseMetadata, Parameters extends EvalParameters, >( - definition: DurableEvalRuntimeDefinition< + definition: WorkflowEvalRuntimeDefinition< Input, Output, Expected, @@ -1674,7 +1693,7 @@ async function materializeCases< >, data: Evaluator["data"], experiment: Experiment | null, -): Promise { +): Promise { const evaluator = definition.evaluator; const iterable = await _internalResolveEvaluatorData( { @@ -1685,7 +1704,7 @@ async function materializeCases< }, experiment, ); - const records: DurableCaseRecord[] = []; + const records: WorkflowCaseRecord[] = []; const seen = new Set(); for await (const datum of iterable) { const caseId = @@ -1696,15 +1715,15 @@ async function materializeCases< : undefined); if (!caseId) { throw new Error( - "Every durable eval case requires id, upsert_id, or caseId", + "Every workflow eval case requires id, upsert_id, or caseId", ); } if (seen.has(caseId)) - throw new Error(`Duplicate durable eval case id: ${caseId}`); + throw new Error(`Duplicate workflow eval case id: ${caseId}`); seen.add(caseId); const trialCount = datum.trialCount ?? evaluator.trialCount ?? 1; if (!Number.isInteger(trialCount) || trialCount < 1) { - throw new Error(`Invalid trialCount for durable eval case ${caseId}`); + throw new Error(`Invalid trialCount for workflow eval case ${caseId}`); } for (let trialIndex = 0; trialIndex < trialCount; trialIndex++) { records.push({ @@ -1729,7 +1748,7 @@ async function materializeCases< return records; } -function taskBatchItem(record: DurableCaseRecord, parameters: JsonValue) { +function taskSubmissionItem(record: WorkflowCaseRecord, parameters: JsonValue) { const datum = record.datum as EvalCase; return { id: record.id, @@ -1742,7 +1761,7 @@ function taskBatchItem(record: DurableCaseRecord, parameters: JsonValue) { }; } -function scorerBatchItem(record: DurableCaseRecord) { +function scorerSubmissionItem(record: WorkflowCaseRecord) { const datum = record.datum as EvalCase; return { id: record.id, @@ -1756,8 +1775,8 @@ function scorerBatchItem(record: DurableCaseRecord) { } async function finishExperiment( - definition: DurableEvalRuntimeDefinition, - state: DurableRunState, + definition: WorkflowEvalRuntimeDefinition, + state: WorkflowRunState, experiment: Experiment | null, ) { const scorerNames = resolveScorers(definition.evaluator.scores ?? []).map( @@ -1821,11 +1840,11 @@ async function finishExperiment( function resolveScorers( scorers: Array< | EvalScorer - | DurableBatchScorer + | WorkflowScorerDefinition >, ) { return scorers.map((scorer, index) => ({ - name: isBatchScorer(scorer) + name: isWorkflowScorer(scorer) ? scorer.name : scorer.name || `scorer_${index}`, scorer, @@ -1833,7 +1852,7 @@ function resolveScorers( } function runKey(projectName: string, evalName: string, runId: string) { - return `durable-eval/v1/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; + return `workflow-eval/v1/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; } function encodedKeyPart(value: string) { @@ -1860,8 +1879,8 @@ function caseRecordKey( return `${key}/cases/${encodedKeyPart(caseId)}/${kind}${suffix ? `/${suffix}` : ""}`; } -function batchRecordKey(key: string, batchId: string) { - return `${key}/batches/${encodedKeyPart(batchId)}`; +function submissionRecordKey(key: string, submissionId: string) { + return `${key}/submissions/${encodedKeyPart(submissionId)}`; } function claimRecordKey(key: string, kind: string, ...parts: string[]) { @@ -1870,7 +1889,7 @@ function claimRecordKey(key: string, kind: string, ...parts: string[]) { } async function claimAction( - store: DurableEvalStore, + store: WorkflowEvalStore, key: string, kind: string, ...parts: string[] @@ -1883,44 +1902,33 @@ async function claimAction( ).created; } -type DurableBatchPlan = Omit< - DurableBatchRecord, - "submissionData" | "externalId" | "status" +type WorkflowSubmissionPlan = Omit< + WorkflowSubmissionRecord, + "submissionData" | "externalId" | "status" | "completionMode" >; -function plannedBatches( - definition: DurableEvalRuntimeDefinition, +function plannedSubmissions( + definition: WorkflowEvalRuntimeDefinition, runId: string, caseIds: string[], ) { const stages: Array<{ kind: "task" | "score"; scorerName?: string }> = []; - if (isBatchTask(definition.evaluator.task)) stages.push({ kind: "task" }); + if (isWorkflowTask(definition.evaluator.task)) stages.push({ kind: "task" }); for (const { name, scorer } of resolveScorers( definition.evaluator.scores ?? [], )) { - if (isBatchScorer(scorer)) { + if (isWorkflowScorer(scorer)) { stages.push({ kind: "score", scorerName: name }); } } - const plans: DurableBatchPlan[] = []; + const plans: WorkflowSubmissionPlan[] = []; for (const { kind, scorerName } of stages) { - const batchSize = - processorForStage(definition, kind, scorerName).batchSize ?? - DEFAULT_BATCH_SIZE; - if (!Number.isInteger(batchSize) || batchSize < 1) { - throw new Error( - `Invalid batchSize for ${scorerName ?? "task"}: ${batchSize}`, - ); - } - for (let offset = 0; offset < caseIds.length; offset += batchSize) { - const itemIds = caseIds.slice(offset, offset + batchSize); + for (const itemId of caseIds) { plans.push({ - id: deterministicId( - stableStringify([runId, kind, scorerName, itemIds]), - ), + id: deterministicId(stableStringify([runId, kind, scorerName, itemId])), kind, scorerName, - itemIds, + itemId, }); } } @@ -1928,8 +1936,8 @@ function plannedBatches( } async function readCaseRecord( - definition: DurableEvalRuntimeDefinition, - store: DurableEvalStore, + definition: WorkflowEvalRuntimeDefinition, + store: WorkflowEvalStore, key: string, id: string, ) { @@ -1946,9 +1954,9 @@ async function readCaseRecord( classificationValues, classificationLogValues, ] = await Promise.all([ - readJson(store, caseRecordKey(key, id, "base")), - readJson(store, caseRecordKey(key, id, "task")), - readJson(store, caseRecordKey(key, id, "task-log")), + readJson(store, caseRecordKey(key, id, "base")), + readJson(store, caseRecordKey(key, id, "task")), + readJson(store, caseRecordKey(key, id, "task-log")), Promise.all( scorers.map(async ({ name }) => ({ name, @@ -1986,7 +1994,7 @@ async function readCaseRecord( })), ), ]); - if (!base) throw new Error(`Durable eval case ${id} is missing`); + if (!base) throw new Error(`Workflow eval case ${id} is missing`); const scores: Record = Object.create(null); for (const { name, value } of scoreValues) { if (value !== undefined) scores[name] = value; @@ -2015,50 +2023,57 @@ async function readCaseRecord( loggedScores, classifications, loggedClassifications, - } satisfies DurableCaseRecord; + } satisfies WorkflowCaseRecord; } async function readRunState( - definition: DurableEvalRuntimeDefinition, - store: DurableEvalStore, + definition: WorkflowEvalRuntimeDefinition, + store: WorkflowEvalStore, key: string, + caseIds?: string[], ) { - const record = await readJson(store, key); + const record = await readJson(store, key); if (!record) return undefined; - const plans = plannedBatches(definition, record.runId, record.caseIds); - const [cases, batchRecords] = await Promise.all([ + const selectedIds = + caseIds ?? + (record.caseCount === 0 + ? [] + : await readJson(store, `${key}/case-ids`)); + if (!selectedIds) + throw new Error(`Workflow eval run ${record.runId} has no case index`); + const plans = plannedSubmissions(definition, record.runId, selectedIds); + const [cases, submissionRecords] = await Promise.all([ Promise.all( - record.caseIds.map((id) => readCaseRecord(definition, store, key, id)), + selectedIds.map((id) => readCaseRecord(definition, store, key, id)), ), Promise.all( plans.map(async ({ id }) => { - return readJson(store, batchRecordKey(key, id)); + return readJson( + store, + submissionRecordKey(key, id), + ); }), ), ]); - const batches = batchRecords.filter( - (value): value is DurableBatchRecord => value !== undefined, + const submissions = submissionRecords.filter( + (value): value is WorkflowSubmissionRecord => value !== undefined, ); - const { caseIds: _caseIds, ...state } = record; - return { ...state, cases, batches }; + return { ...record, cases, submissions }; } async function writeRunRecord( - store: DurableEvalStore, + store: WorkflowEvalStore, key: string, - state: DurableRunState, + state: WorkflowRunState, ) { - const { cases, batches: _batches, ...record } = state; - await writeJson(store, key, { - ...record, - caseIds: cases.map(({ id }) => id), - } satisfies DurableRunRecord); + const { cases: _cases, submissions: _submissions, ...record } = state; + await writeJson(store, key, record); } async function writeCaseBaseRecords( - store: DurableEvalStore, + store: WorkflowEvalStore, key: string, - records: DurableCaseRecord[], + records: WorkflowCaseRecord[], ) { await Promise.all( records.map(({ id, caseId, trialIndex, datum, metadata, tags }) => @@ -2069,15 +2084,15 @@ async function writeCaseBaseRecords( datum, metadata, tags, - } satisfies DurableCaseBaseRecord), + } satisfies WorkflowCaseBaseRecord), ), ); } async function writeCaseRecords( - store: DurableEvalStore, + store: WorkflowEvalStore, key: string, - records: DurableCaseRecord[], + records: WorkflowCaseRecord[], ) { const writes: Promise[] = []; for (const record of records) { @@ -2088,7 +2103,7 @@ async function writeCaseRecords( metadata: record.metadata, tags: record.tags, taskComplete: true, - } satisfies DurableTaskResultRecord), + } satisfies WorkflowTaskResultRecord), ); } if (record.taskLogged) { @@ -2096,7 +2111,7 @@ async function writeCaseRecords( writeJson(store, caseRecordKey(key, record.id, "task-log"), { rootSpan: record.rootSpan, taskLogged: true, - } satisfies DurableTaskLogRecord), + } satisfies WorkflowTaskLogRecord), ); } for (const [name, value] of Object.entries(record.scores)) { @@ -2135,15 +2150,31 @@ async function writeCaseRecords( await Promise.all(writes); } -async function writeBatchRecords( - store: DurableEvalStore, +async function writeSubmissionRecords( + store: WorkflowEvalStore, key: string, - records: DurableBatchRecord[], + records: WorkflowSubmissionRecord[], ) { await Promise.all( - records.map((record) => - writeJson(store, batchRecordKey(key, record.id), record), - ), + records.map(async (record) => { + if (record.externalId !== undefined) { + const locator = await store.getOrSet( + `${key}/external/${encodedKeyPart(record.externalId)}`, + encoder.encode(JSON.stringify(record.id)), + ); + if (JSON.parse(decoder.decode(locator.value)) !== record.id) { + throw new Error(`Duplicate externalId: ${record.externalId}`); + } + } + const progressKey = `${key}/progress/${record.completionMode}`; + await store.addToSet(`${progressKey}/submitted`, record.id); + if (record.status === "complete") { + await store.addToSet(`${progressKey}/complete`, record.id); + } + // Mark completion only after progress is saved, so polling can retry an + // interrupted progress update instead of permanently skipping it. + await writeJson(store, submissionRecordKey(key, record.id), record); + }), ); } @@ -2164,12 +2195,16 @@ function contentVersion(value: Uint8Array) { return (hash >>> 0).toString(16).padStart(8, "0"); } -async function readJson(store: DurableEvalStore, key: string) { +async function readJson(store: WorkflowEvalStore, key: string) { const value = await store.read(key); return value ? (JSON.parse(decoder.decode(value)) as T) : undefined; } -async function writeJson(store: DurableEvalStore, key: string, value: unknown) { +async function writeJson( + store: WorkflowEvalStore, + key: string, + value: unknown, +) { await store.write(key, encoder.encode(stableStringify(value))); } @@ -2197,37 +2232,25 @@ function assertJsonValue(value: unknown, label: string): JsonValue { } } -function resultItemId(value: unknown) { - if ( - typeof value !== "object" || - value === null || - !("id" in value) || - typeof value.id !== "string" - ) { - throw new Error("Batch results must contain a string id"); - } - return value.id; -} - -function isBatchTask( +function isWorkflowTask( value: unknown, -): value is DurableBatchTask { +): value is WorkflowTaskDefinition { return ( typeof value === "object" && value !== null && "kind" in value && - value.kind === BATCH_TASK_KIND + value.kind === WORKFLOW_TASK_KIND ); } -function isBatchScorer( +function isWorkflowScorer( value: unknown, -): value is DurableBatchScorer { +): value is WorkflowScorerDefinition { return ( typeof value === "object" && value !== null && "kind" in value && - value.kind === BATCH_SCORER_KIND + value.kind === WORKFLOW_SCORER_KIND ); }