From fe4c66137b8de20199f14c76cf543ee3c18ebbe8 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 27 Jul 2026 12:24:05 +0000 Subject: [PATCH 01/13] wip feat: Add batch/durable evals api --- .gitignore | 1 + .../durable-eval-webhook/scenario.test.ts | 33 + .../durable-eval-webhook/scenario.ts | 106 + js/README.md | 172 + js/src/cli/functions/load-module.ts | 1 + js/src/cli/index.ts | 182 +- js/src/cli/types.ts | 7 + js/src/cli/util/types.ts | 8 + js/src/durable-eval.test.ts | 876 +++++ js/src/durable-eval.ts | 3387 +++++++++++++++++ js/src/exports.ts | 38 + js/src/framework.ts | 9 + js/src/isomorph.ts | 3 +- js/src/node/config.ts | 1 + 14 files changed, 4822 insertions(+), 2 deletions(-) create mode 100644 e2e/scenarios/durable-eval-webhook/scenario.test.ts create mode 100644 e2e/scenarios/durable-eval-webhook/scenario.ts create mode 100644 js/src/durable-eval.test.ts create mode 100644 js/src/durable-eval.ts diff --git a/.gitignore b/.gitignore index bf65941bd..d6531a16e 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ dist !.aiderignore .pnpm-store **/.bt-tmp +**/.braintrust/evals docker-compose.override.yml Dockerfile.local diff --git a/e2e/scenarios/durable-eval-webhook/scenario.test.ts b/e2e/scenarios/durable-eval-webhook/scenario.test.ts new file mode 100644 index 000000000..4365794c3 --- /dev/null +++ b/e2e/scenarios/durable-eval-webhook/scenario.test.ts @@ -0,0 +1,33 @@ +import { expect, test } from "vitest"; +import { + prepareScenarioDir, + resolveScenarioDir, + withScenarioHarness, +} from "../../helpers/scenario-harness"; +import { findAllSpans } from "../../helpers/trace-selectors"; + +const scenarioDir = await prepareScenarioDir({ + scenarioDir: resolveScenarioDir(import.meta.url), +}); + +test("durable eval collects webhook sub-batches and logs completed rows", async () => { + await withScenarioHarness(async ({ runScenarioDir, testRunEvents }) => { + await runScenarioDir({ scenarioDir }); + + const evalSpans = findAllSpans(testRunEvents(), "eval"); + expect(evalSpans).toHaveLength(3); + expect(evalSpans.map((event) => event.output).sort()).toEqual([2, 4, 6]); + expect( + evalSpans + .map((event) => event.scores) + .sort((left, right) => + JSON.stringify(left).localeCompare(JSON.stringify(right)), + ), + ).toEqual([{ exact: 1 }, { exact: 1 }, { exact: 1 }]); + expect(evalSpans.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) }), + ]); + }); +}); diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts new file mode 100644 index 000000000..ba2f80448 --- /dev/null +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -0,0 +1,106 @@ +import { + BatchTask, + DurableEval, + MemoryDurableEvalStore, + type DurableBatchTaskItem, +} from "braintrust"; +import { + getTestRunId, + runMain, + scopedName, +} from "../../helpers/scenario-runtime"; + +async function main() { + const testRunId = getTestRunId(); + const store = new MemoryDurableEvalStore(); + const jobs = new Map< + string, + DurableBatchTaskItem< + number, + number, + { testRunId: string }, + Record + >[] + >(); + const task = BatchTask< + number, + number, + number, + { testRunId: string }, + Record, + { id: string } + >({ + revision: "task-v1", + batchSize: 2, + maxConcurrentBatches: 2, + async submit(items) { + const id = `provider-${jobs.size + 1}`; + jobs.set(id, items); + return { id }; + }, + completion: { + mode: "webhook", + source: "e2e-provider", + externalId: (handle) => handle.id, + }, + async *collect(handle) { + for (const item of jobs.get(handle.id) ?? []) { + yield { id: item.id, output: item.input * 2 }; + } + }, + }); + const definition = DurableEval( + scopedName("e2e-durable-eval-webhook-project", testRunId), + { + revision: "eval-v1", + experimentName: scopedName( + "e2e-durable-eval-webhook-experiment", + testRunId, + ), + data: [1, 2, 3].map((input) => ({ + id: `case-${input}`, + input, + expected: input * 2, + metadata: { testRunId }, + })), + task, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }, + ); + + const waiting = await definition.run({ + runId: `durable-${testRunId}`, + store, + }); + if ( + waiting.status !== "paused" || + waiting.reason !== "waiting_for_webhook" || + jobs.size !== 2 + ) { + throw new Error("Durable eval did not pause with two webhook batches"); + } + + for (let index = 1; index <= 2; index++) { + const processed = await definition.processBatchResult( + { + eventId: `event-${index}`, + source: "e2e-provider", + externalId: `provider-${index}`, + outcome: { status: "complete" }, + }, + { store }, + ); + if (processed.status !== "processed") { + throw new Error(`Webhook ${index} was not processed`); + } + if (index === 2 && processed.run.status !== "completed") { + throw new Error("Durable eval did not complete after the final webhook"); + } + } +} + +runMain(main); diff --git a/js/README.md b/js/README.md index da6da81a6..ddcf3f7d4 100644 --- a/js/README.md +++ b/js/README.md @@ -44,6 +44,178 @@ async function main() { main().catch(console.error); ``` +## Durable evaluations + +`DurableEval` is an additive evaluation API for work that must survive process +restarts or run through asynchronous provider batch jobs. It checkpoints cases, +task outputs, scores, provider job handles, and attempts outside Braintrust, +while completed results are incrementally merged into a normal Braintrust +experiment. + +Every case needs a stable `id` (or a `caseId` function). Reusing a `runId` +resumes the same input snapshot and never reruns successful work. + +```typescript +import { BatchTask, DurableEval, FileDurableEvalStore } from "braintrust"; + +const store = new FileDurableEvalStore(); +const supportEval = DurableEval("Support bot", { + revision: process.env.GIT_SHA ?? "local", + data: [ + { + id: "password-reset", + input: "How do I reset my password?", + expected: "Open account settings...", + }, + ], + task: BatchTask({ + revision: "generation-v1", + + // Each provider job contains at most 500 eval cases. Up to four jobs may + // be in flight for this task stage at once. + batchSize: 500, + maxConcurrentBatches: 4, + + // Upload/submit an asynchronous provider batch and return a + // JSON-serializable handle. + async submit(items, context) { + const batch = await provider.submit({ + idempotencyKey: context.batchId, + metadata: { durableBatchId: context.batchId }, + items, + }); + return { id: batch.id }; + }, + + completion: { + mode: "webhook", + source: "provider", + externalId: (handle) => handle.id, + }, + + async *collect(handle) { + for await (const item of provider.results(handle.id)) { + yield { + id: item.id, + output: item.output, + }; + } + }, + }), + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], +}); + +const result = await supportEval.run({ + runId: "release-2026-07-27", + store, + deadlineMs: 30 * 60 * 1000, +}); + +if (result.status === "paused") { + await result.resume({ deadlineMs: 30 * 60 * 1000 }); +} +``` + +Use `BatchScorer` for asynchronous batch scoring. Existing per-item tasks and +scorers are also supported and gain checkpointing and retries. The filesystem +store is intended for one machine; distributed workers should share a +compare-and-set-capable `DurableEvalStore`. Launch fixed shards through the API +or the existing CLI: + +```bash +npx braintrust eval evaluation.eval.ts \ + --run-id release-2026-07-27 \ + --shard 0/8 \ + --deadline 6h +``` + +An ordinary error thrown from `submit` is considered ambiguous, because the +provider may have created a job before the connection failed. The run pauses +instead of risking duplicate charges. Throw `DurableEvalNotSubmittedError` only +when it is known that no provider job was created, or implement `recover` to +look up a prior submission using `context.batchId`. + +The `submit`/`completion`/`collect` split is designed for provider-managed batch +APIs, including OpenAI's Batch API. A polling adapter uses +`completion: { mode: "poll", poll }`. A webhook adapter returns while the job is +pending and resumes through `processBatchResult`. `collect` then fetches and +stores the item results before the eval advances to scoring. Every provider +result must carry the durable item `id`. Batch tasks and batch scorers use the +same contract. + +### Webhook processing + +Verify the provider signature against the raw request body before calling the +durable eval definition. The HTTP framework and provider SDK remain outside the +generic durable API: + +```typescript +app.post("/webhooks/provider", rawBodyMiddleware, async (request, response) => { + const event = provider.verifyWebhook(request.rawBody, request.headers); + const batch = await provider.getBatch(event.batchId); + + const result = await supportEval.processBatchResult( + { + eventId: event.id, + source: "provider", + externalId: batch.id, + batchId: batch.metadata?.durableBatchId, + handle: { id: batch.id }, + outcome: + batch.status === "completed" + ? { status: "complete" } + : { + status: "failed", + error: { status: batch.status }, + retryable: false, + }, + // Optional JSON payloads are checkpointed with the deduplicated event. + payload: { type: event.type }, + }, + { + // Web workers and eval workers must use the same durable store. + store, + }, + ); + + response.status(result.status === "pending" ? 202 : 200).end(); +}); +``` + +`eventId` makes delivery idempotent. The SDK indexes a job by its internal +`batchId` before calling `submit`, then adds the provider `externalId` after the +handle is known. An early event is stored in an inbox and attached when that +second index appears. If a worker dies after the provider accepted a batch but +before its handle was checkpointed, include both the internal `batchId` and a +JSON-serializable `handle` in the event to resolve the ambiguous submission. +Without a handle or a successful `recover` callback, the run pauses safely. + +Pure webhook stages return +`{ status: "paused", reason: "waiting_for_webhook" }` instead of holding a +process open. A webhook completion may also define `pollFallback` with +`afterMs`, `poll`, and an optional `intervalMs`. + +All orchestration state lives in the configured `DurableEvalStore`; this API +does not require a Braintrust backend change. + +The definition object exposes `run`, `status`, `retryFailed`, +`resubmitUnknown`, `cancel`, and `processBatchResult`. The corresponding +lifecycle operations are also available through the CLI: + +```bash +npx braintrust eval evaluation.eval.ts --run-id release-2026-07-27 --status +npx braintrust eval evaluation.eval.ts --run-id release-2026-07-27 --retry-failed +npx braintrust eval evaluation.eval.ts --run-id release-2026-07-27 --resubmit-unknown +npx braintrust eval evaluation.eval.ts --run-id release-2026-07-27 --cancel +``` + +`--resubmit-unknown` is intentionally explicit because the original provider +job may exist and resubmitting it can duplicate work and cost. + ## Auto-Instrumentation Braintrust can automatically instrument popular AI SDKs (OpenAI, Anthropic, Vercel AI SDK, and others) to log calls without manual wrapper code. diff --git a/js/src/cli/functions/load-module.ts b/js/src/cli/functions/load-module.ts index 593b45ae7..1ac9a8b8d 100644 --- a/js/src/cli/functions/load-module.ts +++ b/js/src/cli/functions/load-module.ts @@ -26,6 +26,7 @@ export function loadModule({ prompts: [], parameters: [], evaluators: {}, + durableEvaluators: {}, reporters: {}, }; globalThis._lazy_load = true; diff --git a/js/src/cli/index.ts b/js/src/cli/index.ts index ba36d79e3..6d4604352 100755 --- a/js/src/cli/index.ts +++ b/js/src/cli/index.ts @@ -61,6 +61,11 @@ import { } from "./util/debug-logging"; import { pullCommand } from "./util/pull"; import { runDevServer } from "../../dev/server"; +import { + type DurableEvalDefinition, + type DurableEvalOperation, + type DurableEvalRuntimeOptions, +} from "../durable-eval"; // This requires require // https://stackoverflow.com/questions/50822310/how-to-import-package-json-in-typescript @@ -213,6 +218,7 @@ function buildWatchPluginForEvaluator( ): esbuild.Plugin { const evaluators: EvaluatorState = { evaluators: [], + durableEvaluators: [], reporters: {}, }; const plugin = { @@ -412,6 +418,11 @@ interface EvaluatorOpts { jsonl: boolean; filters: Filter[]; progressReporter: ProgressReporter; + durableRunId?: string; + durableShard?: { index: number; count: number }; + durableDeadlineMs?: number; + checkpointDir?: string; + durableOperation?: Exclude; } export function handleBuildFailure({ @@ -466,6 +477,21 @@ function updateEvaluators( reporter: evaluator.reporter, }); } + for (const registration of Object.values( + result.evaluator.durableEvaluators ?? {}, + )) { + evaluators.durableEvaluators.push({ + sourceFile: result.sourceFile, + // Runtime registrations are intentionally generic-erased. + definition: registration.definition as DurableEvalDefinition< + any, + any, + any, + any, + any + >, + }); + } for (const [reporterName, reporter] of Object.entries( result.evaluator.reporters, @@ -528,6 +554,7 @@ export async function buildEvaluators( const evaluators: EvaluatorState = { evaluators: [], + durableEvaluators: [], reporters: {}, }; updateEvaluators(evaluators, buildResults, opts); @@ -548,12 +575,21 @@ async function runOnce( : null; const { evaluators, buildResults } = await buildEvaluators(handles, opts); + if (opts.durableOperation && evaluators.evaluators.length > 0) { + throw new Error( + "Durable lifecycle flags cannot be used with ordinary Eval definitions", + ); + } if (opts.list) { for (const evaluator of evaluators.evaluators) { // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. console.log(evaluator.evaluator.evalName); } + for (const evaluator of evaluators.durableEvaluators) { + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.log(evaluator.definition.evalName); + } return true; } @@ -592,15 +628,46 @@ async function runOnce( } } }); + const durableResultPromises = evaluators.durableEvaluators.map( + async (registration) => { + const options: DurableEvalRuntimeOptions = { + runId: opts.durableRunId, + shard: opts.durableShard, + deadlineMs: opts.durableDeadlineMs, + checkpointDir: opts.checkpointDir, + noSendLogs: opts.noSendLogs, + }; + if (!opts.durableOperation) { + return await registration.definition.run(options); + } + if (!opts.durableRunId) { + throw new Error( + `--${opts.durableOperation.replaceAll("-", "_")} requires --run-id`, + ); + } + const existingOptions = { ...options, runId: opts.durableRunId }; + switch (opts.durableOperation) { + case "status": + return await registration.definition.status(existingOptions); + case "retry-failed": + return await registration.definition.retryFailed(existingOptions); + case "resubmit-unknown": + return await registration.definition.resubmitUnknown(existingOptions); + case "cancel": + return await registration.definition.cancel(existingOptions); + } + }, + ); // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. console.error( styleText( "dim", - `Processing ${styleText("bold", String(resultPromises.length))} evaluator${resultPromises.length === 1 ? "" : "s"}...`, + `Processing ${styleText("bold", String(resultPromises.length + durableResultPromises.length))} evaluator${resultPromises.length + durableResultPromises.length === 1 ? "" : "s"}...`, ), ); const allEvalsResults = await Promise.all(resultPromises); + const allDurableResults = await Promise.all(durableResultPromises); opts.progressReporter.stop(); // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. console.error(""); @@ -655,6 +722,33 @@ async function runOnce( allSuccess = allSuccess && success; } + for (const result of allDurableResults) { + if (opts.jsonl) { + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.log(JSON.stringify(result)); + } else if (result.status === "completed") { + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.error( + `Durable eval ${result.runId} completed (${result.progress.taskSucceeded}/${result.progress.total} tasks)`, + ); + } else { + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.error(`Durable eval ${result.runId} paused: ${result.reason}`); + } + if ( + result.status === "completed" && + (result.failures.tasks > 0 || result.failures.scorers > 0) + ) { + allSuccess = false; + } + if ( + result.status === "paused" && + ["unknown_submission", "provider_unreachable"].includes(result.reason) + ) { + allSuccess = false; + } + } + return allSuccess; } @@ -947,6 +1041,21 @@ async function run(args: RunArgs) { : new BarProgressReporter(), filters: args.filter ? parseFilters(args.filter) : [], list: !!args.list, + durableRunId: args.run_id, + durableShard: args.shard ? parseDurableShard(args.shard) : undefined, + durableDeadlineMs: args.deadline + ? parseDurationMs(args.deadline) + : undefined, + checkpointDir: args.checkpoint_dir, + durableOperation: args.status + ? "status" + : args.retry_failed + ? "retry-failed" + : args.resubmit_unknown + ? "resubmit-unknown" + : args.cancel + ? "cancel" + : undefined, }; if (args.list && args.watch) { @@ -954,6 +1063,23 @@ async function run(args: RunArgs) { console.error(error("Cannot specify both --list and --watch.")); process.exit(1); } + if (args.shard && !args.run_id) { + throw new Error("--shard requires --run-id for durable evals"); + } + const lifecycleFlags = [ + args.status, + args.retry_failed, + args.resubmit_unknown, + args.cancel, + ].filter(Boolean); + if (lifecycleFlags.length > 1) { + throw new Error( + "Specify only one of --status, --retry-failed, --resubmit-unknown, or --cancel", + ); + } + if (lifecycleFlags.length > 0 && !args.run_id) { + throw new Error("Durable lifecycle flags require --run-id"); + } const plugins = evaluatorOpts.watch ? [ @@ -1017,6 +1143,32 @@ async function run(args: RunArgs) { } } +function parseDurableShard(value: string) { + const match = /^(\d+)\/(\d+)$/.exec(value); + if (!match) { + throw new Error(`Invalid --shard value '${value}'; expected INDEX/COUNT`); + } + const shard = { index: Number(match[1]), count: Number(match[2]) }; + if (shard.count < 1 || shard.index < 0 || shard.index >= shard.count) { + throw new Error(`Invalid --shard value '${value}'`); + } + return shard; +} + +function parseDurationMs(value: string) { + const match = /^(\d+(?:\.\d+)?)(ms|s|m|h)?$/.exec(value); + if (!match) { + throw new Error( + `Invalid --deadline value '${value}'; expected e.g. 500ms, 30m, or 6h`, + ); + } + const multipliers = { ms: 1, s: 1_000, m: 60_000, h: 3_600_000 }; + return ( + Number(match[1]) * + multipliers[(match[2] ?? "ms") as keyof typeof multipliers] + ); +} + function addAuthArgs(parser: ArgumentParser) { parser.add_argument("--api-key", { help: "Specify a braintrust api key. If the parameter is not specified, BRAINTRUST_API_KEY or the nearest .env.braintrust file will be used.", @@ -1101,6 +1253,34 @@ async function main() { action: "store_true", help: "Do not show progress bars when processing evaluators.", }); + parser_run.add_argument("--run-id", { + help: "Create or resume a DurableEval run with this stable identifier.", + }); + parser_run.add_argument("--shard", { + help: "Run one DurableEval shard in INDEX/COUNT form, for example 0/8.", + }); + parser_run.add_argument("--deadline", { + help: "Pause DurableEval work after a duration such as 30m or 6h.", + }); + parser_run.add_argument("--checkpoint-dir", { + help: "Override the default .braintrust/evals checkpoint directory.", + }); + parser_run.add_argument("--status", { + action: "store_true", + help: "Inspect a DurableEval run without claiming new work.", + }); + parser_run.add_argument("--retry-failed", { + action: "store_true", + help: "Retry terminal DurableEval failures with the current definition.", + }); + parser_run.add_argument("--resubmit-unknown", { + action: "store_true", + help: "Explicitly resubmit ambiguous provider batches, accepting possible duplicate cost.", + }); + parser_run.add_argument("--cancel", { + action: "store_true", + help: "Cancel active provider batches and terminally cancel unfinished DurableEval work.", + }); parser_run.add_argument("--bundle", { action: "store_true", help: "Experimental (do not use unless you know what you're doing)", diff --git a/js/src/cli/types.ts b/js/src/cli/types.ts index 15184a247..c614c40d0 100644 --- a/js/src/cli/types.ts +++ b/js/src/cli/types.ts @@ -2,6 +2,7 @@ import type * as esbuild from "esbuild"; import type { BaseMetadata } from "../logger"; import type { EvaluatorDef, EvaluatorFile } from "../framework"; import type { ReporterDef } from "../reporters/types"; +import type { DurableEvalDefinition } from "../durable-eval"; export interface BuildSuccess { type: "success"; @@ -34,6 +35,12 @@ export interface EvaluatorState { evaluator: EvaluatorDef; reporter: string | ReporterDef | undefined; }[]; + durableEvaluators: { + sourceFile: string; + // Runtime registration has already passed the public generic boundary. + // eslint-disable-next-line @typescript-eslint/no-explicit-any + definition: DurableEvalDefinition; + }[]; reporters: { [reporter: string]: ReporterDef; }; diff --git a/js/src/cli/util/types.ts b/js/src/cli/util/types.ts index 2909a607e..25532a022 100644 --- a/js/src/cli/util/types.ts +++ b/js/src/cli/util/types.ts @@ -33,6 +33,14 @@ export interface RunArgs extends CommonArgs, AuthArgs, CompileArgs { dev_host: string; dev_port: number; dev_org_name?: string; + run_id?: string; + shard?: string; + deadline?: string; + checkpoint_dir?: string; + status?: boolean; + retry_failed?: boolean; + resubmit_unknown?: boolean; + cancel?: boolean; } export interface BundleArgs extends CommonArgs, AuthArgs, CompileArgs { diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts new file mode 100644 index 000000000..aac780448 --- /dev/null +++ b/js/src/durable-eval.test.ts @@ -0,0 +1,876 @@ +import { describe, expect, test, vi } from "vitest"; +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { + BatchScorer, + BatchTask, + DurableEval, + FileDurableEvalStore, + MemoryDurableEvalStore, + type DurableBatchTaskItem, +} from "./durable-eval"; +import type { EvaluatorFile } from "./framework"; +import { configureNode } from "./node/config"; + +configureNode(); + +describe("DurableEval", () => { + test("filesystem store enforces compare-and-set writes", async () => { + const directory = await mkdtemp(join(tmpdir(), "durable-eval-store-")); + try { + const store = new FileDurableEvalStore(directory); + const first = await store.write( + "runs/one", + new TextEncoder().encode("one"), + { ifAbsent: true }, + ); + expect(first.written).toBe(true); + if (!first.written) throw new Error("initial write failed"); + + await expect( + store.write("runs/one", new TextEncoder().encode("two"), { + ifAbsent: true, + }), + ).resolves.toMatchObject({ written: false }); + await expect( + store.write("runs/one", new TextEncoder().encode("two"), { + ifVersion: "wrong", + }), + ).resolves.toMatchObject({ written: false }); + + const updated = await store.write( + "runs/one", + new TextEncoder().encode("two"), + { ifVersion: first.version }, + ); + expect(updated.written).toBe(true); + expect( + new TextDecoder().decode((await store.read("runs/one"))?.value), + ).toBe("two"); + const keys: string[] = []; + for await (const key of store.list("runs/")) keys.push(key); + expect(keys).toEqual(["runs/one"]); + } finally { + await rm(directory, { recursive: true, force: true }); + } + }); + + test("runs local tasks and scorers and resumes completed work", async () => { + const store = new MemoryDurableEvalStore(); + const task = vi.fn((input: number) => input * 2); + let scorerCalls = 0; + function exact({ + output, + expected, + }: { + output: number; + expected?: number; + }) { + scorerCalls++; + return output === expected ? 1 : 0; + } + const definition = { + revision: "v1", + data: [ + { id: "one", input: 1, expected: 2 }, + { id: "two", input: 2, expected: 4 }, + ], + task, + scores: [exact], + }; + + const durable = DurableEval("durable-local", definition); + const first = await durable.run({ + runId: "run-1", + store, + noSendLogs: true, + }); + expect(first.status).toBe("completed"); + if (first.status !== "completed") throw new Error("run did not complete"); + expect(first.summary.scores.exact?.score).toBe(1); + expect(first.progress).toMatchObject({ + total: 2, + taskSucceeded: 2, + scoreSucceeded: 2, + }); + + const second = await durable.run({ + runId: "run-1", + store, + noSendLogs: true, + }); + expect(second.status).toBe("completed"); + expect(task).toHaveBeenCalledTimes(2); + expect(scorerCalls).toBe(2); + }); + + test("requires stable case identifiers", async () => { + await expect( + DurableEval("missing-ids", { + revision: "v1", + data: [{ input: "hello", expected: "hello" }], + task: (input) => input, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }).run({ + runId: "missing", + store: new MemoryDurableEvalStore(), + noSendLogs: true, + }), + ).rejects.toThrow("must have a non-empty id"); + }); + + test("submits and collects batch tasks and batch scorers", async () => { + const store = new MemoryDurableEvalStore(); + const taskJobs = new Map< + string, + DurableBatchTaskItem>[] + >(); + const task = BatchTask< + number, + number, + number, + void, + Record, + { jobId: string } + >({ + revision: "task-v1", + batchSize: 2, + async submit(items, context) { + taskJobs.set(context.batchId, items); + return { jobId: context.batchId }; + }, + completion: { + mode: "poll", + async poll() { + return { status: "complete" }; + }, + }, + async *collect(handle) { + for (const item of taskJobs.get(handle.jobId) ?? []) { + yield { id: item.id, output: item.input * 2 }; + } + }, + }); + const scoreJobs = new Map< + string, + Array<{ id: string; output: number; expected: number }> + >(); + const scorer = BatchScorer( + { + name: "exact", + revision: "score-v1", + batchSize: 2, + async submit(items, context) { + scoreJobs.set( + context.batchId, + items.map((item) => ({ + id: item.id, + output: item.output, + expected: item.expected, + })), + ); + return { jobId: context.batchId }; + }, + completion: { + mode: "poll", + async poll() { + return { status: "complete" }; + }, + }, + async *collect(handle) { + for (const item of scoreJobs.get(handle.jobId) ?? []) { + yield { + id: item.id, + score: item.output === item.expected ? 1 : 0, + }; + } + }, + }, + ); + + const result = await DurableEval("durable-batches", { + revision: "eval-v1", + data: [ + { id: "one", input: 1, expected: 2 }, + { id: "two", input: 2, expected: 4 }, + { id: "three", input: 3, expected: 6 }, + ], + task, + scores: [scorer], + }).run({ + runId: "batch-run", + store, + noSendLogs: true, + }); + + expect(result.status).toBe("completed"); + if (result.status !== "completed") throw new Error("run did not complete"); + expect(result.summary.scores.exact?.score).toBe(1); + expect(taskJobs).toHaveLength(2); + expect(scoreJobs).toHaveLength(2); + }); + + test("pauses for webhooks and processes bounded sub-batches", async () => { + const store = new MemoryDurableEvalStore(); + const providerJobs = new Map< + string, + DurableBatchTaskItem>[] + >(); + const submittedSizes: number[] = []; + const task = BatchTask< + number, + number, + number, + void, + Record, + { id: string } + >({ + revision: "task-v1", + batchSize: 2, + maxConcurrentBatches: 2, + async submit(items) { + const id = `provider-${providerJobs.size + 1}`; + providerJobs.set(id, items); + submittedSizes.push(items.length); + return { id }; + }, + completion: { + mode: "webhook", + source: "openai", + externalId: (handle) => handle.id, + }, + async *collect(handle) { + for (const item of providerJobs.get(handle.id) ?? []) { + yield { id: item.id, output: item.input * 2 }; + } + }, + }); + const durable = DurableEval("webhook-batches", { + revision: "eval-v1", + data: [ + { id: "one", input: 1, expected: 2 }, + { id: "two", input: 2, expected: 4 }, + { id: "three", input: 3, expected: 6 }, + ], + task, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }); + + const waiting = await durable.run({ + runId: "webhook-run", + store, + noSendLogs: true, + }); + expect(waiting).toMatchObject({ + status: "paused", + reason: "waiting_for_webhook", + }); + expect(submittedSizes).toEqual([2, 1]); + + const first = await durable.processBatchResult( + { + eventId: "event-1", + source: "openai", + externalId: "provider-1", + outcome: { status: "complete" }, + payload: { providerStatus: "completed" }, + }, + { store, noSendLogs: true }, + ); + expect(first).toMatchObject({ + status: "processed", + batchId: expect.any(String), + run: { status: "paused", reason: "waiting_for_webhook" }, + }); + + const second = await durable.processBatchResult( + { + eventId: "event-2", + source: "openai", + externalId: "provider-2", + outcome: { status: "complete" }, + }, + { store, noSendLogs: true }, + ); + expect(second).toMatchObject({ + status: "processed", + run: { + status: "completed", + progress: { taskSucceeded: 3, scoreSucceeded: 3 }, + }, + }); + + await expect( + durable.processBatchResult( + { + eventId: "event-2", + source: "openai", + externalId: "provider-2", + outcome: { status: "complete" }, + }, + { store, noSendLogs: true }, + ), + ).resolves.toMatchObject({ status: "duplicate" }); + + const storedEvents: string[] = []; + for await (const key of store.list("durable-eval/v1/webhooks/events/")) { + const record = await store.read(key); + if (record) storedEvents.push(new TextDecoder().decode(record.value)); + } + expect(storedEvents).toHaveLength(2); + expect(storedEvents.join("\n")).toContain("providerStatus"); + }); + + test("stores an early webhook until the provider handle is indexed", async () => { + const store = new MemoryDurableEvalStore(); + let processEarly: + | (() => Promise<{ status: string; reason?: string }>) + | undefined; + let earlyResult: { status: string; reason?: string } | undefined; + const task = BatchTask< + string, + string, + string, + void, + Record, + { id: string } + >({ + revision: "task-v1", + async submit() { + earlyResult = await processEarly!(); + return { id: "provider-early" }; + }, + completion: { + mode: "webhook", + source: "openai", + externalId: (handle) => handle.id, + }, + async *collect() { + yield { id: "one:trial:0", output: "done" }; + }, + }); + const durable = DurableEval("early-webhook", { + revision: "eval-v1", + data: [{ id: "one", input: "input", expected: "done" }], + task, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }); + processEarly = () => + durable.processBatchResult( + { + eventId: "early-event", + source: "openai", + externalId: "provider-early", + outcome: { status: "complete" }, + }, + { store, noSendLogs: true }, + ); + + const result = await durable.run({ + runId: "early-run", + store, + noSendLogs: true, + }); + expect(earlyResult).toMatchObject({ + status: "pending", + reason: "unmatched", + }); + expect(result.status).toBe("completed"); + }); + + test("uses a webhook handle to resolve an ambiguous submission", async () => { + const store = new MemoryDurableEvalStore(); + let batchId: string | undefined; + const task = BatchTask< + string, + string, + string, + void, + Record, + { id: string } + >({ + revision: "task-v1", + async submit(_items, context) { + batchId = context.batchId; + throw new Error("connection closed after provider accepted the batch"); + }, + completion: { + mode: "webhook", + source: "openai", + externalId: (handle) => handle.id, + }, + async *collect() { + yield { id: "one:trial:0", output: "done" }; + }, + }); + const durable = DurableEval("lost-submit-webhook", { + revision: "eval-v1", + data: [{ id: "one", input: "input", expected: "done" }], + task, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }); + const paused = await durable.run({ + runId: "lost-submit-run", + store, + noSendLogs: true, + }); + expect(paused).toMatchObject({ + status: "paused", + reason: "unknown_submission", + }); + + const processed = await durable.processBatchResult( + { + eventId: "lost-submit-event", + source: "openai", + batchId, + externalId: "provider-lost", + handle: { id: "provider-lost" }, + outcome: { status: "complete" }, + }, + { store, noSendLogs: true }, + ); + expect(processed).toMatchObject({ + status: "processed", + run: { status: "completed" }, + }); + }); + + test("retries result collection when the provider redelivers a webhook", async () => { + const store = new MemoryDurableEvalStore(); + let collectAttempts = 0; + const task = BatchTask< + string, + string, + string, + void, + Record, + { id: string } + >({ + revision: "task-v1", + async submit() { + return { id: "provider-retry" }; + }, + completion: { + mode: "webhook", + source: "openai", + externalId: (handle) => handle.id, + }, + async *collect() { + collectAttempts++; + if (collectAttempts === 1) { + throw new Error("provider result file is temporarily unavailable"); + } + yield { id: "one:trial:0", output: "done" }; + }, + }); + const durable = DurableEval("collect-retry", { + revision: "eval-v1", + data: [{ id: "one", input: "input", expected: "done" }], + task, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }); + await durable.run({ + runId: "collect-retry-run", + store, + noSendLogs: true, + }); + const event = { + eventId: "retry-event", + source: "openai", + externalId: "provider-retry", + outcome: { status: "complete" as const }, + }; + + const unavailable = await durable.processBatchResult(event, { + store, + noSendLogs: true, + }); + expect(unavailable).toMatchObject({ + status: "processed", + run: { status: "paused", reason: "provider_unreachable" }, + }); + + const completed = await durable.processBatchResult(event, { + store, + noSendLogs: true, + }); + expect(completed).toMatchObject({ + status: "processed", + run: { status: "completed" }, + }); + expect(collectAttempts).toBe(2); + }); + + test("does not recover or duplicate a batch while another worker submits it", async () => { + const store = new MemoryDurableEvalStore(); + let submitCount = 0; + let releaseSubmit: () => void = () => undefined; + const submitGate = new Promise((resolve) => { + releaseSubmit = resolve; + }); + const task = BatchTask< + string, + string, + string, + void, + Record, + { jobId: string } + >({ + revision: "task-v1", + async submit(_items, context) { + submitCount++; + await submitGate; + return { jobId: context.batchId }; + }, + completion: { + mode: "poll", + async poll() { + return { status: "complete" }; + }, + }, + async *collect() { + yield { id: "one:trial:0", output: "done" }; + }, + }); + const definition = { + revision: "eval-v1", + data: [{ id: "one", input: "input", expected: "done" }], + task, + scores: [ + function exact({ + output, + expected, + }: { + output: string; + expected?: string; + }) { + return output === expected ? 1 : 0; + }, + ], + }; + + const durable = DurableEval("distributed-submit", definition); + const first = durable.run({ + runId: "distributed-run", + store, + noSendLogs: true, + workerId: "worker-one", + }); + await vi.waitFor(() => expect(submitCount).toBe(1)); + + const second = await durable.run({ + runId: "distributed-run", + store, + noSendLogs: true, + workerId: "worker-two", + deadlineMs: 10, + }); + expect(second.status).toBe("paused"); + expect(submitCount).toBe(1); + + releaseSubmit(); + await expect(first).resolves.toMatchObject({ status: "completed" }); + expect(submitCount).toBe(1); + }); + + test("cancels active provider batches and terminally checkpoints the work", async () => { + const store = new MemoryDurableEvalStore(); + const abort = new AbortController(); + const cancel = vi.fn(async () => undefined); + const task = BatchTask< + string, + string, + string, + void, + Record, + { jobId: string } + >({ + revision: "task-v1", + async submit(_items, context) { + abort.abort(); + return { jobId: context.batchId }; + }, + completion: { + mode: "poll", + async poll() { + return { status: "pending" }; + }, + }, + async *collect() { + // A cancelled batch is never collected. + }, + cancel, + }); + const definition = { + revision: "eval-v1", + data: [{ id: "one", input: "input", expected: "done" }], + task, + scores: [ + function exact({ + output, + expected, + }: { + output: string; + expected?: string; + }) { + return output === expected ? 1 : 0; + }, + ], + }; + + const durable = DurableEval("cancel-batch", definition); + const paused = await durable.run({ + runId: "cancel-run", + store, + noSendLogs: true, + signal: abort.signal, + }); + expect(paused).toMatchObject({ status: "paused", reason: "aborted" }); + + const cancelled = await durable.cancel({ + runId: "cancel-run", + store, + noSendLogs: true, + }); + expect(cancel).toHaveBeenCalledOnce(); + expect(cancelled).toMatchObject({ + status: "completed", + failures: { tasks: 1, scorers: 0 }, + }); + }); + + test("returns a resumable deadline result without resubmitting a batch", async () => { + const store = new MemoryDurableEvalStore(); + let ready = false; + let submits = 0; + const task = BatchTask< + string, + string, + string, + void, + Record, + { jobId: string } + >({ + revision: "task-v1", + async submit(_items, context) { + submits++; + return { jobId: context.batchId }; + }, + completion: { + mode: "poll", + intervalMs: 1, + async poll() { + return ready + ? { status: "complete" as const } + : { status: "pending" as const, retryAfterMs: 1 }; + }, + }, + async *collect() { + yield { id: "one:trial:0", output: "done" }; + }, + }); + + const definition = { + revision: "eval-v1", + data: [{ id: "one", input: "input", expected: "done" }], + task, + scores: [ + function exact({ + output, + expected, + }: { + output: string; + expected?: string; + }) { + return output === expected ? 1 : 0; + }, + ], + }; + const durable = DurableEval("deadline", definition); + const first = await durable.run({ + runId: "deadline-run", + store, + noSendLogs: true, + deadlineMs: 5, + }); + expect(first.status).toBe("paused"); + if (first.status !== "paused") throw new Error("run did not pause"); + expect(first.reason).toBe("deadline"); + + ready = true; + const second = await first.resume({ deadlineMs: 1_000 }); + expect(second.status).toBe("completed"); + expect(submits).toBe(1); + }); + + test("keeps completed work when the definition revision changes", async () => { + const store = new MemoryDurableEvalStore(); + const firstTask = vi.fn((input: number) => input); + await DurableEval("revisions", { + revision: "v1", + data: [{ id: "one", input: 1, expected: 1 }], + task: firstTask, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }).run({ + runId: "revision-run", + store, + noSendLogs: true, + }); + + const changedTask = vi.fn(() => 0); + const result = await DurableEval("revisions", { + revision: "v2", + data: [{ id: "one", input: 1, expected: 1 }], + task: changedTask, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }).run({ + runId: "revision-run", + store, + noSendLogs: true, + }); + + expect(result.status).toBe("completed"); + expect(firstTask).toHaveBeenCalledOnce(); + expect(changedTask).not.toHaveBeenCalled(); + }); + + test("pauses safely when submit may have created a provider job", async () => { + const store = new MemoryDurableEvalStore(); + let submitShouldFail = true; + const definition = { + revision: "v1", + data: [{ id: "one", input: "hello", expected: "hello" }], + task: BatchTask< + string, + string, + string, + void, + Record, + { jobId: string } + >({ + revision: "task-v1", + async submit(_items, context) { + if (submitShouldFail) { + throw new Error("connection closed after request"); + } + return { jobId: context.batchId }; + }, + completion: { + mode: "poll", + async poll() { + return { status: "complete" }; + }, + }, + async *collect() { + yield { id: "one:trial:0", output: "hello" }; + }, + }), + scores: [ + function exact({ + output, + expected, + }: { + output: string; + expected?: string; + }) { + return output === expected ? 1 : 0; + }, + ], + }; + const durable = DurableEval("ambiguous-submit", definition); + const result = await durable.run({ + runId: "ambiguous-run", + store, + noSendLogs: true, + }); + + expect(result.status).toBe("paused"); + if (result.status !== "paused") throw new Error("run did not pause"); + expect(result.reason).toBe("unknown_submission"); + expect(result.progress.unknown).toBe(1); + + const status = await durable.status({ + runId: "ambiguous-run", + store, + noSendLogs: true, + }); + expect(status).toMatchObject({ + status: "paused", + reason: "status_only", + }); + + submitShouldFail = false; + const resumed = await durable.resubmitUnknown({ + runId: "ambiguous-run", + store, + noSendLogs: true, + }); + expect(resumed.status).toBe("completed"); + }); + + test("registers definitions instead of executing during CLI lazy loading", async () => { + const previousEvals = globalThis._evals; + const previousLazy = globalThis._lazy_load; + globalThis._evals = { + functions: [], + prompts: [], + parameters: [], + evaluators: {}, + durableEvaluators: {}, + reporters: {}, + } satisfies EvaluatorFile; + globalThis._lazy_load = true; + const task = vi.fn((input: string) => input); + try { + const definition = DurableEval("lazy-project", { + revision: "v1", + experimentName: "lazy-eval", + data: [{ id: "one", input: "hello", expected: "hello" }], + task, + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }); + expect(globalThis._evals.durableEvaluators?.["lazy-eval"]).toBeDefined(); + expect( + globalThis._evals.durableEvaluators?.["lazy-eval"]?.definition, + ).toBe(definition); + expect(task).not.toHaveBeenCalled(); + } finally { + globalThis._evals = previousEvals; + globalThis._lazy_load = previousLazy; + } + }); +}); diff --git a/js/src/durable-eval.ts b/js/src/durable-eval.ts new file mode 100644 index 000000000..c6025195f --- /dev/null +++ b/js/src/durable-eval.ts @@ -0,0 +1,3387 @@ +import { type Score, SpanTypeAttribute } from "../util/index"; +import iso from "./isomorph"; +import { + type BaseMetadata, + type BraintrustState, + type DefaultMetadataType, + type EvalCase, + type Experiment, + type ExperimentSummary, + NOOP_SPAN, + _internalStartSpanWithInitialMerge, + init as initExperiment, + newId, +} from "./logger"; +import { + type EvalData, + type EvalHooks, + type EvalScorer, + type EvalScorerArgs, + type EvalTask, + type OneOrMoreScores, +} from "./framework"; +import { type EvalParameters, type InferParameters } from "./eval-parameters"; + +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 CHECKPOINT_VERSION = 2; +const DEFAULT_BATCH_SIZE = 1_000; +const DEFAULT_MAX_CONCURRENT_BATCHES = 1; +const JOB_LEASE_MS = 60_000; + +type JsonPrimitive = string | number | boolean | null; +export type JsonValue = + | JsonPrimitive + | JsonValue[] + | { [key: string]: JsonValue }; + +export type DurableEvalWriteCondition = + | { ifAbsent: true } + | { ifVersion: string } + | { unconditional: true }; + +export interface DurableEvalStore { + read( + key: string, + ): Promise<{ value: Uint8Array; version: string } | undefined>; + write( + key: string, + value: Uint8Array, + condition: DurableEvalWriteCondition, + ): Promise< + | { written: true; version: string } + | { written: false; currentVersion?: string } + >; + list(prefix: string): AsyncIterable; +} + +/** + * Throw this from `submit` only when it is known that no provider job was + * created. Other submit errors are treated as ambiguous and pause the run. + */ +export class DurableEvalNotSubmittedError extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options); + this.name = "DurableEvalNotSubmittedError"; + } +} + +export class MemoryDurableEvalStore implements DurableEvalStore { + private readonly values = new Map< + string, + { value: Uint8Array; version: number } + >(); + + async read(key: string) { + const record = this.values.get(key); + return record + ? { value: record.value.slice(), version: String(record.version) } + : undefined; + } + + async write( + key: string, + value: Uint8Array, + condition: DurableEvalWriteCondition, + ) { + const current = this.values.get(key); + if ( + ("ifAbsent" in condition && current) || + ("ifVersion" in condition && + (!current || String(current.version) !== condition.ifVersion)) + ) { + return { + written: false as const, + currentVersion: current ? String(current.version) : undefined, + }; + } + const version = (current?.version ?? 0) + 1; + this.values.set(key, { value: value.slice(), version }); + return { written: true as const, version: String(version) }; + } + + async *list(prefix: string) { + for (const key of [...this.values.keys()].sort()) { + if (key.startsWith(prefix)) { + yield key; + } + } + } +} + +/** + * A filesystem checkpoint store intended for durable, multi-process evals on + * one machine. Distributed deployments should provide a database/object-store + * implementation of {@link DurableEvalStore}. + */ +export class FileDurableEvalStore implements DurableEvalStore { + constructor(public readonly root = ".braintrust/evals") {} + + private assertFilesystem() { + if ( + !iso.pathJoin || + !iso.pathDirname || + !iso.mkdir || + !iso.readFile || + !iso.writeFile || + !iso.readdir || + !iso.stat || + !iso.unlink || + !iso.rename || + !iso.openFile + ) { + throw new Error( + "FileDurableEvalStore is only available in a Node.js filesystem environment", + ); + } + } + + private pathFor(key: string) { + if ( + key.startsWith("/") || + key.split("/").some((part) => part === ".." || part === ".") + ) { + throw new Error(`Invalid durable eval store key: ${key}`); + } + return iso.pathJoin!(this.root, ...key.split("/")); + } + + private async withLock(path: string, fn: () => Promise): Promise { + const lockPath = `${path}.lock`; + await iso.mkdir!(iso.pathDirname!(path), { recursive: true }); + const started = Date.now(); + let handle: { close(): Promise } | undefined; + while (!handle) { + try { + handle = await iso.openFile!(lockPath, "wx"); + } catch (error) { + if (!isErrorCode(error, "EEXIST") || Date.now() - started >= 10_000) { + throw error; + } + await delay(10); + } + } + try { + return await fn(); + } finally { + await handle.close(); + await iso.unlink!(lockPath).catch(() => undefined); + } + } + + async read(key: string) { + this.assertFilesystem(); + const path = this.pathFor(key); + try { + const value = await iso.readFile!(path); + return { value, version: contentVersion(value) }; + } catch (error) { + if (isErrorCode(error, "ENOENT")) return undefined; + throw error; + } + } + + async write( + key: string, + value: Uint8Array, + condition: DurableEvalWriteCondition, + ) { + this.assertFilesystem(); + const path = this.pathFor(key); + return await this.withLock(path, async () => { + const current = await this.read(key); + if ( + ("ifAbsent" in condition && current) || + ("ifVersion" in condition && current?.version !== condition.ifVersion) + ) { + return { + written: false as const, + currentVersion: current?.version, + }; + } + + await iso.mkdir!(iso.pathDirname!(path), { recursive: true }); + const tempPath = `${path}.${newId()}.tmp`; + await iso.writeFile!(tempPath, value); + await iso.rename!(tempPath, path); + return { written: true as const, version: contentVersion(value) }; + }); + } + + async *list(prefix: string) { + this.assertFilesystem(); + const prefixPath = this.pathFor(prefix); + const walk = async function* (path: string): AsyncGenerator { + let entries: string[]; + try { + entries = await iso.readdir!(path); + } catch (error) { + if (isErrorCode(error, "ENOENT")) return; + throw error; + } + for (const entry of entries.sort()) { + if ( + entry.endsWith(".lock") || + entry.endsWith(".tmp") || + entry.includes(".tmp.") + ) { + continue; + } + const child = iso.pathJoin!(path, entry); + const stat = await iso.stat!(child); + if (stat.isDirectory()) { + yield* walk(child); + } else { + yield child; + } + } + }; + + for await (const path of walk(prefixPath)) { + yield path + .slice(this.root.length) + .replace(/^[/\\]+/, "") + .replaceAll("\\", "/"); + } + } +} + +export interface DurableBatchContext { + runId: string; + stage: string; + revision: string; + shard: { index: number; count: number }; + attempt: number; + batchId: string; + itemCount: number; + signal: AbortSignal; +} + +export type DurableBatchPoll = + | { status: "pending"; retryAfterMs?: number } + | { status: "complete" } + | { status: "failed"; error: unknown; retryable?: boolean }; + +export type DurableBatchRecovery = + | { status: "found"; handle: Handle } + | { status: "not_found" } + | { status: "unknown" }; + +export type DurableBatchOutcome = + | { status: "complete" } + | { status: "failed"; error: JsonValue; retryable?: boolean }; + +export type DurableBatchCompletion = + | { + mode: "poll"; + poll( + handle: Handle, + context: DurableBatchContext, + ): Promise; + intervalMs?: number; + } + | { + mode: "webhook"; + source: string; + externalId(handle: Handle, context: DurableBatchContext): string; + pollFallback?: { + afterMs: number; + intervalMs?: number; + poll( + handle: Handle, + context: DurableBatchContext, + ): Promise; + }; + }; + +export interface DurableBatchTaskItem< + Input, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +> { + id: string; + input: Input; + expected: Expected; + metadata: Metadata; + tags: string[] | undefined; + parameters: InferParameters; + trialIndex: number; +} + +export type DurableBatchTaskResult = + | { + id: string; + output: Output; + metadata?: Metadata; + tags?: string[]; + } + | { id: string; error: unknown; retryable?: boolean }; + +export type DurableBatchScorerItem< + Input, + Output, + Expected, + Metadata extends BaseMetadata, +> = EvalScorerArgs & { + id: string; + trialIndex: number; +}; + +export type DurableBatchScorerResult = + | { id: string; score: OneOrMoreScores } + | { id: string; error: unknown; retryable?: boolean }; + +export interface DurableBatchProcessor { + revision: string; + batchSize?: number; + maxConcurrentBatches?: number; + maxAttempts?: number; + submit(items: Item[], context: DurableBatchContext): Promise; + recover?(context: DurableBatchContext): Promise>; + completion: DurableBatchCompletion; + collect( + handle: Handle, + context: DurableBatchContext, + ): AsyncIterable | Promise | AsyncIterable>; + cancel?(handle: Handle, context: DurableBatchContext): Promise; +} + +export interface DurableBatchTask< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, + Handle extends JsonValue, +> extends DurableBatchProcessor< + DurableBatchTaskItem, + DurableBatchTaskResult, + Handle +> { + readonly kind: typeof BATCH_TASK_KIND; +} + +export interface DurableBatchScorer< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Handle extends JsonValue, +> extends DurableBatchProcessor< + DurableBatchScorerItem, + DurableBatchScorerResult, + Handle +> { + readonly kind: typeof BATCH_SCORER_KIND; + name: string; +} + +export function BatchTask< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, + Handle extends JsonValue = JsonValue, +>( + processor: DurableBatchProcessor< + DurableBatchTaskItem, + DurableBatchTaskResult, + Handle + >, +): DurableBatchTask { + return { kind: BATCH_TASK_KIND, ...processor }; +} + +export function BatchScorer< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Handle extends JsonValue = JsonValue, +>( + processor: DurableBatchProcessor< + DurableBatchScorerItem, + DurableBatchScorerResult, + Handle + > & { name: string }, +): DurableBatchScorer { + return { kind: BATCH_SCORER_KIND, ...processor }; +} + +export type DurableEvaluator< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +> = { + revision: string; + data: EvalData; + caseId?: ( + datum: EvalCase, + ) => string | Promise; + task: + | EvalTask + | DurableBatchTask< + Input, + Output, + Expected, + Metadata, + Parameters, + JsonValue + >; + scores: Array< + | EvalScorer + | DurableBatchScorer + >; + parameters?: InferParameters; + experimentName?: string; + description?: string; + metadata?: Record; + tags?: string[]; + trialCount?: number; + projectId?: string; + state?: BraintrustState; +}; + +export interface DurableEvalRuntimeOptions { + runId?: string; + store?: DurableEvalStore; + shard?: { index: number; count: number }; + deadlineMs?: number; + signal?: AbortSignal; + noSendLogs?: boolean; + workerId?: string; + checkpointDir?: string; +} + +export interface DurableEvalExistingRunOptions extends DurableEvalRuntimeOptions { + runId: string; +} + +export type DurableEvalOperation = + | "run" + | "status" + | "retry-failed" + | "resubmit-unknown" + | "cancel"; + +type DurableEvalExecutionOptions = DurableEvalRuntimeOptions & { + operation?: DurableEvalOperation; +}; + +export type DurableBatchResultEvent = { + eventId: string; + source: string; + externalId?: string; + batchId?: string; + handle?: JsonValue; + outcome: DurableBatchOutcome; + payload?: JsonValue; +}; + +export type DurableBatchProcessingResult = + | { + status: "processed"; + batchId: string; + run: DurableEvalResult; + } + | { status: "duplicate"; batchId?: string } + | { + status: "pending"; + reason: "unmatched" | "definition_missing"; + batchId?: string; + }; + +export interface DurableEvalDefinition< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +> { + readonly projectName: string; + readonly evalName: string; + readonly evaluator: DurableEvaluator< + Input, + Output, + Expected, + Metadata, + Parameters + >; + run(options?: DurableEvalRuntimeOptions): Promise; + status(options: DurableEvalExistingRunOptions): Promise; + retryFailed( + options: DurableEvalExistingRunOptions, + ): Promise; + resubmitUnknown( + options: DurableEvalExistingRunOptions, + ): Promise; + cancel(options: DurableEvalExistingRunOptions): Promise; + processBatchResult( + event: DurableBatchResultEvent, + options?: Omit, + ): Promise; +} + +export interface DurableEvalProgress { + total: number; + taskPending: number; + taskSucceeded: number; + taskFailed: number; + scorePending: number; + scoreSucceeded: number; + scoreFailed: number; + unknown: number; +} + +export interface DurableEvalFailureSummary { + tasks: number; + scorers: number; +} + +export type DurableEvalPauseReason = + | "deadline" + | "aborted" + | "shard_complete" + | "waiting_for_webhook" + | "unknown_submission" + | "provider_unreachable" + | "status_only"; + +export type DurableEvalResult = + | { + status: "completed"; + runId: string; + summary: ExperimentSummary; + progress: DurableEvalProgress; + failures: DurableEvalFailureSummary; + } + | { + status: "paused"; + runId: string; + reason: DurableEvalPauseReason; + progress: DurableEvalProgress; + resume( + overrides?: Partial, + ): Promise; + }; + +type SerializedError = { + name?: string; + message: string; + stack?: string; +}; + +type StageState = + | { status: "pending"; attempts: number } + | { + status: "leased"; + attempts: number; + workerId: string; + leaseUntil: number; + } + | { + status: "in_batch"; + attempts: number; + batchId: string; + revision: string; + } + | { + status: "succeeded"; + attempts: number; + revision: string; + value: JsonValue; + } + | { + status: "failed"; + attempts: number; + revision: string; + error: SerializedError; + } + | { + status: "unknown"; + attempts: number; + revision: string; + batchId: string; + }; + +type DurableCaseRecord = { + id: string; + caseId: string; + trialIndex: number; + shard: number; + datum: JsonValue; + task: StageState; + scores: Record; + metadata: JsonValue; + tags?: string[]; + logPending?: boolean; +}; + +type DurableJobRecord = { + id: string; + stage: string; + kind: "task" | "score"; + scorerName?: string; + itemIds: string[]; + attempt: number; + revision: string; + shard: number; + status: + | "preparing" + | "submitting" + | "submitted" + | "complete" + | "failed" + | "unknown"; + workerId?: string; + leaseUntil?: number; + handle?: JsonValue; + external?: { source: string; id: string }; + submittedAt?: number; + outcome?: DurableBatchOutcome; + webhookEventKey?: string; + nextPollAt?: number; + error?: SerializedError; +}; + +type DurableRunManifest = { + schemaVersion: number; + runId: string; + projectName: string; + evalName: string; + revision: string; + shardCount: number; + dataSealed: boolean; + activeScorers: string[]; + experimentName?: string; + createdAt: string; +}; + +type Versioned = { value: T; version: string }; + +type DurableBatchLocator = { + schemaVersion: number; + projectName: string; + evalName: string; + runId: string; + prefix: string; + jobKey: string; + batchId: string; + stage: string; + kind: "task" | "score"; + scorerName?: string; + shard: number; + shardCount: number; +}; + +type DurableWebhookEventRecord = { + schemaVersion: number; + event: DurableBatchResultEvent; + receivedAt: string; + status: "pending" | "applied"; + batchId?: string; +}; + +class DurableEvalDefinitionImpl< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +> implements DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters +> { + readonly evalName: string; + + constructor( + readonly projectName: string, + readonly evaluator: DurableEvaluator< + Input, + Output, + Expected, + Metadata, + Parameters + >, + ) { + this.evalName = evaluator.experimentName ?? projectName; + } + + run(options: DurableEvalRuntimeOptions = {}) { + return runDurableEval(this.projectName, this.evaluator, options); + } + + status(options: DurableEvalExistingRunOptions) { + return runDurableEval(this.projectName, this.evaluator, { + ...options, + operation: "status", + }); + } + + retryFailed(options: DurableEvalExistingRunOptions) { + return runDurableEval(this.projectName, this.evaluator, { + ...options, + operation: "retry-failed", + }); + } + + resubmitUnknown(options: DurableEvalExistingRunOptions) { + return runDurableEval(this.projectName, this.evaluator, { + ...options, + operation: "resubmit-unknown", + }); + } + + cancel(options: DurableEvalExistingRunOptions) { + return runDurableEval(this.projectName, this.evaluator, { + ...options, + operation: "cancel", + }); + } + + processBatchResult( + event: DurableBatchResultEvent, + options: Omit = {}, + ): Promise { + return processDurableBatchResult(this, event, options); + } +} + +export function DurableEval< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +>( + projectName: string, + evaluator: DurableEvaluator, +): DurableEvalDefinition { + const definition = new DurableEvalDefinitionImpl(projectName, evaluator); + if (globalThis._lazy_load) { + globalThis._evals.durableEvaluators ??= {}; + globalThis._evals.durableEvaluators[definition.evalName] = { + definition, + }; + } + return definition; +} + +async function processDurableBatchResult< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + event: DurableBatchResultEvent, + options: Omit, +): Promise { + if (!event.eventId.trim() || !event.source.trim()) { + throw new Error( + "Durable batch webhook events require non-empty eventId and source", + ); + } + if (!event.batchId && !event.externalId) { + throw new Error( + "Durable batch webhook events require batchId or externalId", + ); + } + if (event.externalId !== undefined && !event.externalId.trim()) { + throw new Error("Durable batch webhook externalId must be non-empty"); + } + const normalizedEvent = assertJsonValue( + event, + "durable batch webhook event", + ) as DurableBatchResultEvent; + const store = + options.store ?? + new FileDurableEvalStore(options.checkpointDir ?? ".braintrust/evals"); + const storedEventKey = webhookEventKey(event.source, event.eventId); + const eventRecord: DurableWebhookEventRecord = { + schemaVersion: CHECKPOINT_VERSION, + event: normalizedEvent, + receivedAt: new Date().toISOString(), + status: "pending", + batchId: event.batchId, + }; + const inserted = await writeJson(store, storedEventKey, eventRecord, { + ifAbsent: true, + }); + if (!inserted.written) { + const existing = await readJson( + store, + storedEventKey, + ); + if ( + !existing || + stableStringify(existing.value.event) !== stableStringify(normalizedEvent) + ) { + throw new Error( + `Webhook event ${event.source}/${event.eventId} was reused with different contents`, + ); + } + if (existing.value.status === "applied") { + return { + status: "duplicate", + batchId: existing.value.batchId ?? event.batchId, + }; + } + } + + if (event.externalId) { + await writeJson( + store, + webhookMailboxKey(event.source, event.externalId, storedEventKey), + { eventKey: storedEventKey }, + { ifAbsent: true }, + ); + } + + const internalLocator = event.batchId + ? await readJson( + store, + internalBatchIndexKey(event.batchId), + ) + : undefined; + const externalLocator = event.externalId + ? await readJson( + store, + externalBatchIndexKey(event.source, event.externalId), + ) + : undefined; + if ( + internalLocator && + externalLocator && + internalLocator.value.jobKey !== externalLocator.value.jobKey + ) { + throw new Error( + "Durable batch webhook batchId and externalId resolve to different jobs", + ); + } + const locator = internalLocator?.value ?? externalLocator?.value; + if (!locator) { + return { status: "pending", reason: "unmatched", batchId: event.batchId }; + } + if ( + locator.projectName !== definition.projectName || + locator.evalName !== definition.evalName + ) { + return { + status: "pending", + reason: "definition_missing", + batchId: locator.batchId, + }; + } + + const processor = batchProcessorForLocator(definition.evaluator, locator); + if ( + !processor || + processor.completion.mode !== "webhook" || + processor.completion.source !== event.source + ) { + return { + status: "pending", + reason: "definition_missing", + batchId: locator.batchId, + }; + } + + const attached = await attachWebhookEventToJob( + store, + locator.jobKey, + storedEventKey, + normalizedEvent, + ); + if (attached === "missing") { + return { + status: "pending", + reason: "unmatched", + batchId: locator.batchId, + }; + } + if (attached === "duplicate") { + await markWebhookEventApplied(store, storedEventKey, locator.batchId); + return { status: "duplicate", batchId: locator.batchId }; + } + + const run = await runDurableEval( + definition.projectName, + definition.evaluator, + { + ...options, + runId: locator.runId, + shard: { index: locator.shard, count: locator.shardCount }, + }, + ); + const job = await readJson(store, locator.jobKey); + if ( + job && + (job.value.status === "complete" || job.value.status === "failed") + ) { + await markWebhookEventApplied(store, storedEventKey, locator.batchId); + } + return { status: "processed", batchId: locator.batchId, run }; +} + +async function runDurableEval< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +>( + projectName: string, + evaluator: DurableEvaluator, + options: DurableEvalExecutionOptions = {}, +): Promise { + const runId = options.runId ?? newId(); + const shard = options.shard ?? { index: 0, count: 1 }; + validateShard(shard); + const workerId = options.workerId ?? newId(); + const store = + options.store ?? + new FileDurableEvalStore(options.checkpointDir ?? ".braintrust/evals"); + const evalName = evaluator.experimentName ?? projectName; + const prefix = runPrefix(projectName, evalName, runId); + const manifestKey = `${prefix}/manifest`; + const scorerDefinitions = resolveScorers(evaluator.scores); + const scorerNames = scorerDefinitions.map((definition) => definition.name); + const webhookStages = new Set(); + if ( + isBatchTask(evaluator.task) && + evaluator.task.completion.mode === "webhook" && + evaluator.task.completion.pollFallback === undefined + ) { + webhookStages.add("task"); + } + for (const definition of scorerDefinitions) { + if ( + isBatchScorer(definition.scorer) && + definition.scorer.completion.mode === "webhook" && + definition.scorer.completion.pollFallback === undefined + ) { + webhookStages.add(`score:${definition.name}`); + } + } + if (new Set(scorerNames).size !== scorerNames.length) { + throw new Error("DurableEval scorer names must be unique"); + } + + const existingManifest = await readJson( + store, + manifestKey, + ); + if (options.operation && options.operation !== "run" && !existingManifest) { + throw new Error(`DurableEval run ${runId} does not exist`); + } + let manifest = + existingManifest ?? + (await initializeManifest(store, manifestKey, { + schemaVersion: CHECKPOINT_VERSION, + runId, + projectName, + evalName, + revision: evaluator.revision, + shardCount: shard.count, + dataSealed: false, + activeScorers: scorerNames, + experimentName: evaluator.experimentName, + createdAt: new Date().toISOString(), + })); + if (manifest.value.shardCount !== shard.count) { + throw new Error( + `DurableEval run ${runId} was created with ${manifest.value.shardCount} shards, not ${shard.count}`, + ); + } + + if (!manifest.value.dataSealed) { + await materializeData({ + store, + prefix, + evaluator, + shardCount: shard.count, + }); + } + manifest = await updateManifest(store, manifestKey, (current) => ({ + ...current, + revision: evaluator.revision, + activeScorers: scorerNames, + dataSealed: true, + })); + + const experiment: Experiment | null = options.noSendLogs + ? null + : initExperiment({ + state: evaluator.state, + ...(evaluator.projectId + ? { projectId: evaluator.projectId } + : { project: projectName }), + experiment: manifest.value.experimentName, + update: manifest.value.experimentName !== undefined, + description: evaluator.description, + metadata: evaluator.metadata, + tags: evaluator.tags, + setCurrent: false, + }); + + if (experiment && !manifest.value.experimentName) { + const summary = await experiment.summarize({ summarizeScores: false }); + manifest = await updateManifest(store, manifestKey, (current) => ({ + ...current, + experimentName: summary.experimentName, + })); + } + + await reconcileScorers(store, prefix, scorerNames); + + if (options.operation === "retry-failed") { + await resetStages(store, prefix, scorerNames, "failed"); + } else if (options.operation === "resubmit-unknown") { + await resetStages(store, prefix, scorerNames, "unknown"); + } else if (options.operation === "cancel") { + await cancelRun({ + store, + prefix, + evaluator, + scorerDefinitions, + runId, + shard, + signal: options.signal ?? new AbortController().signal, + }); + } + + if (options.operation === "status" || options.operation === "cancel") { + if (options.operation === "cancel" && experiment) { + await flushPendingLogs({ + store, + prefix, + experiment, + scorerNames, + shard, + runId, + }); + await experiment.flush(); + } + const progress = await collectProgress(store, prefix, scorerNames); + const complete = await isComplete(store, prefix, scorerNames); + if (!complete) { + return pausedResult(runId, "status_only", progress, (overrides = {}) => + runDurableEval(projectName, evaluator, { + ...options, + ...overrides, + operation: "run", + runId, + store, + shard, + }), + ); + } + return { + status: "completed", + runId, + summary: await buildLocalDurableSummary( + store, + prefix, + projectName, + manifest.value.experimentName ?? evalName, + scorerNames, + ), + progress, + failures: { + tasks: progress.taskFailed, + scorers: progress.scoreFailed, + }, + }; + } + + const startedAt = Date.now(); + const deadlineAt = + options.deadlineMs === undefined + ? undefined + : startedAt + Math.max(options.deadlineMs, 0); + const controller = new AbortController(); + const abortHandler = () => controller.abort(); + options.signal?.addEventListener("abort", abortHandler, { once: true }); + + const resume = (overrides: Partial = {}) => + runDurableEval(projectName, evaluator, { + ...options, + ...overrides, + runId, + store, + shard, + operation: "run", + }); + + try { + while (true) { + if (options.signal?.aborted) { + return pausedResult( + runId, + "aborted", + await collectProgress(store, prefix, scorerNames), + resume, + ); + } + if (deadlineAt !== undefined && Date.now() >= deadlineAt) { + return pausedResult( + runId, + "deadline", + await collectProgress(store, prefix, scorerNames), + resume, + ); + } + + let changed = false; + changed = + (await runTaskPass({ + store, + prefix, + projectName, + evalName, + evaluator, + shard, + workerId, + runId, + signal: controller.signal, + })) || changed; + for (const scorer of scorerDefinitions) { + changed = + (await runScorePass({ + store, + prefix, + projectName, + evalName, + scorer, + evaluatorRevision: evaluator.revision, + shard, + workerId, + runId, + signal: controller.signal, + })) || changed; + } + + if (experiment) { + changed = + (await flushPendingLogs({ + store, + prefix, + experiment, + scorerNames, + shard, + runId, + })) || changed; + } + + const progress = await collectProgress(store, prefix, scorerNames); + if (progress.unknown > 0) { + return pausedResult(runId, "unknown_submission", progress, resume); + } + const shardComplete = await isComplete( + store, + prefix, + scorerNames, + shard.index, + ); + const globallyComplete = await isComplete(store, prefix, scorerNames); + if (globallyComplete) { + if (experiment) { + await experiment.flush(); + } + const summary = experiment + ? await experiment.summarize() + : await buildLocalDurableSummary( + store, + prefix, + projectName, + manifest.value.experimentName ?? evalName, + scorerNames, + ); + return { + status: "completed", + runId, + summary, + progress, + failures: { + tasks: progress.taskFailed, + scorers: progress.scoreFailed, + }, + }; + } + if (shardComplete && shard.count > 1) { + return pausedResult(runId, "shard_complete", progress, resume); + } + if (!changed && (await hasProviderErrorJob(store, prefix, shard.index))) { + return pausedResult(runId, "provider_unreachable", progress, resume); + } + if ( + !changed && + (await hasWaitingWebhookJob(store, prefix, shard.index, webhookStages)) + ) { + return pausedResult(runId, "waiting_for_webhook", progress, resume); + } + if (!changed) { + await delay(Math.min(timeRemaining(deadlineAt) ?? 1_000, 1_000)); + } + } + } finally { + options.signal?.removeEventListener("abort", abortHandler); + } +} + +function pausedResult( + runId: string, + reason: DurableEvalPauseReason, + progress: DurableEvalProgress, + resume: ( + overrides?: Partial, + ) => Promise, +): DurableEvalResult { + return { status: "paused", runId, reason, progress, resume }; +} + +function emptyProgress(): DurableEvalProgress { + return { + total: 0, + taskPending: 0, + taskSucceeded: 0, + taskFailed: 0, + scorePending: 0, + scoreSucceeded: 0, + scoreFailed: 0, + unknown: 0, + }; +} + +async function initializeManifest( + store: DurableEvalStore, + key: string, + initial: DurableRunManifest, +): Promise> { + const existing = await readJson(store, key); + if (existing) return existing; + const result = await writeJson(store, key, initial, { ifAbsent: true }); + if (result.written) return { value: initial, version: result.version }; + const raced = await readJson(store, key); + if (!raced) throw new Error("DurableEval manifest initialization failed"); + return raced; +} + +async function updateManifest( + store: DurableEvalStore, + key: string, + update: (manifest: DurableRunManifest) => DurableRunManifest, +) { + while (true) { + const current = await readJson(store, key); + if (!current) throw new Error("DurableEval manifest is missing"); + const next = update(current.value); + const result = await writeJson(store, key, next, { + ifVersion: current.version, + }); + if (result.written) return { value: next, version: result.version }; + } +} + +async function materializeData< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>({ + store, + prefix, + evaluator, + shardCount, +}: { + store: DurableEvalStore; + prefix: string; + evaluator: DurableEvaluator; + shardCount: number; +}) { + const rawData = + typeof evaluator.data === "function" ? evaluator.data() : evaluator.data; + const resolved = rawData instanceof Promise ? await rawData : rawData; + if ( + typeof resolved === "object" && + resolved !== null && + "_type" in resolved + ) { + throw new Error( + "DurableEval does not yet support BaseExperiment data sources", + ); + } + if (!isIterable(resolved) && !isAsyncIterable(resolved)) { + throw new Error("DurableEval data must be an iterable or async iterable"); + } + + const seen = new Set(); + for await (const datum of toAsyncIterable(resolved)) { + const caseId = + datum.id ?? + datum.upsert_id ?? + (evaluator.caseId ? await evaluator.caseId(datum) : undefined); + if (!caseId || typeof caseId !== "string") { + throw new Error( + "Every DurableEval case must have a non-empty id/upsert_id or be resolved by caseId", + ); + } + if (seen.has(caseId)) { + throw new Error(`Duplicate DurableEval 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 DurableEval case ${caseId}`); + } + for (let trialIndex = 0; trialIndex < trialCount; trialIndex++) { + const id = workItemId(caseId, trialIndex); + const record: DurableCaseRecord = { + id, + caseId, + trialIndex, + shard: stableShard(id, shardCount), + datum: assertJsonValue(datum, `case ${caseId}`), + task: { status: "pending", attempts: 0 }, + scores: {}, + metadata: assertJsonValue( + "metadata" in datum ? datum.metadata : {}, + `case ${caseId} metadata`, + ), + tags: datum.tags, + }; + const key = caseKey(prefix, id); + const result = await writeJson(store, key, record, { ifAbsent: true }); + if (!result.written) { + const existing = await readJson(store, key); + if ( + !existing || + stableStringify(existing.value.datum) !== + stableStringify(record.datum) + ) { + throw new Error( + `DurableEval case ${caseId} changed while the run was being prepared`, + ); + } + } + } + } +} + +type ResolvedScorer = { + name: string; + revision: string; + // Runtime orchestration intentionally erases user generics after public + // type-checking at the DurableEval boundary. + scorer: // eslint-disable-next-line @typescript-eslint/no-explicit-any + | EvalScorer + // eslint-disable-next-line @typescript-eslint/no-explicit-any + | DurableBatchScorer; +}; + +function resolveScorers( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + scorers: DurableEvaluator["scores"], +): ResolvedScorer[] { + return scorers.map((scorer) => { + if (isBatchScorer(scorer)) { + if (!scorer.name.trim()) { + throw new Error("DurableEval batch scorers must have a name"); + } + return { name: scorer.name, revision: scorer.revision, scorer }; + } + if (!scorer.name) { + throw new Error("DurableEval scorers must be named functions"); + } + return { name: scorer.name, revision: "unversioned", scorer }; + }); +} + +async function reconcileScorers( + store: DurableEvalStore, + prefix: string, + scorerNames: string[], +) { + for await (const record of listCases(store, prefix)) { + const missing = scorerNames.filter( + (name) => record.value.scores[name] === undefined, + ); + if (!missing.length) continue; + const next = structuredClone(record.value); + for (const name of missing) { + next.scores[name] = { status: "pending", attempts: 0 }; + } + await writeJson(store, caseKey(prefix, record.value.id), next, { + ifVersion: record.version, + }); + } +} + +async function resetStages( + store: DurableEvalStore, + prefix: string, + scorerNames: string[], + status: "failed" | "unknown", +) { + for await (const current of listCases(store, prefix)) { + const next = structuredClone(current.value); + let changed = false; + if (next.task.status === status) { + next.task = { + status: "pending", + attempts: next.task.attempts, + }; + changed = true; + } + for (const name of scorerNames) { + const state = next.scores[name]; + if (state?.status === status) { + next.scores[name] = { + status: "pending", + attempts: state.attempts, + }; + changed = true; + } + } + if (changed) { + await writeJson(store, caseKey(prefix, next.id), next, { + ifVersion: current.version, + }); + } + } +} + +async function cancelRun({ + store, + prefix, + evaluator, + scorerDefinitions, + runId, + shard, + signal, +}: { + store: DurableEvalStore; + prefix: string; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + evaluator: DurableEvaluator; + scorerDefinitions: ResolvedScorer[]; + runId: string; + shard: { index: number; count: number }; + signal: AbortSignal; +}) { + const processors = new Map< + string, + DurableBatchProcessor + >(); + if (isBatchTask(evaluator.task)) { + processors.set( + "task", + evaluator.task as DurableBatchProcessor, + ); + } + for (const definition of scorerDefinitions) { + if (isBatchScorer(definition.scorer)) { + processors.set( + `score:${definition.name}`, + definition.scorer as DurableBatchProcessor, + ); + } + } + + for await (const key of store.list(`${prefix}/jobs/`)) { + const current = await readJson(store, key); + if (!current || current.value.shard !== shard.index) continue; + const job = structuredClone(current.value); + const processor = processors.get(job.stage); + if ( + job.status === "submitted" && + job.handle !== undefined && + processor?.cancel + ) { + await processor.cancel( + job.handle, + batchContext(runId, job.stage, job, shard, signal), + ); + } + if ( + job.status === "preparing" || + job.status === "submitting" || + job.status === "submitted" || + job.status === "unknown" + ) { + job.status = "failed"; + job.error = { name: "CancelledError", message: "Durable eval cancelled" }; + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, key, job, { ifVersion: current.version }); + } + } + + for await (const current of listCases(store, prefix)) { + if (current.value.shard !== shard.index) continue; + const next = structuredClone(current.value); + let changed = false; + if (!isTerminal(next.task)) { + next.task = { + status: "failed", + attempts: next.task.attempts, + revision: evaluator.revision, + error: { + name: "CancelledError", + message: "Durable eval cancelled", + }, + }; + changed = true; + } + if (next.task.status === "succeeded") { + for (const definition of scorerDefinitions) { + const state = next.scores[definition.name]; + if (!isTerminal(state)) { + next.scores[definition.name] = { + status: "failed", + attempts: state?.attempts ?? 0, + revision: definition.revision, + error: { + name: "CancelledError", + message: "Durable eval cancelled", + }, + }; + changed = true; + } + } + } + if (changed) { + next.logPending = true; + await writeJson(store, caseKey(prefix, next.id), next, { + ifVersion: current.version, + }); + } + } +} + +async function runTaskPass({ + store, + prefix, + projectName, + evalName, + evaluator, + shard, + workerId, + runId, + signal, +}: { + store: DurableEvalStore; + prefix: string; + projectName: string; + evalName: string; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + evaluator: DurableEvaluator; + shard: { index: number; count: number }; + workerId: string; + runId: string; + signal: AbortSignal; +}) { + const task = evaluator.task; + if (isBatchTask(task)) { + return await runBatchStage({ + store, + prefix, + projectName, + evalName, + processor: task, + stage: "task", + kind: "task", + shard, + workerId, + runId, + signal, + makeItem: (record) => taskBatchItem(record, evaluator.parameters ?? {}), + applyResult: (record, result, revision, attempt) => + applyTaskResult( + record, + result, + revision, + attempt, + task.maxAttempts ?? 3, + ), + }); + } + const localTask = task as EvalTask< + unknown, + unknown, + unknown, + Record, + EvalParameters + >; + + let changed = false; + for await (const current of listCases(store, prefix)) { + if ( + current.value.shard !== shard.index || + !isClaimable(current.value.task) + ) { + continue; + } + const claimed = structuredClone(current.value); + claimed.task = { + status: "leased", + attempts: current.value.task.attempts, + workerId, + leaseUntil: Date.now() + 60_000, + }; + const claim = await writeJson( + store, + caseKey(prefix, current.value.id), + claimed, + { ifVersion: current.version }, + ); + if (!claim.written) continue; + const datum = current.value.datum as EvalCase< + unknown, + unknown, + BaseMetadata + >; + const attempt = current.value.task.attempts + 1; + const next = structuredClone(current.value); + try { + let metadata = { ...(current.value.metadata as Record) }; + const hooks: EvalHooks< + unknown, + Record, + EvalParameters + > = { + meta: (value) => { + metadata = { ...metadata, ...(value as Record) }; + }, + metadata, + expected: "expected" in datum ? datum.expected : undefined, + span: NOOP_SPAN, + parameters: evaluator.parameters ?? {}, + reportProgress: () => undefined, + trialIndex: current.value.trialIndex, + tags: current.value.tags, + }; + const output = await localTask(datum.input, hooks); + next.metadata = assertJsonValue(hooks.metadata, "task metadata"); + next.tags = hooks.tags; + next.task = { + status: "succeeded", + attempts: attempt, + revision: evaluator.revision, + value: assertJsonValue(output, "task output"), + }; + } catch (error) { + next.task = + attempt < 3 + ? { status: "pending", attempts: attempt } + : { + status: "failed", + attempts: attempt, + revision: evaluator.revision, + error: serializeError(error), + }; + } + next.logPending = true; + const result = await writeJson( + store, + caseKey(prefix, current.value.id), + next, + { ifVersion: claim.version }, + ); + changed = result.written || changed; + } + return changed; +} + +async function runScorePass({ + store, + prefix, + projectName, + evalName, + scorer, + shard, + workerId, + runId, + signal, +}: { + store: DurableEvalStore; + prefix: string; + projectName: string; + evalName: string; + scorer: ResolvedScorer; + evaluatorRevision: string; + shard: { index: number; count: number }; + workerId: string; + runId: string; + signal: AbortSignal; +}) { + if (isBatchScorer(scorer.scorer)) { + const batchScorer = scorer.scorer; + return await runBatchStage({ + store, + prefix, + projectName, + evalName, + processor: batchScorer, + stage: `score:${scorer.name}`, + kind: "score", + scorerName: scorer.name, + shard, + workerId, + runId, + signal, + eligible: (record) => record.task.status === "succeeded", + makeItem: (record) => scorerBatchItem(record), + applyResult: (record, result, revision, attempt) => + applyScoreResult( + record, + scorer.name, + result, + revision, + attempt, + batchScorer.maxAttempts ?? 3, + ), + }); + } + + let changed = false; + for await (const current of listCases(store, prefix)) { + const state = current.value.scores[scorer.name]; + if ( + current.value.shard !== shard.index || + current.value.task.status !== "succeeded" || + !isClaimable(state) + ) { + continue; + } + const claimed = structuredClone(current.value); + claimed.scores[scorer.name] = { + status: "leased", + attempts: state.attempts, + workerId, + leaseUntil: Date.now() + 60_000, + }; + const claim = await writeJson( + store, + caseKey(prefix, current.value.id), + claimed, + { ifVersion: current.version }, + ); + if (!claim.written) continue; + const datum = current.value.datum as EvalCase< + unknown, + unknown, + BaseMetadata + >; + const attempt = state.attempts + 1; + const next = structuredClone(current.value); + try { + const raw = await scorer.scorer({ + input: datum.input, + expected: "expected" in datum ? datum.expected : undefined, + metadata: current.value.metadata as BaseMetadata, + output: current.value.task.value, + }); + next.scores[scorer.name] = { + status: "succeeded", + attempts: attempt, + revision: scorer.revision, + value: assertJsonValue( + normalizeScores(raw, scorer.name), + `scorer ${scorer.name} output`, + ), + }; + } catch (error) { + next.scores[scorer.name] = + attempt < 3 + ? { status: "pending", attempts: attempt } + : { + status: "failed", + attempts: attempt, + revision: scorer.revision, + error: serializeError(error), + }; + } + next.logPending = true; + const result = await writeJson( + store, + caseKey(prefix, current.value.id), + next, + { ifVersion: claim.version }, + ); + changed = result.written || changed; + } + return changed; +} + +async function runBatchStage({ + store, + prefix, + projectName, + evalName, + processor, + stage, + kind, + scorerName, + shard, + workerId, + runId, + signal, + eligible = () => true, + makeItem, + applyResult, +}: { + store: DurableEvalStore; + prefix: string; + projectName: string; + evalName: string; + processor: DurableBatchProcessor; + stage: string; + kind: "task" | "score"; + scorerName?: string; + shard: { index: number; count: number }; + workerId: string; + runId: string; + signal: AbortSignal; + eligible?: (record: DurableCaseRecord) => boolean; + makeItem: (record: DurableCaseRecord) => Item; + applyResult: ( + record: DurableCaseRecord, + result: Result, + revision: string, + attempt: number, + ) => DurableCaseRecord; +}) { + const batchSize = processor.batchSize ?? DEFAULT_BATCH_SIZE; + const maxConcurrentBatches = + processor.maxConcurrentBatches ?? DEFAULT_MAX_CONCURRENT_BATCHES; + if (!Number.isInteger(batchSize) || batchSize < 1) { + throw new Error(`Invalid batchSize for durable stage ${stage}`); + } + if (!Number.isInteger(maxConcurrentBatches) || maxConcurrentBatches < 1) { + throw new Error(`Invalid maxConcurrentBatches for durable stage ${stage}`); + } + + let changed = false; + let activeJobs = 0; + for await (const current of listJobs(store, prefix, stage)) { + let job = current.value; + if (job.shard !== shard.index) { + continue; + } + if (job.status === "preparing") { + activeJobs++; + if ((job.leaseUntil ?? 0) > Date.now()) continue; + const preparationError = new Error( + "Batch preparation lease expired before provider submission", + ); + await retryBatchItems({ + store, + prefix, + job, + scorerName, + processor, + error: preparationError, + retryable: true, + }); + job.status = "failed"; + job.error = serializeError(preparationError); + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: current.version, + }); + activeJobs--; + changed = true; + continue; + } + if (!["submitting", "submitted"].includes(job.status)) continue; + + const completion = processor.completion; + const fallback = + completion.mode === "webhook" ? completion.pollFallback : undefined; + const fallbackReady = + fallback !== undefined && + Date.now() >= (job.submittedAt ?? Date.now()) + fallback.afterMs; + const needsProviderCheck = + job.status === "submitting" || + job.outcome !== undefined || + completion.mode === "poll" || + fallbackReady; + + if (!needsProviderCheck) { + activeJobs++; + continue; + } + if ((job.leaseUntil ?? 0) > Date.now()) { + activeJobs++; + continue; + } + + const claimedJob = { + ...job, + workerId, + leaseUntil: Date.now() + JOB_LEASE_MS, + }; + const jobClaim = await writeJson(store, current.key, claimedJob, { + ifVersion: current.version, + }); + if (!jobClaim.written) continue; + job = claimedJob; + let jobVersion = current.version; + jobVersion = jobClaim.version; + const context = batchContext(runId, stage, job, shard, signal); + let handle = job.handle as Handle | undefined; + if (!handle) { + const recovery = processor.recover + ? await processor.recover(context) + : { status: "unknown" as const }; + if (recovery.status === "found") { + handle = recovery.handle; + const external = externalBatchReference(processor, handle, context); + const recovered = await persistSubmittedJob({ + store, + jobKey: current.key, + handle: assertJsonValue(handle, `${stage} batch handle`), + external, + submittedAt: job.submittedAt ?? Date.now(), + }); + job = recovered.value; + jobVersion = recovered.version; + const locator = batchLocator({ + projectName, + evalName, + runId, + prefix, + jobKey: current.key, + job, + shardCount: shard.count, + }); + if (external) { + await registerExternalBatchLocator(store, external, locator); + } + await attachPendingWebhookEvents(store, current.key, job); + const latest = await readJson(store, current.key); + if (latest) { + job = latest.value; + jobVersion = latest.version; + } + changed = true; + } else if (recovery.status === "not_found") { + await retryBatchItems({ + store, + prefix, + job, + scorerName, + processor, + error: new Error("Provider confirmed batch was not submitted"), + retryable: true, + }); + job.status = "failed"; + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: jobVersion, + }); + changed = true; + continue; + } else { + await markBatchUnknown(store, prefix, job, scorerName); + job.status = "unknown"; + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: jobVersion, + }); + changed = true; + continue; + } + } + + let poll: DurableBatchPoll | undefined = job.outcome; + if ( + job.nextPollAt && + job.nextPollAt > Date.now() && + (!fallbackReady || job.outcome !== undefined) + ) { + activeJobs++; + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: jobVersion, + }); + continue; + } + if (!poll) { + const poller = + completion.mode === "poll" + ? completion + : fallbackReady + ? fallback + : undefined; + if (!poller) { + activeJobs++; + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: jobVersion, + }); + continue; + } + try { + poll = await poller.poll(handle, context); + delete job.error; + } catch (error) { + activeJobs++; + job.nextPollAt = Date.now() + (poller.intervalMs ?? 10_000); + job.error = serializeError(error); + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: jobVersion, + }); + continue; + } + } + if (poll.status === "pending") { + activeJobs++; + delete job.error; + job.nextPollAt = + Date.now() + + (poll.retryAfterMs ?? + (completion.mode === "poll" + ? completion.intervalMs + : fallback?.intervalMs) ?? + 10_000); + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: jobVersion, + }); + continue; + } + if (poll.status === "failed") { + await retryBatchItems({ + store, + prefix, + job, + scorerName, + processor, + error: poll.error, + retryable: poll.retryable ?? false, + }); + job.status = "failed"; + job.error = serializeError(poll.error); + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: jobVersion, + }); + await markWebhookApplied(store, job); + changed = true; + continue; + } + + try { + delete job.error; + const seen = new Set(); + const collected = await processor.collect(handle, context); + for await (const result of toAsyncIterable(collected)) { + const resultId = resultItemId(result); + if (seen.has(resultId) || !job.itemIds.includes(resultId)) { + throw new Error( + `Batch stage ${stage} returned an unknown or duplicate item id ${resultId}`, + ); + } + seen.add(resultId); + const caseRecord = await readJson( + store, + caseKey(prefix, resultId), + ); + if ( + !caseRecord || + !stageBelongsToJob(caseRecord.value, job.kind, job.id, scorerName) + ) { + continue; + } + const next = applyResult( + structuredClone(caseRecord.value), + result, + job.revision, + job.attempt, + ); + next.logPending = true; + await writeJson(store, caseKey(prefix, resultId), next, { + ifVersion: caseRecord.version, + }); + } + for (const itemId of job.itemIds) { + if (seen.has(itemId)) continue; + const caseRecord = await readJson( + store, + caseKey(prefix, itemId), + ); + if ( + !caseRecord || + !stageBelongsToJob(caseRecord.value, job.kind, job.id, scorerName) + ) { + continue; + } + const missing = { + id: itemId, + error: new Error(`Batch stage ${stage} returned no result`), + retryable: true, + } as Result; + const next = applyResult( + structuredClone(caseRecord.value), + missing, + job.revision, + job.attempt, + ); + next.logPending = true; + await writeJson(store, caseKey(prefix, itemId), next, { + ifVersion: caseRecord.version, + }); + } + } catch (error) { + activeJobs++; + job.error = serializeError(error); + job.nextPollAt = Date.now() + 10_000; + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: jobVersion, + }); + changed = true; + continue; + } + job.status = "complete"; + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, current.key, job, { + ifVersion: jobVersion, + }); + await markWebhookApplied(store, job); + changed = true; + } + + if (activeJobs >= maxConcurrentBatches) { + return changed; + } + const ready: Versioned[] = []; + let attempts: number | undefined; + for await (const record of listCases(store, prefix)) { + if (record.value.shard !== shard.index || !eligible(record.value)) continue; + const state = + kind === "task" ? record.value.task : record.value.scores[scorerName!]; + if (isClaimable(state)) { + attempts ??= state.attempts; + if (state.attempts !== attempts) continue; + ready.push(record); + if (ready.length >= batchSize) break; + } + } + if (!ready.length) return changed; + + const attempt = (attempts ?? 0) + 1; + const batchId = deterministicId( + `${projectName}:${evalName}:${runId}:${stage}:${ready + .map((record) => record.value.id) + .join(",")}:${attempt}`, + ); + let job: DurableJobRecord = { + id: batchId, + stage, + kind, + scorerName, + itemIds: ready.map((record) => record.value.id), + attempt, + revision: processor.revision, + shard: shard.index, + status: "preparing", + workerId, + leaseUntil: Date.now() + JOB_LEASE_MS, + }; + const jobKey = `${prefix}/jobs/${encodeURIComponent(stage)}/${batchId}`; + const created = await writeJson(store, jobKey, job, { ifAbsent: true }); + if (!created.written) return changed; + await registerBatchLocator( + store, + batchLocator({ + projectName, + evalName, + runId, + prefix, + jobKey, + job, + shardCount: shard.count, + }), + ); + + const claimed: Versioned[] = []; + for (const record of ready) { + const next = structuredClone(record.value); + const state: StageState = { + status: "in_batch", + attempts: attempt, + batchId, + revision: processor.revision, + }; + if (kind === "task") next.task = state; + else next.scores[scorerName!] = state; + const claim = await writeJson( + store, + caseKey(prefix, record.value.id), + next, + { + ifVersion: record.version, + }, + ); + if (claim.written) claimed.push(record); + } + job.itemIds = claimed.map((record) => record.value.id); + if (!claimed.length) { + job.status = "failed"; + job.error = serializeError(new Error("Batch lost all item claims")); + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, jobKey, job, { + ifVersion: created.version, + }); + return changed; + } + + job.status = "submitting"; + const prepared = await writeJson(store, jobKey, job, { + ifVersion: created.version, + }); + if (!prepared.written) return changed; + + const context = batchContext(runId, stage, job, shard, signal); + const items = claimed.map((record) => makeItem(record.value)); + try { + const handle = await processor.submit(items, context); + const external = externalBatchReference(processor, handle, context); + const submitted = await persistSubmittedJob({ + store, + jobKey, + handle: assertJsonValue(handle, `${stage} batch handle`), + external, + submittedAt: Date.now(), + }); + job = submitted.value; + const locator = batchLocator({ + projectName, + evalName, + runId, + prefix, + jobKey, + job, + shardCount: shard.count, + }); + if (external) { + await registerExternalBatchLocator(store, external, locator); + } + await attachPendingWebhookEvents(store, jobKey, job); + } catch (error) { + if (error instanceof DurableEvalNotSubmittedError) { + await retryBatchItems({ + store, + prefix, + job, + scorerName, + processor, + error, + retryable: true, + }); + job.status = "failed"; + job.error = serializeError(error); + } else { + await markBatchUnknown(store, prefix, job, scorerName); + job.status = "unknown"; + job.error = serializeError(error); + } + delete job.workerId; + delete job.leaseUntil; + await writeJson(store, jobKey, job, { + ifVersion: prepared.version, + }); + } + return true; +} + +function taskBatchItem( + record: DurableCaseRecord, + parameters: Record, +): DurableBatchTaskItem { + const datum = record.datum as EvalCase; + return { + id: record.id, + input: datum.input, + expected: "expected" in datum ? datum.expected : undefined, + metadata: record.metadata as BaseMetadata, + tags: record.tags, + parameters, + trialIndex: record.trialIndex, + }; +} + +function scorerBatchItem( + record: DurableCaseRecord, +): DurableBatchScorerItem> { + const datum = record.datum as EvalCase; + if (record.task.status !== "succeeded") { + throw new Error("Cannot score a task that has not succeeded"); + } + return { + id: record.id, + input: datum.input, + output: record.task.value, + expected: "expected" in datum ? datum.expected : undefined, + metadata: record.metadata as Record, + trialIndex: record.trialIndex, + }; +} + +function applyTaskResult( + record: DurableCaseRecord, + result: DurableBatchTaskResult, + revision: string, + attempt: number, + maxAttempts: number, +) { + if ("error" in result) { + record.task = + (result.retryable ?? false) && attempt < maxAttempts + ? { status: "pending", attempts: attempt } + : { + status: "failed", + attempts: attempt, + revision, + error: serializeError(result.error), + }; + } else { + record.task = { + status: "succeeded", + attempts: attempt, + revision, + value: assertJsonValue(result.output, "batch task output"), + }; + if (result.metadata !== undefined) { + record.metadata = assertJsonValue(result.metadata, "batch task metadata"); + } + if (result.tags !== undefined) record.tags = result.tags; + } + return record; +} + +function applyScoreResult( + record: DurableCaseRecord, + scorerName: string, + result: DurableBatchScorerResult, + revision: string, + attempt: number, + maxAttempts: number, +) { + if ("error" in result) { + record.scores[scorerName] = + (result.retryable ?? false) && attempt < maxAttempts + ? { status: "pending", attempts: attempt } + : { + status: "failed", + attempts: attempt, + revision, + error: serializeError(result.error), + }; + } else { + record.scores[scorerName] = { + status: "succeeded", + attempts: attempt, + revision, + value: assertJsonValue( + normalizeScores(result.score, scorerName), + `batch scorer ${scorerName} output`, + ), + }; + } + return record; +} + +async function retryBatchItems({ + store, + prefix, + job, + scorerName, + processor, + error, + retryable, +}: { + store: DurableEvalStore; + prefix: string; + job: DurableJobRecord; + scorerName?: string; + processor: DurableBatchProcessor; + error: unknown; + retryable: boolean; +}) { + for (const itemId of job.itemIds) { + const current = await readJson( + store, + caseKey(prefix, itemId), + ); + if (!current) continue; + const next = structuredClone(current.value); + if (!stageBelongsToJob(next, job.kind, job.id, scorerName)) { + continue; + } + const state = + retryable && job.attempt < (processor.maxAttempts ?? 3) + ? ({ status: "pending", attempts: job.attempt } satisfies StageState) + : ({ + status: "failed", + attempts: job.attempt, + revision: job.revision, + error: serializeError(error), + } satisfies StageState); + if (job.kind === "task") next.task = state; + else next.scores[scorerName!] = state; + next.logPending = true; + await writeJson(store, caseKey(prefix, itemId), next, { + ifVersion: current.version, + }); + } +} + +async function markBatchUnknown( + store: DurableEvalStore, + prefix: string, + job: DurableJobRecord, + scorerName?: string, +) { + for (const itemId of job.itemIds) { + const current = await readJson( + store, + caseKey(prefix, itemId), + ); + if (!current) continue; + const next = structuredClone(current.value); + if (!stageBelongsToJob(next, job.kind, job.id, scorerName)) continue; + const state: StageState = { + status: "unknown", + attempts: job.attempt, + revision: job.revision, + batchId: job.id, + }; + if (job.kind === "task") next.task = state; + else next.scores[scorerName!] = state; + await writeJson(store, caseKey(prefix, itemId), next, { + ifVersion: current.version, + }); + } +} + +function stageBelongsToJob( + record: DurableCaseRecord, + kind: "task" | "score", + batchId: string, + scorerName?: string, +) { + const state = kind === "task" ? record.task : record.scores[scorerName ?? ""]; + return ( + (state?.status === "in_batch" || state?.status === "unknown") && + state.batchId === batchId + ); +} + +function batchLocator({ + projectName, + evalName, + runId, + prefix, + jobKey, + job, + shardCount, +}: { + projectName: string; + evalName: string; + runId: string; + prefix: string; + jobKey: string; + job: DurableJobRecord; + shardCount: number; +}): DurableBatchLocator { + return { + schemaVersion: CHECKPOINT_VERSION, + projectName, + evalName, + runId, + prefix, + jobKey, + batchId: job.id, + stage: job.stage, + kind: job.kind, + scorerName: job.scorerName, + shard: job.shard, + shardCount, + }; +} + +async function registerBatchLocator( + store: DurableEvalStore, + locator: DurableBatchLocator, +) { + await registerLocator(store, internalBatchIndexKey(locator.batchId), locator); +} + +async function registerExternalBatchLocator( + store: DurableEvalStore, + external: { source: string; id: string }, + locator: DurableBatchLocator, +) { + await registerLocator( + store, + externalBatchIndexKey(external.source, external.id), + locator, + ); +} + +async function registerLocator( + store: DurableEvalStore, + key: string, + locator: DurableBatchLocator, +) { + const inserted = await writeJson(store, key, locator, { ifAbsent: true }); + if (inserted.written) return; + const existing = await readJson(store, key); + if (!existing || existing.value.jobKey !== locator.jobKey) { + throw new Error( + `Durable batch index collision for batch ${locator.batchId}`, + ); + } +} + +function externalBatchReference( + processor: DurableBatchProcessor, + handle: Handle, + context: DurableBatchContext, +) { + if (processor.completion.mode !== "webhook") return undefined; + if (!processor.completion.source.trim()) { + throw new Error("Durable batch webhook source must be non-empty"); + } + const id = processor.completion.externalId(handle, context); + if (!id.trim()) { + throw new Error("Durable batch webhook externalId must be non-empty"); + } + return { source: processor.completion.source, id }; +} + +async function persistSubmittedJob({ + store, + jobKey, + handle, + external, + submittedAt, +}: { + store: DurableEvalStore; + jobKey: string; + handle: JsonValue; + external?: { source: string; id: string }; + submittedAt: number; +}): Promise> { + while (true) { + const current = await readJson(store, jobKey); + if (!current) { + throw new Error("Durable batch job disappeared during submission"); + } + if ( + current.value.handle !== undefined && + stableStringify(current.value.handle) !== stableStringify(handle) + ) { + throw new Error( + `Durable batch ${current.value.id} recovered with a different handle`, + ); + } + if ( + current.value.external && + external && + (current.value.external.source !== external.source || + current.value.external.id !== external.id) + ) { + throw new Error( + `Durable batch ${current.value.id} resolved to a different external job`, + ); + } + if ( + current.value.status === "complete" || + current.value.status === "failed" + ) { + return current; + } + const next: DurableJobRecord = { + ...current.value, + status: "submitted", + handle, + external: external ?? current.value.external, + submittedAt: current.value.submittedAt ?? submittedAt, + nextPollAt: Date.now(), + }; + delete next.workerId; + delete next.leaseUntil; + const written = await writeJson(store, jobKey, next, { + ifVersion: current.version, + }); + if (written.written) { + return { value: next, version: written.version }; + } + } +} + +function batchProcessorForLocator< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + evaluator: DurableEvaluator, + locator: DurableBatchLocator, +): DurableBatchProcessor | undefined { + if (locator.kind === "task") { + return isBatchTask(evaluator.task) + ? (evaluator.task as DurableBatchProcessor) + : undefined; + } + const scorer = evaluator.scores.find( + (candidate) => + isBatchScorer(candidate) && candidate.name === locator.scorerName, + ); + return scorer as + | DurableBatchProcessor + | undefined; +} + +async function attachWebhookEventToJob( + store: DurableEvalStore, + jobKey: string, + storedEventKey: string, + event: DurableBatchResultEvent, +): Promise<"attached" | "duplicate" | "missing"> { + while (true) { + const current = await readJson(store, jobKey); + if (!current) return "missing"; + if ( + current.value.status === "complete" || + current.value.status === "failed" + ) { + return "duplicate"; + } + if (current.value.webhookEventKey === storedEventKey) { + const next = { + ...current.value, + nextPollAt: Date.now(), + }; + delete next.error; + delete next.workerId; + delete next.leaseUntil; + const written = await writeJson(store, jobKey, next, { + ifVersion: current.version, + }); + if (written.written) return "attached"; + continue; + } + if (current.value.webhookEventKey) return "duplicate"; + if ( + event.handle !== undefined && + current.value.handle !== undefined && + stableStringify(event.handle) !== stableStringify(current.value.handle) + ) { + throw new Error( + `Webhook handle does not match durable batch ${current.value.id}`, + ); + } + if ( + event.externalId && + current.value.external && + (current.value.external.source !== event.source || + current.value.external.id !== event.externalId) + ) { + throw new Error( + `Webhook externalId does not match durable batch ${current.value.id}`, + ); + } + const handle = event.handle ?? current.value.handle; + const next: DurableJobRecord = { + ...current.value, + handle, + external: event.externalId + ? { source: event.source, id: event.externalId } + : current.value.external, + outcome: event.outcome, + webhookEventKey: storedEventKey, + nextPollAt: Date.now(), + ...(handle !== undefined + ? { + status: "submitted", + submittedAt: current.value.submittedAt ?? Date.now(), + } + : {}), + }; + delete next.workerId; + delete next.leaseUntil; + const written = await writeJson(store, jobKey, next, { + ifVersion: current.version, + }); + if (written.written) return "attached"; + } +} + +async function attachPendingWebhookEvents( + store: DurableEvalStore, + jobKey: string, + job: DurableJobRecord, +) { + if (!job.external) return; + for await (const pointerKey of store.list( + webhookMailboxPrefix(job.external.source, job.external.id), + )) { + const pointer = await readJson<{ eventKey: string }>(store, pointerKey); + if (!pointer) continue; + const event = await readJson( + store, + pointer.value.eventKey, + ); + if (!event || event.value.status === "applied") continue; + const attached = await attachWebhookEventToJob( + store, + jobKey, + pointer.value.eventKey, + event.value.event, + ); + if (attached !== "attached") { + await markWebhookEventApplied(store, pointer.value.eventKey, job.id); + } + } +} + +async function markWebhookApplied( + store: DurableEvalStore, + job: DurableJobRecord, +) { + if (job.webhookEventKey) { + await markWebhookEventApplied(store, job.webhookEventKey, job.id); + } +} + +async function markWebhookEventApplied( + store: DurableEvalStore, + eventKey: string, + batchId: string, +) { + while (true) { + const current = await readJson(store, eventKey); + if (!current || current.value.status === "applied") return; + const next: DurableWebhookEventRecord = { + ...current.value, + status: "applied", + batchId, + }; + const written = await writeJson(store, eventKey, next, { + ifVersion: current.version, + }); + if (written.written) return; + } +} + +async function hasWaitingWebhookJob( + store: DurableEvalStore, + prefix: string, + shard: number, + webhookStages: Set, +) { + for await (const key of store.list(`${prefix}/jobs/`)) { + const job = await readJson(store, key); + if ( + job?.value.shard === shard && + job.value.status === "submitted" && + job.value.outcome === undefined && + webhookStages.has(job.value.stage) + ) { + return true; + } + } + return false; +} + +async function hasProviderErrorJob( + store: DurableEvalStore, + prefix: string, + shard: number, +) { + for await (const key of store.list(`${prefix}/jobs/`)) { + const job = await readJson(store, key); + if ( + job?.value.shard === shard && + job.value.status === "submitted" && + job.value.error !== undefined + ) { + return true; + } + } + return false; +} + +async function flushPendingLogs({ + store, + prefix, + experiment, + scorerNames, + shard, + runId, +}: { + store: DurableEvalStore; + prefix: string; + experiment: Experiment; + scorerNames: string[]; + shard: { index: number; count: number }; + runId: string; +}) { + let changed = false; + for await (const current of listCases(store, prefix)) { + if (current.value.shard !== shard.index || !current.value.logPending) { + continue; + } + await logDurableCase(experiment, current.value, scorerNames, runId); + await experiment.flush(); + const latest = await readJson( + store, + caseKey(prefix, current.value.id), + ); + if (latest) { + const next = structuredClone(latest.value); + next.logPending = false; + const result = await writeJson( + store, + caseKey(prefix, current.value.id), + next, + { ifVersion: latest.version }, + ); + changed = result.written || changed; + } + } + return changed; +} + +async function logDurableCase( + experiment: Experiment, + record: DurableCaseRecord, + scorerNames: string[], + runId: string, +) { + const datum = record.datum as EvalCase; + const rootSpanId = deterministicId(`${runId}:${record.id}:root`); + const root = _internalStartSpanWithInitialMerge({ + parent: await experiment.export(), + spanId: rootSpanId, + name: "eval", + spanAttributes: { type: SpanTypeAttribute.EVAL }, + event: { + id: deterministicId(`${runId}:${record.id}:row`), + input: datum.input, + expected: "expected" in datum ? datum.expected : undefined, + metadata: { + ...(record.metadata as Record), + durable_eval: { + run_id: runId, + case_id: record.caseId, + trial_index: record.trialIndex, + }, + }, + tags: record.tags, + ...(record.task.status === "succeeded" + ? { output: record.task.value } + : record.task.status === "failed" + ? { error: record.task.error.message } + : {}), + scores: collectedScores(record, scorerNames), + }, + state: evaluatorState(experiment), + }); + const parent = await root.export(); + if (record.task.status === "succeeded" || record.task.status === "failed") { + const task = _internalStartSpanWithInitialMerge({ + parent, + spanId: deterministicId(`${runId}:${record.id}:task-span`), + name: "task", + spanAttributes: { type: SpanTypeAttribute.TASK }, + event: { + id: deterministicId(`${runId}:${record.id}:task-row`), + input: datum.input, + metadata: { + durable_eval: { + revision: record.task.revision, + attempts: record.task.attempts, + }, + }, + ...(record.task.status === "succeeded" + ? { output: record.task.value } + : { error: record.task.error.message }), + }, + state: evaluatorState(experiment), + }); + task.end(); + } + for (const scorerName of scorerNames) { + const state = record.scores[scorerName]; + if (state?.status !== "succeeded" && state?.status !== "failed") continue; + const scorer = _internalStartSpanWithInitialMerge({ + parent, + spanId: deterministicId(`${runId}:${record.id}:score:${scorerName}:span`), + name: scorerName, + spanAttributes: { type: SpanTypeAttribute.SCORE }, + event: { + id: deterministicId(`${runId}:${record.id}:score:${scorerName}:row`), + input: { + input: datum.input, + output: + record.task.status === "succeeded" ? record.task.value : undefined, + expected: "expected" in datum ? datum.expected : undefined, + }, + metadata: { + durable_eval: { + revision: state.revision, + attempts: state.attempts, + }, + }, + ...(state.status === "succeeded" + ? { + output: state.value, + scores: state.value as Record, + } + : { error: state.error.message }), + }, + state: evaluatorState(experiment), + }); + scorer.end(); + } + root.end(); +} + +function evaluatorState(experiment: Experiment) { + return experiment.loggingState; +} + +function collectedScores(record: DurableCaseRecord, scorerNames: string[]) { + const scores: Record = {}; + for (const name of scorerNames) { + const state = record.scores[name]; + if (state?.status === "succeeded") { + Object.assign(scores, state.value); + } + } + return scores; +} + +async function collectProgress( + store: DurableEvalStore, + prefix: string, + scorerNames: string[], +): Promise { + const progress: DurableEvalProgress = { + ...emptyProgress(), + }; + for await (const record of listCases(store, prefix)) { + progress.total++; + countStage(record.value.task, progress, "task"); + if (record.value.task.status === "succeeded") { + for (const name of scorerNames) { + countStage(record.value.scores[name], progress, "score"); + } + } + } + return progress; +} + +function countStage( + state: StageState | undefined, + progress: DurableEvalProgress, + kind: "task" | "score", +) { + if (!state || ["pending", "leased", "in_batch"].includes(state.status)) { + if (kind === "task") progress.taskPending++; + else progress.scorePending++; + } else if (state.status === "succeeded") { + if (kind === "task") progress.taskSucceeded++; + else progress.scoreSucceeded++; + } else if (state.status === "failed") { + if (kind === "task") progress.taskFailed++; + else progress.scoreFailed++; + } else if (state.status === "unknown") { + progress.unknown++; + } +} + +async function isComplete( + store: DurableEvalStore, + prefix: string, + scorerNames: string[], + shard?: number, +) { + for await (const record of listCases(store, prefix)) { + if (shard !== undefined && record.value.shard !== shard) continue; + if (!isTerminal(record.value.task)) return false; + if (record.value.task.status === "succeeded") { + for (const name of scorerNames) { + if (!isTerminal(record.value.scores[name])) return false; + } + } + } + return true; +} + +function isTerminal(state: StageState | undefined) { + return state?.status === "succeeded" || state?.status === "failed"; +} + +function isClaimable(state: StageState | undefined) { + return ( + state?.status === "pending" || + (state?.status === "leased" && state.leaseUntil <= Date.now()) + ); +} + +async function buildLocalDurableSummary( + store: DurableEvalStore, + prefix: string, + projectName: string, + experimentName: string, + scorerNames: string[], +): Promise { + const totals: Record = {}; + for await (const record of listCases(store, prefix)) { + for (const [name, score] of Object.entries( + collectedScores(record.value, scorerNames), + )) { + if (score === null) continue; + const current = totals[name] ?? { total: 0, count: 0 }; + current.total += score; + current.count++; + totals[name] = current; + } + } + return { + projectName, + experimentName, + scores: Object.fromEntries( + Object.entries(totals).map(([name, value]) => [ + name, + { + name, + score: value.total / value.count, + improvements: 0, + regressions: 0, + }, + ]), + ), + }; +} + +function normalizeScores( + value: OneOrMoreScores, + defaultName: string, +): Record { + if (value === null) return { [defaultName]: null }; + if (typeof value === "number") return { [defaultName]: value }; + const values = Array.isArray(value) ? value : [value]; + return Object.fromEntries( + values.map((score: Score) => [score.name ?? defaultName, score.score]), + ); +} + +function batchContext( + runId: string, + stage: string, + job: DurableJobRecord, + shard: { index: number; count: number }, + signal: AbortSignal, +): DurableBatchContext { + return { + runId, + stage, + revision: job.revision, + shard, + attempt: job.attempt, + batchId: job.id, + itemCount: job.itemIds.length, + signal, + }; +} + +function resultItemId(value: unknown): string { + 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( + value: unknown, +): value is DurableBatchTask< + unknown, + unknown, + unknown, + BaseMetadata, + EvalParameters, + JsonValue +> { + return ( + typeof value === "object" && + value !== null && + "kind" in value && + value.kind === BATCH_TASK_KIND + ); +} + +function isBatchScorer( + value: unknown, +): value is DurableBatchScorer< + unknown, + unknown, + unknown, + BaseMetadata, + JsonValue +> { + return ( + typeof value === "object" && + value !== null && + "kind" in value && + value.kind === BATCH_SCORER_KIND + ); +} + +async function* listCases( + store: DurableEvalStore, + prefix: string, +): AsyncGenerator> { + for await (const key of store.list(`${prefix}/cases/`)) { + const record = await readJson(store, key); + if (record) yield record; + } +} + +async function* listJobs( + store: DurableEvalStore, + prefix: string, + stage: string, +): AsyncGenerator & { key: string }> { + for await (const key of store.list( + `${prefix}/jobs/${encodeURIComponent(stage)}/`, + )) { + const record = await readJson(store, key); + if (record) yield { ...record, key }; + } +} + +function runPrefix(projectName: string, evalName: string, runId: string) { + return `durable-eval/v1/${contentVersion( + encoder.encode(`${projectName}\0${evalName}\0${runId}`), + )}`; +} + +function internalBatchIndexKey(batchId: string) { + return `durable-eval/v1/indices/batches/${contentVersion( + encoder.encode(batchId), + )}`; +} + +function externalBatchIndexKey(source: string, externalId: string) { + return `durable-eval/v1/indices/external/${contentVersion( + encoder.encode(`${source}\0${externalId}`), + )}`; +} + +function webhookEventKey(source: string, eventId: string) { + return `durable-eval/v1/webhooks/events/${contentVersion( + encoder.encode(`${source}\0${eventId}`), + )}`; +} + +function webhookMailboxPrefix(source: string, externalId: string) { + return `durable-eval/v1/webhooks/mailboxes/${contentVersion( + encoder.encode(`${source}\0${externalId}`), + )}/`; +} + +function webhookMailboxKey( + source: string, + externalId: string, + eventKey: string, +) { + return `${webhookMailboxPrefix(source, externalId)}${contentVersion( + encoder.encode(eventKey), + )}`; +} + +function caseKey(prefix: string, id: string) { + return `${prefix}/cases/${encodeURIComponent(id)}`; +} + +function workItemId(caseId: string, trialIndex: number) { + return `${caseId}:trial:${trialIndex}`; +} + +function stableShard(id: string, count: number) { + const hash = contentVersion(encoder.encode(id)); + return Number.parseInt(hash.slice(0, 8), 16) % count; +} + +function deterministicId(value: string) { + const hex = contentVersion(encoder.encode(value)) + .padEnd(32, "0") + .slice(0, 32); + return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-5${hex.slice(13, 16)}-a${hex.slice(17, 20)}-${hex.slice(20, 32)}`; +} + +function contentVersion(value: Uint8Array) { + if (iso.hash) return iso.hash(decoder.decode(value)); + let hash = 2166136261; + for (const byte of value) { + hash ^= byte; + hash = Math.imul(hash, 16777619); + } + return (hash >>> 0).toString(16).padStart(8, "0"); +} + +async function readJson( + store: DurableEvalStore, + key: string, +): Promise | undefined> { + const record = await store.read(key); + if (!record) return undefined; + return { + value: JSON.parse(decoder.decode(record.value)) as T, + version: record.version, + }; +} + +async function writeJson( + store: DurableEvalStore, + key: string, + value: T, + condition: DurableEvalWriteCondition, +) { + return await store.write( + key, + encoder.encode(stableStringify(value)), + condition, + ); +} + +function stableStringify(value: unknown): string { + return JSON.stringify(value, (_key, nested) => { + if (nested && typeof nested === "object" && !Array.isArray(nested)) { + return Object.fromEntries( + Object.entries(nested).sort(([left], [right]) => + left.localeCompare(right), + ), + ); + } + return nested; + }); +} + +function assertJsonValue(value: unknown, label: string): JsonValue { + try { + const serialized = JSON.stringify(value); + if (serialized === undefined) { + throw new Error("value serializes to undefined"); + } + return JSON.parse(serialized) as JsonValue; + } catch (error) { + throw new Error(`${label} must be JSON serializable`, { cause: error }); + } +} + +function serializeError(error: unknown): SerializedError { + if (error instanceof Error) { + return { name: error.name, message: error.message, stack: error.stack }; + } + return { message: String(error) }; +} + +function isErrorCode(error: unknown, code: string) { + return ( + typeof error === "object" && + error !== null && + "code" in error && + error.code === code + ); +} + +function validateShard(shard: { index: number; count: number }) { + if ( + !Number.isInteger(shard.index) || + !Number.isInteger(shard.count) || + shard.count < 1 || + shard.index < 0 || + shard.index >= shard.count + ) { + throw new Error(`Invalid DurableEval shard ${shard.index}/${shard.count}`); + } +} + +function isIterable(value: unknown): value is Iterable { + return ( + typeof value === "object" && value !== null && Symbol.iterator in value + ); +} + +function isAsyncIterable(value: unknown): value is AsyncIterable { + return ( + typeof value === "object" && value !== null && Symbol.asyncIterator in value + ); +} + +async function* toAsyncIterable( + value: Iterable | AsyncIterable, +): AsyncGenerator { + if (isAsyncIterable(value)) { + for await (const item of value) yield item; + } else { + for (const item of value) yield item; + } +} + +function delay(ms: number) { + return new Promise((resolve) => setTimeout(resolve, Math.max(ms, 0))); +} + +function timeRemaining(deadlineAt: number | undefined) { + return deadlineAt === undefined + ? undefined + : Math.max(deadlineAt - Date.now(), 0); +} diff --git a/js/src/exports.ts b/js/src/exports.ts index 36a38946a..bee995e8e 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -262,6 +262,44 @@ export { defaultErrorScoreHandler, } from "./framework"; +export type { + DurableBatchCompletion, + DurableBatchContext, + DurableBatchOutcome, + DurableBatchPoll, + DurableBatchProcessingResult, + DurableBatchProcessor, + DurableBatchRecovery, + DurableBatchResultEvent, + DurableBatchScorer, + DurableBatchScorerItem, + DurableBatchScorerResult, + DurableBatchTask, + DurableBatchTaskItem, + DurableBatchTaskResult, + DurableEvalDefinition, + DurableEvalExistingRunOptions, + DurableEvalFailureSummary, + DurableEvalOperation, + DurableEvalPauseReason, + DurableEvalProgress, + DurableEvalResult, + DurableEvalRuntimeOptions, + DurableEvalStore, + DurableEvalWriteCondition, + DurableEvaluator, + JsonValue, +} from "./durable-eval"; + +export { + BatchScorer, + BatchTask, + DurableEval, + DurableEvalNotSubmittedError, + FileDurableEvalStore, + MemoryDurableEvalStore, +} from "./durable-eval"; + export { agentAssertionScorer } from "./agent-assertions"; export { DatasetPipeline } from "./dataset-pipeline"; diff --git a/js/src/framework.ts b/js/src/framework.ts index 8628eb0e9..8461799d7 100644 --- a/js/src/framework.ts +++ b/js/src/framework.ts @@ -456,6 +456,14 @@ export type EvaluatorFile = { reporter?: ReporterDef | string; }; }; + durableEvaluators?: Record< + string, + { + // Kept opaque here to avoid coupling the existing Eval framework to the + // additive durable evaluator's generic surface. + definition: unknown; + } + >; reporters: { [reporterName: string]: ReporterDef }; }; @@ -559,6 +567,7 @@ globalThis._evals = { prompts: [], parameters: [], evaluators: {}, + durableEvaluators: {}, reporters: {}, }; diff --git a/js/src/isomorph.ts b/js/src/isomorph.ts index 05e27d066..cfa19e217 100644 --- a/js/src/isomorph.ts +++ b/js/src/isomorph.ts @@ -266,9 +266,10 @@ interface Common { path: string, opts?: { recursive?: boolean }, ) => Promise; - writeFile?: (filename: string, data: string) => Promise; + writeFile?: (filename: string, data: string | Uint8Array) => Promise; readFile?: (filename: string) => Promise; readdir?: (path: string) => Promise; + rename?: (oldPath: string, newPath: string) => Promise; utimes?: (path: string, atime: Date, mtime: Date) => Promise; unlink?: (path: string) => Promise; // eslint-disable-next-line @typescript-eslint/no-explicit-any diff --git a/js/src/node/config.ts b/js/src/node/config.ts index f9743837a..19c266ea1 100644 --- a/js/src/node/config.ts +++ b/js/src/node/config.ts @@ -128,6 +128,7 @@ export function configureNode() { iso.writeFile = fs.writeFile; iso.readFile = fs.readFile; iso.readdir = fs.readdir; + iso.rename = fs.rename; iso.stat = fs.stat; iso.statSync = fsSync.statSync; iso.utimes = fs.utimes; From 42b25233089626879bc68f48506ab4bb1ec72cec Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 27 Jul 2026 12:49:08 +0000 Subject: [PATCH 02/13] Update PR #2297 --- .changeset/durable-batches-evaluate.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/durable-batches-evaluate.md diff --git a/.changeset/durable-batches-evaluate.md b/.changeset/durable-batches-evaluate.md new file mode 100644 index 000000000..a0486161b --- /dev/null +++ b/.changeset/durable-batches-evaluate.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Add batch/durable evals api From 2505af07552f27df75e619346724c5cc306363b4 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 27 Jul 2026 15:51:38 +0000 Subject: [PATCH 03/13] Update PR #2297 --- .../durable-eval-webhook/scenario.test.ts | 19 +- .../durable-eval-webhook/scenario.ts | 32 +- js/src/cli/index.ts | 92 +- js/src/durable-eval.test.ts | 387 ++++- js/src/durable-eval.ts | 1408 ++++++++++++----- 5 files changed, 1524 insertions(+), 414 deletions(-) diff --git a/e2e/scenarios/durable-eval-webhook/scenario.test.ts b/e2e/scenarios/durable-eval-webhook/scenario.test.ts index 4365794c3..3522a8786 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.test.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.test.ts @@ -15,19 +15,30 @@ test("durable eval collects webhook sub-batches and logs completed rows", async await runScenarioDir({ scenarioDir }); const evalSpans = findAllSpans(testRunEvents(), "eval"); - expect(evalSpans).toHaveLength(3); - expect(evalSpans.map((event) => event.output).sort()).toEqual([2, 4, 6]); + const webhookSpans = evalSpans.filter( + (event) => event.metadata?.kind === "webhook", + ); + expect(webhookSpans).toHaveLength(3); + expect(webhookSpans.map((event) => event.output).sort()).toEqual([2, 4, 6]); expect( - evalSpans + webhookSpans .map((event) => event.scores) .sort((left, right) => JSON.stringify(left).localeCompare(JSON.stringify(right)), ), ).toEqual([{ exact: 1 }, { exact: 1 }, { exact: 1 }]); - expect(evalSpans.map((event) => event.metadata?.durable_eval)).toEqual([ + 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 shardedSpans = evalSpans.filter( + (event) => event.metadata?.kind === "sharded", + ); + expect(shardedSpans).toHaveLength(4); + expect(new Set(shardedSpans.map((event) => event.experimentId)).size).toBe( + 1, + ); }); }); diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts index ba2f80448..8b5d7cfb4 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -18,7 +18,7 @@ async function main() { DurableBatchTaskItem< number, number, - { testRunId: string }, + { testRunId: string; kind: string }, Record >[] >(); @@ -26,7 +26,7 @@ async function main() { number, number, number, - { testRunId: string }, + { testRunId: string; kind: string }, Record, { id: string } >({ @@ -61,7 +61,7 @@ async function main() { id: `case-${input}`, input, expected: input * 2, - metadata: { testRunId }, + metadata: { testRunId, kind: "webhook" }, })), task, scores: [ @@ -101,6 +101,32 @@ async function main() { throw new Error("Durable eval did not complete after the final webhook"); } } + + const sharded = DurableEval( + scopedName("e2e-durable-eval-sharded-project", testRunId), + { + revision: "eval-v1", + data: [1, 2, 3, 4].map((input) => ({ + id: `shard-case-${input}`, + input, + metadata: { testRunId, kind: "sharded" }, + })), + task: (input) => input * 10, + scores: [], + }, + ); + await Promise.all([ + sharded.run({ + runId: `sharded-${testRunId}`, + shard: { index: 0, count: 2 }, + store, + }), + sharded.run({ + runId: `sharded-${testRunId}`, + shard: { index: 1, count: 2 }, + store, + }), + ]); } runMain(main); diff --git a/js/src/cli/index.ts b/js/src/cli/index.ts index 6d4604352..3dac2d388 100755 --- a/js/src/cli/index.ts +++ b/js/src/cli/index.ts @@ -249,6 +249,9 @@ function buildWatchPluginForEvaluator( evaluators.evaluators = evaluators.evaluators.filter( (e) => e.sourceFile !== inFile, ); + evaluators.durableEvaluators = evaluators.durableEvaluators.filter( + (e) => e.sourceFile !== inFile, + ); // Update the evaluators and reporters for (const evaluator of Object.values(evalResult.evaluators)) { @@ -264,6 +267,21 @@ function buildWatchPluginForEvaluator( reporter: evaluator.reporter, }); } + for (const registration of Object.values( + evalResult.durableEvaluators ?? {}, + )) { + evaluators.durableEvaluators.push({ + sourceFile: inFile, + // Runtime registrations are intentionally generic-erased. + definition: registration.definition as DurableEvalDefinition< + any, + any, + any, + any, + any + >, + }); + } for (const [reporterName, reporter] of Object.entries( evalResult.reporters, )) { @@ -310,6 +328,20 @@ function buildWatchPluginForEvaluator( addReport(evalReports, resolvedReporter, report); } + for (const registration of evaluators.durableEvaluators.filter( + (candidate) => candidate.sourceFile === inFile, + )) { + const result = await runDurableEvaluator( + registration.definition, + opts, + ); + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.error( + result.status === "completed" + ? `Durable eval ${result.runId} completed` + : `Durable eval ${result.runId} paused: ${result.reason}`, + ); + } for (const [reporterName, { reporter, results }] of Object.entries( evalReports, @@ -561,6 +593,36 @@ export async function buildEvaluators( return { evaluators, buildResults }; } +async function runDurableEvaluator( + definition: DurableEvalDefinition, + opts: EvaluatorOpts, +) { + const options: DurableEvalRuntimeOptions = { + runId: opts.durableRunId, + shard: opts.durableShard, + deadlineMs: opts.durableDeadlineMs, + checkpointDir: opts.checkpointDir, + noSendLogs: opts.noSendLogs, + }; + if (!opts.durableOperation) { + return await definition.run(options); + } + if (!opts.durableRunId) { + throw new Error(`--${opts.durableOperation} requires --run-id`); + } + const existingOptions = { ...options, runId: opts.durableRunId }; + switch (opts.durableOperation) { + case "status": + return await definition.status(existingOptions); + case "retry-failed": + return await definition.retryFailed(existingOptions); + case "resubmit-unknown": + return await definition.resubmitUnknown(existingOptions); + case "cancel": + return await definition.cancel(existingOptions); + } +} + async function runOnce( handles: Record, opts: EvaluatorOpts, @@ -629,34 +691,8 @@ async function runOnce( } }); const durableResultPromises = evaluators.durableEvaluators.map( - async (registration) => { - const options: DurableEvalRuntimeOptions = { - runId: opts.durableRunId, - shard: opts.durableShard, - deadlineMs: opts.durableDeadlineMs, - checkpointDir: opts.checkpointDir, - noSendLogs: opts.noSendLogs, - }; - if (!opts.durableOperation) { - return await registration.definition.run(options); - } - if (!opts.durableRunId) { - throw new Error( - `--${opts.durableOperation.replaceAll("-", "_")} requires --run-id`, - ); - } - const existingOptions = { ...options, runId: opts.durableRunId }; - switch (opts.durableOperation) { - case "status": - return await registration.definition.status(existingOptions); - case "retry-failed": - return await registration.definition.retryFailed(existingOptions); - case "resubmit-unknown": - return await registration.definition.resubmitUnknown(existingOptions); - case "cancel": - return await registration.definition.cancel(existingOptions); - } - }, + async (registration) => + await runDurableEvaluator(registration.definition, opts), ); // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts index aac780448..43dc05c30 100644 --- a/js/src/durable-eval.test.ts +++ b/js/src/durable-eval.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test, vi } from "vitest"; -import { mkdtemp, rm } from "node:fs/promises"; +import { mkdtemp, rm, utimes, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { @@ -9,6 +9,8 @@ import { FileDurableEvalStore, MemoryDurableEvalStore, type DurableBatchTaskItem, + type DurableEvalStore, + type DurableEvalWriteCondition, } from "./durable-eval"; import type { EvaluatorFile } from "./framework"; import { configureNode } from "./node/config"; @@ -56,6 +58,33 @@ describe("DurableEval", () => { } }); + test("filesystem store recovers stale locks and rejects Windows traversal keys", async () => { + const directory = await mkdtemp(join(tmpdir(), "durable-eval-locks-")); + try { + const store = new FileDurableEvalStore(directory); + const value = new TextEncoder().encode("one"); + const first = await store.write("runs/one", value, { ifAbsent: true }); + expect(first.written).toBe(true); + + const lockPath = join(directory, "runs", "one.lock"); + await writeFile(lockPath, "orphaned"); + const staleTime = new Date(Date.now() - 60_000); + await utimes(lockPath, staleTime, staleTime); + const current = await store.read("runs/one"); + await expect( + store.write("runs/one", new TextEncoder().encode("two"), { + ifVersion: current!.version, + }), + ).resolves.toMatchObject({ written: true }); + + await expect( + store.write("..\\outside", value, { ifAbsent: true }), + ).rejects.toThrow("Invalid durable eval store key"); + } finally { + await rm(directory, { recursive: true, force: true }); + } + }); + test("runs local tasks and scorers and resumes completed work", async () => { const store = new MemoryDurableEvalStore(); const task = vi.fn((input: number) => input * 2); @@ -105,6 +134,31 @@ describe("DurableEval", () => { expect(scorerCalls).toBe(2); }); + test("persists metadata written through hooks.meta", async () => { + const result = await DurableEval("metadata-hooks", { + revision: "v1", + data: [{ id: "one", input: "hello" }], + task: (_input, hooks) => { + hooks.meta({ source: "deprecated-hook" }); + return hooks.metadata.source; + }, + scores: [ + function exact({ output }) { + return output === "deprecated-hook" ? 1 : 0; + }, + ], + }).run({ + runId: "metadata-run", + store: new MemoryDurableEvalStore(), + noSendLogs: true, + }); + + expect(result).toMatchObject({ + status: "completed", + summary: { scores: { exact: { score: 1 } } }, + }); + }); + test("requires stable case identifiers", async () => { await expect( DurableEval("missing-ids", { @@ -215,6 +269,84 @@ describe("DurableEval", () => { expect(scoreJobs).toHaveLength(2); }); + test("retries raced case writes before completing a batch job", async () => { + class FailOneResultWriteStore implements DurableEvalStore { + readonly inner = new MemoryDurableEvalStore(); + failed = false; + + read(key: string) { + return this.inner.read(key); + } + + async write( + key: string, + value: Uint8Array, + condition: DurableEvalWriteCondition, + ) { + const serialized = new TextDecoder().decode(value); + if ( + !this.failed && + key.includes("/cases/") && + serialized.includes('"status":"succeeded"') + ) { + this.failed = true; + return { + written: false as const, + currentVersion: (await this.inner.read(key))?.version, + }; + } + return this.inner.write(key, value, condition); + } + + list(prefix: string) { + return this.inner.list(prefix); + } + } + + const store = new FailOneResultWriteStore(); + const result = await DurableEval("cas-batch-results", { + revision: "eval-v1", + data: [{ id: "one", input: 1, expected: 2 }], + task: BatchTask< + number, + number, + number, + void, + Record, + { id: string } + >({ + revision: "task-v1", + async submit() { + return { id: "provider-job" }; + }, + completion: { + mode: "poll", + async poll() { + return { status: "complete" }; + }, + }, + async *collect() { + yield { id: "one:trial:0", output: 2 }; + }, + }), + scores: [ + function exact({ output, expected }) { + return output === expected ? 1 : 0; + }, + ], + }).run({ + runId: "cas-run", + store, + noSendLogs: true, + }); + + expect(store.failed).toBe(true); + expect(result).toMatchObject({ + status: "completed", + progress: { taskSucceeded: 1, scoreSucceeded: 1 }, + }); + }); + test("pauses for webhooks and processes bounded sub-batches", async () => { const store = new MemoryDurableEvalStore(); const providerJobs = new Map< @@ -595,6 +727,76 @@ describe("DurableEval", () => { expect(submitCount).toBe(1); }); + test("renews provider leases while polling exceeds the lease duration", async () => { + vi.useFakeTimers(); + try { + const store = new MemoryDurableEvalStore(); + let pollCalls = 0; + let releasePoll: () => void = () => undefined; + const pollGate = new Promise((resolve) => { + releasePoll = resolve; + }); + const durable = DurableEval("long-provider-poll", { + revision: "eval-v1", + data: [{ id: "one", input: "input" }], + task: BatchTask< + string, + string, + void, + void, + Record, + { id: string } + >({ + revision: "task-v1", + async submit() { + return { id: "long-job" }; + }, + completion: { + mode: "poll", + async poll() { + pollCalls++; + await pollGate; + return { status: "complete" }; + }, + }, + async *collect() { + yield { id: "one:trial:0", output: "done" }; + }, + }), + scores: [], + }); + + const first = durable.run({ + runId: "long-poll-run", + store, + noSendLogs: true, + workerId: "worker-one", + }); + while (pollCalls === 0) await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(61_000); + + const second = durable.run({ + runId: "long-poll-run", + store, + noSendLogs: true, + workerId: "worker-two", + deadlineMs: 5, + }); + await vi.advanceTimersByTimeAsync(10); + await expect(second).resolves.toMatchObject({ + status: "paused", + reason: "deadline", + }); + expect(pollCalls).toBe(1); + + releasePoll(); + await vi.advanceTimersByTimeAsync(0); + await expect(first).resolves.toMatchObject({ status: "completed" }); + } finally { + vi.useRealTimers(); + } + }); + test("cancels active provider batches and terminally checkpoints the work", async () => { const store = new MemoryDurableEvalStore(); const abort = new AbortController(); @@ -609,12 +811,12 @@ describe("DurableEval", () => { >({ revision: "task-v1", async submit(_items, context) { - abort.abort(); return { jobId: context.batchId }; }, completion: { mode: "poll", async poll() { + abort.abort(); return { status: "pending" }; }, }, @@ -725,6 +927,141 @@ describe("DurableEval", () => { expect(submits).toBe(1); }); + test("deadline aborts an active provider submission", async () => { + const neverSubmitted = new Promise<{ id: string }>(() => undefined); + const durable = DurableEval("submit-deadline", { + revision: "eval-v1", + data: [{ id: "one", input: "input" }], + task: BatchTask< + string, + string, + void, + void, + Record, + { id: string } + >({ + revision: "task-v1", + async submit() { + return await neverSubmitted; + }, + completion: { + mode: "poll", + async poll() { + return { status: "pending" }; + }, + }, + async *collect() { + // The submission never returns a handle. + }, + }), + scores: [], + }); + + const started = Date.now(); + const result = await durable.run({ + runId: "submit-deadline-run", + store: new MemoryDurableEvalStore(), + noSendLogs: true, + deadlineMs: 10, + }); + expect(result).toMatchObject({ status: "paused", reason: "deadline" }); + expect(Date.now() - started).toBeLessThan(1_000); + }); + + test("reserves one deterministic experiment name across concurrent shards", async () => { + const store = new MemoryDurableEvalStore(); + const durable = DurableEval("experiment-reservation", { + revision: "v1", + data: [ + { id: "one", input: 1 }, + { id: "two", input: 2 }, + { id: "three", input: 3 }, + { id: "four", input: 4 }, + ], + task: (input) => input, + scores: [], + }); + + await Promise.all([ + durable.run({ + runId: "shared-run", + shard: { index: 0, count: 2 }, + store, + noSendLogs: true, + }), + durable.run({ + runId: "shared-run", + shard: { index: 1, count: 2 }, + store, + noSendLogs: true, + }), + ]); + + const manifests: string[] = []; + for await (const key of store.list("durable-eval/v1/")) { + if (!key.endsWith("/manifest")) continue; + const record = await store.read(key); + if (record) manifests.push(new TextDecoder().decode(record.value)); + } + expect(manifests).toHaveLength(1); + expect(JSON.parse(manifests[0])).toMatchObject({ + experimentName: "experiment-reservation-shared-run", + shardCount: 2, + }); + + await expect( + durable.status({ + runId: "shared-run", + store, + noSendLogs: true, + }), + ).resolves.toMatchObject({ status: "completed" }); + }); + + test("applies lifecycle retries to every shard when no shard is specified", async () => { + const store = new MemoryDurableEvalStore(); + let succeed = false; + const durable = DurableEval("sharded-lifecycle", { + revision: "v1", + data: [1, 2, 3, 4].map((input) => ({ + id: `case-${input}`, + input, + })), + task: (input) => { + if (!succeed) throw new Error("try again later"); + return input; + }, + scores: [], + }); + + await Promise.all([ + durable.run({ + runId: "sharded-lifecycle-run", + shard: { index: 0, count: 2 }, + store, + noSendLogs: true, + }), + durable.run({ + runId: "sharded-lifecycle-run", + shard: { index: 1, count: 2 }, + store, + noSendLogs: true, + }), + ]); + + succeed = true; + await expect( + durable.retryFailed({ + runId: "sharded-lifecycle-run", + store, + noSendLogs: true, + }), + ).resolves.toMatchObject({ + status: "completed", + progress: { taskSucceeded: 4, taskFailed: 0 }, + }); + }); + test("keeps completed work when the definition revision changes", async () => { const store = new MemoryDurableEvalStore(); const firstTask = vi.fn((input: number) => input); @@ -764,6 +1101,52 @@ describe("DurableEval", () => { expect(changedTask).not.toHaveBeenCalled(); }); + test("removes inactive scorer state when a definition changes", async () => { + const store = new MemoryDurableEvalStore(); + await DurableEval("remove-scorer", { + revision: "v1", + data: [{ id: "one", input: 1 }], + task: (input) => input, + scores: [ + function oldScore() { + return 1; + }, + ], + }).run({ + runId: "remove-scorer-run", + store, + noSendLogs: true, + }); + + const result = await DurableEval("remove-scorer", { + revision: "v2", + data: [{ id: "one", input: 1 }], + task: (input) => input, + scores: [], + }).run({ + runId: "remove-scorer-run", + store, + noSendLogs: true, + }); + expect(result).toMatchObject({ + status: "completed", + summary: { scores: {} }, + }); + + const cases: Array> = []; + for await (const key of store.list("durable-eval/v1/")) { + if (!key.includes("/cases/")) continue; + const record = await store.read(key); + if (record) + cases.push(JSON.parse(new TextDecoder().decode(record.value))); + } + expect(cases).toHaveLength(1); + expect(cases[0]).toMatchObject({ + scores: {}, + removedScores: ["oldScore"], + }); + }); + test("pauses safely when submit may have created a provider job", async () => { const store = new MemoryDurableEvalStore(); let submitShouldFail = true; diff --git a/js/src/durable-eval.ts b/js/src/durable-eval.ts index c6025195f..c0806f143 100644 --- a/js/src/durable-eval.ts +++ b/js/src/durable-eval.ts @@ -30,6 +30,9 @@ const CHECKPOINT_VERSION = 2; const DEFAULT_BATCH_SIZE = 1_000; const DEFAULT_MAX_CONCURRENT_BATCHES = 1; const JOB_LEASE_MS = 60_000; +const LEASE_HEARTBEAT_MS = JOB_LEASE_MS / 3; +const FILE_LOCK_STALE_MS = 10_000; +const FILE_LOCK_WAIT_MS = 15_000; type JsonPrimitive = string | number | boolean | null; export type JsonValue = @@ -141,6 +144,7 @@ export class FileDurableEvalStore implements DurableEvalStore { private pathFor(key: string) { if ( key.startsWith("/") || + key.includes("\\") || key.split("/").some((part) => part === ".." || part === ".") ) { throw new Error(`Invalid durable eval store key: ${key}`); @@ -150,16 +154,43 @@ export class FileDurableEvalStore implements DurableEvalStore { private async withLock(path: string, fn: () => Promise): Promise { const lockPath = `${path}.lock`; + const lockContents = stableStringify({ + owner: newId(), + createdAt: new Date().toISOString(), + }); await iso.mkdir!(iso.pathDirname!(path), { recursive: true }); const started = Date.now(); - let handle: { close(): Promise } | undefined; + let handle: + | { close(): Promise; writeFile(value: string): Promise } + | undefined; while (!handle) { try { - handle = await iso.openFile!(lockPath, "wx"); + const candidate = await iso.openFile!(lockPath, "wx"); + try { + await candidate.writeFile(lockContents); + handle = candidate; + } catch (error) { + await candidate.close(); + await iso.unlink!(lockPath).catch(() => undefined); + throw error; + } } catch (error) { - if (!isErrorCode(error, "EEXIST") || Date.now() - started >= 10_000) { + if (!isErrorCode(error, "EEXIST")) { throw error; } + try { + const lockStat = await iso.stat!(lockPath); + if (Date.now() - lockStat.mtimeMs >= FILE_LOCK_STALE_MS) { + await iso.unlink!(lockPath).catch((unlinkError) => { + if (!isErrorCode(unlinkError, "ENOENT")) throw unlinkError; + }); + continue; + } + } catch (statError) { + if (!isErrorCode(statError, "ENOENT")) throw statError; + continue; + } + if (Date.now() - started >= FILE_LOCK_WAIT_MS) throw error; await delay(10); } } @@ -167,7 +198,14 @@ export class FileDurableEvalStore implements DurableEvalStore { return await fn(); } finally { await handle.close(); - await iso.unlink!(lockPath).catch(() => undefined); + try { + const currentLock = await iso.readFile!(lockPath); + if (decoder.decode(currentLock) === lockContents) { + await iso.unlink!(lockPath).catch(() => undefined); + } + } catch (error) { + if (!isErrorCode(error, "ENOENT")) throw error; + } } } @@ -621,6 +659,7 @@ type DurableCaseRecord = { metadata: JsonValue; tags?: string[]; logPending?: boolean; + removedScores?: string[]; }; type DurableJobRecord = { @@ -663,6 +702,15 @@ type DurableRunManifest = { createdAt: string; }; +type DurableExperimentInitRecord = { + schemaVersion: number; + experimentName: string; + status: "initializing" | "ready"; + workerId?: string; + leaseUntil?: number; + experimentId?: string; +}; + type Versioned = { value: T; version: string }; type DurableBatchLocator = { @@ -920,6 +968,9 @@ async function processDurableBatchResult< await markWebhookEventApplied(store, storedEventKey, locator.batchId); return { status: "duplicate", batchId: locator.batchId }; } + if (attached === "in_progress") { + return { status: "duplicate", batchId: locator.batchId }; + } const run = await runDurableEval( definition.projectName, @@ -952,8 +1003,6 @@ async function runDurableEval< options: DurableEvalExecutionOptions = {}, ): Promise { const runId = options.runId ?? newId(); - const shard = options.shard ?? { index: 0, count: 1 }; - validateShard(shard); const workerId = options.workerId ?? newId(); const store = options.store ?? @@ -988,6 +1037,11 @@ async function runDurableEval< store, manifestKey, ); + const shard = options.shard ?? { + index: 0, + count: existingManifest?.value.shardCount ?? 1, + }; + validateShard(shard); if (options.operation && options.operation !== "run" && !existingManifest) { throw new Error(`DurableEval run ${runId} does not exist`); } @@ -1002,7 +1056,7 @@ async function runDurableEval< shardCount: shard.count, dataSealed: false, activeScorers: scorerNames, - experimentName: evaluator.experimentName, + experimentName: evaluator.experimentName ?? `${evalName}-${runId}`, createdAt: new Date().toISOString(), })); if (manifest.value.shardCount !== shard.count) { @@ -1010,6 +1064,44 @@ async function runDurableEval< `DurableEval run ${runId} was created with ${manifest.value.shardCount} shards, not ${shard.count}`, ); } + if ( + options.shard === undefined && + manifest.value.shardCount > 1 && + (options.operation === "retry-failed" || + options.operation === "resubmit-unknown" || + options.operation === "cancel") + ) { + let lastResult: DurableEvalResult | undefined; + let pausedLifecycleResult: DurableEvalResult | undefined; + for (let index = 0; index < manifest.value.shardCount; index++) { + const result = await runDurableEval(projectName, evaluator, { + ...options, + runId, + store, + shard: { index, count: manifest.value.shardCount }, + }); + lastResult = result; + if ( + result.status === "paused" && + result.reason !== "shard_complete" && + result.reason !== "status_only" && + !pausedLifecycleResult + ) { + pausedLifecycleResult = result; + } + } + if (!lastResult) throw new Error("DurableEval sharded lifecycle failed"); + return pausedLifecycleResult ?? lastResult; + } + if (!manifest.value.experimentName) { + manifest = await updateManifest(store, manifestKey, (current) => ({ + ...current, + experimentName: + current.experimentName ?? + evaluator.experimentName ?? + `${evalName}-${runId}`, + })); + } if (!manifest.value.dataSealed) { await materializeData({ @@ -1028,33 +1120,32 @@ async function runDurableEval< const experiment: Experiment | null = options.noSendLogs ? null - : initExperiment({ - state: evaluator.state, - ...(evaluator.projectId - ? { projectId: evaluator.projectId } - : { project: projectName }), - experiment: manifest.value.experimentName, - update: manifest.value.experimentName !== undefined, - description: evaluator.description, - metadata: evaluator.metadata, - tags: evaluator.tags, - setCurrent: false, + : await initializeDurableExperiment({ + store, + key: `${prefix}/experiment`, + workerId, + experimentName: manifest.value.experimentName!, + create: () => + initExperiment({ + state: evaluator.state, + ...(evaluator.projectId + ? { projectId: evaluator.projectId } + : { project: projectName }), + experiment: manifest.value.experimentName, + update: true, + description: evaluator.description, + metadata: evaluator.metadata, + tags: evaluator.tags, + setCurrent: false, + }), }); - if (experiment && !manifest.value.experimentName) { - const summary = await experiment.summarize({ summarizeScores: false }); - manifest = await updateManifest(store, manifestKey, (current) => ({ - ...current, - experimentName: summary.experimentName, - })); - } - await reconcileScorers(store, prefix, scorerNames); if (options.operation === "retry-failed") { - await resetStages(store, prefix, scorerNames, "failed"); + await resetStages(store, prefix, scorerNames, "failed", shard.index); } else if (options.operation === "resubmit-unknown") { - await resetStages(store, prefix, scorerNames, "unknown"); + await resetStages(store, prefix, scorerNames, "unknown", shard.index); } else if (options.operation === "cancel") { await cancelRun({ store, @@ -1119,6 +1210,13 @@ async function runDurableEval< const controller = new AbortController(); const abortHandler = () => controller.abort(); options.signal?.addEventListener("abort", abortHandler, { once: true }); + const deadlineTimer = + deadlineAt === undefined + ? undefined + : setTimeout( + () => controller.abort(), + Math.max(deadlineAt - Date.now(), 0), + ); const resume = (overrides: Partial = {}) => runDurableEval(projectName, evaluator, { @@ -1170,7 +1268,6 @@ async function runDurableEval< projectName, evalName, scorer, - evaluatorRevision: evaluator.revision, shard, workerId, runId, @@ -1191,6 +1288,14 @@ async function runDurableEval< } const progress = await collectProgress(store, prefix, scorerNames); + if (controller.signal.aborted) { + return pausedResult( + runId, + options.signal?.aborted ? "aborted" : "deadline", + progress, + resume, + ); + } if (progress.unknown > 0) { return pausedResult(runId, "unknown_submission", progress, resume); } @@ -1242,6 +1347,7 @@ async function runDurableEval< } } } finally { + if (deadlineTimer) clearTimeout(deadlineTimer); options.signal?.removeEventListener("abort", abortHandler); } } @@ -1300,6 +1406,123 @@ async function updateManifest( } } +async function initializeDurableExperiment({ + store, + key, + workerId, + experimentName, + create, +}: { + store: DurableEvalStore; + key: string; + workerId: string; + experimentName: string; + create: () => Experiment; +}): Promise { + while (true) { + const current = await readJson(store, key); + if (current?.value.status === "ready") { + const experiment = create(); + const experimentId = await experiment.id; + if ( + current.value.experimentId && + current.value.experimentId !== experimentId + ) { + throw new Error( + `DurableEval experiment ${experimentName} resolved to multiple experiment IDs`, + ); + } + return experiment; + } + + const now = Date.now(); + if (!current || (current.value.leaseUntil ?? 0) <= now) { + const claimed: DurableExperimentInitRecord = { + schemaVersion: CHECKPOINT_VERSION, + experimentName, + status: "initializing", + workerId, + leaseUntil: now + JOB_LEASE_MS, + }; + const write = await writeJson( + store, + key, + claimed, + current ? { ifVersion: current.version } : { ifAbsent: true }, + ); + if (write.written) { + const heartbeat = startLeaseHeartbeat(async () => { + let renewed = false; + await updateExperimentInitRecord(store, key, (record) => { + if ( + record.status !== "initializing" || + record.workerId !== workerId + ) { + return undefined; + } + renewed = true; + return { ...record, leaseUntil: Date.now() + JOB_LEASE_MS }; + }); + return renewed; + }); + try { + const experiment = create(); + const experimentId = await experiment.id; + const finalized = await updateExperimentInitRecord( + store, + key, + (record) => + record.status === "initializing" && record.workerId === workerId + ? { + schemaVersion: CHECKPOINT_VERSION, + experimentName, + status: "ready", + experimentId, + } + : undefined, + ); + if (!finalized) { + throw new Error( + `DurableEval lost the experiment initialization lease for ${experimentName}`, + ); + } + return experiment; + } catch (error) { + await updateExperimentInitRecord(store, key, (record) => + record.status === "initializing" && record.workerId === workerId + ? { ...record, leaseUntil: 0 } + : undefined, + ); + throw error; + } finally { + await heartbeat.stop(); + } + } + } + + await delay(25); + } +} + +async function updateExperimentInitRecord( + store: DurableEvalStore, + key: string, + update: ( + record: DurableExperimentInitRecord, + ) => DurableExperimentInitRecord | undefined, +) { + while (true) { + const current = await readJson(store, key); + if (!current) return false; + const next = update(structuredClone(current.value)); + if (!next) return false; + const written = await writeJson(store, key, next, { + ifVersion: current.version, + }); + if (written.written) return true; + } +} + async function materializeData< Input, Output, @@ -1421,17 +1644,59 @@ async function reconcileScorers( prefix: string, scorerNames: string[], ) { + const active = new Set(scorerNames); for await (const record of listCases(store, prefix)) { - const missing = scorerNames.filter( - (name) => record.value.scores[name] === undefined, - ); - if (!missing.length) continue; - const next = structuredClone(record.value); - for (const name of missing) { - next.scores[name] = { status: "pending", attempts: 0 }; - } - await writeJson(store, caseKey(prefix, record.value.id), next, { - ifVersion: record.version, + await updateCaseRecord(store, prefix, record.value.id, (next) => { + const missing = scorerNames.filter( + (name) => next.scores[name] === undefined, + ); + const removed = Object.keys(next.scores).filter( + (name) => !active.has(name), + ); + if (!missing.length && !removed.length) return undefined; + for (const name of missing) { + next.scores[name] = { status: "pending", attempts: 0 }; + } + const removedScoreKeys = removed.flatMap((name) => { + const state = next.scores[name]; + return state?.status === "succeeded" + ? typeof state.value === "object" && + state.value !== null && + !Array.isArray(state.value) + ? Object.keys(state.value) + : [name] + : []; + }); + for (const name of removed) { + delete next.scores[name]; + } + if (removedScoreKeys.length) { + next.removedScores = [ + ...new Set([...(next.removedScores ?? []), ...removedScoreKeys]), + ]; + next.logPending = true; + } + return next; + }); + } + for await (const key of store.list(`${prefix}/jobs/`)) { + await updateJobRecord(store, key, (job) => { + if ( + !job.stage.startsWith("score:") || + active.has(job.stage.slice("score:".length)) || + job.status === "complete" || + job.status === "failed" + ) { + return undefined; + } + job.status = "failed"; + job.error = { + name: "ScorerRemovedError", + message: `Scorer ${job.stage.slice("score:".length)} was removed`, + }; + delete job.workerId; + delete job.leaseUntil; + return job; }); } } @@ -1441,32 +1706,31 @@ async function resetStages( prefix: string, scorerNames: string[], status: "failed" | "unknown", + shard: number, ) { for await (const current of listCases(store, prefix)) { - const next = structuredClone(current.value); - let changed = false; - if (next.task.status === status) { - next.task = { - status: "pending", - attempts: next.task.attempts, - }; - changed = true; - } - for (const name of scorerNames) { - const state = next.scores[name]; - if (state?.status === status) { - next.scores[name] = { + if (current.value.shard !== shard) continue; + await updateCaseRecord(store, prefix, current.value.id, (next) => { + let changed = false; + if (next.task.status === status) { + next.task = { status: "pending", - attempts: state.attempts, + attempts: next.task.attempts, }; changed = true; } - } - if (changed) { - await writeJson(store, caseKey(prefix, next.id), next, { - ifVersion: current.version, - }); - } + for (const name of scorerNames) { + const state = next.scores[name]; + if (state?.status === status) { + next.scores[name] = { + status: "pending", + attempts: state.attempts, + }; + changed = true; + } + } + return changed ? next : undefined; + }); } } @@ -1522,59 +1786,63 @@ async function cancelRun({ batchContext(runId, job.stage, job, shard, signal), ); } - if ( - job.status === "preparing" || - job.status === "submitting" || - job.status === "submitted" || - job.status === "unknown" - ) { - job.status = "failed"; - job.error = { name: "CancelledError", message: "Durable eval cancelled" }; - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, key, job, { ifVersion: current.version }); - } + await updateJobRecord(store, key, (latest) => { + if ( + latest.status !== "preparing" && + latest.status !== "submitting" && + latest.status !== "submitted" && + latest.status !== "unknown" + ) { + return undefined; + } + latest.status = "failed"; + latest.error = { + name: "CancelledError", + message: "Durable eval cancelled", + }; + delete latest.workerId; + delete latest.leaseUntil; + return latest; + }); } for await (const current of listCases(store, prefix)) { if (current.value.shard !== shard.index) continue; - const next = structuredClone(current.value); - let changed = false; - if (!isTerminal(next.task)) { - next.task = { - status: "failed", - attempts: next.task.attempts, - revision: evaluator.revision, - error: { - name: "CancelledError", - message: "Durable eval cancelled", - }, - }; - changed = true; - } - if (next.task.status === "succeeded") { - for (const definition of scorerDefinitions) { - const state = next.scores[definition.name]; - if (!isTerminal(state)) { - next.scores[definition.name] = { - status: "failed", - attempts: state?.attempts ?? 0, - revision: definition.revision, - error: { - name: "CancelledError", - message: "Durable eval cancelled", - }, - }; - changed = true; + await updateCaseRecord(store, prefix, current.value.id, (next) => { + let changed = false; + if (!isTerminal(next.task)) { + next.task = { + status: "failed", + attempts: next.task.attempts, + revision: evaluator.revision, + error: { + name: "CancelledError", + message: "Durable eval cancelled", + }, + }; + changed = true; + } + if (next.task.status === "succeeded") { + for (const definition of scorerDefinitions) { + const state = next.scores[definition.name]; + if (!isTerminal(state)) { + next.scores[definition.name] = { + status: "failed", + attempts: state?.attempts ?? 0, + revision: definition.revision, + error: { + name: "CancelledError", + message: "Durable eval cancelled", + }, + }; + changed = true; + } } } - } - if (changed) { + if (!changed) return undefined; next.logPending = true; - await writeJson(store, caseKey(prefix, next.id), next, { - ifVersion: current.version, - }); - } + return next; + }); } } @@ -1646,7 +1914,7 @@ async function runTaskPass({ status: "leased", attempts: current.value.task.attempts, workerId, - leaseUntil: Date.now() + 60_000, + leaseUntil: Date.now() + JOB_LEASE_MS, }; const claim = await writeJson( store, @@ -1655,22 +1923,38 @@ async function runTaskPass({ { ifVersion: current.version }, ); if (!claim.written) continue; + const heartbeat = startCaseLeaseHeartbeat({ + store, + prefix, + id: current.value.id, + kind: "task", + workerId, + }); const datum = current.value.datum as EvalCase< unknown, unknown, BaseMetadata >; const attempt = current.value.task.attempts + 1; - const next = structuredClone(current.value); + let taskResult: + | { + status: "succeeded"; + output: JsonValue; + metadata: JsonValue; + tags?: string[]; + } + | { status: "failed"; error: unknown }; try { - let metadata = { ...(current.value.metadata as Record) }; + const metadata = { + ...(current.value.metadata as Record), + }; const hooks: EvalHooks< unknown, Record, EvalParameters > = { meta: (value) => { - metadata = { ...metadata, ...(value as Record) }; + Object.assign(metadata, value); }, metadata, expected: "expected" in datum ? datum.expected : undefined, @@ -1680,34 +1964,57 @@ async function runTaskPass({ trialIndex: current.value.trialIndex, tags: current.value.tags, }; - const output = await localTask(datum.input, hooks); - next.metadata = assertJsonValue(hooks.metadata, "task metadata"); - next.tags = hooks.tags; - next.task = { + const output = await awaitWithSignal( + Promise.resolve().then(() => localTask(datum.input, hooks)), + signal, + ); + taskResult = { status: "succeeded", - attempts: attempt, - revision: evaluator.revision, - value: assertJsonValue(output, "task output"), + output: assertJsonValue(output, "task output"), + metadata: assertJsonValue(hooks.metadata, "task metadata"), + tags: hooks.tags, }; } catch (error) { - next.task = - attempt < 3 - ? { status: "pending", attempts: attempt } - : { - status: "failed", - attempts: attempt, - revision: evaluator.revision, - error: serializeError(error), - }; + taskResult = { status: "failed", error }; + } finally { + await heartbeat.stop(); } - next.logPending = true; - const result = await writeJson( - store, - caseKey(prefix, current.value.id), - next, - { ifVersion: claim.version }, - ); - changed = result.written || changed; + changed = + (await updateCaseRecord(store, prefix, current.value.id, (next) => { + if (next.task.status !== "leased" || next.task.workerId !== workerId) { + return undefined; + } + if (signal.aborted) { + next.task = { + status: "pending", + attempts: next.task.attempts, + }; + return next; + } + if (taskResult.status === "succeeded") { + next.metadata = taskResult.metadata; + next.tags = taskResult.tags; + next.task = { + status: "succeeded", + attempts: attempt, + revision: evaluator.revision, + value: taskResult.output, + }; + } else { + next.task = + attempt < 3 + ? { status: "pending", attempts: attempt } + : { + status: "failed", + attempts: attempt, + revision: evaluator.revision, + error: serializeError(taskResult.error), + }; + } + next.logPending = true; + return next; + })) || changed; + if (signal.aborted) return changed; } return changed; } @@ -1728,7 +2035,6 @@ async function runScorePass({ projectName: string; evalName: string; scorer: ResolvedScorer; - evaluatorRevision: string; shard: { index: number; count: number }; workerId: string; runId: string; @@ -1763,6 +2069,12 @@ async function runScorePass({ }); } + const localScorer = scorer.scorer as EvalScorer< + unknown, + unknown, + unknown, + Record + >; let changed = false; for await (const current of listCases(store, prefix)) { const state = current.value.scores[scorer.name]; @@ -1778,7 +2090,7 @@ async function runScorePass({ status: "leased", attempts: state.attempts, workerId, - leaseUntil: Date.now() + 60_000, + leaseUntil: Date.now() + JOB_LEASE_MS, }; const claim = await writeJson( store, @@ -1787,48 +2099,81 @@ async function runScorePass({ { ifVersion: current.version }, ); if (!claim.written) continue; + const heartbeat = startCaseLeaseHeartbeat({ + store, + prefix, + id: current.value.id, + kind: "score", + scorerName: scorer.name, + workerId, + }); const datum = current.value.datum as EvalCase< unknown, unknown, BaseMetadata >; const attempt = state.attempts + 1; - const next = structuredClone(current.value); + const taskOutput = current.value.task.value; + let scoreResult: + | { status: "succeeded"; value: JsonValue } + | { status: "failed"; error: unknown }; try { - const raw = await scorer.scorer({ - input: datum.input, - expected: "expected" in datum ? datum.expected : undefined, - metadata: current.value.metadata as BaseMetadata, - output: current.value.task.value, - }); - next.scores[scorer.name] = { + const raw = await awaitWithSignal( + Promise.resolve().then(() => + localScorer({ + input: datum.input, + expected: "expected" in datum ? datum.expected : undefined, + metadata: current.value.metadata as Record, + output: taskOutput, + }), + ), + signal, + ); + scoreResult = { status: "succeeded", - attempts: attempt, - revision: scorer.revision, value: assertJsonValue( normalizeScores(raw, scorer.name), `scorer ${scorer.name} output`, ), }; } catch (error) { - next.scores[scorer.name] = - attempt < 3 - ? { status: "pending", attempts: attempt } - : { - status: "failed", - attempts: attempt, - revision: scorer.revision, - error: serializeError(error), - }; + scoreResult = { status: "failed", error }; + } finally { + await heartbeat.stop(); } - next.logPending = true; - const result = await writeJson( - store, - caseKey(prefix, current.value.id), - next, - { ifVersion: claim.version }, - ); - changed = result.written || changed; + changed = + (await updateCaseRecord(store, prefix, current.value.id, (next) => { + const latest = next.scores[scorer.name]; + if (latest?.status !== "leased" || latest.workerId !== workerId) { + return undefined; + } + if (signal.aborted) { + next.scores[scorer.name] = { + status: "pending", + attempts: latest.attempts, + }; + return next; + } + next.scores[scorer.name] = + scoreResult.status === "succeeded" + ? { + status: "succeeded", + attempts: attempt, + revision: scorer.revision, + value: scoreResult.value, + } + : attempt < 3 + ? { status: "pending", attempts: attempt } + : { + status: "failed", + attempts: attempt, + revision: scorer.revision, + error: serializeError(scoreResult.error), + }; + next.logPending = true; + return next; + })) || changed; + if (signal.aborted) return changed; } return changed; } @@ -1891,25 +2236,40 @@ async function runBatchStage({ if (job.status === "preparing") { activeJobs++; if ((job.leaseUntil ?? 0) > Date.now()) continue; + const takeover = { + ...job, + workerId, + leaseUntil: Date.now() + JOB_LEASE_MS, + }; + const takeoverResult = await writeJson(store, current.key, takeover, { + ifVersion: current.version, + }); + if (!takeoverResult.written) continue; + job = takeover; const preparationError = new Error( "Batch preparation lease expired before provider submission", ); - await retryBatchItems({ - store, - prefix, - job, - scorerName, - processor, - error: preparationError, - retryable: true, - }); - job.status = "failed"; - job.error = serializeError(preparationError); - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: current.version, + const failed = await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId || latest.status !== "preparing") { + return undefined; + } + latest.status = "failed"; + latest.error = serializeError(preparationError); + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); + if (failed?.written && failed.value.status === "failed") { + await retryBatchItems({ + store, + prefix, + job, + scorerName, + processor, + error: preparationError, + retryable: true, + }); + } activeJobs--; changed = true; continue; @@ -1947,26 +2307,50 @@ async function runBatchStage({ }); if (!jobClaim.written) continue; job = claimedJob; - let jobVersion = current.version; - jobVersion = jobClaim.version; const context = batchContext(runId, stage, job, shard, signal); let handle = job.handle as Handle | undefined; if (!handle) { - const recovery = processor.recover - ? await processor.recover(context) - : { status: "unknown" as const }; + let recovery: DurableBatchRecovery; + try { + recovery = processor.recover + ? await withJobLeaseHeartbeat({ + store, + key: current.key, + workerId, + signal, + operation: () => processor.recover!(context), + }) + : { status: "unknown" }; + } catch (error) { + await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId) return undefined; + latest.error = serializeError(error); + latest.nextPollAt = Date.now() + 10_000; + delete latest.workerId; + delete latest.leaseUntil; + return latest; + }); + changed = true; + continue; + } + const latest = await readJson(store, current.key); + if (!latest || latest.value.workerId !== workerId) { + changed = true; + continue; + } + job = latest.value; if (recovery.status === "found") { handle = recovery.handle; const external = externalBatchReference(processor, handle, context); const recovered = await persistSubmittedJob({ store, jobKey: current.key, + workerId, handle: assertJsonValue(handle, `${stage} batch handle`), external, submittedAt: job.submittedAt ?? Date.now(), }); job = recovered.value; - jobVersion = recovered.version; const locator = batchLocator({ projectName, evalName, @@ -1983,35 +2367,41 @@ async function runBatchStage({ const latest = await readJson(store, current.key); if (latest) { job = latest.value; - jobVersion = latest.version; } changed = true; } else if (recovery.status === "not_found") { - await retryBatchItems({ - store, - prefix, - job, - scorerName, - processor, - error: new Error("Provider confirmed batch was not submitted"), - retryable: true, - }); - job.status = "failed"; - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: jobVersion, + const error = new Error("Provider confirmed batch was not submitted"); + const failed = await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId) return undefined; + latest.status = "failed"; + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); + if (failed?.written && failed.value.status === "failed") { + await retryBatchItems({ + store, + prefix, + job, + scorerName, + processor, + error, + retryable: true, + }); + } changed = true; continue; } else { - await markBatchUnknown(store, prefix, job, scorerName); - job.status = "unknown"; - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: jobVersion, + const unknown = await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId) return undefined; + latest.status = "unknown"; + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); + if (unknown?.written) { + await markBatchUnknown(store, prefix, job, scorerName); + } changed = true; continue; } @@ -2024,10 +2414,11 @@ async function runBatchStage({ (!fallbackReady || job.outcome !== undefined) ) { activeJobs++; - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: jobVersion, + await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId) return undefined; + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); continue; } @@ -2040,72 +2431,98 @@ async function runBatchStage({ : undefined; if (!poller) { activeJobs++; - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: jobVersion, + await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId) return undefined; + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); continue; } try { - poll = await poller.poll(handle, context); + poll = await withJobLeaseHeartbeat({ + store, + key: current.key, + workerId, + signal, + operation: () => poller.poll(handle, context), + }); + const latest = await readJson(store, current.key); + if (!latest || latest.value.workerId !== workerId) { + changed = true; + continue; + } + job = latest.value; delete job.error; } catch (error) { activeJobs++; - job.nextPollAt = Date.now() + (poller.intervalMs ?? 10_000); - job.error = serializeError(error); - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: jobVersion, + await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId) return undefined; + latest.nextPollAt = Date.now() + (poller.intervalMs ?? 10_000); + latest.error = serializeError(error); + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); continue; } } if (poll.status === "pending") { activeJobs++; - delete job.error; - job.nextPollAt = - Date.now() + - (poll.retryAfterMs ?? - (completion.mode === "poll" - ? completion.intervalMs - : fallback?.intervalMs) ?? - 10_000); - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: jobVersion, + await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId) return undefined; + delete latest.error; + latest.nextPollAt = + Date.now() + + (poll.retryAfterMs ?? + (completion.mode === "poll" + ? completion.intervalMs + : fallback?.intervalMs) ?? + 10_000); + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); continue; } if (poll.status === "failed") { - await retryBatchItems({ - store, - prefix, - job, - scorerName, - processor, - error: poll.error, - retryable: poll.retryable ?? false, - }); - job.status = "failed"; - job.error = serializeError(poll.error); - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: jobVersion, + const failed = await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId) return undefined; + latest.status = "failed"; + latest.error = serializeError(poll.error); + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); - await markWebhookApplied(store, job); + if (failed?.written && failed.value.status === "failed") { + await retryBatchItems({ + store, + prefix, + job, + scorerName, + processor, + error: poll.error, + retryable: poll.retryable ?? false, + }); + await markWebhookApplied(store, failed.value); + } changed = true; continue; } + const collectionHeartbeat = startJobLeaseHeartbeat( + store, + current.key, + workerId, + ); try { delete job.error; const seen = new Set(); - const collected = await processor.collect(handle, context); - for await (const result of toAsyncIterable(collected)) { + const collected = await awaitWithSignal( + Promise.resolve().then(() => processor.collect(handle, context)), + signal, + ); + for await (const result of toAbortableAsyncIterable(collected, signal)) { const resultId = resultItemId(result); if (seen.has(resultId) || !job.itemIds.includes(resultId)) { throw new Error( @@ -2113,74 +2530,63 @@ async function runBatchStage({ ); } seen.add(resultId); - const caseRecord = await readJson( - store, - caseKey(prefix, resultId), - ); - if ( - !caseRecord || - !stageBelongsToJob(caseRecord.value, job.kind, job.id, scorerName) - ) { - continue; - } - const next = applyResult( - structuredClone(caseRecord.value), - result, - job.revision, - job.attempt, - ); - next.logPending = true; - await writeJson(store, caseKey(prefix, resultId), next, { - ifVersion: caseRecord.version, + await updateCaseRecord(store, prefix, resultId, (record) => { + if (!stageBelongsToJob(record, job.kind, job.id, scorerName)) { + return undefined; + } + const next = applyResult(record, result, job.revision, job.attempt); + next.logPending = true; + return next; }); } for (const itemId of job.itemIds) { if (seen.has(itemId)) continue; - const caseRecord = await readJson( - store, - caseKey(prefix, itemId), - ); - if ( - !caseRecord || - !stageBelongsToJob(caseRecord.value, job.kind, job.id, scorerName) - ) { - continue; - } const missing = { id: itemId, error: new Error(`Batch stage ${stage} returned no result`), retryable: true, } as Result; - const next = applyResult( - structuredClone(caseRecord.value), - missing, - job.revision, - job.attempt, - ); - next.logPending = true; - await writeJson(store, caseKey(prefix, itemId), next, { - ifVersion: caseRecord.version, + await updateCaseRecord(store, prefix, itemId, (record) => { + if (!stageBelongsToJob(record, job.kind, job.id, scorerName)) { + return undefined; + } + const next = applyResult(record, missing, job.revision, job.attempt); + next.logPending = true; + return next; }); } } catch (error) { activeJobs++; - job.error = serializeError(error); - job.nextPollAt = Date.now() + 10_000; - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: jobVersion, + await updateJobRecord(store, current.key, (latest) => { + if (latest.workerId !== workerId) return undefined; + latest.error = serializeError(error); + latest.nextPollAt = Date.now() + 10_000; + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); changed = true; continue; + } finally { + await collectionHeartbeat.stop(); } - job.status = "complete"; - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, current.key, job, { - ifVersion: jobVersion, + const completed = await updateJobRecord(store, current.key, (latest) => { + if ( + latest.id !== job.id || + latest.status !== "submitted" || + latest.workerId !== workerId + ) { + return undefined; + } + latest.status = "complete"; + delete latest.workerId; + delete latest.leaseUntil; + delete latest.error; + return latest; }); - await markWebhookApplied(store, job); + if (completed?.value.status !== "complete") continue; + job = completed.value; + await markWebhookApplied(store, completed.value); changed = true; } @@ -2237,8 +2643,10 @@ async function runBatchStage({ }), ); + const preparationHeartbeat = startJobLeaseHeartbeat(store, jobKey, workerId); const claimed: Versioned[] = []; for (const record of ready) { + if (signal.aborted) break; const next = structuredClone(record.value); const state: StageState = { status: "in_batch", @@ -2259,31 +2667,82 @@ async function runBatchStage({ if (claim.written) claimed.push(record); } job.itemIds = claimed.map((record) => record.value.id); + if (signal.aborted) { + await preparationHeartbeat.stop(); + const failed = await updateJobRecord(store, jobKey, (latest) => { + if (latest.workerId !== workerId || latest.status !== "preparing") { + return undefined; + } + latest.itemIds = job.itemIds; + latest.status = "failed"; + latest.error = serializeError(new DurableEvalAbortError()); + delete latest.workerId; + delete latest.leaseUntil; + return latest; + }); + if (failed?.written) { + await retryBatchItems({ + store, + prefix, + job, + scorerName, + processor, + error: new DurableEvalAbortError(), + retryable: true, + }); + } + return true; + } if (!claimed.length) { - job.status = "failed"; - job.error = serializeError(new Error("Batch lost all item claims")); - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, jobKey, job, { - ifVersion: created.version, + await preparationHeartbeat.stop(); + await updateJobRecord(store, jobKey, (latest) => { + if (latest.workerId !== workerId || latest.status !== "preparing") { + return undefined; + } + latest.status = "failed"; + latest.error = serializeError(new Error("Batch lost all item claims")); + delete latest.workerId; + delete latest.leaseUntil; + return latest; }); return changed; } - job.status = "submitting"; - const prepared = await writeJson(store, jobKey, job, { - ifVersion: created.version, + const prepared = await updateJobRecord(store, jobKey, (latest) => { + if (latest.workerId !== workerId || latest.status !== "preparing") { + return undefined; + } + latest.itemIds = job.itemIds; + latest.status = "submitting"; + latest.leaseUntil = Date.now() + JOB_LEASE_MS; + return latest; }); - if (!prepared.written) return changed; + await preparationHeartbeat.stop(); + if ( + !prepared || + !prepared.written || + prepared.value.status !== "submitting" || + prepared.value.workerId !== workerId + ) { + return changed; + } + job = prepared.value; const context = batchContext(runId, stage, job, shard, signal); const items = claimed.map((record) => makeItem(record.value)); try { - const handle = await processor.submit(items, context); + const handle = await withJobLeaseHeartbeat({ + store, + key: jobKey, + workerId, + signal, + operation: () => processor.submit(items, context), + }); const external = externalBatchReference(processor, handle, context); const submitted = await persistSubmittedJob({ store, jobKey, + workerId, handle: assertJsonValue(handle, `${stage} batch handle`), external, submittedAt: Date.now(), @@ -2303,7 +2762,29 @@ async function runBatchStage({ } await attachPendingWebhookEvents(store, jobKey, job); } catch (error) { - if (error instanceof DurableEvalNotSubmittedError) { + const definitelyNotSubmitted = + error instanceof DurableEvalNotSubmittedError; + const terminalStatus = definitelyNotSubmitted ? "failed" : "unknown"; + const updated = await updateJobRecord(store, jobKey, (latest) => { + if ( + latest.workerId !== workerId || + latest.status === "complete" || + latest.status === "failed" || + latest.webhookEventKey + ) { + return undefined; + } + latest.status = terminalStatus; + latest.error = serializeError(error); + delete latest.workerId; + delete latest.leaseUntil; + return latest; + }); + if ( + updated?.written && + updated.value.status === "failed" && + definitelyNotSubmitted + ) { await retryBatchItems({ store, prefix, @@ -2313,18 +2794,13 @@ async function runBatchStage({ error, retryable: true, }); - job.status = "failed"; - job.error = serializeError(error); - } else { + } else if ( + updated?.written && + updated.value.status === "unknown" && + !definitelyNotSubmitted + ) { await markBatchUnknown(store, prefix, job, scorerName); - job.status = "unknown"; - job.error = serializeError(error); } - delete job.workerId; - delete job.leaseUntil; - await writeJson(store, jobKey, job, { - ifVersion: prepared.version, - }); } return true; } @@ -2444,29 +2920,23 @@ async function retryBatchItems({ retryable: boolean; }) { for (const itemId of job.itemIds) { - const current = await readJson( - store, - caseKey(prefix, itemId), - ); - if (!current) continue; - const next = structuredClone(current.value); - if (!stageBelongsToJob(next, job.kind, job.id, scorerName)) { - continue; - } - const state = - retryable && job.attempt < (processor.maxAttempts ?? 3) - ? ({ status: "pending", attempts: job.attempt } satisfies StageState) - : ({ - status: "failed", - attempts: job.attempt, - revision: job.revision, - error: serializeError(error), - } satisfies StageState); - if (job.kind === "task") next.task = state; - else next.scores[scorerName!] = state; - next.logPending = true; - await writeJson(store, caseKey(prefix, itemId), next, { - ifVersion: current.version, + await updateCaseRecord(store, prefix, itemId, (record) => { + if (!stageBelongsToJob(record, job.kind, job.id, scorerName)) { + return undefined; + } + const state = + retryable && job.attempt < (processor.maxAttempts ?? 3) + ? ({ status: "pending", attempts: job.attempt } satisfies StageState) + : ({ + status: "failed", + attempts: job.attempt, + revision: job.revision, + error: serializeError(error), + } satisfies StageState); + if (job.kind === "task") record.task = state; + else record.scores[scorerName!] = state; + record.logPending = true; + return record; }); } } @@ -2478,23 +2948,19 @@ async function markBatchUnknown( scorerName?: string, ) { for (const itemId of job.itemIds) { - const current = await readJson( - store, - caseKey(prefix, itemId), - ); - if (!current) continue; - const next = structuredClone(current.value); - if (!stageBelongsToJob(next, job.kind, job.id, scorerName)) continue; - const state: StageState = { - status: "unknown", - attempts: job.attempt, - revision: job.revision, - batchId: job.id, - }; - if (job.kind === "task") next.task = state; - else next.scores[scorerName!] = state; - await writeJson(store, caseKey(prefix, itemId), next, { - ifVersion: current.version, + await updateCaseRecord(store, prefix, itemId, (record) => { + if (!stageBelongsToJob(record, job.kind, job.id, scorerName)) { + return undefined; + } + const state: StageState = { + status: "unknown", + attempts: job.attempt, + revision: job.revision, + batchId: job.id, + }; + if (job.kind === "task") record.task = state; + else record.scores[scorerName!] = state; + return record; }); } } @@ -2598,12 +3064,14 @@ function externalBatchReference( async function persistSubmittedJob({ store, jobKey, + workerId, handle, external, submittedAt, }: { store: DurableEvalStore; jobKey: string; + workerId: string; handle: JsonValue; external?: { source: string; id: string }; submittedAt: number; @@ -2613,6 +3081,9 @@ async function persistSubmittedJob({ if (!current) { throw new Error("Durable batch job disappeared during submission"); } + if (current.value.workerId !== workerId) { + throw new Error("Durable batch submission lease was lost"); + } if ( current.value.handle !== undefined && stableStringify(current.value.handle) !== stableStringify(handle) @@ -2685,7 +3156,7 @@ async function attachWebhookEventToJob( jobKey: string, storedEventKey: string, event: DurableBatchResultEvent, -): Promise<"attached" | "duplicate" | "missing"> { +): Promise<"attached" | "duplicate" | "in_progress" | "missing"> { while (true) { const current = await readJson(store, jobKey); if (!current) return "missing"; @@ -2696,6 +3167,12 @@ async function attachWebhookEventToJob( return "duplicate"; } if (current.value.webhookEventKey === storedEventKey) { + if ( + current.value.workerId && + (current.value.leaseUntil ?? 0) > Date.now() + ) { + return "in_progress"; + } const next = { ...current.value, nextPollAt: Date.now(), @@ -2777,7 +3254,7 @@ async function attachPendingWebhookEvents( pointer.value.eventKey, event.value.event, ); - if (attached !== "attached") { + if (attached === "duplicate") { await markWebhookEventApplied(store, pointer.value.eventKey, job.id); } } @@ -2822,7 +3299,7 @@ async function hasWaitingWebhookJob( const job = await readJson(store, key); if ( job?.value.shard === shard && - job.value.status === "submitted" && + (job.value.status === "submitting" || job.value.status === "submitted") && job.value.outcome === undefined && webhookStages.has(job.value.stage) ) { @@ -2841,7 +3318,7 @@ async function hasProviderErrorJob( const job = await readJson(store, key); if ( job?.value.shard === shard && - job.value.status === "submitted" && + (job.value.status === "submitting" || job.value.status === "submitted") && job.value.error !== undefined ) { return true; @@ -2872,21 +3349,16 @@ async function flushPendingLogs({ } await logDurableCase(experiment, current.value, scorerNames, runId); await experiment.flush(); - const latest = await readJson( + const next = structuredClone(current.value); + next.logPending = false; + delete next.removedScores; + const result = await writeJson( store, caseKey(prefix, current.value.id), + next, + { ifVersion: current.version }, ); - if (latest) { - const next = structuredClone(latest.value); - next.logPending = false; - const result = await writeJson( - store, - caseKey(prefix, current.value.id), - next, - { ifVersion: latest.version }, - ); - changed = result.written || changed; - } + changed = result.written || changed; } return changed; } @@ -2992,6 +3464,9 @@ function evaluatorState(experiment: Experiment) { function collectedScores(record: DurableCaseRecord, scorerNames: string[]) { const scores: Record = {}; + for (const name of record.removedScores ?? []) { + scores[name] = null; + } for (const name of scorerNames) { const state = record.scores[name]; if (state?.status === "succeeded") { @@ -3301,6 +3776,166 @@ async function writeJson( ); } +async function updateCaseRecord( + store: DurableEvalStore, + prefix: string, + id: string, + update: (record: DurableCaseRecord) => DurableCaseRecord | undefined, +) { + while (true) { + const current = await readJson( + store, + caseKey(prefix, id), + ); + if (!current) return false; + const next = update(structuredClone(current.value)); + if (!next) return false; + const written = await writeJson(store, caseKey(prefix, id), next, { + ifVersion: current.version, + }); + if (written.written) return true; + } +} + +async function updateJobRecord( + store: DurableEvalStore, + key: string, + update: (record: DurableJobRecord) => DurableJobRecord | undefined, +) { + while (true) { + const current = await readJson(store, key); + if (!current) return undefined; + const next = update(structuredClone(current.value)); + if (!next) return { ...current, written: false as const }; + const written = await writeJson(store, key, next, { + ifVersion: current.version, + }); + if (written.written) { + return { value: next, version: written.version, written: true as const }; + } + } +} + +function startCaseLeaseHeartbeat({ + store, + prefix, + id, + kind, + scorerName, + workerId, +}: { + store: DurableEvalStore; + prefix: string; + id: string; + kind: "task" | "score"; + scorerName?: string; + workerId: string; +}) { + return startLeaseHeartbeat(async () => { + const updated = await updateCaseRecord(store, prefix, id, (record) => { + const state = + kind === "task" ? record.task : record.scores[scorerName ?? ""]; + if (state?.status !== "leased" || state.workerId !== workerId) { + return undefined; + } + state.leaseUntil = Date.now() + JOB_LEASE_MS; + return record; + }); + return updated; + }); +} + +function startJobLeaseHeartbeat( + store: DurableEvalStore, + key: string, + workerId: string, +) { + return startLeaseHeartbeat(async () => { + const updated = await updateJobRecord(store, key, (job) => { + if ( + job.workerId !== workerId || + (job.status !== "preparing" && + job.status !== "submitting" && + job.status !== "submitted") + ) { + return undefined; + } + job.leaseUntil = Date.now() + JOB_LEASE_MS; + return job; + }); + return updated?.value.workerId === workerId; + }); +} + +async function withJobLeaseHeartbeat({ + store, + key, + workerId, + signal, + operation, +}: { + store: DurableEvalStore; + key: string; + workerId: string; + signal: AbortSignal; + operation: () => Promise; +}) { + if (signal.aborted) throw new DurableEvalAbortError(); + const heartbeat = startJobLeaseHeartbeat(store, key, workerId); + try { + return await awaitWithSignal(Promise.resolve().then(operation), signal); + } finally { + await heartbeat.stop(); + } +} + +function startLeaseHeartbeat(renew: () => Promise) { + let stopped = false; + let timer: ReturnType | undefined; + let pending = Promise.resolve(); + const schedule = () => { + if (stopped) return; + timer = setTimeout(() => { + pending = renew() + .then((owned) => { + if (!owned) stopped = true; + }) + .finally(schedule); + }, LEASE_HEARTBEAT_MS); + }; + schedule(); + return { + async stop() { + stopped = true; + if (timer) clearTimeout(timer); + await pending; + }, + }; +} + +function awaitWithSignal( + operation: Promise, + signal: AbortSignal, +): Promise { + if (signal.aborted) { + return Promise.reject(new DurableEvalAbortError()); + } + return new Promise((resolve, reject) => { + const abort = () => reject(new DurableEvalAbortError()); + signal.addEventListener("abort", abort, { once: true }); + operation.then(resolve, reject).finally(() => { + signal.removeEventListener("abort", abort); + }); + }); +} + +class DurableEvalAbortError extends Error { + constructor() { + super("Durable eval operation was aborted"); + this.name = "AbortError"; + } +} + function stableStringify(value: unknown): string { return JSON.stringify(value, (_key, nested) => { if (nested && typeof nested === "object" && !Array.isArray(nested)) { @@ -3376,6 +4011,25 @@ async function* toAsyncIterable( } } +async function* toAbortableAsyncIterable( + value: Iterable | AsyncIterable, + signal: AbortSignal, +): AsyncGenerator { + if (isAsyncIterable(value)) { + const iterator = value[Symbol.asyncIterator](); + while (true) { + const next = await awaitWithSignal(iterator.next(), signal); + if (next.done) return; + yield next.value; + } + } else { + for (const item of value) { + if (signal.aborted) throw new DurableEvalAbortError(); + yield item; + } + } +} + function delay(ms: number) { return new Promise((resolve) => setTimeout(resolve, Math.max(ms, 0))); } From beea0e8363109d3f2d96a090dccef99fd57274b6 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Tue, 28 Jul 2026 10:40:37 +0000 Subject: [PATCH 04/13] Update PR #2297 --- js/src/durable-eval.test.ts | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts index 43dc05c30..0cfd1a043 100644 --- a/js/src/durable-eval.test.ts +++ b/js/src/durable-eval.test.ts @@ -137,7 +137,13 @@ describe("DurableEval", () => { test("persists metadata written through hooks.meta", async () => { const result = await DurableEval("metadata-hooks", { revision: "v1", - data: [{ id: "one", input: "hello" }], + data: [ + { + id: "one", + input: "hello", + metadata: { source: "initial" }, + }, + ], task: (_input, hooks) => { hooks.meta({ source: "deprecated-hook" }); return hooks.metadata.source; From 4fc69497398232771cd501ef724a1d2ad3369fb0 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Wed, 29 Jul 2026 11:25:16 +0000 Subject: [PATCH 05/13] Update PR #2297 --- .../durable-eval-webhook/scenario.ts | 60 ++++++++++++++----- js/README.md | 20 +++---- js/src/exports.ts | 31 +--------- 3 files changed, 54 insertions(+), 57 deletions(-) diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts index 8b5d7cfb4..8987970af 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -1,27 +1,55 @@ -import { - BatchTask, - DurableEval, - MemoryDurableEvalStore, - type DurableBatchTaskItem, -} from "braintrust"; +import { BatchTask, DurableEval, type DurableEvalStore } from "braintrust"; import { getTestRunId, runMain, scopedName, } from "../../helpers/scenario-runtime"; -async function main() { - const testRunId = getTestRunId(); - const store = new MemoryDurableEvalStore(); - const jobs = new Map< +class MemoryStore implements DurableEvalStore { + private readonly values = new Map< string, - DurableBatchTaskItem< - number, - number, - { testRunId: string; kind: string }, - Record - >[] + { value: Uint8Array; version: number } >(); + + async read(key: string) { + const record = this.values.get(key); + return record + ? { value: record.value.slice(), version: String(record.version) } + : undefined; + } + + async write( + key: string, + value: Uint8Array, + condition: Parameters[2], + ) { + const current = this.values.get(key); + if ( + ("ifAbsent" in condition && current) || + ("ifVersion" in condition && + (!current || String(current.version) !== condition.ifVersion)) + ) { + return { + written: false as const, + currentVersion: current ? String(current.version) : undefined, + }; + } + const version = (current?.version ?? 0) + 1; + this.values.set(key, { value: value.slice(), version }); + return { written: true as const, version: String(version) }; + } + + async *list(prefix: string) { + for (const key of [...this.values.keys()].sort()) { + if (key.startsWith(prefix)) yield key; + } + } +} + +async function main() { + const testRunId = getTestRunId(); + const store = new MemoryStore(); + const jobs = new Map>(); const task = BatchTask< number, number, diff --git a/js/README.md b/js/README.md index ddcf3f7d4..920d11881 100644 --- a/js/README.md +++ b/js/README.md @@ -56,9 +56,9 @@ Every case needs a stable `id` (or a `caseId` function). Reusing a `runId` resumes the same input snapshot and never reruns successful work. ```typescript -import { BatchTask, DurableEval, FileDurableEvalStore } from "braintrust"; +import { BatchTask, DurableEval, type DurableEvalStore } from "braintrust"; -const store = new FileDurableEvalStore(); +const store: DurableEvalStore = checkpointStore; const supportEval = DurableEval("Support bot", { revision: process.env.GIT_SHA ?? "local", data: [ @@ -121,10 +121,9 @@ if (result.status === "paused") { ``` Use `BatchScorer` for asynchronous batch scoring. Existing per-item tasks and -scorers are also supported and gain checkpointing and retries. The filesystem -store is intended for one machine; distributed workers should share a -compare-and-set-capable `DurableEvalStore`. Launch fixed shards through the API -or the existing CLI: +scorers are also supported and gain checkpointing and retries. Supply a +compare-and-set-capable `DurableEvalStore`; distributed workers must share the +same store. Launch fixed shards through the API or the existing CLI: ```bash npx braintrust eval evaluation.eval.ts \ @@ -133,11 +132,10 @@ npx braintrust eval evaluation.eval.ts \ --deadline 6h ``` -An ordinary error thrown from `submit` is considered ambiguous, because the -provider may have created a job before the connection failed. The run pauses -instead of risking duplicate charges. Throw `DurableEvalNotSubmittedError` only -when it is known that no provider job was created, or implement `recover` to -look up a prior submission using `context.batchId`. +An error thrown from `submit` is considered ambiguous, because the provider may +have created a job before the connection failed. The run pauses instead of +risking duplicate charges. Implement `recover` to look up a prior submission +using `context.batchId`. The `submit`/`completion`/`collect` split is designed for provider-managed batch APIs, including OpenAI's Batch API. A polling adapter uses diff --git a/js/src/exports.ts b/js/src/exports.ts index bee995e8e..0e9415917 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -262,42 +262,13 @@ export { defaultErrorScoreHandler, } from "./framework"; -export type { - DurableBatchCompletion, - DurableBatchContext, - DurableBatchOutcome, - DurableBatchPoll, - DurableBatchProcessingResult, - DurableBatchProcessor, - DurableBatchRecovery, - DurableBatchResultEvent, - DurableBatchScorer, - DurableBatchScorerItem, - DurableBatchScorerResult, - DurableBatchTask, - DurableBatchTaskItem, - DurableBatchTaskResult, - DurableEvalDefinition, - DurableEvalExistingRunOptions, - DurableEvalFailureSummary, - DurableEvalOperation, - DurableEvalPauseReason, - DurableEvalProgress, - DurableEvalResult, - DurableEvalRuntimeOptions, - DurableEvalStore, - DurableEvalWriteCondition, - DurableEvaluator, - JsonValue, -} from "./durable-eval"; +export type { DurableEvalStore } from "./durable-eval"; export { BatchScorer, BatchTask, DurableEval, DurableEvalNotSubmittedError, - FileDurableEvalStore, - MemoryDurableEvalStore, } from "./durable-eval"; export { agentAssertionScorer } from "./agent-assertions"; From 6237c9ebc117b847d3cee67d72297705a3b9093d Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Wed, 29 Jul 2026 16:23:47 +0000 Subject: [PATCH 06/13] Update PR #2297 --- js/README.md | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/js/README.md b/js/README.md index 920d11881..e0d202d99 100644 --- a/js/README.md +++ b/js/README.md @@ -60,6 +60,8 @@ import { BatchTask, DurableEval, type DurableEvalStore } from "braintrust"; const store: DurableEvalStore = checkpointStore; const supportEval = DurableEval("Support bot", { + // Records which eval definition produced each checkpointed result. Update it + // when eval-level data or orchestration changes; completed cases stay reused. revision: process.env.GIT_SHA ?? "local", data: [ { @@ -69,6 +71,8 @@ const supportEval = DurableEval("Support bot", { }, ], task: BatchTask({ + // Batch stages have their own revision so results retain the version of the + // prompt, model, and task code that produced them. revision: "generation-v1", // Each provider job contains at most 500 eval cases. Up to four jobs may @@ -87,6 +91,10 @@ const supportEval = DurableEval("Support bot", { return { id: batch.id }; }, + // Completion controls how DurableEval learns that the asynchronous job is + // ready. "webhook" pauses until processBatchResult receives an event; + // "poll" calls a supplied poll function. Webhooks may also define a + // pollFallback for delayed or missing events. completion: { mode: "webhook", source: "provider", @@ -147,21 +155,24 @@ same contract. ### Webhook processing -Verify the provider signature against the raw request body before calling the -durable eval definition. The HTTP framework and provider SDK remain outside the -generic durable API: +Pass a provider batch's terminal state to the durable eval definition: ```typescript -app.post("/webhooks/provider", rawBodyMiddleware, async (request, response) => { - const event = provider.verifyWebhook(request.rawBody, request.headers); +app.post("/webhooks/provider", async (request, response) => { + const event = request.body; const batch = await provider.getBatch(event.batchId); const result = await supportEval.processBatchResult( { + // Uniquely identifies this webhook event so repeated delivery is safe. eventId: event.id, + // Namespaces provider job IDs when multiple webhook sources share a store. source: "provider", + // Matches the event to the provider job handle saved during submission. externalId: batch.id, + // Matches events that arrive before the provider handle is checkpointed. batchId: batch.metadata?.durableBatchId, + // Supplies the provider handle when recovering an early or ambiguous event. handle: { id: batch.id }, outcome: batch.status === "completed" @@ -169,6 +180,8 @@ app.post("/webhooks/provider", rawBodyMiddleware, async (request, response) => { : { status: "failed", error: { status: batch.status }, + // false makes this failure terminal; true requeues its items until + // the batch stage reaches maxAttempts. retryable: false, }, // Optional JSON payloads are checkpointed with the deduplicated event. From 365ace2286aef4c53260dbc95a135b09622b0fcf Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:16:58 +0000 Subject: [PATCH 07/13] Update PR #2297 --- .../ai-sdk-harness-instrumentation/scenario.test.ts | 11 ++++++++++- knip.jsonc | 3 +++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/e2e/scenarios/ai-sdk-harness-instrumentation/scenario.test.ts b/e2e/scenarios/ai-sdk-harness-instrumentation/scenario.test.ts index a03599324..0ad7e0525 100644 --- a/e2e/scenarios/ai-sdk-harness-instrumentation/scenario.test.ts +++ b/e2e/scenarios/ai-sdk-harness-instrumentation/scenario.test.ts @@ -245,7 +245,16 @@ describe.sequential("HarnessAgent instrumentation variants", () => { ); const harnessSpans = findAllSpans(events, "harness"); expect(harnessSpans).toHaveLength(4); - const bashSpans = findAllSpans(events, "bash"); + // The harness may issue additional bash calls while coordinating a + // suspended turn. Assert only the two commands requested from the + // agent; coordination calls are not part of this contract. + const bashSpans = findAllSpans(events, "bash").filter((span) => { + const input = String(span.input); + return ( + input.includes("printf GENERATE_OK") || + input.includes("printf STREAM_OK") + ); + }); expect(bashSpans).toHaveLength(2); for (const bashSpan of bashSpans) { expect(bashSpan.span.type).toBe("tool"); diff --git a/knip.jsonc b/knip.jsonc index 9752cb9a2..7337fa278 100644 --- a/knip.jsonc +++ b/knip.jsonc @@ -10,6 +10,9 @@ ], "ignoreIssues": { "**/generated_types.ts": ["exports", "types"], + // These support the inferred signatures of the intentionally small public + // DurableEval API and must remain exported for declaration bundling. + "js/src/durable-eval.ts": ["types"], }, "workspaces": { "dev-packages/seinfeld": { From 11669b429393ce05a1c27b16030a2f174c24ad16 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Tue, 4 Aug 2026 11:30:41 +0000 Subject: [PATCH 08/13] Update PR #2297 --- .../durable-eval-webhook/scenario.test.ts | 8 - .../durable-eval-webhook/scenario.ts | 159 +- js/README.md | 235 +- js/src/cli/functions/load-module.ts | 1 - js/src/cli/index.ts | 218 +- js/src/cli/types.ts | 7 - js/src/cli/util/types.ts | 8 - js/src/durable-eval.test.ts | 1450 ++---- js/src/durable-eval.ts | 4538 ++++------------- js/src/exports.ts | 7 +- js/src/framework.ts | 9 - 11 files changed, 1511 insertions(+), 5129 deletions(-) diff --git a/e2e/scenarios/durable-eval-webhook/scenario.test.ts b/e2e/scenarios/durable-eval-webhook/scenario.test.ts index 3522a8786..2a5efd08f 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.test.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.test.ts @@ -32,13 +32,5 @@ test("durable eval collects webhook sub-batches and logs completed rows", async expect.objectContaining({ run_id: expect.any(String) }), expect.objectContaining({ run_id: expect.any(String) }), ]); - - const shardedSpans = evalSpans.filter( - (event) => event.metadata?.kind === "sharded", - ); - expect(shardedSpans).toHaveLength(4); - expect(new Set(shardedSpans.map((event) => event.experimentId)).size).toBe( - 1, - ); }); }); diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts index 8987970af..37e3a8a4c 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -6,43 +6,14 @@ import { } from "../../helpers/scenario-runtime"; class MemoryStore implements DurableEvalStore { - private readonly values = new Map< - string, - { value: Uint8Array; version: number } - >(); + private readonly values = new Map(); async read(key: string) { - const record = this.values.get(key); - return record - ? { value: record.value.slice(), version: String(record.version) } - : undefined; + return this.values.get(key)?.slice(); } - async write( - key: string, - value: Uint8Array, - condition: Parameters[2], - ) { - const current = this.values.get(key); - if ( - ("ifAbsent" in condition && current) || - ("ifVersion" in condition && - (!current || String(current.version) !== condition.ifVersion)) - ) { - return { - written: false as const, - currentVersion: current ? String(current.version) : undefined, - }; - } - const version = (current?.version ?? 0) + 1; - this.values.set(key, { value: value.slice(), version }); - return { written: true as const, version: String(version) }; - } - - async *list(prefix: string) { - for (const key of [...this.values.keys()].sort()) { - if (key.startsWith(prefix)) yield key; - } + async write(key: string, value: Uint8Array) { + this.values.set(key, value.slice()); } } @@ -55,32 +26,54 @@ async function main() { number, number, { testRunId: string; kind: string }, - Record, - { id: string } + Record >({ - revision: "task-v1", - batchSize: 2, - maxConcurrentBatches: 2, - async submit(items) { - const id = `provider-${jobs.size + 1}`; - jobs.set(id, items); - return { id }; - }, - completion: { - mode: "webhook", - source: "e2e-provider", - externalId: (handle) => handle.id, - }, - async *collect(handle) { - for (const item of jobs.get(handle.id) ?? []) { - yield { id: item.id, output: item.input * 2 }; - } + workflow(workflow) { + const generated = workflow.batch("generate", { + batchSize: 2, + input: (item) => item.input, + async submit(items) { + const id = `generate-${jobs.size + 1}`; + jobs.set(id, items); + return { id }; + }, + completion: { + mode: "webhook", + externalId: (handle) => handle.id, + }, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: item.input * 2, + })); + }, + }); + return workflow.batch("finalize", { + needs: { generated }, + input: (_item, { generated }) => generated, + batchSize: 2, + async submit(items) { + const id = `finalize-${jobs.size + 1}`; + jobs.set(id, items); + return { id }; + }, + completion: { + mode: "webhook", + externalId: (handle) => handle.id, + }, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: item.input, + })); + }, + }); }, }); const definition = DurableEval( scopedName("e2e-durable-eval-webhook-project", testRunId), { - revision: "eval-v1", + store, experimentName: scopedName( "e2e-durable-eval-webhook-experiment", testRunId, @@ -100,61 +93,23 @@ async function main() { }, ); - const waiting = await definition.run({ + const waiting = await definition.start({ runId: `durable-${testRunId}`, - store, }); - if ( - waiting.status !== "paused" || - waiting.reason !== "waiting_for_webhook" || - jobs.size !== 2 - ) { + if (waiting.status !== "waiting" || jobs.size !== 2) { throw new Error("Durable eval did not pause with two webhook batches"); } - for (let index = 1; index <= 2; index++) { - const processed = await definition.processBatchResult( - { - eventId: `event-${index}`, - source: "e2e-provider", - externalId: `provider-${index}`, - outcome: { status: "complete" }, - }, - { store }, - ); - if (processed.status !== "processed") { - throw new Error(`Webhook ${index} was not processed`); - } - if (index === 2 && processed.run.status !== "completed") { - throw new Error("Durable eval did not complete after the final webhook"); - } + let processed = waiting; + const completedJobs = new Set(); + while (completedJobs.size < jobs.size || processed.status !== "completed") { + const externalId = [...jobs.keys()].find((id) => !completedJobs.has(id)); + if (!externalId) throw new Error("Durable eval stopped before completion"); + completedJobs.add(externalId); + processed = await definition.processBatchResult({ + externalId, + }); } - - const sharded = DurableEval( - scopedName("e2e-durable-eval-sharded-project", testRunId), - { - revision: "eval-v1", - data: [1, 2, 3, 4].map((input) => ({ - id: `shard-case-${input}`, - input, - metadata: { testRunId, kind: "sharded" }, - })), - task: (input) => input * 10, - scores: [], - }, - ); - await Promise.all([ - sharded.run({ - runId: `sharded-${testRunId}`, - shard: { index: 0, count: 2 }, - store, - }), - sharded.run({ - runId: `sharded-${testRunId}`, - shard: { index: 1, count: 2 }, - store, - }), - ]); } runMain(main); diff --git a/js/README.md b/js/README.md index e0d202d99..610894ccf 100644 --- a/js/README.md +++ b/js/README.md @@ -46,23 +46,20 @@ main().catch(console.error); ## Durable evaluations -`DurableEval` is an additive evaluation API for work that must survive process -restarts or run through asynchronous provider batch jobs. It checkpoints cases, -task outputs, scores, provider job handles, and attempts outside Braintrust, -while completed results are incrementally merged into a normal Braintrust -experiment. +`DurableEval` 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 +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). Reusing a `runId` -resumes the same input snapshot and never reruns successful work. +Every case needs a stable `id` (or a `caseId` function). ```typescript import { BatchTask, DurableEval, type DurableEvalStore } from "braintrust"; const store: DurableEvalStore = checkpointStore; const supportEval = DurableEval("Support bot", { - // Records which eval definition produced each checkpointed result. Update it - // when eval-level data or orchestration changes; completed cases stay reused. - revision: process.env.GIT_SHA ?? "local", + store, data: [ { id: "password-reset", @@ -71,17 +68,10 @@ const supportEval = DurableEval("Support bot", { }, ], task: BatchTask({ - // Batch stages have their own revision so results retain the version of the - // prompt, model, and task code that produced them. - revision: "generation-v1", - - // Each provider job contains at most 500 eval cases. Up to four jobs may - // be in flight for this task stage at once. + // Each provider job contains at most 500 eval cases. batchSize: 500, - maxConcurrentBatches: 4, - // Upload/submit an asynchronous provider batch and return a - // JSON-serializable handle. + // Submit one sub-batch and return a JSON-serializable provider handle. async submit(items, context) { const batch = await provider.submit({ idempotencyKey: context.batchId, @@ -91,23 +81,18 @@ const supportEval = DurableEval("Support bot", { return { id: batch.id }; }, - // Completion controls how DurableEval learns that the asynchronous job is - // ready. "webhook" pauses until processBatchResult receives an event; - // "poll" calls a supplied poll function. Webhooks may also define a - // pollFallback for delayed or missing events. completion: { + // "webhook" waits for processBatchResult(). Use "poll" with a poll() + // callback when the provider does not send completion events. mode: "webhook", - source: "provider", externalId: (handle) => handle.id, }, - async *collect(handle) { - for await (const item of provider.results(handle.id)) { - yield { - id: item.id, - output: item.output, - }; - } + async collect(handle) { + return (await provider.results(handle.id)).map((item) => ({ + id: item.id, + output: item.output, + })); }, }), scores: [ @@ -117,115 +102,131 @@ const supportEval = DurableEval("Support bot", { ], }); -const result = await supportEval.run({ +const result = await supportEval.start({ runId: "release-2026-07-27", - store, - deadlineMs: 30 * 60 * 1000, }); - -if (result.status === "paused") { - await result.resume({ deadlineMs: 30 * 60 * 1000 }); -} ``` -Use `BatchScorer` for asynchronous batch scoring. Existing per-item tasks and -scorers are also supported and gain checkpointing and retries. Supply a -compare-and-set-capable `DurableEvalStore`; distributed workers must share the -same store. Launch fixed shards through the API or the existing CLI: +`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. -```bash -npx braintrust eval evaluation.eval.ts \ - --run-id release-2026-07-27 \ - --shard 0/8 \ - --deadline 6h +### Polling + +Polling adapters report the provider's current status through `completion`: + +```typescript +completion: { + mode: "poll", + async poll(handle) { + const batch = await provider.getBatch(handle.id); + if (batch.status === "completed") return { status: "complete" }; + if (batch.status === "failed") { + return { status: "failed", error: batch.error }; + } + return { status: "pending" }; + }, +}, ``` -An error thrown from `submit` is considered ambiguous, because the provider may -have created a job before the connection failed. The run pauses instead of -risking duplicate charges. Implement `recover` to look up a prior submission -using `context.batchId`. +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: -The `submit`/`completion`/`collect` split is designed for provider-managed batch -APIs, including OpenAI's Batch API. A polling adapter uses -`completion: { mode: "poll", poll }`. A webhook adapter returns while the job is -pending and resumes through `processBatchResult`. `collect` then fetches and -stores the item results before the eval advances to scoring. Every provider -result must carry the durable item `id`. Batch tasks and batch scorers use the -same contract. +```typescript +const result = await supportEval.poll({ + runId: "release-2026-07-27", +}); -### Webhook processing +if (result.status === "waiting" && result.pending.poll > 0) { + scheduleAnotherPoll(); +} +``` -Pass a provider batch's terminal state to the durable eval definition: +`start()`, `poll()`, and `processBatchResult()` return the current eval status. +A waiting result includes the number of submitted batches using each completion +mode: ```typescript -app.post("/webhooks/provider", async (request, response) => { - const event = request.body; - const batch = await provider.getBatch(event.batchId); +{ + status: "waiting", + runId: "release-2026-07-27", + pending: { poll: 2, webhook: 1 }, +} +``` - const result = await supportEval.processBatchResult( - { - // Uniquely identifies this webhook event so repeated delivery is safe. - eventId: event.id, - // Namespaces provider job IDs when multiple webhook sources share a store. - source: "provider", - // Matches the event to the provider job handle saved during submission. - externalId: batch.id, - // Matches events that arrive before the provider handle is checkpointed. - batchId: batch.metadata?.durableBatchId, - // Supplies the provider handle when recovering an early or ambiguous event. - handle: { id: batch.id }, - outcome: - batch.status === "completed" - ? { status: "complete" } - : { - status: "failed", - error: { status: batch.status }, - // false makes this failure terminal; true requeues its items until - // the batch stage reaches maxAttempts. - retryable: false, - }, - // Optional JSON payloads are checkpointed with the deduplicated event. - payload: { type: event.type }, - }, - { - // Web workers and eval workers must use the same durable store. - store, - }, - ); +Use `status()` to read the same information without polling providers, +collecting results, or advancing the workflow: - response.status(result.status === "pending" ? 202 : 200).end(); +```typescript +const status = await supportEval.status({ + runId: "release-2026-07-27", }); ``` -`eventId` makes delivery idempotent. The SDK indexes a job by its internal -`batchId` before calling `submit`, then adds the provider `externalId` after the -handle is known. An early event is stored in an inbox and attached when that -second index appears. If a worker dies after the provider accepted a batch but -before its handle was checkpointed, include both the internal `batchId` and a -JSON-serializable `handle` in the event to resolve the ambiguous submission. -Without a handle or a successful `recover` callback, the run pauses safely. +Completed statuses have zero pending batches and include the saved experiment +summary. They can be read repeatedly without logging the eval again. -Pure webhook stages return -`{ status: "paused", reason: "waiting_for_webhook" }` instead of holding a -process open. A webhook completion may also define `pollFallback` with -`afterMs`, `poll`, and an optional `intervalMs`. +### Multi-stage workflows -All orchestration state lives in the configured `DurableEvalStore`; this API -does not require a Braintrust backend change. +The direct `BatchTask({ submit, completion, collect })` form remains the +one-batch shorthand. Use `workflow` when a task or scorer requires multiple +provider batch operations: -The definition object exposes `run`, `status`, `retryFailed`, -`resubmitUnknown`, `cancel`, and `processBatchResult`. The corresponding -lifecycle operations are also available through the CLI: +```typescript +task: BatchTask({ + workflow(workflow) { + const draft = workflow.batch("draft", { + input: ({ input }) => ({ prompt: input }), + batchSize: 500, + submit: submitDraftBatch, + completion: draftCompletion, + collect: collectDraftBatch, + }); + + return workflow.batch("revise", { + needs: { draft }, + input: ({ input }, { draft }) => ({ original: input, draft }), + batchSize: 500, + submit: submitRevisionBatch, + completion: revisionCompletion, + collect: collectRevisionBatch, + }); + }, +}), +``` -```bash -npx braintrust eval evaluation.eval.ts --run-id release-2026-07-27 --status -npx braintrust eval evaluation.eval.ts --run-id release-2026-07-27 --retry-failed -npx braintrust eval evaluation.eval.ts --run-id release-2026-07-27 --resubmit-unknown -npx braintrust eval evaluation.eval.ts --run-id release-2026-07-27 --cancel +Every named batch is a persisted workflow node. `needs` can express sequential +operations, parallel branches, and joins. The returned node supplies the final +task output or scorer result. + +### 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 result = await supportEval.processBatchResult({ + // The provider's batch ID. DurableEval saved it from submit()'s handle. + 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 handle. + batchId: batch.metadata?.durableBatchId, + }); + + response.status(result.status === "waiting" ? 202 : 200).end(); +}); ``` -`--resubmit-unknown` is intentionally explicit because the original provider -job may exist and resubmitting it can duplicate work and cost. +The method accepts either `externalId` or `batchId`. The stored batch locator +identifies the task or scorer workflow node, whose `collect()` results are +stored before the eval advances. Webhook idempotency and provider failure +handling remain application responsibilities for now. ## Auto-Instrumentation diff --git a/js/src/cli/functions/load-module.ts b/js/src/cli/functions/load-module.ts index 1ac9a8b8d..593b45ae7 100644 --- a/js/src/cli/functions/load-module.ts +++ b/js/src/cli/functions/load-module.ts @@ -26,7 +26,6 @@ export function loadModule({ prompts: [], parameters: [], evaluators: {}, - durableEvaluators: {}, reporters: {}, }; globalThis._lazy_load = true; diff --git a/js/src/cli/index.ts b/js/src/cli/index.ts index 3dac2d388..ba36d79e3 100755 --- a/js/src/cli/index.ts +++ b/js/src/cli/index.ts @@ -61,11 +61,6 @@ import { } from "./util/debug-logging"; import { pullCommand } from "./util/pull"; import { runDevServer } from "../../dev/server"; -import { - type DurableEvalDefinition, - type DurableEvalOperation, - type DurableEvalRuntimeOptions, -} from "../durable-eval"; // This requires require // https://stackoverflow.com/questions/50822310/how-to-import-package-json-in-typescript @@ -218,7 +213,6 @@ function buildWatchPluginForEvaluator( ): esbuild.Plugin { const evaluators: EvaluatorState = { evaluators: [], - durableEvaluators: [], reporters: {}, }; const plugin = { @@ -249,9 +243,6 @@ function buildWatchPluginForEvaluator( evaluators.evaluators = evaluators.evaluators.filter( (e) => e.sourceFile !== inFile, ); - evaluators.durableEvaluators = evaluators.durableEvaluators.filter( - (e) => e.sourceFile !== inFile, - ); // Update the evaluators and reporters for (const evaluator of Object.values(evalResult.evaluators)) { @@ -267,21 +258,6 @@ function buildWatchPluginForEvaluator( reporter: evaluator.reporter, }); } - for (const registration of Object.values( - evalResult.durableEvaluators ?? {}, - )) { - evaluators.durableEvaluators.push({ - sourceFile: inFile, - // Runtime registrations are intentionally generic-erased. - definition: registration.definition as DurableEvalDefinition< - any, - any, - any, - any, - any - >, - }); - } for (const [reporterName, reporter] of Object.entries( evalResult.reporters, )) { @@ -328,20 +304,6 @@ function buildWatchPluginForEvaluator( addReport(evalReports, resolvedReporter, report); } - for (const registration of evaluators.durableEvaluators.filter( - (candidate) => candidate.sourceFile === inFile, - )) { - const result = await runDurableEvaluator( - registration.definition, - opts, - ); - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error( - result.status === "completed" - ? `Durable eval ${result.runId} completed` - : `Durable eval ${result.runId} paused: ${result.reason}`, - ); - } for (const [reporterName, { reporter, results }] of Object.entries( evalReports, @@ -450,11 +412,6 @@ interface EvaluatorOpts { jsonl: boolean; filters: Filter[]; progressReporter: ProgressReporter; - durableRunId?: string; - durableShard?: { index: number; count: number }; - durableDeadlineMs?: number; - checkpointDir?: string; - durableOperation?: Exclude; } export function handleBuildFailure({ @@ -509,21 +466,6 @@ function updateEvaluators( reporter: evaluator.reporter, }); } - for (const registration of Object.values( - result.evaluator.durableEvaluators ?? {}, - )) { - evaluators.durableEvaluators.push({ - sourceFile: result.sourceFile, - // Runtime registrations are intentionally generic-erased. - definition: registration.definition as DurableEvalDefinition< - any, - any, - any, - any, - any - >, - }); - } for (const [reporterName, reporter] of Object.entries( result.evaluator.reporters, @@ -586,43 +528,12 @@ export async function buildEvaluators( const evaluators: EvaluatorState = { evaluators: [], - durableEvaluators: [], reporters: {}, }; updateEvaluators(evaluators, buildResults, opts); return { evaluators, buildResults }; } -async function runDurableEvaluator( - definition: DurableEvalDefinition, - opts: EvaluatorOpts, -) { - const options: DurableEvalRuntimeOptions = { - runId: opts.durableRunId, - shard: opts.durableShard, - deadlineMs: opts.durableDeadlineMs, - checkpointDir: opts.checkpointDir, - noSendLogs: opts.noSendLogs, - }; - if (!opts.durableOperation) { - return await definition.run(options); - } - if (!opts.durableRunId) { - throw new Error(`--${opts.durableOperation} requires --run-id`); - } - const existingOptions = { ...options, runId: opts.durableRunId }; - switch (opts.durableOperation) { - case "status": - return await definition.status(existingOptions); - case "retry-failed": - return await definition.retryFailed(existingOptions); - case "resubmit-unknown": - return await definition.resubmitUnknown(existingOptions); - case "cancel": - return await definition.cancel(existingOptions); - } -} - async function runOnce( handles: Record, opts: EvaluatorOpts, @@ -637,21 +548,12 @@ async function runOnce( : null; const { evaluators, buildResults } = await buildEvaluators(handles, opts); - if (opts.durableOperation && evaluators.evaluators.length > 0) { - throw new Error( - "Durable lifecycle flags cannot be used with ordinary Eval definitions", - ); - } if (opts.list) { for (const evaluator of evaluators.evaluators) { // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. console.log(evaluator.evaluator.evalName); } - for (const evaluator of evaluators.durableEvaluators) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.log(evaluator.definition.evalName); - } return true; } @@ -690,20 +592,15 @@ async function runOnce( } } }); - const durableResultPromises = evaluators.durableEvaluators.map( - async (registration) => - await runDurableEvaluator(registration.definition, opts), - ); // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. console.error( styleText( "dim", - `Processing ${styleText("bold", String(resultPromises.length + durableResultPromises.length))} evaluator${resultPromises.length + durableResultPromises.length === 1 ? "" : "s"}...`, + `Processing ${styleText("bold", String(resultPromises.length))} evaluator${resultPromises.length === 1 ? "" : "s"}...`, ), ); const allEvalsResults = await Promise.all(resultPromises); - const allDurableResults = await Promise.all(durableResultPromises); opts.progressReporter.stop(); // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. console.error(""); @@ -758,33 +655,6 @@ async function runOnce( allSuccess = allSuccess && success; } - for (const result of allDurableResults) { - if (opts.jsonl) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.log(JSON.stringify(result)); - } else if (result.status === "completed") { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error( - `Durable eval ${result.runId} completed (${result.progress.taskSucceeded}/${result.progress.total} tasks)`, - ); - } else { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error(`Durable eval ${result.runId} paused: ${result.reason}`); - } - if ( - result.status === "completed" && - (result.failures.tasks > 0 || result.failures.scorers > 0) - ) { - allSuccess = false; - } - if ( - result.status === "paused" && - ["unknown_submission", "provider_unreachable"].includes(result.reason) - ) { - allSuccess = false; - } - } - return allSuccess; } @@ -1077,21 +947,6 @@ async function run(args: RunArgs) { : new BarProgressReporter(), filters: args.filter ? parseFilters(args.filter) : [], list: !!args.list, - durableRunId: args.run_id, - durableShard: args.shard ? parseDurableShard(args.shard) : undefined, - durableDeadlineMs: args.deadline - ? parseDurationMs(args.deadline) - : undefined, - checkpointDir: args.checkpoint_dir, - durableOperation: args.status - ? "status" - : args.retry_failed - ? "retry-failed" - : args.resubmit_unknown - ? "resubmit-unknown" - : args.cancel - ? "cancel" - : undefined, }; if (args.list && args.watch) { @@ -1099,23 +954,6 @@ async function run(args: RunArgs) { console.error(error("Cannot specify both --list and --watch.")); process.exit(1); } - if (args.shard && !args.run_id) { - throw new Error("--shard requires --run-id for durable evals"); - } - const lifecycleFlags = [ - args.status, - args.retry_failed, - args.resubmit_unknown, - args.cancel, - ].filter(Boolean); - if (lifecycleFlags.length > 1) { - throw new Error( - "Specify only one of --status, --retry-failed, --resubmit-unknown, or --cancel", - ); - } - if (lifecycleFlags.length > 0 && !args.run_id) { - throw new Error("Durable lifecycle flags require --run-id"); - } const plugins = evaluatorOpts.watch ? [ @@ -1179,32 +1017,6 @@ async function run(args: RunArgs) { } } -function parseDurableShard(value: string) { - const match = /^(\d+)\/(\d+)$/.exec(value); - if (!match) { - throw new Error(`Invalid --shard value '${value}'; expected INDEX/COUNT`); - } - const shard = { index: Number(match[1]), count: Number(match[2]) }; - if (shard.count < 1 || shard.index < 0 || shard.index >= shard.count) { - throw new Error(`Invalid --shard value '${value}'`); - } - return shard; -} - -function parseDurationMs(value: string) { - const match = /^(\d+(?:\.\d+)?)(ms|s|m|h)?$/.exec(value); - if (!match) { - throw new Error( - `Invalid --deadline value '${value}'; expected e.g. 500ms, 30m, or 6h`, - ); - } - const multipliers = { ms: 1, s: 1_000, m: 60_000, h: 3_600_000 }; - return ( - Number(match[1]) * - multipliers[(match[2] ?? "ms") as keyof typeof multipliers] - ); -} - function addAuthArgs(parser: ArgumentParser) { parser.add_argument("--api-key", { help: "Specify a braintrust api key. If the parameter is not specified, BRAINTRUST_API_KEY or the nearest .env.braintrust file will be used.", @@ -1289,34 +1101,6 @@ async function main() { action: "store_true", help: "Do not show progress bars when processing evaluators.", }); - parser_run.add_argument("--run-id", { - help: "Create or resume a DurableEval run with this stable identifier.", - }); - parser_run.add_argument("--shard", { - help: "Run one DurableEval shard in INDEX/COUNT form, for example 0/8.", - }); - parser_run.add_argument("--deadline", { - help: "Pause DurableEval work after a duration such as 30m or 6h.", - }); - parser_run.add_argument("--checkpoint-dir", { - help: "Override the default .braintrust/evals checkpoint directory.", - }); - parser_run.add_argument("--status", { - action: "store_true", - help: "Inspect a DurableEval run without claiming new work.", - }); - parser_run.add_argument("--retry-failed", { - action: "store_true", - help: "Retry terminal DurableEval failures with the current definition.", - }); - parser_run.add_argument("--resubmit-unknown", { - action: "store_true", - help: "Explicitly resubmit ambiguous provider batches, accepting possible duplicate cost.", - }); - parser_run.add_argument("--cancel", { - action: "store_true", - help: "Cancel active provider batches and terminally cancel unfinished DurableEval work.", - }); parser_run.add_argument("--bundle", { action: "store_true", help: "Experimental (do not use unless you know what you're doing)", diff --git a/js/src/cli/types.ts b/js/src/cli/types.ts index c614c40d0..15184a247 100644 --- a/js/src/cli/types.ts +++ b/js/src/cli/types.ts @@ -2,7 +2,6 @@ import type * as esbuild from "esbuild"; import type { BaseMetadata } from "../logger"; import type { EvaluatorDef, EvaluatorFile } from "../framework"; import type { ReporterDef } from "../reporters/types"; -import type { DurableEvalDefinition } from "../durable-eval"; export interface BuildSuccess { type: "success"; @@ -35,12 +34,6 @@ export interface EvaluatorState { evaluator: EvaluatorDef; reporter: string | ReporterDef | undefined; }[]; - durableEvaluators: { - sourceFile: string; - // Runtime registration has already passed the public generic boundary. - // eslint-disable-next-line @typescript-eslint/no-explicit-any - definition: DurableEvalDefinition; - }[]; reporters: { [reporter: string]: ReporterDef; }; diff --git a/js/src/cli/util/types.ts b/js/src/cli/util/types.ts index 25532a022..2909a607e 100644 --- a/js/src/cli/util/types.ts +++ b/js/src/cli/util/types.ts @@ -33,14 +33,6 @@ export interface RunArgs extends CommonArgs, AuthArgs, CompileArgs { dev_host: string; dev_port: number; dev_org_name?: string; - run_id?: string; - shard?: string; - deadline?: string; - checkpoint_dir?: string; - status?: boolean; - retry_failed?: boolean; - resubmit_unknown?: boolean; - cancel?: boolean; } export interface BundleArgs extends CommonArgs, AuthArgs, CompileArgs { diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts index 0cfd1a043..3c4f70f1d 100644 --- a/js/src/durable-eval.test.ts +++ b/js/src/durable-eval.test.ts @@ -1,161 +1,45 @@ import { describe, expect, test, vi } from "vitest"; -import { mkdtemp, rm, utimes, writeFile } from "node:fs/promises"; -import { tmpdir } from "node:os"; -import { join } from "node:path"; +import { configureNode } from "./node/config"; import { BatchScorer, BatchTask, DurableEval, - FileDurableEvalStore, - MemoryDurableEvalStore, + type DurableBatchScorerItem, type DurableBatchTaskItem, type DurableEvalStore, - type DurableEvalWriteCondition, } from "./durable-eval"; -import type { EvaluatorFile } from "./framework"; -import { configureNode } from "./node/config"; configureNode(); -describe("DurableEval", () => { - test("filesystem store enforces compare-and-set writes", async () => { - const directory = await mkdtemp(join(tmpdir(), "durable-eval-store-")); - try { - const store = new FileDurableEvalStore(directory); - const first = await store.write( - "runs/one", - new TextEncoder().encode("one"), - { ifAbsent: true }, - ); - expect(first.written).toBe(true); - if (!first.written) throw new Error("initial write failed"); - - await expect( - store.write("runs/one", new TextEncoder().encode("two"), { - ifAbsent: true, - }), - ).resolves.toMatchObject({ written: false }); - await expect( - store.write("runs/one", new TextEncoder().encode("two"), { - ifVersion: "wrong", - }), - ).resolves.toMatchObject({ written: false }); +class MemoryStore implements DurableEvalStore { + private readonly values = new Map(); - const updated = await store.write( - "runs/one", - new TextEncoder().encode("two"), - { ifVersion: first.version }, - ); - expect(updated.written).toBe(true); - expect( - new TextDecoder().decode((await store.read("runs/one"))?.value), - ).toBe("two"); - const keys: string[] = []; - for await (const key of store.list("runs/")) keys.push(key); - expect(keys).toEqual(["runs/one"]); - } finally { - await rm(directory, { recursive: true, force: true }); - } - }); + async read(key: string) { + return this.values.get(key)?.slice(); + } - test("filesystem store recovers stale locks and rejects Windows traversal keys", async () => { - const directory = await mkdtemp(join(tmpdir(), "durable-eval-locks-")); - try { - const store = new FileDurableEvalStore(directory); - const value = new TextEncoder().encode("one"); - const first = await store.write("runs/one", value, { ifAbsent: true }); - expect(first.written).toBe(true); + async write(key: string, value: Uint8Array) { + this.values.set(key, value.slice()); + } +} - const lockPath = join(directory, "runs", "one.lock"); - await writeFile(lockPath, "orphaned"); - const staleTime = new Date(Date.now() - 60_000); - await utimes(lockPath, staleTime, staleTime); - const current = await store.read("runs/one"); - await expect( - store.write("runs/one", new TextEncoder().encode("two"), { - ifVersion: current!.version, - }), - ).resolves.toMatchObject({ written: true }); - - await expect( - store.write("..\\outside", value, { ifAbsent: true }), - ).rejects.toThrow("Invalid durable eval store key"); - } finally { - await rm(directory, { recursive: true, force: true }); - } - }); - - test("runs local tasks and scorers and resumes completed work", async () => { - const store = new MemoryDurableEvalStore(); +describe("DurableEval", () => { + test("runs ordinary tasks and scorers", async () => { const task = vi.fn((input: number) => input * 2); - let scorerCalls = 0; - function exact({ - output, - expected, - }: { - output: number; - expected?: number; - }) { - scorerCalls++; - return output === expected ? 1 : 0; - } - const definition = { - revision: "v1", + const result = await DurableEval("local", { + store: new MemoryStore(), data: [ { id: "one", input: 1, expected: 2 }, { id: "two", input: 2, expected: 4 }, ], task, - scores: [exact], - }; - - const durable = DurableEval("durable-local", definition); - const first = await durable.run({ - runId: "run-1", - store, - noSendLogs: true, - }); - expect(first.status).toBe("completed"); - if (first.status !== "completed") throw new Error("run did not complete"); - expect(first.summary.scores.exact?.score).toBe(1); - expect(first.progress).toMatchObject({ - total: 2, - taskSucceeded: 2, - scoreSucceeded: 2, - }); - - const second = await durable.run({ - runId: "run-1", - store, - noSendLogs: true, - }); - expect(second.status).toBe("completed"); - expect(task).toHaveBeenCalledTimes(2); - expect(scorerCalls).toBe(2); - }); - - test("persists metadata written through hooks.meta", async () => { - const result = await DurableEval("metadata-hooks", { - revision: "v1", - data: [ - { - id: "one", - input: "hello", - metadata: { source: "initial" }, - }, - ], - task: (_input, hooks) => { - hooks.meta({ source: "deprecated-hook" }); - return hooks.metadata.source; - }, scores: [ - function exact({ output }) { - return output === "deprecated-hook" ? 1 : 0; + function exact({ output, expected }) { + return output === expected ? 1 : 0; }, ], - }).run({ - runId: "metadata-run", - store: new MemoryDurableEvalStore(), + }).start({ + runId: "local-run", noSendLogs: true, }); @@ -163,203 +47,135 @@ describe("DurableEval", () => { status: "completed", summary: { scores: { exact: { score: 1 } } }, }); + expect(task).toHaveBeenCalledTimes(2); }); - test("requires stable case identifiers", async () => { - await expect( - DurableEval("missing-ids", { - revision: "v1", - data: [{ input: "hello", expected: "hello" }], - task: (input) => input, - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }).run({ - runId: "missing", - store: new MemoryDurableEvalStore(), - noSendLogs: true, - }), - ).rejects.toThrow("must have a non-empty id"); - }); - - test("submits and collects batch tasks and batch scorers", async () => { - const store = new MemoryDurableEvalStore(); + 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 = BatchTask< number, number, number, void, Record, - { jobId: string } + { id: string } >({ - revision: "task-v1", batchSize: 2, - async submit(items, context) { - taskJobs.set(context.batchId, items); - return { jobId: context.batchId }; + async submit(items) { + const id = `task-${taskJobs.size + 1}`; + taskJobs.set(id, items); + return { id }; }, completion: { mode: "poll", - async poll() { - return { status: "complete" }; - }, + poll: taskPoll, }, - async *collect(handle) { - for (const item of taskJobs.get(handle.jobId) ?? []) { - yield { id: item.id, output: item.input * 2 }; - } + async collect(handle) { + return (taskJobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: item.input * 2, + })); }, }); - const scoreJobs = new Map< - string, - Array<{ id: string; output: number; expected: number }> - >(); - const scorer = BatchScorer( - { - name: "exact", - revision: "score-v1", - batchSize: 2, - async submit(items, context) { - scoreJobs.set( - context.batchId, - items.map((item) => ({ - id: item.id, - output: item.output, - expected: item.expected, - })), - ); - return { jobId: context.batchId }; - }, - completion: { - mode: "poll", - async poll() { - return { status: "complete" }; - }, - }, - async *collect(handle) { - for (const item of scoreJobs.get(handle.jobId) ?? []) { - yield { - id: item.id, - score: item.output === item.expected ? 1 : 0, - }; - } - }, + const scorer = BatchScorer({ + 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(handle) { + return (scoreJobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + score: item.output === item.expected ? 1 : 0, + })); + }, + }); - const result = await DurableEval("durable-batches", { - revision: "eval-v1", - data: [ - { id: "one", input: 1, expected: 2 }, - { id: "two", input: 2, expected: 4 }, - { id: "three", input: 3, expected: 6 }, - ], + const store = new MemoryStore(); + const durable = DurableEval("polling-batches", { + store, + data: [1, 2, 3].map((input) => ({ + id: `case-${input}`, + input, + expected: input * 2, + })), task, scores: [scorer], - }).run({ - runId: "batch-run", - store, - noSendLogs: true, }); + const options = { + runId: "polling-run", + noSendLogs: true, + }; - expect(result.status).toBe("completed"); - if (result.status !== "completed") throw new Error("run did not complete"); - expect(result.summary.scores.exact?.score).toBe(1); - expect(taskJobs).toHaveLength(2); - expect(scoreJobs).toHaveLength(2); - }); - - test("retries raced case writes before completing a batch job", async () => { - class FailOneResultWriteStore implements DurableEvalStore { - readonly inner = new MemoryDurableEvalStore(); - failed = false; - - read(key: string) { - return this.inner.read(key); - } - - async write( - key: string, - value: Uint8Array, - condition: DurableEvalWriteCondition, - ) { - const serialized = new TextDecoder().decode(value); - if ( - !this.failed && - key.includes("/cases/") && - serialized.includes('"status":"succeeded"') - ) { - this.failed = true; - return { - written: false as const, - currentVersion: (await this.inner.read(key))?.version, - }; - } - return this.inner.write(key, value, condition); - } - - list(prefix: string) { - return this.inner.list(prefix); - } - } + await expect(durable.start(options)).resolves.toEqual({ + status: "waiting", + runId: "polling-run", + pending: { poll: 2, webhook: 0 }, + }); + expect(taskJobs.size).toBe(2); + expect(scoreJobs.size).toBe(0); + await expect(durable.status(options)).resolves.toEqual({ + status: "waiting", + runId: "polling-run", + pending: { poll: 2, webhook: 0 }, + }); + expect(taskJobs.size).toBe(2); + expect(taskPoll).not.toHaveBeenCalled(); - const store = new FailOneResultWriteStore(); - const result = await DurableEval("cas-batch-results", { - revision: "eval-v1", - data: [{ id: "one", input: 1, expected: 2 }], - task: BatchTask< - number, - number, - number, - void, - Record, - { id: string } - >({ - revision: "task-v1", - async submit() { - return { id: "provider-job" }; - }, - completion: { - mode: "poll", - async poll() { - return { status: "complete" }; - }, - }, - async *collect() { - yield { id: "one:trial:0", output: 2 }; - }, - }), - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }).run({ - runId: "cas-run", - store, - noSendLogs: true, + await expect(durable.poll(options)).resolves.toEqual({ + status: "waiting", + runId: "polling-run", + 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(store.failed).toBe(true); expect(result).toMatchObject({ status: "completed", - progress: { taskSucceeded: 1, scoreSucceeded: 1 }, + 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("pauses for webhooks and processes bounded sub-batches", async () => { - const store = new MemoryDurableEvalStore(); - const providerJobs = new Map< + test("processes task and scorer webhook batches through one method", async () => { + const store = new MemoryStore(); + const taskJobs = new Map< string, DurableBatchTaskItem>[] >(); - const submittedSizes: number[] = []; + const scoreJobs = new Map< + string, + DurableBatchScorerItem[] + >(); const task = BatchTask< number, number, @@ -368,898 +184,258 @@ describe("DurableEval", () => { Record, { id: string } >({ - revision: "task-v1", batchSize: 2, - maxConcurrentBatches: 2, async submit(items) { - const id = `provider-${providerJobs.size + 1}`; - providerJobs.set(id, items); - submittedSizes.push(items.length); + const id = `task-provider-${taskJobs.size + 1}`; + taskJobs.set(id, items); return { id }; }, completion: { mode: "webhook", - source: "openai", externalId: (handle) => handle.id, }, - async *collect(handle) { - for (const item of providerJobs.get(handle.id) ?? []) { - yield { id: item.id, output: item.input * 2 }; - } + async collect(handle) { + return (taskJobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: item.input * 2, + })); }, }); - const durable = DurableEval("webhook-batches", { - revision: "eval-v1", - data: [ - { id: "one", input: 1, expected: 2 }, - { id: "two", input: 2, expected: 4 }, - { id: "three", input: 3, expected: 6 }, - ], - task, - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }); - - const waiting = await durable.run({ - runId: "webhook-run", - store, - noSendLogs: true, - }); - expect(waiting).toMatchObject({ - status: "paused", - reason: "waiting_for_webhook", - }); - expect(submittedSizes).toEqual([2, 1]); - - const first = await durable.processBatchResult( - { - eventId: "event-1", - source: "openai", - externalId: "provider-1", - outcome: { status: "complete" }, - payload: { providerStatus: "completed" }, - }, - { store, noSendLogs: true }, - ); - expect(first).toMatchObject({ - status: "processed", - batchId: expect.any(String), - run: { status: "paused", reason: "waiting_for_webhook" }, - }); - - const second = await durable.processBatchResult( - { - eventId: "event-2", - source: "openai", - externalId: "provider-2", - outcome: { status: "complete" }, - }, - { store, noSendLogs: true }, - ); - expect(second).toMatchObject({ - status: "processed", - run: { - status: "completed", - progress: { taskSucceeded: 3, scoreSucceeded: 3 }, - }, - }); - - await expect( - durable.processBatchResult( - { - eventId: "event-2", - source: "openai", - externalId: "provider-2", - outcome: { status: "complete" }, - }, - { store, noSendLogs: true }, - ), - ).resolves.toMatchObject({ status: "duplicate" }); - - const storedEvents: string[] = []; - for await (const key of store.list("durable-eval/v1/webhooks/events/")) { - const record = await store.read(key); - if (record) storedEvents.push(new TextDecoder().decode(record.value)); - } - expect(storedEvents).toHaveLength(2); - expect(storedEvents.join("\n")).toContain("providerStatus"); - }); - - test("stores an early webhook until the provider handle is indexed", async () => { - const store = new MemoryDurableEvalStore(); - let processEarly: - | (() => Promise<{ status: string; reason?: string }>) - | undefined; - let earlyResult: { status: string; reason?: string } | undefined; - const task = BatchTask< - string, - string, - string, - void, - Record, - { id: string } - >({ - revision: "task-v1", - async submit() { - earlyResult = await processEarly!(); - return { id: "provider-early" }; - }, - completion: { - mode: "webhook", - source: "openai", - externalId: (handle) => handle.id, - }, - async *collect() { - yield { id: "one:trial:0", output: "done" }; - }, - }); - const durable = DurableEval("early-webhook", { - revision: "eval-v1", - data: [{ id: "one", input: "input", expected: "done" }], - task, - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }); - processEarly = () => - durable.processBatchResult( - { - eventId: "early-event", - source: "openai", - externalId: "provider-early", - outcome: { status: "complete" }, - }, - { store, noSendLogs: true }, - ); - - const result = await durable.run({ - runId: "early-run", - store, - noSendLogs: true, - }); - expect(earlyResult).toMatchObject({ - status: "pending", - reason: "unmatched", - }); - expect(result.status).toBe("completed"); - }); - - test("uses a webhook handle to resolve an ambiguous submission", async () => { - const store = new MemoryDurableEvalStore(); - let batchId: string | undefined; - const task = BatchTask< - string, - string, - string, - void, - Record, - { id: string } - >({ - revision: "task-v1", - async submit(_items, context) { - batchId = context.batchId; - throw new Error("connection closed after provider accepted the batch"); + const scorer = BatchScorer({ + name: "exact", + batchSize: 2, + async submit(items) { + const id = `score-provider-${scoreJobs.size + 1}`; + scoreJobs.set(id, items); + return { id }; }, completion: { mode: "webhook", - source: "openai", externalId: (handle) => handle.id, }, - async *collect() { - yield { id: "one:trial:0", output: "done" }; + async collect(handle) { + return (scoreJobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + score: item.output === item.expected ? 1 : 0, + })); }, }); - const durable = DurableEval("lost-submit-webhook", { - revision: "eval-v1", - data: [{ id: "one", input: "input", expected: "done" }], - task, - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }); - const paused = await durable.run({ - runId: "lost-submit-run", + const durable = DurableEval("webhook-batches", { store, - noSendLogs: true, - }); - expect(paused).toMatchObject({ - status: "paused", - reason: "unknown_submission", - }); - - const processed = await durable.processBatchResult( - { - eventId: "lost-submit-event", - source: "openai", - batchId, - externalId: "provider-lost", - handle: { id: "provider-lost" }, - outcome: { status: "complete" }, - }, - { store, noSendLogs: true }, - ); - expect(processed).toMatchObject({ - status: "processed", - run: { status: "completed" }, - }); - }); - - test("retries result collection when the provider redelivers a webhook", async () => { - const store = new MemoryDurableEvalStore(); - let collectAttempts = 0; - const task = BatchTask< - string, - string, - string, - void, - Record, - { id: string } - >({ - revision: "task-v1", - async submit() { - return { id: "provider-retry" }; - }, - completion: { - mode: "webhook", - source: "openai", - externalId: (handle) => handle.id, - }, - async *collect() { - collectAttempts++; - if (collectAttempts === 1) { - throw new Error("provider result file is temporarily unavailable"); - } - yield { id: "one:trial:0", output: "done" }; - }, - }); - const durable = DurableEval("collect-retry", { - revision: "eval-v1", - data: [{ id: "one", input: "input", expected: "done" }], + data: [1, 2, 3].map((input) => ({ + id: `case-${input}`, + input, + expected: input * 2, + })), task, - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }); - await durable.run({ - runId: "collect-retry-run", - store, - noSendLogs: true, - }); - const event = { - eventId: "retry-event", - source: "openai", - externalId: "provider-retry", - outcome: { status: "complete" as const }, - }; - - const unavailable = await durable.processBatchResult(event, { - store, - noSendLogs: true, - }); - expect(unavailable).toMatchObject({ - status: "processed", - run: { status: "paused", reason: "provider_unreachable" }, - }); - - const completed = await durable.processBatchResult(event, { - store, - noSendLogs: true, - }); - expect(completed).toMatchObject({ - status: "processed", - run: { status: "completed" }, - }); - expect(collectAttempts).toBe(2); - }); - - test("does not recover or duplicate a batch while another worker submits it", async () => { - const store = new MemoryDurableEvalStore(); - let submitCount = 0; - let releaseSubmit: () => void = () => undefined; - const submitGate = new Promise((resolve) => { - releaseSubmit = resolve; - }); - const task = BatchTask< - string, - string, - string, - void, - Record, - { jobId: string } - >({ - revision: "task-v1", - async submit(_items, context) { - submitCount++; - await submitGate; - return { jobId: context.batchId }; - }, - completion: { - mode: "poll", - async poll() { - return { status: "complete" }; - }, - }, - async *collect() { - yield { id: "one:trial:0", output: "done" }; - }, + scores: [scorer], }); - const definition = { - revision: "eval-v1", - data: [{ id: "one", input: "input", expected: "done" }], - task, - scores: [ - function exact({ - output, - expected, - }: { - output: string; - expected?: string; - }) { - return output === expected ? 1 : 0; - }, - ], - }; - const durable = DurableEval("distributed-submit", definition); - const first = durable.run({ - runId: "distributed-run", - store, - noSendLogs: true, - workerId: "worker-one", + await expect( + durable.start({ runId: "webhook-run", noSendLogs: true }), + ).resolves.toEqual({ + status: "waiting", + runId: "webhook-run", + pending: { poll: 0, webhook: 2 }, }); - await vi.waitFor(() => expect(submitCount).toBe(1)); + expect(taskJobs.size).toBe(2); - const second = await durable.run({ - runId: "distributed-run", - store, - noSendLogs: true, - workerId: "worker-two", - deadlineMs: 10, + await expect( + durable.start({ runId: "webhook-run", noSendLogs: true }), + ).resolves.toEqual({ + status: "waiting", + runId: "webhook-run", + pending: { poll: 0, webhook: 2 }, }); - expect(second.status).toBe("paused"); - expect(submitCount).toBe(1); - - releaseSubmit(); - await expect(first).resolves.toMatchObject({ status: "completed" }); - expect(submitCount).toBe(1); - }); + expect(taskJobs.size).toBe(2); - test("renews provider leases while polling exceeds the lease duration", async () => { - vi.useFakeTimers(); - try { - const store = new MemoryDurableEvalStore(); - let pollCalls = 0; - let releasePoll: () => void = () => undefined; - const pollGate = new Promise((resolve) => { - releasePoll = resolve; - }); - const durable = DurableEval("long-provider-poll", { - revision: "eval-v1", - data: [{ id: "one", input: "input" }], - task: BatchTask< - string, - string, - void, - void, - Record, - { id: string } - >({ - revision: "task-v1", - async submit() { - return { id: "long-job" }; - }, - completion: { - mode: "poll", - async poll() { - pollCalls++; - await pollGate; - return { status: "complete" }; - }, - }, - async *collect() { - yield { id: "one:trial:0", output: "done" }; - }, - }), - scores: [], - }); - - const first = durable.run({ - runId: "long-poll-run", - store, - noSendLogs: true, - workerId: "worker-one", - }); - while (pollCalls === 0) await vi.advanceTimersByTimeAsync(0); - await vi.advanceTimersByTimeAsync(61_000); - - const second = durable.run({ - runId: "long-poll-run", - store, - noSendLogs: true, - workerId: "worker-two", - deadlineMs: 5, - }); - await vi.advanceTimersByTimeAsync(10); - await expect(second).resolves.toMatchObject({ - status: "paused", - reason: "deadline", + const taskIds = [...taskJobs.keys()]; + for (const [index, externalId] of taskIds.entries()) { + const result = await durable.processBatchResult( + { externalId }, + { noSendLogs: true }, + ); + expect(result).toMatchObject({ + status: "waiting", + pending: { poll: 0, webhook: index === 0 ? 1 : 2 }, }); - expect(pollCalls).toBe(1); - - releasePoll(); - await vi.advanceTimersByTimeAsync(0); - await expect(first).resolves.toMatchObject({ status: "completed" }); - } finally { - vi.useRealTimers(); } - }); - - test("cancels active provider batches and terminally checkpoints the work", async () => { - const store = new MemoryDurableEvalStore(); - const abort = new AbortController(); - const cancel = vi.fn(async () => undefined); - const task = BatchTask< - string, - string, - string, - void, - Record, - { jobId: string } - >({ - revision: "task-v1", - async submit(_items, context) { - return { jobId: context.batchId }; - }, - completion: { - mode: "poll", - async poll() { - abort.abort(); - return { status: "pending" }; - }, - }, - async *collect() { - // A cancelled batch is never collected. - }, - cancel, - }); - const definition = { - revision: "eval-v1", - data: [{ id: "one", input: "input", expected: "done" }], - task, - scores: [ - function exact({ - output, - expected, - }: { - output: string; - expected?: string; - }) { - return output === expected ? 1 : 0; - }, - ], - }; - - const durable = DurableEval("cancel-batch", definition); - const paused = await durable.run({ - runId: "cancel-run", - store, - noSendLogs: true, - signal: abort.signal, - }); - expect(paused).toMatchObject({ status: "paused", reason: "aborted" }); - - const cancelled = await durable.cancel({ - runId: "cancel-run", - store, - noSendLogs: true, - }); - expect(cancel).toHaveBeenCalledOnce(); - expect(cancelled).toMatchObject({ + expect(scoreJobs.size).toBe(2); + + const scoreIds = [...scoreJobs.keys()]; + let result; + for (const externalId of scoreIds) { + result = await durable.processBatchResult( + { externalId }, + { noSendLogs: true }, + ); + } + expect(result).toMatchObject({ status: "completed", - failures: { tasks: 1, scorers: 0 }, + pending: { poll: 0, webhook: 0 }, + summary: { scores: { exact: { score: 1 } } }, }); }); - test("returns a resumable deadline result without resubmitting a batch", async () => { - const store = new MemoryDurableEvalStore(); - let ready = false; - let submits = 0; - const task = BatchTask< - string, - string, - string, - void, - Record, - { jobId: string } - >({ - revision: "task-v1", - async submit(_items, context) { - submits++; - return { jobId: context.batchId }; - }, - completion: { - mode: "poll", - intervalMs: 1, - async poll() { - return ready - ? { status: "complete" as const } - : { status: "pending" as const, retryAfterMs: 1 }; - }, - }, - async *collect() { - yield { id: "one:trial:0", output: "done" }; + test("runs multi-stage task and scorer workflows", async () => { + const jobs = new Map>(); + const submit = async ( + prefix: string, + items: Array<{ id: string; input: unknown }>, + ) => { + const id = `${prefix}-${jobs.size + 1}`; + jobs.set(id, items); + return { id }; + }; + const completion = { + mode: "poll" as const, + async poll() { + return { status: "complete" as const }; }, - }); - - const definition = { - revision: "eval-v1", - data: [{ id: "one", input: "input", expected: "done" }], - task, - scores: [ - function exact({ - output, - expected, - }: { - output: string; - expected?: string; - }) { - return output === expected ? 1 : 0; - }, - ], }; - const durable = DurableEval("deadline", definition); - const first = await durable.run({ - runId: "deadline-run", - store, - noSendLogs: true, - deadlineMs: 5, - }); - expect(first.status).toBe("paused"); - if (first.status !== "paused") throw new Error("run did not pause"); - expect(first.reason).toBe("deadline"); - ready = true; - const second = await first.resume({ deadlineMs: 1_000 }); - expect(second.status).toBe("completed"); - expect(submits).toBe(1); - }); - - test("deadline aborts an active provider submission", async () => { - const neverSubmitted = new Promise<{ id: string }>(() => undefined); - const durable = DurableEval("submit-deadline", { - revision: "eval-v1", - data: [{ id: "one", input: "input" }], - task: BatchTask< - string, - string, - void, - void, - Record, - { id: string } - >({ - revision: "task-v1", - async submit() { - return await neverSubmitted; + const task = BatchTask>( + { + workflow(w) { + const doubled = w.batch("double", { + batchSize: 2, + input: (item) => item.input, + submit: (items) => submit("double", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) * 2, + })); + }, + }); + const incremented = w.batch("increment", { + needs: { doubled }, + input: (_item, outputs) => outputs.doubled, + submit: (items) => submit("increment", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) + 1, + })); + }, + }); + const decremented = w.batch("decrement", { + needs: { doubled }, + input: (_item, outputs) => outputs.doubled, + submit: (items) => submit("decrement", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) - 1, + })); + }, + }); + return w.batch("combine", { + needs: { incremented, decremented }, + input: (_item, outputs) => outputs, + submit: (items) => submit("combine", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => { + const value = item.input as { + incremented: number; + decremented: number; + }; + return { + id: item.id, + output: (value.incremented + value.decremented) / 2, + }; + }); + }, + }); }, - completion: { - mode: "poll", - async poll() { - return { status: "pending" }; + }, + ); + const scorer = BatchScorer({ + name: "exact", + workflow(w) { + const comparison = w.batch("compare", { + input: (item) => ({ + output: item.output, + expected: item.expected, + }), + submit: (items) => submit("compare", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => { + const value = item.input as { + output: number; + expected: number; + }; + return { + id: item.id, + output: value.output === value.expected, + }; + }); }, - }, - async *collect() { - // The submission never returns a handle. - }, - }), - scores: [], - }); - - const started = Date.now(); - const result = await durable.run({ - runId: "submit-deadline-run", - store: new MemoryDurableEvalStore(), - noSendLogs: true, - deadlineMs: 10, - }); - expect(result).toMatchObject({ status: "paused", reason: "deadline" }); - expect(Date.now() - started).toBeLessThan(1_000); - }); - - test("reserves one deterministic experiment name across concurrent shards", async () => { - const store = new MemoryDurableEvalStore(); - const durable = DurableEval("experiment-reservation", { - revision: "v1", - data: [ - { id: "one", input: 1 }, - { id: "two", input: 2 }, - { id: "three", input: 3 }, - { id: "four", input: 4 }, - ], - task: (input) => input, - scores: [], - }); - - await Promise.all([ - durable.run({ - runId: "shared-run", - shard: { index: 0, count: 2 }, - store, - noSendLogs: true, - }), - durable.run({ - runId: "shared-run", - shard: { index: 1, count: 2 }, - store, - noSendLogs: true, - }), - ]); - - const manifests: string[] = []; - for await (const key of store.list("durable-eval/v1/")) { - if (!key.endsWith("/manifest")) continue; - const record = await store.read(key); - if (record) manifests.push(new TextDecoder().decode(record.value)); - } - expect(manifests).toHaveLength(1); - expect(JSON.parse(manifests[0])).toMatchObject({ - experimentName: "experiment-reservation-shared-run", - shardCount: 2, + }); + return w.batch("score", { + needs: { comparison }, + input: (_item, outputs) => outputs.comparison, + submit: (items) => submit("score", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: item.input ? 1 : 0, + })); + }, + }); + }, }); - - await expect( - durable.status({ - runId: "shared-run", - store, - noSendLogs: true, - }), - ).resolves.toMatchObject({ status: "completed" }); - }); - - test("applies lifecycle retries to every shard when no shard is specified", async () => { - const store = new MemoryDurableEvalStore(); - let succeed = false; - const durable = DurableEval("sharded-lifecycle", { - revision: "v1", - data: [1, 2, 3, 4].map((input) => ({ + const store = new MemoryStore(); + const durable = DurableEval("workflow", { + store, + data: [1, 2, 3].map((input) => ({ id: `case-${input}`, input, + expected: input * 2, })), - task: (input) => { - if (!succeed) throw new Error("try again later"); - return input; - }, - scores: [], - }); - - await Promise.all([ - durable.run({ - runId: "sharded-lifecycle-run", - shard: { index: 0, count: 2 }, - store, - noSendLogs: true, - }), - durable.run({ - runId: "sharded-lifecycle-run", - shard: { index: 1, count: 2 }, - store, - noSendLogs: true, - }), - ]); - - succeed = true; - await expect( - durable.retryFailed({ - runId: "sharded-lifecycle-run", - store, - noSendLogs: true, - }), - ).resolves.toMatchObject({ - status: "completed", - progress: { taskSucceeded: 4, taskFailed: 0 }, - }); - }); - - test("keeps completed work when the definition revision changes", async () => { - const store = new MemoryDurableEvalStore(); - const firstTask = vi.fn((input: number) => input); - await DurableEval("revisions", { - revision: "v1", - data: [{ id: "one", input: 1, expected: 1 }], - task: firstTask, - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }).run({ - runId: "revision-run", - store, - noSendLogs: true, - }); - - const changedTask = vi.fn(() => 0); - const result = await DurableEval("revisions", { - revision: "v2", - data: [{ id: "one", input: 1, expected: 1 }], - task: changedTask, - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }).run({ - runId: "revision-run", - store, - noSendLogs: true, + task, + scores: [scorer], }); - - expect(result.status).toBe("completed"); - expect(firstTask).toHaveBeenCalledOnce(); - expect(changedTask).not.toHaveBeenCalled(); - }); - - test("removes inactive scorer state when a definition changes", async () => { - const store = new MemoryDurableEvalStore(); - await DurableEval("remove-scorer", { - revision: "v1", - data: [{ id: "one", input: 1 }], - task: (input) => input, - scores: [ - function oldScore() { - return 1; - }, - ], - }).run({ - runId: "remove-scorer-run", - store, + const options = { + runId: "workflow-run", noSendLogs: true, - }); + }; - const result = await DurableEval("remove-scorer", { - revision: "v2", - data: [{ id: "one", input: 1 }], - task: (input) => input, - scores: [], - }).run({ - runId: "remove-scorer-run", - store, - noSendLogs: true, + await expect(durable.start(options)).resolves.toMatchObject({ + status: "waiting", }); - expect(result).toMatchObject({ - status: "completed", - summary: { scores: {} }, + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "waiting", }); - - const cases: Array> = []; - for await (const key of store.list("durable-eval/v1/")) { - if (!key.includes("/cases/")) continue; - const record = await store.read(key); - if (record) - cases.push(JSON.parse(new TextDecoder().decode(record.value))); - } - expect(cases).toHaveLength(1); - expect(cases[0]).toMatchObject({ - scores: {}, - removedScores: ["oldScore"], - }); - }); - - test("pauses safely when submit may have created a provider job", async () => { - const store = new MemoryDurableEvalStore(); - let submitShouldFail = true; - const definition = { - revision: "v1", - data: [{ id: "one", input: "hello", expected: "hello" }], - task: BatchTask< - string, - string, - string, - void, - Record, - { jobId: string } - >({ - revision: "task-v1", - async submit(_items, context) { - if (submitShouldFail) { - throw new Error("connection closed after request"); - } - return { jobId: context.batchId }; - }, - completion: { - mode: "poll", - async poll() { - return { status: "complete" }; - }, - }, - async *collect() { - yield { id: "one:trial:0", output: "hello" }; - }, - }), - scores: [ - function exact({ - output, - expected, - }: { - output: string; - expected?: string; - }) { - return output === expected ? 1 : 0; - }, - ], - }; - const durable = DurableEval("ambiguous-submit", definition); - const result = await durable.run({ - runId: "ambiguous-run", - store, - noSendLogs: true, + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "waiting", }); - - expect(result.status).toBe("paused"); - if (result.status !== "paused") throw new Error("run did not pause"); - expect(result.reason).toBe("unknown_submission"); - expect(result.progress.unknown).toBe(1); - - const status = await durable.status({ - runId: "ambiguous-run", - store, - noSendLogs: true, + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "waiting", }); - expect(status).toMatchObject({ - status: "paused", - reason: "status_only", + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "waiting", }); - - submitShouldFail = false; - const resumed = await durable.resubmitUnknown({ - runId: "ambiguous-run", - store, - noSendLogs: true, + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "completed", + summary: { scores: { exact: { score: 1 } } }, }); - expect(resumed.status).toBe("completed"); }); - test("registers definitions instead of executing during CLI lazy loading", async () => { - const previousEvals = globalThis._evals; - const previousLazy = globalThis._lazy_load; - globalThis._evals = { - functions: [], - prompts: [], - parameters: [], - evaluators: {}, - durableEvaluators: {}, - reporters: {}, - } satisfies EvaluatorFile; - globalThis._lazy_load = true; - const task = vi.fn((input: string) => input); - try { - const definition = DurableEval("lazy-project", { - revision: "v1", - experimentName: "lazy-eval", - data: [{ id: "one", input: "hello", expected: "hello" }], - task, - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], - }); - expect(globalThis._evals.durableEvaluators?.["lazy-eval"]).toBeDefined(); - expect( - globalThis._evals.durableEvaluators?.["lazy-eval"]?.definition, - ).toBe(definition); - expect(task).not.toHaveBeenCalled(); - } finally { - globalThis._evals = previousEvals; - globalThis._lazy_load = previousLazy; - } + test("requires stable case ids", async () => { + await expect( + DurableEval("missing-ids", { + store: new MemoryStore(), + data: [{ input: "hello" }], + task: (input) => input, + scores: [], + }).start({ noSendLogs: true }), + ).rejects.toThrow("requires id, upsert_id, or caseId"); }); }); diff --git a/js/src/durable-eval.ts b/js/src/durable-eval.ts index c0806f143..160de1893 100644 --- a/js/src/durable-eval.ts +++ b/js/src/durable-eval.ts @@ -1,4 +1,13 @@ import { type Score, SpanTypeAttribute } from "../util/index"; +import { type EvalParameters, type InferParameters } from "./eval-parameters"; +import { + type EvalData, + type EvalHooks, + type EvalScorer, + type EvalScorerArgs, + type EvalTask, + type OneOrMoreScores, +} from "./framework"; import iso from "./isomorph"; import { type BaseMetadata, @@ -8,19 +17,9 @@ import { type Experiment, type ExperimentSummary, NOOP_SPAN, - _internalStartSpanWithInitialMerge, init as initExperiment, newId, } from "./logger"; -import { - type EvalData, - type EvalHooks, - type EvalScorer, - type EvalScorerArgs, - type EvalTask, - type OneOrMoreScores, -} from "./framework"; -import { type EvalParameters, type InferParameters } from "./eval-parameters"; const encoder = new TextEncoder(); const decoder = new TextDecoder(); @@ -28,11 +27,6 @@ const BATCH_TASK_KIND = "braintrust.durable.batch-task"; const BATCH_SCORER_KIND = "braintrust.durable.batch-scorer"; const CHECKPOINT_VERSION = 2; const DEFAULT_BATCH_SIZE = 1_000; -const DEFAULT_MAX_CONCURRENT_BATCHES = 1; -const JOB_LEASE_MS = 60_000; -const LEASE_HEARTBEAT_MS = JOB_LEASE_MS / 3; -const FILE_LOCK_STALE_MS = 10_000; -const FILE_LOCK_WAIT_MS = 15_000; type JsonPrimitive = string | number | boolean | null; export type JsonValue = @@ -40,276 +34,24 @@ export type JsonValue = | JsonValue[] | { [key: string]: JsonValue }; -export type DurableEvalWriteCondition = - | { ifAbsent: true } - | { ifVersion: string } - | { unconditional: true }; - -export interface DurableEvalStore { - read( - key: string, - ): Promise<{ value: Uint8Array; version: string } | undefined>; - write( - key: string, - value: Uint8Array, - condition: DurableEvalWriteCondition, - ): Promise< - | { written: true; version: string } - | { written: false; currentVersion?: string } - >; - list(prefix: string): AsyncIterable; -} - -/** - * Throw this from `submit` only when it is known that no provider job was - * created. Other submit errors are treated as ambiguous and pause the run. - */ -export class DurableEvalNotSubmittedError extends Error { - constructor(message: string, options?: ErrorOptions) { - super(message, options); - this.name = "DurableEvalNotSubmittedError"; - } -} - -export class MemoryDurableEvalStore implements DurableEvalStore { - private readonly values = new Map< - string, - { value: Uint8Array; version: number } - >(); - - async read(key: string) { - const record = this.values.get(key); - return record - ? { value: record.value.slice(), version: String(record.version) } - : undefined; - } - - async write( - key: string, - value: Uint8Array, - condition: DurableEvalWriteCondition, - ) { - const current = this.values.get(key); - if ( - ("ifAbsent" in condition && current) || - ("ifVersion" in condition && - (!current || String(current.version) !== condition.ifVersion)) - ) { - return { - written: false as const, - currentVersion: current ? String(current.version) : undefined, - }; - } - const version = (current?.version ?? 0) + 1; - this.values.set(key, { value: value.slice(), version }); - return { written: true as const, version: String(version) }; - } - - async *list(prefix: string) { - for (const key of [...this.values.keys()].sort()) { - if (key.startsWith(prefix)) { - yield key; - } - } - } -} - /** - * A filesystem checkpoint store intended for durable, multi-process evals on - * one machine. Distributed deployments should provide a database/object-store - * implementation of {@link DurableEvalStore}. + * Minimal persistence used to reconnect provider webhooks with submitted + * batches. DurableEval does not require any Braintrust backend changes. */ -export class FileDurableEvalStore implements DurableEvalStore { - constructor(public readonly root = ".braintrust/evals") {} - - private assertFilesystem() { - if ( - !iso.pathJoin || - !iso.pathDirname || - !iso.mkdir || - !iso.readFile || - !iso.writeFile || - !iso.readdir || - !iso.stat || - !iso.unlink || - !iso.rename || - !iso.openFile - ) { - throw new Error( - "FileDurableEvalStore is only available in a Node.js filesystem environment", - ); - } - } - - private pathFor(key: string) { - if ( - key.startsWith("/") || - key.includes("\\") || - key.split("/").some((part) => part === ".." || part === ".") - ) { - throw new Error(`Invalid durable eval store key: ${key}`); - } - return iso.pathJoin!(this.root, ...key.split("/")); - } - - private async withLock(path: string, fn: () => Promise): Promise { - const lockPath = `${path}.lock`; - const lockContents = stableStringify({ - owner: newId(), - createdAt: new Date().toISOString(), - }); - await iso.mkdir!(iso.pathDirname!(path), { recursive: true }); - const started = Date.now(); - let handle: - | { close(): Promise; writeFile(value: string): Promise } - | undefined; - while (!handle) { - try { - const candidate = await iso.openFile!(lockPath, "wx"); - try { - await candidate.writeFile(lockContents); - handle = candidate; - } catch (error) { - await candidate.close(); - await iso.unlink!(lockPath).catch(() => undefined); - throw error; - } - } catch (error) { - if (!isErrorCode(error, "EEXIST")) { - throw error; - } - try { - const lockStat = await iso.stat!(lockPath); - if (Date.now() - lockStat.mtimeMs >= FILE_LOCK_STALE_MS) { - await iso.unlink!(lockPath).catch((unlinkError) => { - if (!isErrorCode(unlinkError, "ENOENT")) throw unlinkError; - }); - continue; - } - } catch (statError) { - if (!isErrorCode(statError, "ENOENT")) throw statError; - continue; - } - if (Date.now() - started >= FILE_LOCK_WAIT_MS) throw error; - await delay(10); - } - } - try { - return await fn(); - } finally { - await handle.close(); - try { - const currentLock = await iso.readFile!(lockPath); - if (decoder.decode(currentLock) === lockContents) { - await iso.unlink!(lockPath).catch(() => undefined); - } - } catch (error) { - if (!isErrorCode(error, "ENOENT")) throw error; - } - } - } - - async read(key: string) { - this.assertFilesystem(); - const path = this.pathFor(key); - try { - const value = await iso.readFile!(path); - return { value, version: contentVersion(value) }; - } catch (error) { - if (isErrorCode(error, "ENOENT")) return undefined; - throw error; - } - } - - async write( - key: string, - value: Uint8Array, - condition: DurableEvalWriteCondition, - ) { - this.assertFilesystem(); - const path = this.pathFor(key); - return await this.withLock(path, async () => { - const current = await this.read(key); - if ( - ("ifAbsent" in condition && current) || - ("ifVersion" in condition && current?.version !== condition.ifVersion) - ) { - return { - written: false as const, - currentVersion: current?.version, - }; - } - - await iso.mkdir!(iso.pathDirname!(path), { recursive: true }); - const tempPath = `${path}.${newId()}.tmp`; - await iso.writeFile!(tempPath, value); - await iso.rename!(tempPath, path); - return { written: true as const, version: contentVersion(value) }; - }); - } - - async *list(prefix: string) { - this.assertFilesystem(); - const prefixPath = this.pathFor(prefix); - const walk = async function* (path: string): AsyncGenerator { - let entries: string[]; - try { - entries = await iso.readdir!(path); - } catch (error) { - if (isErrorCode(error, "ENOENT")) return; - throw error; - } - for (const entry of entries.sort()) { - if ( - entry.endsWith(".lock") || - entry.endsWith(".tmp") || - entry.includes(".tmp.") - ) { - continue; - } - const child = iso.pathJoin!(path, entry); - const stat = await iso.stat!(child); - if (stat.isDirectory()) { - yield* walk(child); - } else { - yield child; - } - } - }; - - for await (const path of walk(prefixPath)) { - yield path - .slice(this.root.length) - .replace(/^[/\\]+/, "") - .replaceAll("\\", "/"); - } - } +export interface DurableEvalStore { + read(key: string): Promise; + write(key: string, value: Uint8Array): Promise; } export interface DurableBatchContext { runId: string; - stage: string; - revision: string; - shard: { index: number; count: number }; - attempt: number; batchId: string; - itemCount: number; - signal: AbortSignal; } export type DurableBatchPoll = - | { status: "pending"; retryAfterMs?: number } + | { status: "pending" } | { status: "complete" } - | { status: "failed"; error: unknown; retryable?: boolean }; - -export type DurableBatchRecovery = - | { status: "found"; handle: Handle } - | { status: "not_found" } - | { status: "unknown" }; - -export type DurableBatchOutcome = - | { status: "complete" } - | { status: "failed"; error: JsonValue; retryable?: boolean }; + | { status: "failed"; error: unknown }; export type DurableBatchCompletion = | { @@ -318,20 +60,10 @@ export type DurableBatchCompletion = handle: Handle, context: DurableBatchContext, ): Promise; - intervalMs?: number; } | { mode: "webhook"; - source: string; externalId(handle: Handle, context: DurableBatchContext): string; - pollFallback?: { - afterMs: number; - intervalMs?: number; - poll( - handle: Handle, - context: DurableBatchContext, - ): Promise; - }; }; export interface DurableBatchTaskItem< @@ -356,7 +88,7 @@ export type DurableBatchTaskResult = metadata?: Metadata; tags?: string[]; } - | { id: string; error: unknown; retryable?: boolean }; + | { id: string; error: unknown }; export type DurableBatchScorerItem< Input, @@ -370,23 +102,85 @@ export type DurableBatchScorerItem< export type DurableBatchScorerResult = | { id: string; score: OneOrMoreScores } - | { id: string; error: unknown; retryable?: boolean }; + | { id: string; error: unknown }; export interface DurableBatchProcessor { - revision: string; batchSize?: number; - maxConcurrentBatches?: number; - maxAttempts?: number; submit(items: Item[], context: DurableBatchContext): Promise; - recover?(context: DurableBatchContext): Promise>; completion: DurableBatchCompletion; - collect( - handle: Handle, - context: DurableBatchContext, - ): AsyncIterable | Promise | AsyncIterable>; - cancel?(handle: Handle, context: DurableBatchContext): Promise; + collect(handle: Handle, context: DurableBatchContext): Promise; +} + +export interface DurableWorkflowBatchItem { + id: string; + input: Input; +} + +export type DurableWorkflowBatchResult = + | { + id: string; + output: Output; + metadata?: Metadata; + tags?: string[]; + } + | { id: string; error: unknown }; + +const WORKFLOW_NODE_OUTPUT: unique symbol = Symbol("DurableWorkflowNodeOutput"); + +export interface DurableWorkflowNode { + readonly [WORKFLOW_NODE_OUTPUT]: Output; +} + +type DurableWorkflowNodeMap = Record>; + +type DurableWorkflowNodeOutputs = { + [Name in keyof Nodes]: Nodes[Name] extends DurableWorkflowNode + ? Output + : never; +}; + +export interface DurableWorkflowBuilder< + RootItem, + Metadata extends BaseMetadata, +> { + batch< + Output, + Needs extends DurableWorkflowNodeMap = Record, + Input = RootItem, + Handle extends JsonValue = JsonValue, + >( + name: string, + processor: DurableBatchProcessor< + DurableWorkflowBatchItem, + DurableWorkflowBatchResult, + Handle + > & { + needs?: Needs; + input?: ( + item: RootItem, + outputs: DurableWorkflowNodeOutputs, + ) => Input; + }, + ): DurableWorkflowNode; } +type DurableWorkflowNodeDefinition = { + name: string; + needs: Record; + item: ( + rootItem: unknown, + outputs: Record, + id: string, + ) => unknown; + processor: DurableBatchProcessor; + result: (result: any) => unknown; +}; + +type DurableWorkflowDefinition = { + nodes: DurableWorkflowNodeDefinition[]; + outputNode: string; +}; + export interface DurableBatchTask< Input, Output, @@ -394,12 +188,14 @@ export interface DurableBatchTask< Metadata extends BaseMetadata, Parameters extends EvalParameters, Handle extends JsonValue, -> extends DurableBatchProcessor< - DurableBatchTaskItem, - DurableBatchTaskResult, - Handle > { readonly kind: typeof BATCH_TASK_KIND; + readonly processor?: DurableBatchProcessor< + DurableBatchTaskItem, + DurableBatchTaskResult, + Handle + >; + readonly workflow?: DurableWorkflowDefinition; } export interface DurableBatchScorer< @@ -408,13 +204,15 @@ export interface DurableBatchScorer< Expected, Metadata extends BaseMetadata, Handle extends JsonValue, -> extends DurableBatchProcessor< - DurableBatchScorerItem, - DurableBatchScorerResult, - Handle > { readonly kind: typeof BATCH_SCORER_KIND; name: string; + readonly processor?: DurableBatchProcessor< + DurableBatchScorerItem, + DurableBatchScorerResult, + Handle + >; + readonly workflow?: DurableWorkflowDefinition; } export function BatchTask< @@ -430,8 +228,50 @@ export function BatchTask< DurableBatchTaskResult, Handle >, -): DurableBatchTask { - return { kind: BATCH_TASK_KIND, ...processor }; +): DurableBatchTask; +export function BatchTask< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +>(config: { + workflow( + builder: DurableWorkflowBuilder< + DurableBatchTaskItem, + Metadata + >, + ): DurableWorkflowNode; +}): DurableBatchTask; +export function BatchTask( + config: + | DurableBatchProcessor + | { + workflow( + builder: DurableWorkflowBuilder, + ): DurableWorkflowNode; + }, +): DurableBatchTask< + unknown, + unknown, + unknown, + BaseMetadata, + EvalParameters, + JsonValue +> { + return { + kind: BATCH_TASK_KIND, + ...("workflow" in config + ? { workflow: buildWorkflow(config.workflow) } + : { processor: config }), + } as DurableBatchTask< + unknown, + unknown, + unknown, + BaseMetadata, + EvalParameters, + JsonValue + >; } export function BatchScorer< @@ -446,8 +286,106 @@ export function BatchScorer< DurableBatchScorerResult, Handle > & { name: string }, -): DurableBatchScorer { - return { kind: BATCH_SCORER_KIND, ...processor }; +): DurableBatchScorer; +export function BatchScorer< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, +>(config: { + name: string; + workflow( + builder: DurableWorkflowBuilder< + DurableBatchScorerItem, + Metadata + >, + ): DurableWorkflowNode; +}): DurableBatchScorer; +export function BatchScorer( + config: + | (DurableBatchProcessor & { name: string }) + | { + name: string; + workflow( + builder: DurableWorkflowBuilder, + ): DurableWorkflowNode; + }, +): DurableBatchScorer { + return { + kind: BATCH_SCORER_KIND, + name: config.name, + ...("workflow" in config + ? { workflow: buildWorkflow(config.workflow) } + : { processor: config }), + } as DurableBatchScorer; +} + +function buildWorkflow( + define: ( + builder: DurableWorkflowBuilder, + ) => DurableWorkflowNode, +): DurableWorkflowDefinition { + const nodes: DurableWorkflowNodeDefinition[] = []; + const nodeNames = new Map, string>(); + const builder: DurableWorkflowBuilder = { + batch(name, config) { + if (!name.trim()) throw new Error("Workflow batch names cannot be empty"); + if (nodes.some((node) => node.name === name)) { + throw new Error(`Duplicate workflow batch name: ${name}`); + } + const needs = Object.fromEntries( + Object.entries(config.needs ?? {}).map(([alias, dependency]) => { + const dependencyName = nodeNames.get(dependency); + if (!dependencyName) { + throw new Error( + `Workflow batch ${name} depends on an unknown or later batch`, + ); + } + return [alias, dependencyName]; + }), + ); + const { input, needs: _needs, ...processor } = config; + const handle = {} as DurableWorkflowNode; + nodeNames.set(handle, name); + nodes.push({ + name, + needs, + item(rootItem, outputs, id) { + return { + id, + input: input + ? input(rootItem as RootItem, outputs as never) + : rootItem, + }; + }, + processor: processor as DurableBatchProcessor, + result: (result) => result.output, + }); + return handle as never; + }, + }; + const output = define(builder); + const outputNode = nodeNames.get(output); + if (!outputNode) { + throw new Error("A batch workflow must return one of its batch nodes"); + } + const reachable = new Set(); + const visit = (name: string) => { + if (reachable.has(name)) return; + reachable.add(name); + const node = nodes.find((candidate) => candidate.name === name)!; + Object.values(node.needs).forEach(visit); + }; + visit(outputNode); + const unused = nodes.filter((node) => !reachable.has(node.name)); + if (unused.length > 0) { + throw new Error( + `Batch workflow contains nodes that do not contribute to its output: ${unused + .map((node) => node.name) + .join(", ")}`, + ); + } + return { nodes, outputNode }; } export type DurableEvaluator< @@ -457,7 +395,7 @@ export type DurableEvaluator< Metadata extends BaseMetadata = DefaultMetadataType, Parameters extends EvalParameters = EvalParameters, > = { - revision: string; + store: DurableEvalStore; data: EvalData; caseId?: ( datum: EvalCase, @@ -488,51 +426,31 @@ export type DurableEvaluator< export interface DurableEvalRuntimeOptions { runId?: string; - store?: DurableEvalStore; - shard?: { index: number; count: number }; - deadlineMs?: number; - signal?: AbortSignal; noSendLogs?: boolean; - workerId?: string; - checkpointDir?: string; -} - -export interface DurableEvalExistingRunOptions extends DurableEvalRuntimeOptions { - runId: string; } -export type DurableEvalOperation = - | "run" - | "status" - | "retry-failed" - | "resubmit-unknown" - | "cancel"; - -type DurableEvalExecutionOptions = DurableEvalRuntimeOptions & { - operation?: DurableEvalOperation; -}; - -export type DurableBatchResultEvent = { - eventId: string; - source: string; - externalId?: string; +export type DurableBatchResult = { batchId?: string; - handle?: JsonValue; - outcome: DurableBatchOutcome; - payload?: JsonValue; + externalId?: string; }; -export type DurableBatchProcessingResult = +export type DurableEvalResult = | { - status: "processed"; - batchId: string; - run: DurableEvalResult; + status: "waiting"; + runId: string; + pending: { + poll: number; + webhook: number; + }; } - | { status: "duplicate"; batchId?: string } | { - status: "pending"; - reason: "unmatched" | "definition_missing"; - batchId?: string; + status: "completed"; + runId: string; + pending: { + poll: 0; + webhook: 0; + }; + summary: ExperimentSummary; }; export interface DurableEvalDefinition< @@ -551,189 +469,59 @@ export interface DurableEvalDefinition< Metadata, Parameters >; - run(options?: DurableEvalRuntimeOptions): Promise; - status(options: DurableEvalExistingRunOptions): Promise; - retryFailed( - options: DurableEvalExistingRunOptions, + start(options?: DurableEvalRuntimeOptions): Promise; + status( + options: DurableEvalRuntimeOptions & { runId: string }, ): Promise; - resubmitUnknown( - options: DurableEvalExistingRunOptions, + poll( + options: DurableEvalRuntimeOptions & { runId: string }, ): Promise; - cancel(options: DurableEvalExistingRunOptions): Promise; processBatchResult( - event: DurableBatchResultEvent, - options?: Omit, - ): Promise; -} - -export interface DurableEvalProgress { - total: number; - taskPending: number; - taskSucceeded: number; - taskFailed: number; - scorePending: number; - scoreSucceeded: number; - scoreFailed: number; - unknown: number; -} - -export interface DurableEvalFailureSummary { - tasks: number; - scorers: number; + result: DurableBatchResult, + options?: Omit, + ): Promise; } -export type DurableEvalPauseReason = - | "deadline" - | "aborted" - | "shard_complete" - | "waiting_for_webhook" - | "unknown_submission" - | "provider_unreachable" - | "status_only"; - -export type DurableEvalResult = - | { - status: "completed"; - runId: string; - summary: ExperimentSummary; - progress: DurableEvalProgress; - failures: DurableEvalFailureSummary; - } - | { - status: "paused"; - runId: string; - reason: DurableEvalPauseReason; - progress: DurableEvalProgress; - resume( - overrides?: Partial, - ): Promise; - }; - -type SerializedError = { - name?: string; - message: string; - stack?: string; -}; - -type StageState = - | { status: "pending"; attempts: number } - | { - status: "leased"; - attempts: number; - workerId: string; - leaseUntil: number; - } - | { - status: "in_batch"; - attempts: number; - batchId: string; - revision: string; - } - | { - status: "succeeded"; - attempts: number; - revision: string; - value: JsonValue; - } - | { - status: "failed"; - attempts: number; - revision: string; - error: SerializedError; - } - | { - status: "unknown"; - attempts: number; - revision: string; - batchId: string; - }; - type DurableCaseRecord = { id: string; caseId: string; trialIndex: number; - shard: number; datum: JsonValue; - task: StageState; - scores: Record; metadata: JsonValue; tags?: string[]; - logPending?: boolean; - removedScores?: string[]; + taskComplete: boolean; + output?: JsonValue; + taskNodeOutputs: Record; + scores: Record; + scoreNodeOutputs: Record>; }; -type DurableJobRecord = { +type DurableBatchRecord = { id: string; - stage: string; kind: "task" | "score"; scorerName?: string; + nodeName: string; itemIds: string[]; - attempt: number; - revision: string; - shard: number; - status: - | "preparing" - | "submitting" - | "submitted" - | "complete" - | "failed" - | "unknown"; - workerId?: string; - leaseUntil?: number; - handle?: JsonValue; - external?: { source: string; id: string }; - submittedAt?: number; - outcome?: DurableBatchOutcome; - webhookEventKey?: string; - nextPollAt?: number; - error?: SerializedError; + handle: JsonValue; + externalId?: string; + status: "submitted" | "complete"; }; -type DurableRunManifest = { +type DurableRunState = { schemaVersion: number; runId: string; projectName: string; evalName: string; - revision: string; - shardCount: number; - dataSealed: boolean; - activeScorers: string[]; - experimentName?: string; - createdAt: string; -}; - -type DurableExperimentInitRecord = { - schemaVersion: number; experimentName: string; - status: "initializing" | "ready"; - workerId?: string; - leaseUntil?: number; - experimentId?: string; + status: "running" | "completed"; + summary?: ExperimentSummary; + cases: DurableCaseRecord[]; + batches: DurableBatchRecord[]; }; -type Versioned = { value: T; version: string }; - type DurableBatchLocator = { - schemaVersion: number; - projectName: string; - evalName: string; - runId: string; - prefix: string; - jobKey: string; + runKey: string; batchId: string; - stage: string; - kind: "task" | "score"; - scorerName?: string; - shard: number; - shardCount: number; -}; - -type DurableWebhookEventRecord = { - schemaVersion: number; - event: DurableBatchResultEvent; - receivedAt: string; - status: "pending" | "applied"; - batchId?: string; }; class DurableEvalDefinitionImpl< @@ -764,43 +552,27 @@ class DurableEvalDefinitionImpl< this.evalName = evaluator.experimentName ?? projectName; } - run(options: DurableEvalRuntimeOptions = {}) { - return runDurableEval(this.projectName, this.evaluator, options); - } - - status(options: DurableEvalExistingRunOptions) { - return runDurableEval(this.projectName, this.evaluator, { - ...options, - operation: "status", - }); - } - - retryFailed(options: DurableEvalExistingRunOptions) { - return runDurableEval(this.projectName, this.evaluator, { - ...options, - operation: "retry-failed", - }); + start(options: DurableEvalRuntimeOptions = {}): Promise { + return startDurableEval(this, options); } - resubmitUnknown(options: DurableEvalExistingRunOptions) { - return runDurableEval(this.projectName, this.evaluator, { - ...options, - operation: "resubmit-unknown", - }); + status( + options: DurableEvalRuntimeOptions & { runId: string }, + ): Promise { + return getDurableEvalStatus(this, options); } - cancel(options: DurableEvalExistingRunOptions) { - return runDurableEval(this.projectName, this.evaluator, { - ...options, - operation: "cancel", - }); + poll( + options: DurableEvalRuntimeOptions & { runId: string }, + ): Promise { + return pollDurableEval(this, options); } processBatchResult( - event: DurableBatchResultEvent, - options: Omit = {}, - ): Promise { - return processDurableBatchResult(this, event, options); + result: DurableBatchResult, + options: Omit = {}, + ): Promise { + return processDurableBatchResult(this, result, options); } } @@ -814,17 +586,10 @@ export function DurableEval< projectName: string, evaluator: DurableEvaluator, ): DurableEvalDefinition { - const definition = new DurableEvalDefinitionImpl(projectName, evaluator); - if (globalThis._lazy_load) { - globalThis._evals.durableEvaluators ??= {}; - globalThis._evals.durableEvaluators[definition.evalName] = { - definition, - }; - } - return definition; + return new DurableEvalDefinitionImpl(projectName, evaluator); } -async function processDurableBatchResult< +async function startDurableEval< Input, Output, Expected, @@ -838,2548 +603,693 @@ async function processDurableBatchResult< Metadata, Parameters >, - event: DurableBatchResultEvent, - options: Omit, -): Promise { - if (!event.eventId.trim() || !event.source.trim()) { - throw new Error( - "Durable batch webhook events require non-empty eventId and source", - ); - } - if (!event.batchId && !event.externalId) { - throw new Error( - "Durable batch webhook events require batchId or externalId", - ); - } - if (event.externalId !== undefined && !event.externalId.trim()) { - throw new Error("Durable batch webhook externalId must be non-empty"); - } - const normalizedEvent = assertJsonValue( - event, - "durable batch webhook event", - ) as DurableBatchResultEvent; - const store = - options.store ?? - new FileDurableEvalStore(options.checkpointDir ?? ".braintrust/evals"); - const storedEventKey = webhookEventKey(event.source, event.eventId); - const eventRecord: DurableWebhookEventRecord = { - schemaVersion: CHECKPOINT_VERSION, - event: normalizedEvent, - receivedAt: new Date().toISOString(), - status: "pending", - batchId: event.batchId, - }; - const inserted = await writeJson(store, storedEventKey, eventRecord, { - ifAbsent: true, - }); - if (!inserted.written) { - const existing = await readJson( - store, - storedEventKey, - ); - if ( - !existing || - stableStringify(existing.value.event) !== stableStringify(normalizedEvent) - ) { - throw new Error( - `Webhook event ${event.source}/${event.eventId} was reused with different contents`, - ); - } - if (existing.value.status === "applied") { - return { - status: "duplicate", - batchId: existing.value.batchId ?? event.batchId, - }; - } + options: DurableEvalRuntimeOptions, +): Promise { + const store = definition.evaluator.store; + const runId = options.runId ?? newId(); + const key = runKey(definition.projectName, definition.evalName, runId); + let state = await readJson(store, key); + if (!state) { + state = { + schemaVersion: CHECKPOINT_VERSION, + runId, + projectName: definition.projectName, + evalName: definition.evalName, + experimentName: + definition.evaluator.experimentName ?? + `${definition.evalName}-${runId}`, + status: "running", + cases: await materializeCases(definition.evaluator), + batches: [], + }; + await writeJson(store, key, state); } + return advanceDurableEval(definition, state, store, key, options); +} - if (event.externalId) { - await writeJson( - store, - webhookMailboxKey(event.source, event.externalId, storedEventKey), - { eventKey: storedEventKey }, - { ifAbsent: true }, - ); - } - - const internalLocator = event.batchId - ? await readJson( - store, - internalBatchIndexKey(event.batchId), - ) - : undefined; - const externalLocator = event.externalId - ? await readJson( - store, - externalBatchIndexKey(event.source, event.externalId), - ) - : undefined; - if ( - internalLocator && - externalLocator && - internalLocator.value.jobKey !== externalLocator.value.jobKey - ) { - throw new Error( - "Durable batch webhook batchId and externalId resolve to different jobs", - ); - } - const locator = internalLocator?.value ?? externalLocator?.value; - if (!locator) { - return { status: "pending", reason: "unmatched", batchId: event.batchId }; - } - if ( - locator.projectName !== definition.projectName || - locator.evalName !== definition.evalName - ) { - return { - status: "pending", - reason: "definition_missing", - batchId: locator.batchId, - }; - } - - const processor = batchProcessorForLocator(definition.evaluator, locator); - if ( - !processor || - processor.completion.mode !== "webhook" || - processor.completion.source !== event.source - ) { - return { - status: "pending", - reason: "definition_missing", - batchId: locator.batchId, - }; - } - - const attached = await attachWebhookEventToJob( - store, - locator.jobKey, - storedEventKey, - normalizedEvent, - ); - if (attached === "missing") { - return { - status: "pending", - reason: "unmatched", - batchId: locator.batchId, - }; - } - if (attached === "duplicate") { - await markWebhookEventApplied(store, storedEventKey, locator.batchId); - return { status: "duplicate", batchId: locator.batchId }; - } - if (attached === "in_progress") { - return { status: "duplicate", batchId: locator.batchId }; - } - - const run = await runDurableEval( - definition.projectName, - definition.evaluator, - { - ...options, - runId: locator.runId, - shard: { index: locator.shard, count: locator.shardCount }, - }, - ); - const job = await readJson(store, locator.jobKey); - if ( - job && - (job.value.status === "complete" || job.value.status === "failed") - ) { - await markWebhookEventApplied(store, storedEventKey, locator.batchId); - } - return { status: "processed", batchId: locator.batchId, run }; -} - -async function runDurableEval< +async function getDurableEvalStatus< Input, Output, - Expected = void, - Metadata extends BaseMetadata = DefaultMetadataType, - Parameters extends EvalParameters = EvalParameters, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, >( - projectName: string, - evaluator: DurableEvaluator, - options: DurableEvalExecutionOptions = {}, + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + options: DurableEvalRuntimeOptions & { runId: string }, ): Promise { - const runId = options.runId ?? newId(); - const workerId = options.workerId ?? newId(); - const store = - options.store ?? - new FileDurableEvalStore(options.checkpointDir ?? ".braintrust/evals"); - const evalName = evaluator.experimentName ?? projectName; - const prefix = runPrefix(projectName, evalName, runId); - const manifestKey = `${prefix}/manifest`; - const scorerDefinitions = resolveScorers(evaluator.scores); - const scorerNames = scorerDefinitions.map((definition) => definition.name); - const webhookStages = new Set(); - if ( - isBatchTask(evaluator.task) && - evaluator.task.completion.mode === "webhook" && - evaluator.task.completion.pollFallback === undefined - ) { - webhookStages.add("task"); - } - for (const definition of scorerDefinitions) { - if ( - isBatchScorer(definition.scorer) && - definition.scorer.completion.mode === "webhook" && - definition.scorer.completion.pollFallback === undefined - ) { - webhookStages.add(`score:${definition.name}`); - } - } - if (new Set(scorerNames).size !== scorerNames.length) { - throw new Error("DurableEval scorer names must be unique"); - } - - const existingManifest = await readJson( + const store = definition.evaluator.store; + const state = await readJson( store, - manifestKey, + runKey(definition.projectName, definition.evalName, options.runId), ); - const shard = options.shard ?? { - index: 0, - count: existingManifest?.value.shardCount ?? 1, - }; - validateShard(shard); - if (options.operation && options.operation !== "run" && !existingManifest) { - throw new Error(`DurableEval run ${runId} does not exist`); - } - let manifest = - existingManifest ?? - (await initializeManifest(store, manifestKey, { - schemaVersion: CHECKPOINT_VERSION, - runId, - projectName, - evalName, - revision: evaluator.revision, - shardCount: shard.count, - dataSealed: false, - activeScorers: scorerNames, - experimentName: evaluator.experimentName ?? `${evalName}-${runId}`, - createdAt: new Date().toISOString(), - })); - if (manifest.value.shardCount !== shard.count) { - throw new Error( - `DurableEval run ${runId} was created with ${manifest.value.shardCount} shards, not ${shard.count}`, - ); - } - if ( - options.shard === undefined && - manifest.value.shardCount > 1 && - (options.operation === "retry-failed" || - options.operation === "resubmit-unknown" || - options.operation === "cancel") - ) { - let lastResult: DurableEvalResult | undefined; - let pausedLifecycleResult: DurableEvalResult | undefined; - for (let index = 0; index < manifest.value.shardCount; index++) { - const result = await runDurableEval(projectName, evaluator, { - ...options, - runId, - store, - shard: { index, count: manifest.value.shardCount }, - }); - lastResult = result; - if ( - result.status === "paused" && - result.reason !== "shard_complete" && - result.reason !== "status_only" && - !pausedLifecycleResult - ) { - pausedLifecycleResult = result; - } - } - if (!lastResult) throw new Error("DurableEval sharded lifecycle failed"); - return pausedLifecycleResult ?? lastResult; - } - if (!manifest.value.experimentName) { - manifest = await updateManifest(store, manifestKey, (current) => ({ - ...current, - experimentName: - current.experimentName ?? - evaluator.experimentName ?? - `${evalName}-${runId}`, - })); - } - - if (!manifest.value.dataSealed) { - await materializeData({ - store, - prefix, - evaluator, - shardCount: shard.count, - }); - } - manifest = await updateManifest(store, manifestKey, (current) => ({ - ...current, - revision: evaluator.revision, - activeScorers: scorerNames, - dataSealed: true, - })); - - const experiment: Experiment | null = options.noSendLogs - ? null - : await initializeDurableExperiment({ - store, - key: `${prefix}/experiment`, - workerId, - experimentName: manifest.value.experimentName!, - create: () => - initExperiment({ - state: evaluator.state, - ...(evaluator.projectId - ? { projectId: evaluator.projectId } - : { project: projectName }), - experiment: manifest.value.experimentName, - update: true, - description: evaluator.description, - metadata: evaluator.metadata, - tags: evaluator.tags, - setCurrent: false, - }), - }); - - await reconcileScorers(store, prefix, scorerNames); - - if (options.operation === "retry-failed") { - await resetStages(store, prefix, scorerNames, "failed", shard.index); - } else if (options.operation === "resubmit-unknown") { - await resetStages(store, prefix, scorerNames, "unknown", shard.index); - } else if (options.operation === "cancel") { - await cancelRun({ - store, - prefix, - evaluator, - scorerDefinitions, - runId, - shard, - signal: options.signal ?? new AbortController().signal, - }); - } - - if (options.operation === "status" || options.operation === "cancel") { - if (options.operation === "cancel" && experiment) { - await flushPendingLogs({ - store, - prefix, - experiment, - scorerNames, - shard, - runId, - }); - await experiment.flush(); - } - const progress = await collectProgress(store, prefix, scorerNames); - const complete = await isComplete(store, prefix, scorerNames); - if (!complete) { - return pausedResult(runId, "status_only", progress, (overrides = {}) => - runDurableEval(projectName, evaluator, { - ...options, - ...overrides, - operation: "run", - runId, - store, - shard, - }), - ); - } - return { - status: "completed", - runId, - summary: await buildLocalDurableSummary( - store, - prefix, - projectName, - manifest.value.experimentName ?? evalName, - scorerNames, - ), - progress, - failures: { - tasks: progress.taskFailed, - scorers: progress.scoreFailed, - }, - }; - } - - const startedAt = Date.now(); - const deadlineAt = - options.deadlineMs === undefined - ? undefined - : startedAt + Math.max(options.deadlineMs, 0); - const controller = new AbortController(); - const abortHandler = () => controller.abort(); - options.signal?.addEventListener("abort", abortHandler, { once: true }); - const deadlineTimer = - deadlineAt === undefined - ? undefined - : setTimeout( - () => controller.abort(), - Math.max(deadlineAt - Date.now(), 0), - ); - - const resume = (overrides: Partial = {}) => - runDurableEval(projectName, evaluator, { - ...options, - ...overrides, - runId, - store, - shard, - operation: "run", - }); - - try { - while (true) { - if (options.signal?.aborted) { - return pausedResult( - runId, - "aborted", - await collectProgress(store, prefix, scorerNames), - resume, - ); - } - if (deadlineAt !== undefined && Date.now() >= deadlineAt) { - return pausedResult( - runId, - "deadline", - await collectProgress(store, prefix, scorerNames), - resume, - ); - } - - let changed = false; - changed = - (await runTaskPass({ - store, - prefix, - projectName, - evalName, - evaluator, - shard, - workerId, - runId, - signal: controller.signal, - })) || changed; - for (const scorer of scorerDefinitions) { - changed = - (await runScorePass({ - store, - prefix, - projectName, - evalName, - scorer, - shard, - workerId, - runId, - signal: controller.signal, - })) || changed; - } - - if (experiment) { - changed = - (await flushPendingLogs({ - store, - prefix, - experiment, - scorerNames, - shard, - runId, - })) || changed; - } - - const progress = await collectProgress(store, prefix, scorerNames); - if (controller.signal.aborted) { - return pausedResult( - runId, - options.signal?.aborted ? "aborted" : "deadline", - progress, - resume, - ); - } - if (progress.unknown > 0) { - return pausedResult(runId, "unknown_submission", progress, resume); - } - const shardComplete = await isComplete( - store, - prefix, - scorerNames, - shard.index, - ); - const globallyComplete = await isComplete(store, prefix, scorerNames); - if (globallyComplete) { - if (experiment) { - await experiment.flush(); - } - const summary = experiment - ? await experiment.summarize() - : await buildLocalDurableSummary( - store, - prefix, - projectName, - manifest.value.experimentName ?? evalName, - scorerNames, - ); - return { - status: "completed", - runId, - summary, - progress, - failures: { - tasks: progress.taskFailed, - scorers: progress.scoreFailed, - }, - }; - } - if (shardComplete && shard.count > 1) { - return pausedResult(runId, "shard_complete", progress, resume); - } - if (!changed && (await hasProviderErrorJob(store, prefix, shard.index))) { - return pausedResult(runId, "provider_unreachable", progress, resume); - } - if ( - !changed && - (await hasWaitingWebhookJob(store, prefix, shard.index, webhookStages)) - ) { - return pausedResult(runId, "waiting_for_webhook", progress, resume); - } - if (!changed) { - await delay(Math.min(timeRemaining(deadlineAt) ?? 1_000, 1_000)); - } - } - } finally { - if (deadlineTimer) clearTimeout(deadlineTimer); - options.signal?.removeEventListener("abort", abortHandler); - } -} - -function pausedResult( - runId: string, - reason: DurableEvalPauseReason, - progress: DurableEvalProgress, - resume: ( - overrides?: Partial, - ) => Promise, -): DurableEvalResult { - return { status: "paused", runId, reason, progress, resume }; -} - -function emptyProgress(): DurableEvalProgress { - return { - total: 0, - taskPending: 0, - taskSucceeded: 0, - taskFailed: 0, - scorePending: 0, - scoreSucceeded: 0, - scoreFailed: 0, - unknown: 0, - }; -} - -async function initializeManifest( - store: DurableEvalStore, - key: string, - initial: DurableRunManifest, -): Promise> { - const existing = await readJson(store, key); - if (existing) return existing; - const result = await writeJson(store, key, initial, { ifAbsent: true }); - if (result.written) return { value: initial, version: result.version }; - const raced = await readJson(store, key); - if (!raced) throw new Error("DurableEval manifest initialization failed"); - return raced; + if (!state) throw new Error(`DurableEval run ${options.runId} is missing`); + return currentStatus(definition, state); } -async function updateManifest( - store: DurableEvalStore, - key: string, - update: (manifest: DurableRunManifest) => DurableRunManifest, -) { - while (true) { - const current = await readJson(store, key); - if (!current) throw new Error("DurableEval manifest is missing"); - const next = update(current.value); - const result = await writeJson(store, key, next, { - ifVersion: current.version, - }); - if (result.written) return { value: next, version: result.version }; - } -} - -async function initializeDurableExperiment({ - store, - key, - workerId, - experimentName, - create, -}: { - store: DurableEvalStore; - key: string; - workerId: string; - experimentName: string; - create: () => Experiment; -}): Promise { - while (true) { - const current = await readJson(store, key); - if (current?.value.status === "ready") { - const experiment = create(); - const experimentId = await experiment.id; - if ( - current.value.experimentId && - current.value.experimentId !== experimentId - ) { - throw new Error( - `DurableEval experiment ${experimentName} resolved to multiple experiment IDs`, - ); - } - return experiment; - } - - const now = Date.now(); - if (!current || (current.value.leaseUntil ?? 0) <= now) { - const claimed: DurableExperimentInitRecord = { - schemaVersion: CHECKPOINT_VERSION, - experimentName, - status: "initializing", - workerId, - leaseUntil: now + JOB_LEASE_MS, - }; - const write = await writeJson( - store, - key, - claimed, - current ? { ifVersion: current.version } : { ifAbsent: true }, - ); - if (write.written) { - const heartbeat = startLeaseHeartbeat(async () => { - let renewed = false; - await updateExperimentInitRecord(store, key, (record) => { - if ( - record.status !== "initializing" || - record.workerId !== workerId - ) { - return undefined; - } - renewed = true; - return { ...record, leaseUntil: Date.now() + JOB_LEASE_MS }; - }); - return renewed; - }); - try { - const experiment = create(); - const experimentId = await experiment.id; - const finalized = await updateExperimentInitRecord( - store, - key, - (record) => - record.status === "initializing" && record.workerId === workerId - ? { - schemaVersion: CHECKPOINT_VERSION, - experimentName, - status: "ready", - experimentId, - } - : undefined, - ); - if (!finalized) { - throw new Error( - `DurableEval lost the experiment initialization lease for ${experimentName}`, - ); - } - return experiment; - } catch (error) { - await updateExperimentInitRecord(store, key, (record) => - record.status === "initializing" && record.workerId === workerId - ? { ...record, leaseUntil: 0 } - : undefined, - ); - throw error; - } finally { - await heartbeat.stop(); - } - } - } - - await delay(25); - } -} - -async function updateExperimentInitRecord( - store: DurableEvalStore, - key: string, - update: ( - record: DurableExperimentInitRecord, - ) => DurableExperimentInitRecord | undefined, -) { - while (true) { - const current = await readJson(store, key); - if (!current) return false; - const next = update(structuredClone(current.value)); - if (!next) return false; - const written = await writeJson(store, key, next, { - ifVersion: current.version, - }); - if (written.written) return true; - } -} - -async function materializeData< +async function processDurableBatchResult< Input, Output, Expected, Metadata extends BaseMetadata, Parameters extends EvalParameters, ->({ - store, - prefix, - evaluator, - shardCount, -}: { - store: DurableEvalStore; - prefix: string; - evaluator: DurableEvaluator; - shardCount: number; -}) { - const rawData = - typeof evaluator.data === "function" ? evaluator.data() : evaluator.data; - const resolved = rawData instanceof Promise ? await rawData : rawData; - if ( - typeof resolved === "object" && - resolved !== null && - "_type" in resolved - ) { - throw new Error( - "DurableEval does not yet support BaseExperiment data sources", - ); - } - if (!isIterable(resolved) && !isAsyncIterable(resolved)) { - throw new Error("DurableEval data must be an iterable or async iterable"); - } - - const seen = new Set(); - for await (const datum of toAsyncIterable(resolved)) { - const caseId = - datum.id ?? - datum.upsert_id ?? - (evaluator.caseId ? await evaluator.caseId(datum) : undefined); - if (!caseId || typeof caseId !== "string") { - throw new Error( - "Every DurableEval case must have a non-empty id/upsert_id or be resolved by caseId", - ); - } - if (seen.has(caseId)) { - throw new Error(`Duplicate DurableEval 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 DurableEval case ${caseId}`); - } - for (let trialIndex = 0; trialIndex < trialCount; trialIndex++) { - const id = workItemId(caseId, trialIndex); - const record: DurableCaseRecord = { - id, - caseId, - trialIndex, - shard: stableShard(id, shardCount), - datum: assertJsonValue(datum, `case ${caseId}`), - task: { status: "pending", attempts: 0 }, - scores: {}, - metadata: assertJsonValue( - "metadata" in datum ? datum.metadata : {}, - `case ${caseId} metadata`, - ), - tags: datum.tags, - }; - const key = caseKey(prefix, id); - const result = await writeJson(store, key, record, { ifAbsent: true }); - if (!result.written) { - const existing = await readJson(store, key); - if ( - !existing || - stableStringify(existing.value.datum) !== - stableStringify(record.datum) - ) { - throw new Error( - `DurableEval case ${caseId} changed while the run was being prepared`, - ); - } - } - } - } -} - -type ResolvedScorer = { - name: string; - revision: string; - // Runtime orchestration intentionally erases user generics after public - // type-checking at the DurableEval boundary. - scorer: // eslint-disable-next-line @typescript-eslint/no-explicit-any - | EvalScorer - // eslint-disable-next-line @typescript-eslint/no-explicit-any - | DurableBatchScorer; -}; - -function resolveScorers( - // eslint-disable-next-line @typescript-eslint/no-explicit-any - scorers: DurableEvaluator["scores"], -): ResolvedScorer[] { - return scorers.map((scorer) => { - if (isBatchScorer(scorer)) { - if (!scorer.name.trim()) { - throw new Error("DurableEval batch scorers must have a name"); - } - return { name: scorer.name, revision: scorer.revision, scorer }; - } - if (!scorer.name) { - throw new Error("DurableEval scorers must be named functions"); - } - return { name: scorer.name, revision: "unversioned", scorer }; - }); -} - -async function reconcileScorers( - store: DurableEvalStore, - prefix: string, - scorerNames: string[], -) { - const active = new Set(scorerNames); - for await (const record of listCases(store, prefix)) { - await updateCaseRecord(store, prefix, record.value.id, (next) => { - const missing = scorerNames.filter( - (name) => next.scores[name] === undefined, - ); - const removed = Object.keys(next.scores).filter( - (name) => !active.has(name), - ); - if (!missing.length && !removed.length) return undefined; - for (const name of missing) { - next.scores[name] = { status: "pending", attempts: 0 }; - } - const removedScoreKeys = removed.flatMap((name) => { - const state = next.scores[name]; - return state?.status === "succeeded" - ? typeof state.value === "object" && - state.value !== null && - !Array.isArray(state.value) - ? Object.keys(state.value) - : [name] - : []; - }); - for (const name of removed) { - delete next.scores[name]; - } - if (removedScoreKeys.length) { - next.removedScores = [ - ...new Set([...(next.removedScores ?? []), ...removedScoreKeys]), - ]; - next.logPending = true; - } - return next; - }); - } - for await (const key of store.list(`${prefix}/jobs/`)) { - await updateJobRecord(store, key, (job) => { - if ( - !job.stage.startsWith("score:") || - active.has(job.stage.slice("score:".length)) || - job.status === "complete" || - job.status === "failed" - ) { - return undefined; - } - job.status = "failed"; - job.error = { - name: "ScorerRemovedError", - message: `Scorer ${job.stage.slice("score:".length)} was removed`, - }; - delete job.workerId; - delete job.leaseUntil; - return job; - }); - } -} - -async function resetStages( - store: DurableEvalStore, - prefix: string, - scorerNames: string[], - status: "failed" | "unknown", - shard: number, -) { - for await (const current of listCases(store, prefix)) { - if (current.value.shard !== shard) continue; - await updateCaseRecord(store, prefix, current.value.id, (next) => { - let changed = false; - if (next.task.status === status) { - next.task = { - status: "pending", - attempts: next.task.attempts, - }; - changed = true; - } - for (const name of scorerNames) { - const state = next.scores[name]; - if (state?.status === status) { - next.scores[name] = { - status: "pending", - attempts: state.attempts, - }; - changed = true; - } - } - return changed ? next : undefined; - }); - } -} - -async function cancelRun({ - store, - prefix, - evaluator, - scorerDefinitions, - runId, - shard, - signal, -}: { - store: DurableEvalStore; - prefix: string; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - evaluator: DurableEvaluator; - scorerDefinitions: ResolvedScorer[]; - runId: string; - shard: { index: number; count: number }; - signal: AbortSignal; -}) { - const processors = new Map< - string, - DurableBatchProcessor - >(); - if (isBatchTask(evaluator.task)) { - processors.set( - "task", - evaluator.task as DurableBatchProcessor, - ); - } - for (const definition of scorerDefinitions) { - if (isBatchScorer(definition.scorer)) { - processors.set( - `score:${definition.name}`, - definition.scorer as DurableBatchProcessor, - ); - } - } - - for await (const key of store.list(`${prefix}/jobs/`)) { - const current = await readJson(store, key); - if (!current || current.value.shard !== shard.index) continue; - const job = structuredClone(current.value); - const processor = processors.get(job.stage); - if ( - job.status === "submitted" && - job.handle !== undefined && - processor?.cancel - ) { - await processor.cancel( - job.handle, - batchContext(runId, job.stage, job, shard, signal), - ); - } - await updateJobRecord(store, key, (latest) => { - if ( - latest.status !== "preparing" && - latest.status !== "submitting" && - latest.status !== "submitted" && - latest.status !== "unknown" - ) { - return undefined; - } - latest.status = "failed"; - latest.error = { - name: "CancelledError", - message: "Durable eval cancelled", - }; - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - } - - for await (const current of listCases(store, prefix)) { - if (current.value.shard !== shard.index) continue; - await updateCaseRecord(store, prefix, current.value.id, (next) => { - let changed = false; - if (!isTerminal(next.task)) { - next.task = { - status: "failed", - attempts: next.task.attempts, - revision: evaluator.revision, - error: { - name: "CancelledError", - message: "Durable eval cancelled", - }, - }; - changed = true; - } - if (next.task.status === "succeeded") { - for (const definition of scorerDefinitions) { - const state = next.scores[definition.name]; - if (!isTerminal(state)) { - next.scores[definition.name] = { - status: "failed", - attempts: state?.attempts ?? 0, - revision: definition.revision, - error: { - name: "CancelledError", - message: "Durable eval cancelled", - }, - }; - changed = true; - } - } - } - if (!changed) return undefined; - next.logPending = true; - return next; - }); - } -} - -async function runTaskPass({ - store, - prefix, - projectName, - evalName, - evaluator, - shard, - workerId, - runId, - signal, -}: { - store: DurableEvalStore; - prefix: string; - projectName: string; - evalName: string; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - evaluator: DurableEvaluator; - shard: { index: number; count: number }; - workerId: string; - runId: string; - signal: AbortSignal; -}) { - const task = evaluator.task; - if (isBatchTask(task)) { - return await runBatchStage({ - store, - prefix, - projectName, - evalName, - processor: task, - stage: "task", - kind: "task", - shard, - workerId, - runId, - signal, - makeItem: (record) => taskBatchItem(record, evaluator.parameters ?? {}), - applyResult: (record, result, revision, attempt) => - applyTaskResult( - record, - result, - revision, - attempt, - task.maxAttempts ?? 3, - ), - }); - } - const localTask = task as EvalTask< - unknown, - unknown, - unknown, - Record, - EvalParameters - >; - - let changed = false; - for await (const current of listCases(store, prefix)) { - if ( - current.value.shard !== shard.index || - !isClaimable(current.value.task) - ) { - continue; - } - const claimed = structuredClone(current.value); - claimed.task = { - status: "leased", - attempts: current.value.task.attempts, - workerId, - leaseUntil: Date.now() + JOB_LEASE_MS, - }; - const claim = await writeJson( - store, - caseKey(prefix, current.value.id), - claimed, - { ifVersion: current.version }, - ); - if (!claim.written) continue; - const heartbeat = startCaseLeaseHeartbeat({ - store, - prefix, - id: current.value.id, - kind: "task", - workerId, - }); - const datum = current.value.datum as EvalCase< - unknown, - unknown, - BaseMetadata - >; - const attempt = current.value.task.attempts + 1; - let taskResult: - | { - status: "succeeded"; - output: JsonValue; - metadata: JsonValue; - tags?: string[]; - } - | { status: "failed"; error: unknown }; - try { - const metadata = { - ...(current.value.metadata as Record), - }; - const hooks: EvalHooks< - unknown, - Record, - EvalParameters - > = { - meta: (value) => { - Object.assign(metadata, value); - }, - metadata, - expected: "expected" in datum ? datum.expected : undefined, - span: NOOP_SPAN, - parameters: evaluator.parameters ?? {}, - reportProgress: () => undefined, - trialIndex: current.value.trialIndex, - tags: current.value.tags, - }; - const output = await awaitWithSignal( - Promise.resolve().then(() => localTask(datum.input, hooks)), - signal, - ); - taskResult = { - status: "succeeded", - output: assertJsonValue(output, "task output"), - metadata: assertJsonValue(hooks.metadata, "task metadata"), - tags: hooks.tags, - }; - } catch (error) { - taskResult = { status: "failed", error }; - } finally { - await heartbeat.stop(); - } - changed = - (await updateCaseRecord(store, prefix, current.value.id, (next) => { - if (next.task.status !== "leased" || next.task.workerId !== workerId) { - return undefined; - } - if (signal.aborted) { - next.task = { - status: "pending", - attempts: next.task.attempts, - }; - return next; - } - if (taskResult.status === "succeeded") { - next.metadata = taskResult.metadata; - next.tags = taskResult.tags; - next.task = { - status: "succeeded", - attempts: attempt, - revision: evaluator.revision, - value: taskResult.output, - }; - } else { - next.task = - attempt < 3 - ? { status: "pending", attempts: attempt } - : { - status: "failed", - attempts: attempt, - revision: evaluator.revision, - error: serializeError(taskResult.error), - }; - } - next.logPending = true; - return next; - })) || changed; - if (signal.aborted) return changed; - } - return changed; -} - -async function runScorePass({ - store, - prefix, - projectName, - evalName, - scorer, - shard, - workerId, - runId, - signal, -}: { - store: DurableEvalStore; - prefix: string; - projectName: string; - evalName: string; - scorer: ResolvedScorer; - shard: { index: number; count: number }; - workerId: string; - runId: string; - signal: AbortSignal; -}) { - if (isBatchScorer(scorer.scorer)) { - const batchScorer = scorer.scorer; - return await runBatchStage({ - store, - prefix, - projectName, - evalName, - processor: batchScorer, - stage: `score:${scorer.name}`, - kind: "score", - scorerName: scorer.name, - shard, - workerId, - runId, - signal, - eligible: (record) => record.task.status === "succeeded", - makeItem: (record) => scorerBatchItem(record), - applyResult: (record, result, revision, attempt) => - applyScoreResult( - record, - scorer.name, - result, - revision, - attempt, - batchScorer.maxAttempts ?? 3, - ), - }); - } - - const localScorer = scorer.scorer as EvalScorer< - unknown, - unknown, - unknown, - Record - >; - let changed = false; - for await (const current of listCases(store, prefix)) { - const state = current.value.scores[scorer.name]; - if ( - current.value.shard !== shard.index || - current.value.task.status !== "succeeded" || - !isClaimable(state) - ) { - continue; - } - const claimed = structuredClone(current.value); - claimed.scores[scorer.name] = { - status: "leased", - attempts: state.attempts, - workerId, - leaseUntil: Date.now() + JOB_LEASE_MS, - }; - const claim = await writeJson( - store, - caseKey(prefix, current.value.id), - claimed, - { ifVersion: current.version }, - ); - if (!claim.written) continue; - const heartbeat = startCaseLeaseHeartbeat({ - store, - prefix, - id: current.value.id, - kind: "score", - scorerName: scorer.name, - workerId, - }); - const datum = current.value.datum as EvalCase< - unknown, - unknown, - BaseMetadata - >; - const attempt = state.attempts + 1; - const taskOutput = current.value.task.value; - let scoreResult: - | { status: "succeeded"; value: JsonValue } - | { status: "failed"; error: unknown }; - try { - const raw = await awaitWithSignal( - Promise.resolve().then(() => - localScorer({ - input: datum.input, - expected: "expected" in datum ? datum.expected : undefined, - metadata: current.value.metadata as Record, - output: taskOutput, - }), - ), - signal, - ); - scoreResult = { - status: "succeeded", - value: assertJsonValue( - normalizeScores(raw, scorer.name), - `scorer ${scorer.name} output`, - ), - }; - } catch (error) { - scoreResult = { status: "failed", error }; - } finally { - await heartbeat.stop(); - } - changed = - (await updateCaseRecord(store, prefix, current.value.id, (next) => { - const latest = next.scores[scorer.name]; - if (latest?.status !== "leased" || latest.workerId !== workerId) { - return undefined; - } - if (signal.aborted) { - next.scores[scorer.name] = { - status: "pending", - attempts: latest.attempts, - }; - return next; - } - next.scores[scorer.name] = - scoreResult.status === "succeeded" - ? { - status: "succeeded", - attempts: attempt, - revision: scorer.revision, - value: scoreResult.value, - } - : attempt < 3 - ? { status: "pending", attempts: attempt } - : { - status: "failed", - attempts: attempt, - revision: scorer.revision, - error: serializeError(scoreResult.error), - }; - next.logPending = true; - return next; - })) || changed; - if (signal.aborted) return changed; - } - return changed; -} - -async function runBatchStage({ - store, - prefix, - projectName, - evalName, - processor, - stage, - kind, - scorerName, - shard, - workerId, - runId, - signal, - eligible = () => true, - makeItem, - applyResult, -}: { - store: DurableEvalStore; - prefix: string; - projectName: string; - evalName: string; - processor: DurableBatchProcessor; - stage: string; - kind: "task" | "score"; - scorerName?: string; - shard: { index: number; count: number }; - workerId: string; - runId: string; - signal: AbortSignal; - eligible?: (record: DurableCaseRecord) => boolean; - makeItem: (record: DurableCaseRecord) => Item; - applyResult: ( - record: DurableCaseRecord, - result: Result, - revision: string, - attempt: number, - ) => DurableCaseRecord; -}) { - const batchSize = processor.batchSize ?? DEFAULT_BATCH_SIZE; - const maxConcurrentBatches = - processor.maxConcurrentBatches ?? DEFAULT_MAX_CONCURRENT_BATCHES; - if (!Number.isInteger(batchSize) || batchSize < 1) { - throw new Error(`Invalid batchSize for durable stage ${stage}`); - } - if (!Number.isInteger(maxConcurrentBatches) || maxConcurrentBatches < 1) { - throw new Error(`Invalid maxConcurrentBatches for durable stage ${stage}`); - } - - let changed = false; - let activeJobs = 0; - for await (const current of listJobs(store, prefix, stage)) { - let job = current.value; - if (job.shard !== shard.index) { - continue; - } - if (job.status === "preparing") { - activeJobs++; - if ((job.leaseUntil ?? 0) > Date.now()) continue; - const takeover = { - ...job, - workerId, - leaseUntil: Date.now() + JOB_LEASE_MS, - }; - const takeoverResult = await writeJson(store, current.key, takeover, { - ifVersion: current.version, - }); - if (!takeoverResult.written) continue; - job = takeover; - const preparationError = new Error( - "Batch preparation lease expired before provider submission", - ); - const failed = await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId || latest.status !== "preparing") { - return undefined; - } - latest.status = "failed"; - latest.error = serializeError(preparationError); - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - if (failed?.written && failed.value.status === "failed") { - await retryBatchItems({ - store, - prefix, - job, - scorerName, - processor, - error: preparationError, - retryable: true, - }); - } - activeJobs--; - changed = true; - continue; - } - if (!["submitting", "submitted"].includes(job.status)) continue; - - const completion = processor.completion; - const fallback = - completion.mode === "webhook" ? completion.pollFallback : undefined; - const fallbackReady = - fallback !== undefined && - Date.now() >= (job.submittedAt ?? Date.now()) + fallback.afterMs; - const needsProviderCheck = - job.status === "submitting" || - job.outcome !== undefined || - completion.mode === "poll" || - fallbackReady; - - if (!needsProviderCheck) { - activeJobs++; - continue; - } - if ((job.leaseUntil ?? 0) > Date.now()) { - activeJobs++; - continue; - } - - const claimedJob = { - ...job, - workerId, - leaseUntil: Date.now() + JOB_LEASE_MS, - }; - const jobClaim = await writeJson(store, current.key, claimedJob, { - ifVersion: current.version, - }); - if (!jobClaim.written) continue; - job = claimedJob; - const context = batchContext(runId, stage, job, shard, signal); - let handle = job.handle as Handle | undefined; - if (!handle) { - let recovery: DurableBatchRecovery; - try { - recovery = processor.recover - ? await withJobLeaseHeartbeat({ - store, - key: current.key, - workerId, - signal, - operation: () => processor.recover!(context), - }) - : { status: "unknown" }; - } catch (error) { - await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId) return undefined; - latest.error = serializeError(error); - latest.nextPollAt = Date.now() + 10_000; - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - changed = true; - continue; - } - const latest = await readJson(store, current.key); - if (!latest || latest.value.workerId !== workerId) { - changed = true; - continue; - } - job = latest.value; - if (recovery.status === "found") { - handle = recovery.handle; - const external = externalBatchReference(processor, handle, context); - const recovered = await persistSubmittedJob({ - store, - jobKey: current.key, - workerId, - handle: assertJsonValue(handle, `${stage} batch handle`), - external, - submittedAt: job.submittedAt ?? Date.now(), - }); - job = recovered.value; - const locator = batchLocator({ - projectName, - evalName, - runId, - prefix, - jobKey: current.key, - job, - shardCount: shard.count, - }); - if (external) { - await registerExternalBatchLocator(store, external, locator); - } - await attachPendingWebhookEvents(store, current.key, job); - const latest = await readJson(store, current.key); - if (latest) { - job = latest.value; - } - changed = true; - } else if (recovery.status === "not_found") { - const error = new Error("Provider confirmed batch was not submitted"); - const failed = await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId) return undefined; - latest.status = "failed"; - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - if (failed?.written && failed.value.status === "failed") { - await retryBatchItems({ - store, - prefix, - job, - scorerName, - processor, - error, - retryable: true, - }); - } - changed = true; - continue; - } else { - const unknown = await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId) return undefined; - latest.status = "unknown"; - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - if (unknown?.written) { - await markBatchUnknown(store, prefix, job, scorerName); - } - changed = true; - continue; - } - } - - let poll: DurableBatchPoll | undefined = job.outcome; - if ( - job.nextPollAt && - job.nextPollAt > Date.now() && - (!fallbackReady || job.outcome !== undefined) - ) { - activeJobs++; - await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId) return undefined; - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - continue; - } - if (!poll) { - const poller = - completion.mode === "poll" - ? completion - : fallbackReady - ? fallback - : undefined; - if (!poller) { - activeJobs++; - await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId) return undefined; - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - continue; - } - try { - poll = await withJobLeaseHeartbeat({ - store, - key: current.key, - workerId, - signal, - operation: () => poller.poll(handle, context), - }); - const latest = await readJson(store, current.key); - if (!latest || latest.value.workerId !== workerId) { - changed = true; - continue; - } - job = latest.value; - delete job.error; - } catch (error) { - activeJobs++; - await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId) return undefined; - latest.nextPollAt = Date.now() + (poller.intervalMs ?? 10_000); - latest.error = serializeError(error); - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - continue; - } - } - if (poll.status === "pending") { - activeJobs++; - await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId) return undefined; - delete latest.error; - latest.nextPollAt = - Date.now() + - (poll.retryAfterMs ?? - (completion.mode === "poll" - ? completion.intervalMs - : fallback?.intervalMs) ?? - 10_000); - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - continue; - } - if (poll.status === "failed") { - const failed = await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId) return undefined; - latest.status = "failed"; - latest.error = serializeError(poll.error); - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - if (failed?.written && failed.value.status === "failed") { - await retryBatchItems({ - store, - prefix, - job, - scorerName, - processor, - error: poll.error, - retryable: poll.retryable ?? false, - }); - await markWebhookApplied(store, failed.value); - } - changed = true; - continue; - } - - const collectionHeartbeat = startJobLeaseHeartbeat( - store, - current.key, - workerId, - ); - try { - delete job.error; - const seen = new Set(); - const collected = await awaitWithSignal( - Promise.resolve().then(() => processor.collect(handle, context)), - signal, - ); - for await (const result of toAbortableAsyncIterable(collected, signal)) { - const resultId = resultItemId(result); - if (seen.has(resultId) || !job.itemIds.includes(resultId)) { - throw new Error( - `Batch stage ${stage} returned an unknown or duplicate item id ${resultId}`, - ); - } - seen.add(resultId); - await updateCaseRecord(store, prefix, resultId, (record) => { - if (!stageBelongsToJob(record, job.kind, job.id, scorerName)) { - return undefined; - } - const next = applyResult(record, result, job.revision, job.attempt); - next.logPending = true; - return next; - }); - } - for (const itemId of job.itemIds) { - if (seen.has(itemId)) continue; - const missing = { - id: itemId, - error: new Error(`Batch stage ${stage} returned no result`), - retryable: true, - } as Result; - await updateCaseRecord(store, prefix, itemId, (record) => { - if (!stageBelongsToJob(record, job.kind, job.id, scorerName)) { - return undefined; - } - const next = applyResult(record, missing, job.revision, job.attempt); - next.logPending = true; - return next; - }); - } - } catch (error) { - activeJobs++; - await updateJobRecord(store, current.key, (latest) => { - if (latest.workerId !== workerId) return undefined; - latest.error = serializeError(error); - latest.nextPollAt = Date.now() + 10_000; - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - changed = true; - continue; - } finally { - await collectionHeartbeat.stop(); - } - const completed = await updateJobRecord(store, current.key, (latest) => { - if ( - latest.id !== job.id || - latest.status !== "submitted" || - latest.workerId !== workerId - ) { - return undefined; - } - latest.status = "complete"; - delete latest.workerId; - delete latest.leaseUntil; - delete latest.error; - return latest; - }); - if (completed?.value.status !== "complete") continue; - job = completed.value; - await markWebhookApplied(store, completed.value); - changed = true; - } - - if (activeJobs >= maxConcurrentBatches) { - return changed; - } - const ready: Versioned[] = []; - let attempts: number | undefined; - for await (const record of listCases(store, prefix)) { - if (record.value.shard !== shard.index || !eligible(record.value)) continue; - const state = - kind === "task" ? record.value.task : record.value.scores[scorerName!]; - if (isClaimable(state)) { - attempts ??= state.attempts; - if (state.attempts !== attempts) continue; - ready.push(record); - if (ready.length >= batchSize) break; - } - } - if (!ready.length) return changed; - - const attempt = (attempts ?? 0) + 1; - const batchId = deterministicId( - `${projectName}:${evalName}:${runId}:${stage}:${ready - .map((record) => record.value.id) - .join(",")}:${attempt}`, - ); - let job: DurableJobRecord = { - id: batchId, - stage, - kind, - scorerName, - itemIds: ready.map((record) => record.value.id), - attempt, - revision: processor.revision, - shard: shard.index, - status: "preparing", - workerId, - leaseUntil: Date.now() + JOB_LEASE_MS, - }; - const jobKey = `${prefix}/jobs/${encodeURIComponent(stage)}/${batchId}`; - const created = await writeJson(store, jobKey, job, { ifAbsent: true }); - if (!created.written) return changed; - await registerBatchLocator( - store, - batchLocator({ - projectName, - evalName, - runId, - prefix, - jobKey, - job, - shardCount: shard.count, - }), - ); - - const preparationHeartbeat = startJobLeaseHeartbeat(store, jobKey, workerId); - const claimed: Versioned[] = []; - for (const record of ready) { - if (signal.aborted) break; - const next = structuredClone(record.value); - const state: StageState = { - status: "in_batch", - attempts: attempt, - batchId, - revision: processor.revision, - }; - if (kind === "task") next.task = state; - else next.scores[scorerName!] = state; - const claim = await writeJson( - store, - caseKey(prefix, record.value.id), - next, - { - ifVersion: record.version, - }, - ); - if (claim.written) claimed.push(record); - } - job.itemIds = claimed.map((record) => record.value.id); - if (signal.aborted) { - await preparationHeartbeat.stop(); - const failed = await updateJobRecord(store, jobKey, (latest) => { - if (latest.workerId !== workerId || latest.status !== "preparing") { - return undefined; - } - latest.itemIds = job.itemIds; - latest.status = "failed"; - latest.error = serializeError(new DurableEvalAbortError()); - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - if (failed?.written) { - await retryBatchItems({ - store, - prefix, - job, - scorerName, - processor, - error: new DurableEvalAbortError(), - retryable: true, - }); - } - return true; - } - if (!claimed.length) { - await preparationHeartbeat.stop(); - await updateJobRecord(store, jobKey, (latest) => { - if (latest.workerId !== workerId || latest.status !== "preparing") { - return undefined; - } - latest.status = "failed"; - latest.error = serializeError(new Error("Batch lost all item claims")); - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - return changed; - } - - const prepared = await updateJobRecord(store, jobKey, (latest) => { - if (latest.workerId !== workerId || latest.status !== "preparing") { - return undefined; - } - latest.itemIds = job.itemIds; - latest.status = "submitting"; - latest.leaseUntil = Date.now() + JOB_LEASE_MS; - return latest; - }); - await preparationHeartbeat.stop(); - if ( - !prepared || - !prepared.written || - prepared.value.status !== "submitting" || - prepared.value.workerId !== workerId - ) { - return changed; - } - job = prepared.value; - - const context = batchContext(runId, stage, job, shard, signal); - const items = claimed.map((record) => makeItem(record.value)); - try { - const handle = await withJobLeaseHeartbeat({ - store, - key: jobKey, - workerId, - signal, - operation: () => processor.submit(items, context), - }); - const external = externalBatchReference(processor, handle, context); - const submitted = await persistSubmittedJob({ - store, - jobKey, - workerId, - handle: assertJsonValue(handle, `${stage} batch handle`), - external, - submittedAt: Date.now(), - }); - job = submitted.value; - const locator = batchLocator({ - projectName, - evalName, - runId, - prefix, - jobKey, - job, - shardCount: shard.count, - }); - if (external) { - await registerExternalBatchLocator(store, external, locator); - } - await attachPendingWebhookEvents(store, jobKey, job); - } catch (error) { - const definitelyNotSubmitted = - error instanceof DurableEvalNotSubmittedError; - const terminalStatus = definitelyNotSubmitted ? "failed" : "unknown"; - const updated = await updateJobRecord(store, jobKey, (latest) => { - if ( - latest.workerId !== workerId || - latest.status === "complete" || - latest.status === "failed" || - latest.webhookEventKey - ) { - return undefined; - } - latest.status = terminalStatus; - latest.error = serializeError(error); - delete latest.workerId; - delete latest.leaseUntil; - return latest; - }); - if ( - updated?.written && - updated.value.status === "failed" && - definitelyNotSubmitted - ) { - await retryBatchItems({ - store, - prefix, - job, - scorerName, - processor, - error, - retryable: true, - }); - } else if ( - updated?.written && - updated.value.status === "unknown" && - !definitelyNotSubmitted - ) { - await markBatchUnknown(store, prefix, job, scorerName); - } - } - return true; -} - -function taskBatchItem( - record: DurableCaseRecord, - parameters: Record, -): DurableBatchTaskItem { - const datum = record.datum as EvalCase; - return { - id: record.id, - input: datum.input, - expected: "expected" in datum ? datum.expected : undefined, - metadata: record.metadata as BaseMetadata, - tags: record.tags, - parameters, - trialIndex: record.trialIndex, - }; -} - -function scorerBatchItem( - record: DurableCaseRecord, -): DurableBatchScorerItem> { - const datum = record.datum as EvalCase; - if (record.task.status !== "succeeded") { - throw new Error("Cannot score a task that has not succeeded"); - } - return { - id: record.id, - input: datum.input, - output: record.task.value, - expected: "expected" in datum ? datum.expected : undefined, - metadata: record.metadata as Record, - trialIndex: record.trialIndex, - }; -} - -function applyTaskResult( - record: DurableCaseRecord, - result: DurableBatchTaskResult, - revision: string, - attempt: number, - maxAttempts: number, -) { - if ("error" in result) { - record.task = - (result.retryable ?? false) && attempt < maxAttempts - ? { status: "pending", attempts: attempt } - : { - status: "failed", - attempts: attempt, - revision, - error: serializeError(result.error), - }; - } else { - record.task = { - status: "succeeded", - attempts: attempt, - revision, - value: assertJsonValue(result.output, "batch task output"), - }; - if (result.metadata !== undefined) { - record.metadata = assertJsonValue(result.metadata, "batch task metadata"); - } - if (result.tags !== undefined) record.tags = result.tags; - } - return record; -} - -function applyScoreResult( - record: DurableCaseRecord, - scorerName: string, - result: DurableBatchScorerResult, - revision: string, - attempt: number, - maxAttempts: number, -) { - if ("error" in result) { - record.scores[scorerName] = - (result.retryable ?? false) && attempt < maxAttempts - ? { status: "pending", attempts: attempt } - : { - status: "failed", - attempts: attempt, - revision, - error: serializeError(result.error), - }; - } else { - record.scores[scorerName] = { - status: "succeeded", - attempts: attempt, - revision, - value: assertJsonValue( - normalizeScores(result.score, scorerName), - `batch scorer ${scorerName} output`, - ), - }; +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + result: DurableBatchResult, + options: Omit, +): Promise { + if (!result.batchId && !result.externalId) { + throw new Error("Batch results require batchId or externalId"); } - return record; -} - -async function retryBatchItems({ - store, - prefix, - job, - scorerName, - processor, - error, - retryable, -}: { - store: DurableEvalStore; - prefix: string; - job: DurableJobRecord; - scorerName?: string; - processor: DurableBatchProcessor; - error: unknown; - retryable: boolean; -}) { - for (const itemId of job.itemIds) { - await updateCaseRecord(store, prefix, itemId, (record) => { - if (!stageBelongsToJob(record, job.kind, job.id, scorerName)) { - return undefined; - } - const state = - retryable && job.attempt < (processor.maxAttempts ?? 3) - ? ({ status: "pending", attempts: job.attempt } satisfies StageState) - : ({ - status: "failed", - attempts: job.attempt, - revision: job.revision, - error: serializeError(error), - } satisfies StageState); - if (job.kind === "task") record.task = state; - else record.scores[scorerName!] = state; - record.logPending = true; - return record; - }); + const store = definition.evaluator.store; + const locator = await locateBatch( + store, + definition.projectName, + definition.evalName, + result, + ); + if (!locator) { + throw new Error("No submitted batch matches this result"); } -} - -async function markBatchUnknown( - store: DurableEvalStore, - prefix: string, - job: DurableJobRecord, - scorerName?: string, -) { - for (const itemId of job.itemIds) { - await updateCaseRecord(store, prefix, itemId, (record) => { - if (!stageBelongsToJob(record, job.kind, job.id, scorerName)) { - return undefined; - } - const state: StageState = { - status: "unknown", - attempts: job.attempt, - revision: job.revision, - batchId: job.id, - }; - if (job.kind === "task") record.task = state; - else record.scores[scorerName!] = state; - return record; - }); + const state = await readJson(store, locator.runKey); + if (!state) + throw new Error(`DurableEval run for batch ${locator.batchId} is missing`); + const batch = state.batches.find( + (candidate) => candidate.id === locator.batchId, + ); + if (!batch) throw new Error(`Batch ${locator.batchId} is missing`); + if (batch.status !== "complete") { + await collectBatch(definition, state, batch); + batch.status = "complete"; + await writeJson(store, locator.runKey, state); } + return advanceDurableEval(definition, state, store, locator.runKey, options); } -function stageBelongsToJob( - record: DurableCaseRecord, - kind: "task" | "score", - batchId: string, - scorerName?: string, -) { - const state = kind === "task" ? record.task : record.scores[scorerName ?? ""]; - return ( - (state?.status === "in_batch" || state?.status === "unknown") && - state.batchId === batchId +async function pollDurableEval< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + options: DurableEvalRuntimeOptions & { runId: string }, +): Promise { + const store = definition.evaluator.store; + const key = runKey( + definition.projectName, + definition.evalName, + options.runId, ); -} - -function batchLocator({ - projectName, - evalName, - runId, - prefix, - jobKey, - job, - shardCount, -}: { - projectName: string; - evalName: string; - runId: string; - prefix: string; - jobKey: string; - job: DurableJobRecord; - shardCount: number; -}): DurableBatchLocator { - return { - schemaVersion: CHECKPOINT_VERSION, - projectName, - evalName, - runId, - prefix, - jobKey, - batchId: job.id, - stage: job.stage, - kind: job.kind, - scorerName: job.scorerName, - shard: job.shard, - shardCount, - }; -} - -async function registerBatchLocator( - store: DurableEvalStore, - locator: DurableBatchLocator, -) { - await registerLocator(store, internalBatchIndexKey(locator.batchId), locator); -} + const state = await readJson(store, key); + if (!state) throw new Error(`DurableEval run ${options.runId} is missing`); -async function registerExternalBatchLocator( - store: DurableEvalStore, - external: { source: string; id: string }, - locator: DurableBatchLocator, -) { - await registerLocator( - store, - externalBatchIndexKey(external.source, external.id), - locator, + const batches = state.batches.filter((batch) => { + if (batch.status === "complete") return false; + return processorForBatch(definition, batch).completion.mode === "poll"; + }); + const results = await Promise.all( + batches.map(async (batch) => ({ + batch, + result: await ( + processorForBatch(definition, batch).completion as Extract< + DurableBatchCompletion, + { mode: "poll" } + > + ).poll(batch.handle, { + runId: state.runId, + batchId: batch.id, + }), + })), ); + for (const { batch, result } of results) { + if (result.status === "failed") throw asError(result.error); + if (result.status !== "complete") continue; + await collectBatch(definition, state, batch); + batch.status = "complete"; + await writeJson(store, key, state); + } + return advanceDurableEval(definition, state, store, key, options); } -async function registerLocator( +async function advanceDurableEval< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + state: DurableRunState, store: DurableEvalStore, key: string, - locator: DurableBatchLocator, -) { - const inserted = await writeJson(store, key, locator, { ifAbsent: true }); - if (inserted.written) return; - const existing = await readJson(store, key); - if (!existing || existing.value.jobKey !== locator.jobKey) { - throw new Error( - `Durable batch index collision for batch ${locator.batchId}`, - ); + options: Omit, +): Promise { + if (state.status === "completed") return currentStatus(definition, state); + await runTaskStage(definition, state, store, key); + if (state.cases.some((record) => !record.taskComplete)) { + return currentStatus(definition, state); } -} -function externalBatchReference( - processor: DurableBatchProcessor, - handle: Handle, - context: DurableBatchContext, -) { - if (processor.completion.mode !== "webhook") return undefined; - if (!processor.completion.source.trim()) { - throw new Error("Durable batch webhook source must be non-empty"); - } - const id = processor.completion.externalId(handle, context); - if (!id.trim()) { - throw new Error("Durable batch webhook externalId must be non-empty"); + await runScoreStages(definition, state, store, key); + const scorerNames = resolveScorers(definition.evaluator.scores).map( + ({ name }) => name, + ); + if ( + state.cases.some((record) => + scorerNames.some((name) => !(name in record.scores)), + ) + ) { + return currentStatus(definition, state); } - return { source: processor.completion.source, id }; + + state.summary = await finishExperiment(definition, state, options.noSendLogs); + state.status = "completed"; + await writeJson(store, key, state); + return currentStatus(definition, state); } -async function persistSubmittedJob({ - store, - jobKey, - workerId, - handle, - external, - submittedAt, -}: { - store: DurableEvalStore; - jobKey: string; - workerId: string; - handle: JsonValue; - external?: { source: string; id: string }; - submittedAt: number; -}): Promise> { - while (true) { - const current = await readJson(store, jobKey); - if (!current) { - throw new Error("Durable batch job disappeared during submission"); - } - if (current.value.workerId !== workerId) { - throw new Error("Durable batch submission lease was lost"); - } - if ( - current.value.handle !== undefined && - stableStringify(current.value.handle) !== stableStringify(handle) - ) { - throw new Error( - `Durable batch ${current.value.id} recovered with a different handle`, - ); - } - if ( - current.value.external && - external && - (current.value.external.source !== external.source || - current.value.external.id !== external.id) - ) { - throw new Error( - `Durable batch ${current.value.id} resolved to a different external job`, - ); - } - if ( - current.value.status === "complete" || - current.value.status === "failed" - ) { - return current; +function currentStatus( + definition: DurableEvalDefinition, + state: DurableRunState, +): DurableEvalResult { + if (state.status === "completed") { + if (!state.summary) { + throw new Error(`DurableEval run ${state.runId} has no saved summary`); } - const next: DurableJobRecord = { - ...current.value, - status: "submitted", - handle, - external: external ?? current.value.external, - submittedAt: current.value.submittedAt ?? submittedAt, - nextPollAt: Date.now(), + return { + status: "completed", + runId: state.runId, + pending: { poll: 0, webhook: 0 }, + summary: state.summary, }; - delete next.workerId; - delete next.leaseUntil; - const written = await writeJson(store, jobKey, next, { - ifVersion: current.version, - }); - if (written.written) { - return { value: next, version: written.version }; - } } + const pending = { poll: 0, webhook: 0 }; + for (const batch of state.batches) { + if (batch.status === "complete") continue; + pending[processorForBatch(definition, batch).completion.mode]++; + } + return { status: "waiting", runId: state.runId, pending }; } -function batchProcessorForLocator< +async function runTaskStage< Input, Output, Expected, Metadata extends BaseMetadata, Parameters extends EvalParameters, >( - evaluator: DurableEvaluator, - locator: DurableBatchLocator, -): DurableBatchProcessor | undefined { - if (locator.kind === "task") { - return isBatchTask(evaluator.task) - ? (evaluator.task as DurableBatchProcessor) - : undefined; + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + state: DurableRunState, + store: DurableEvalStore, + key: string, +) { + if (isBatchTask(definition.evaluator.task)) { + await ensureWorkflowBatches(definition, state, store, key, "task"); + return; + } + + const task = definition.evaluator.task as EvalTask< + Input, + Output, + Expected, + Metadata, + Parameters + >; + for (const record of state.cases) { + if (record.taskComplete) continue; + const datum = record.datum as EvalCase; + const metadata = { ...(record.metadata as Record) }; + const hooks: EvalHooks = { + meta(value) { + Object.assign(metadata, value); + }, + metadata: metadata as EvalHooks< + Expected, + Metadata, + Parameters + >["metadata"], + expected: ("expected" in datum ? datum.expected : undefined) as Expected, + span: NOOP_SPAN, + parameters: (definition.evaluator.parameters ?? + {}) as InferParameters, + reportProgress: () => undefined, + trialIndex: record.trialIndex, + tags: record.tags, + }; + record.output = assertJsonValue( + await task(datum.input, hooks), + `task output for ${record.caseId}`, + ); + record.metadata = assertJsonValue(hooks.metadata, "task metadata"); + record.tags = hooks.tags; + record.taskComplete = true; + await writeJson(store, key, state); } - const scorer = evaluator.scores.find( - (candidate) => - isBatchScorer(candidate) && candidate.name === locator.scorerName, - ); - return scorer as - | DurableBatchProcessor - | undefined; } -async function attachWebhookEventToJob( +async function runScoreStages< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + state: DurableRunState, store: DurableEvalStore, - jobKey: string, - storedEventKey: string, - event: DurableBatchResultEvent, -): Promise<"attached" | "duplicate" | "in_progress" | "missing"> { - while (true) { - const current = await readJson(store, jobKey); - if (!current) return "missing"; - if ( - current.value.status === "complete" || - current.value.status === "failed" - ) { - return "duplicate"; - } - if (current.value.webhookEventKey === storedEventKey) { - if ( - current.value.workerId && - (current.value.leaseUntil ?? 0) > Date.now() - ) { - return "in_progress"; - } - const next = { - ...current.value, - nextPollAt: Date.now(), - }; - delete next.error; - delete next.workerId; - delete next.leaseUntil; - const written = await writeJson(store, jobKey, next, { - ifVersion: current.version, - }); - if (written.written) return "attached"; + key: string, +) { + const scorers = resolveScorers(definition.evaluator.scores); + for (const { name, scorer } of scorers) { + if (isBatchScorer(scorer)) { + await ensureWorkflowBatches(definition, state, store, key, "score", name); continue; } - if (current.value.webhookEventKey) return "duplicate"; - if ( - event.handle !== undefined && - current.value.handle !== undefined && - stableStringify(event.handle) !== stableStringify(current.value.handle) - ) { - throw new Error( - `Webhook handle does not match durable batch ${current.value.id}`, + for (const record of state.cases) { + if (name in record.scores) continue; + const datum = record.datum as EvalCase; + const value = await ( + scorer as EvalScorer + )({ + input: datum.input, + ...(datum.tags ? { tags: datum.tags } : {}), + ...(datum.id ? { id: datum.id } : {}), + ...(datum.upsert_id ? { upsert_id: datum.upsert_id } : {}), + ...(datum.trialCount ? { trialCount: datum.trialCount } : {}), + ...("expected" in datum ? { expected: datum.expected } : {}), + metadata: record.metadata as Metadata, + output: record.output as Output, + } as unknown as EvalScorerArgs); + record.scores[name] = assertJsonValue( + normalizeScores(value, name), + `scorer ${name} output`, ); + await writeJson(store, key, state); } - if ( - event.externalId && - current.value.external && - (current.value.external.source !== event.source || - current.value.external.id !== event.externalId) - ) { + } +} + +async function ensureWorkflowBatches( + definition: DurableEvalDefinition, + state: DurableRunState, + store: DurableEvalStore, + key: string, + kind: "task" | "score", + scorerName?: string, +) { + const workflow = workflowForStage(definition, kind, scorerName); + for (const node of workflow.nodes) { + const batchSize = node.processor.batchSize ?? DEFAULT_BATCH_SIZE; + if (!Number.isInteger(batchSize) || batchSize < 1) { throw new Error( - `Webhook externalId does not match durable batch ${current.value.id}`, + `Invalid batchSize for ${scorerName ? `${scorerName}.${node.name}` : node.name}`, ); } - const handle = event.handle ?? current.value.handle; - const next: DurableJobRecord = { - ...current.value, - handle, - external: event.externalId - ? { source: event.source, id: event.externalId } - : current.value.external, - outcome: event.outcome, - webhookEventKey: storedEventKey, - nextPollAt: Date.now(), - ...(handle !== undefined - ? { - status: "submitted", - submittedAt: current.value.submittedAt ?? Date.now(), - } - : {}), - }; - delete next.workerId; - delete next.leaseUntil; - const written = await writeJson(store, jobKey, next, { - ifVersion: current.version, + const assigned = new Set( + state.batches + .filter( + (batch) => + batch.kind === kind && + batch.scorerName === scorerName && + batch.nodeName === node.name, + ) + .flatMap((batch) => batch.itemIds), + ); + const eligible = state.cases.filter((record) => { + const outputs = nodeOutputsFor(record, kind, scorerName); + return ( + !assigned.has(record.id) && + !(node.name in outputs) && + Object.values(node.needs).every((dependency) => dependency in outputs) + ); }); - if (written.written) return "attached"; + for (let offset = 0; offset < eligible.length; offset += batchSize) { + const records = eligible.slice(offset, offset + batchSize); + const batchId = newId(); + const context = { runId: state.runId, batchId }; + const items = records.map((record) => + itemForNode(definition, record, kind, scorerName, node), + ); + const handle = assertJsonValue( + await node.processor.submit(items, context), + `handle for batch ${batchId}`, + ); + const externalId = + node.processor.completion.mode === "webhook" + ? node.processor.completion.externalId(handle, context) + : undefined; + if (externalId !== undefined && !externalId.trim()) { + throw new Error(`Batch ${batchId} produced an empty externalId`); + } + const batch: DurableBatchRecord = { + id: batchId, + kind, + scorerName, + nodeName: node.name, + itemIds: records.map((record) => record.id), + handle, + externalId, + status: "submitted", + }; + state.batches.push(batch); + await writeJson(store, key, state); + await indexBatch(store, definition, batch, key); + } } } -async function attachPendingWebhookEvents( - store: DurableEvalStore, - jobKey: string, - job: DurableJobRecord, +async function collectBatch( + definition: DurableEvalDefinition, + state: DurableRunState, + batch: DurableBatchRecord, ) { - if (!job.external) return; - for await (const pointerKey of store.list( - webhookMailboxPrefix(job.external.source, job.external.id), - )) { - const pointer = await readJson<{ eventKey: string }>(store, pointerKey); - if (!pointer) continue; - const event = await readJson( - store, - pointer.value.eventKey, + const workflow = workflowForStage(definition, batch.kind, batch.scorerName); + const node = workflow.nodes.find( + (candidate) => candidate.name === batch.nodeName, + ); + if (!node) + throw new Error( + `Definition no longer contains batch node ${batch.nodeName}`, ); - if (!event || event.value.status === "applied") continue; - const attached = await attachWebhookEventToJob( - store, - jobKey, - pointer.value.eventKey, - event.value.event, + const processor = node.processor; + const context = { runId: state.runId, batchId: batch.id }; + const results = await processor.collect(batch.handle, context); + if (!Array.isArray(results)) { + throw new Error(`collect for batch ${batch.id} must return an array`); + } + const expectedIds = new Set(batch.itemIds); + const seen = new Set(); + 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); + if ("error" in result) throw asError(result.error); + const record = state.cases.find((candidate) => candidate.id === id)!; + const output = assertJsonValue( + node.result(result), + `output for ${batch.nodeName} item ${id}`, ); - if (attached === "duplicate") { - await markWebhookEventApplied(store, pointer.value.eventKey, job.id); + nodeOutputsFor(record, batch.kind, batch.scorerName)[batch.nodeName] = + output; + if (batch.nodeName !== workflow.outputNode) continue; + if (batch.kind === "task") { + record.output = output; + if ("metadata" in result && result.metadata !== undefined) { + record.metadata = assertJsonValue( + 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( + normalizeScores(output as OneOrMoreScores, batch.scorerName!), + `score for ${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(", ")}`, + ); + } } -async function markWebhookApplied( - store: DurableEvalStore, - job: DurableJobRecord, -) { - if (job.webhookEventKey) { - await markWebhookEventApplied(store, job.webhookEventKey, job.id); +function processorForBatch( + definition: DurableEvalDefinition, + batch: DurableBatchRecord, +): DurableBatchProcessor { + const node = workflowForStage( + definition, + batch.kind, + batch.scorerName, + ).nodes.find((candidate) => candidate.name === batch.nodeName); + if (!node) + throw new Error( + `Definition no longer contains batch node ${batch.nodeName}`, + ); + return node.processor; +} + +function workflowForStage( + definition: DurableEvalDefinition, + kind: "task" | "score", + scorerName?: string, +): DurableWorkflowDefinition { + let stage: + | DurableBatchTask + | DurableBatchScorer; + if (kind === "task") { + if (!isBatchTask(definition.evaluator.task)) { + throw new Error("Definition no longer contains the batch task"); + } + stage = definition.evaluator.task; + } else { + const scorer = resolveScorers(definition.evaluator.scores).find( + ({ name }) => name === scorerName, + )?.scorer; + if (!isBatchScorer(scorer)) { + throw new Error(`Definition no longer contains scorer ${scorerName}`); + } + stage = scorer; + } + if (stage.workflow) return stage.workflow; + if (!stage.processor) { + throw new Error("Batch definition has neither a processor nor a workflow"); } + return { + outputNode: "$batch", + nodes: [ + { + name: "$batch", + needs: {}, + item: (rootItem) => rootItem, + processor: stage.processor as DurableBatchProcessor< + any, + any, + JsonValue + >, + result: + kind === "task" + ? (result) => result.output + : (result) => result.score, + }, + ], + }; } -async function markWebhookEventApplied( - store: DurableEvalStore, - eventKey: string, - batchId: string, +function itemForNode( + definition: DurableEvalDefinition, + record: DurableCaseRecord, + kind: "task" | "score", + scorerName: string | undefined, + node: DurableWorkflowNodeDefinition, ) { - while (true) { - const current = await readJson(store, eventKey); - if (!current || current.value.status === "applied") return; - const next: DurableWebhookEventRecord = { - ...current.value, - status: "applied", - batchId, - }; - const written = await writeJson(store, eventKey, next, { - ifVersion: current.version, - }); - if (written.written) return; - } + const rootItem = + kind === "task" + ? taskBatchItem(record, definition.evaluator.parameters ?? {}) + : scorerBatchItem(record); + const outputs = nodeOutputsFor(record, kind, scorerName); + return node.item( + rootItem, + Object.fromEntries( + Object.entries(node.needs).map(([alias, dependency]) => [ + alias, + outputs[dependency], + ]), + ), + record.id, + ); } -async function hasWaitingWebhookJob( - store: DurableEvalStore, - prefix: string, - shard: number, - webhookStages: Set, +function nodeOutputsFor( + record: DurableCaseRecord, + kind: "task" | "score", + scorerName?: string, ) { - for await (const key of store.list(`${prefix}/jobs/`)) { - const job = await readJson(store, key); - if ( - job?.value.shard === shard && - (job.value.status === "submitting" || job.value.status === "submitted") && - job.value.outcome === undefined && - webhookStages.has(job.value.stage) - ) { - return true; + if (kind === "task") return record.taskNodeOutputs; + return (record.scoreNodeOutputs[scorerName!] ??= {}); +} + +async function materializeCases( + evaluator: DurableEvaluator, +): Promise { + const raw = + typeof evaluator.data === "function" ? evaluator.data() : evaluator.data; + const data = raw instanceof Promise ? await raw : raw; + if (typeof data === "object" && data !== null && "_type" in data) { + throw new Error("DurableEval does not support BaseExperiment data sources"); + } + if (!isIterable(data) && !isAsyncIterable(data)) { + throw new Error("DurableEval data must be iterable"); + } + const records: DurableCaseRecord[] = []; + const seen = new Set(); + for await (const datum of toAsyncIterable(data)) { + const caseId = + datum.id ?? + datum.upsert_id ?? + (evaluator.caseId ? await evaluator.caseId(datum) : undefined); + if (!caseId) { + throw new Error( + "Every DurableEval case requires id, upsert_id, or caseId", + ); + } + if (seen.has(caseId)) + throw new Error(`Duplicate DurableEval 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 DurableEval case ${caseId}`); + } + for (let trialIndex = 0; trialIndex < trialCount; trialIndex++) { + records.push({ + id: `${caseId}:trial:${trialIndex}`, + caseId, + trialIndex, + datum: assertJsonValue(datum, `case ${caseId}`), + metadata: assertJsonValue( + "metadata" in datum ? datum.metadata : {}, + `metadata for ${caseId}`, + ), + tags: datum.tags, + taskComplete: false, + taskNodeOutputs: {}, + scores: {}, + scoreNodeOutputs: {}, + }); } } - return false; + return records; +} + +function taskBatchItem(record: DurableCaseRecord, parameters: JsonValue) { + const datum = record.datum as EvalCase; + return { + id: record.id, + input: datum.input, + expected: "expected" in datum ? datum.expected : undefined, + metadata: record.metadata, + tags: record.tags, + parameters, + trialIndex: record.trialIndex, + }; +} + +function scorerBatchItem(record: DurableCaseRecord) { + const datum = record.datum as EvalCase; + return { + id: record.id, + input: datum.input, + output: record.output, + expected: "expected" in datum ? datum.expected : undefined, + metadata: record.metadata, + tags: record.tags, + trialIndex: record.trialIndex, + }; } -async function hasProviderErrorJob( - store: DurableEvalStore, - prefix: string, - shard: number, +async function finishExperiment( + definition: DurableEvalDefinition, + state: DurableRunState, + noSendLogs = false, ) { - for await (const key of store.list(`${prefix}/jobs/`)) { - const job = await readJson(store, key); - if ( - job?.value.shard === shard && - (job.value.status === "submitting" || job.value.status === "submitted") && - job.value.error !== undefined - ) { - return true; - } - } - return false; -} - -async function flushPendingLogs({ - store, - prefix, - experiment, - scorerNames, - shard, - runId, -}: { - store: DurableEvalStore; - prefix: string; - experiment: Experiment; - scorerNames: string[]; - shard: { index: number; count: number }; - runId: string; -}) { - let changed = false; - for await (const current of listCases(store, prefix)) { - if (current.value.shard !== shard.index || !current.value.logPending) { - continue; - } - await logDurableCase(experiment, current.value, scorerNames, runId); - await experiment.flush(); - const next = structuredClone(current.value); - next.logPending = false; - delete next.removedScores; - const result = await writeJson( - store, - caseKey(prefix, current.value.id), - next, - { ifVersion: current.version }, + const scorerNames = resolveScorers(definition.evaluator.scores).map( + ({ name }) => name, + ); + if (noSendLogs) { + return buildLocalSummary( + state, + definition.projectName, + state.experimentName, + scorerNames, ); - changed = result.written || changed; } - return changed; + const experiment = initExperiment({ + state: definition.evaluator.state, + ...(definition.evaluator.projectId + ? { projectId: definition.evaluator.projectId } + : { project: definition.projectName }), + experiment: state.experimentName, + update: true, + description: definition.evaluator.description, + metadata: definition.evaluator.metadata, + tags: definition.evaluator.tags, + setCurrent: false, + }); + for (const record of state.cases) { + logCase(experiment, state.runId, record, scorerNames); + } + await experiment.flush(); + return await experiment.summarize(); } -async function logDurableCase( +function logCase( experiment: Experiment, + runId: string, record: DurableCaseRecord, scorerNames: string[], - runId: string, ) { const datum = record.datum as EvalCase; - const rootSpanId = deterministicId(`${runId}:${record.id}:root`); - const root = _internalStartSpanWithInitialMerge({ - parent: await experiment.export(), - spanId: rootSpanId, + const scores = Object.assign( + {}, + ...scorerNames.map((name) => record.scores[name] ?? {}), + ) as Record; + const span = experiment.startSpan({ name: "eval", + spanId: deterministicId(`${runId}:${record.id}:span`), spanAttributes: { type: SpanTypeAttribute.EVAL }, event: { id: deterministicId(`${runId}:${record.id}:row`), input: datum.input, expected: "expected" in datum ? datum.expected : undefined, + output: record.output, + scores, metadata: { ...(record.metadata as Record), durable_eval: { @@ -3389,178 +1299,28 @@ async function logDurableCase( }, }, tags: record.tags, - ...(record.task.status === "succeeded" - ? { output: record.task.value } - : record.task.status === "failed" - ? { error: record.task.error.message } - : {}), - scores: collectedScores(record, scorerNames), }, - state: evaluatorState(experiment), }); - const parent = await root.export(); - if (record.task.status === "succeeded" || record.task.status === "failed") { - const task = _internalStartSpanWithInitialMerge({ - parent, - spanId: deterministicId(`${runId}:${record.id}:task-span`), - name: "task", - spanAttributes: { type: SpanTypeAttribute.TASK }, - event: { - id: deterministicId(`${runId}:${record.id}:task-row`), - input: datum.input, - metadata: { - durable_eval: { - revision: record.task.revision, - attempts: record.task.attempts, - }, - }, - ...(record.task.status === "succeeded" - ? { output: record.task.value } - : { error: record.task.error.message }), - }, - state: evaluatorState(experiment), - }); - task.end(); - } - for (const scorerName of scorerNames) { - const state = record.scores[scorerName]; - if (state?.status !== "succeeded" && state?.status !== "failed") continue; - const scorer = _internalStartSpanWithInitialMerge({ - parent, - spanId: deterministicId(`${runId}:${record.id}:score:${scorerName}:span`), - name: scorerName, - spanAttributes: { type: SpanTypeAttribute.SCORE }, - event: { - id: deterministicId(`${runId}:${record.id}:score:${scorerName}:row`), - input: { - input: datum.input, - output: - record.task.status === "succeeded" ? record.task.value : undefined, - expected: "expected" in datum ? datum.expected : undefined, - }, - metadata: { - durable_eval: { - revision: state.revision, - attempts: state.attempts, - }, - }, - ...(state.status === "succeeded" - ? { - output: state.value, - scores: state.value as Record, - } - : { error: state.error.message }), - }, - state: evaluatorState(experiment), - }); - scorer.end(); - } - root.end(); -} - -function evaluatorState(experiment: Experiment) { - return experiment.loggingState; -} - -function collectedScores(record: DurableCaseRecord, scorerNames: string[]) { - const scores: Record = {}; - for (const name of record.removedScores ?? []) { - scores[name] = null; - } - for (const name of scorerNames) { - const state = record.scores[name]; - if (state?.status === "succeeded") { - Object.assign(scores, state.value); - } - } - return scores; -} - -async function collectProgress( - store: DurableEvalStore, - prefix: string, - scorerNames: string[], -): Promise { - const progress: DurableEvalProgress = { - ...emptyProgress(), - }; - for await (const record of listCases(store, prefix)) { - progress.total++; - countStage(record.value.task, progress, "task"); - if (record.value.task.status === "succeeded") { - for (const name of scorerNames) { - countStage(record.value.scores[name], progress, "score"); - } - } - } - return progress; -} - -function countStage( - state: StageState | undefined, - progress: DurableEvalProgress, - kind: "task" | "score", -) { - if (!state || ["pending", "leased", "in_batch"].includes(state.status)) { - if (kind === "task") progress.taskPending++; - else progress.scorePending++; - } else if (state.status === "succeeded") { - if (kind === "task") progress.taskSucceeded++; - else progress.scoreSucceeded++; - } else if (state.status === "failed") { - if (kind === "task") progress.taskFailed++; - else progress.scoreFailed++; - } else if (state.status === "unknown") { - progress.unknown++; - } -} - -async function isComplete( - store: DurableEvalStore, - prefix: string, - scorerNames: string[], - shard?: number, -) { - for await (const record of listCases(store, prefix)) { - if (shard !== undefined && record.value.shard !== shard) continue; - if (!isTerminal(record.value.task)) return false; - if (record.value.task.status === "succeeded") { - for (const name of scorerNames) { - if (!isTerminal(record.value.scores[name])) return false; - } - } - } - return true; -} - -function isTerminal(state: StageState | undefined) { - return state?.status === "succeeded" || state?.status === "failed"; -} - -function isClaimable(state: StageState | undefined) { - return ( - state?.status === "pending" || - (state?.status === "leased" && state.leaseUntil <= Date.now()) - ); + span.end(); } -async function buildLocalDurableSummary( - store: DurableEvalStore, - prefix: string, +function buildLocalSummary( + state: DurableRunState, projectName: string, experimentName: string, scorerNames: string[], -): Promise { +): ExperimentSummary { const totals: Record = {}; - for await (const record of listCases(store, prefix)) { - for (const [name, score] of Object.entries( - collectedScores(record.value, scorerNames), - )) { - if (score === null) continue; - const current = totals[name] ?? { total: 0, count: 0 }; - current.total += score; - current.count++; - totals[name] = current; + for (const record of state.cases) { + for (const scorerName of scorerNames) { + const scores = record.scores[scorerName] as Record; + for (const [name, score] of Object.entries(scores)) { + if (score === null) continue; + const total = totals[name] ?? { total: 0, count: 0 }; + total.total += score; + total.count++; + totals[name] = total; + } } } return { @@ -3580,10 +1340,21 @@ async function buildLocalDurableSummary( }; } -function normalizeScores( - value: OneOrMoreScores, - defaultName: string, -): Record { +function resolveScorers( + scorers: Array< + | EvalScorer + | DurableBatchScorer + >, +) { + return scorers.map((scorer, index) => ({ + name: isBatchScorer(scorer) + ? scorer.name + : scorer.name || `scorer_${index}`, + scorer, + })); +} + +function normalizeScores(value: OneOrMoreScores, defaultName: string) { if (value === null) return { [defaultName]: null }; if (typeof value === "number") return { [defaultName]: value }; const values = Array.isArray(value) ? value : [value]; @@ -3592,146 +1363,77 @@ function normalizeScores( ); } -function batchContext( - runId: string, - stage: string, - job: DurableJobRecord, - shard: { index: number; count: number }, - signal: AbortSignal, -): DurableBatchContext { - return { - runId, - stage, - revision: job.revision, - shard, - attempt: job.attempt, - batchId: job.id, - itemCount: job.itemIds.length, - signal, - }; -} - -function resultItemId(value: unknown): string { - 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( - value: unknown, -): value is DurableBatchTask< - unknown, - unknown, - unknown, - BaseMetadata, - EvalParameters, - JsonValue -> { - return ( - typeof value === "object" && - value !== null && - "kind" in value && - value.kind === BATCH_TASK_KIND - ); -} - -function isBatchScorer( - value: unknown, -): value is DurableBatchScorer< - unknown, - unknown, - unknown, - BaseMetadata, - JsonValue -> { - return ( - typeof value === "object" && - value !== null && - "kind" in value && - value.kind === BATCH_SCORER_KIND - ); -} - -async function* listCases( +async function indexBatch( store: DurableEvalStore, - prefix: string, -): AsyncGenerator> { - for await (const key of store.list(`${prefix}/cases/`)) { - const record = await readJson(store, key); - if (record) yield record; + definition: DurableEvalDefinition, + batch: DurableBatchRecord, + key: string, +) { + const locator = { runKey: key, batchId: batch.id }; + await writeJson( + store, + batchIndexKey( + definition.projectName, + definition.evalName, + "batch", + batch.id, + ), + locator, + ); + if (batch.externalId) { + await writeJson( + store, + batchIndexKey( + definition.projectName, + definition.evalName, + "external", + batch.externalId, + ), + locator, + ); } } -async function* listJobs( +async function locateBatch( store: DurableEvalStore, - prefix: string, - stage: string, -): AsyncGenerator & { key: string }> { - for await (const key of store.list( - `${prefix}/jobs/${encodeURIComponent(stage)}/`, - )) { - const record = await readJson(store, key); - if (record) yield { ...record, key }; + projectName: string, + evalName: string, + result: DurableBatchResult, +) { + const byBatch = result.batchId + ? await readJson( + store, + batchIndexKey(projectName, evalName, "batch", result.batchId), + ) + : undefined; + const byExternal = result.externalId + ? await readJson( + store, + batchIndexKey(projectName, evalName, "external", result.externalId), + ) + : undefined; + if ( + byBatch && + byExternal && + (byBatch.runKey !== byExternal.runKey || + byBatch.batchId !== byExternal.batchId) + ) { + throw new Error("batchId and externalId identify different batches"); } + return byBatch ?? byExternal; } -function runPrefix(projectName: string, evalName: string, runId: string) { - return `durable-eval/v1/${contentVersion( - encoder.encode(`${projectName}\0${evalName}\0${runId}`), - )}`; +function runKey(projectName: string, evalName: string, runId: string) { + return `durable-eval/v2/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; } -function internalBatchIndexKey(batchId: string) { - return `durable-eval/v1/indices/batches/${contentVersion( - encoder.encode(batchId), - )}`; -} - -function externalBatchIndexKey(source: string, externalId: string) { - return `durable-eval/v1/indices/external/${contentVersion( - encoder.encode(`${source}\0${externalId}`), - )}`; -} - -function webhookEventKey(source: string, eventId: string) { - return `durable-eval/v1/webhooks/events/${contentVersion( - encoder.encode(`${source}\0${eventId}`), - )}`; -} - -function webhookMailboxPrefix(source: string, externalId: string) { - return `durable-eval/v1/webhooks/mailboxes/${contentVersion( - encoder.encode(`${source}\0${externalId}`), - )}/`; -} - -function webhookMailboxKey( - source: string, - externalId: string, - eventKey: string, +function batchIndexKey( + projectName: string, + evalName: string, + type: "batch" | "external", + id: string, ) { - return `${webhookMailboxPrefix(source, externalId)}${contentVersion( - encoder.encode(eventKey), - )}`; -} - -function caseKey(prefix: string, id: string) { - return `${prefix}/cases/${encodeURIComponent(id)}`; -} - -function workItemId(caseId: string, trialIndex: number) { - return `${caseId}:trial:${trialIndex}`; -} - -function stableShard(id: string, count: number) { - const hash = contentVersion(encoder.encode(id)); - return Number.parseInt(hash.slice(0, 8), 16) % count; + return `durable-eval/v2/index/${contentVersion(encoder.encode(`${projectName}\0${evalName}`))}/${type}/${contentVersion(encoder.encode(id))}`; } function deterministicId(value: string) { @@ -3751,192 +1453,16 @@ function contentVersion(value: Uint8Array) { return (hash >>> 0).toString(16).padStart(8, "0"); } -async function readJson( - store: DurableEvalStore, - key: string, -): Promise | undefined> { - const record = await store.read(key); - if (!record) return undefined; - return { - value: JSON.parse(decoder.decode(record.value)) as T, - version: record.version, - }; -} - -async function writeJson( - store: DurableEvalStore, - key: string, - value: T, - condition: DurableEvalWriteCondition, -) { - return await store.write( - key, - encoder.encode(stableStringify(value)), - condition, - ); -} - -async function updateCaseRecord( - store: DurableEvalStore, - prefix: string, - id: string, - update: (record: DurableCaseRecord) => DurableCaseRecord | undefined, -) { - while (true) { - const current = await readJson( - store, - caseKey(prefix, id), - ); - if (!current) return false; - const next = update(structuredClone(current.value)); - if (!next) return false; - const written = await writeJson(store, caseKey(prefix, id), next, { - ifVersion: current.version, - }); - if (written.written) return true; - } -} - -async function updateJobRecord( - store: DurableEvalStore, - key: string, - update: (record: DurableJobRecord) => DurableJobRecord | undefined, -) { - while (true) { - const current = await readJson(store, key); - if (!current) return undefined; - const next = update(structuredClone(current.value)); - if (!next) return { ...current, written: false as const }; - const written = await writeJson(store, key, next, { - ifVersion: current.version, - }); - if (written.written) { - return { value: next, version: written.version, written: true as const }; - } - } -} - -function startCaseLeaseHeartbeat({ - store, - prefix, - id, - kind, - scorerName, - workerId, -}: { - store: DurableEvalStore; - prefix: string; - id: string; - kind: "task" | "score"; - scorerName?: string; - workerId: string; -}) { - return startLeaseHeartbeat(async () => { - const updated = await updateCaseRecord(store, prefix, id, (record) => { - const state = - kind === "task" ? record.task : record.scores[scorerName ?? ""]; - if (state?.status !== "leased" || state.workerId !== workerId) { - return undefined; - } - state.leaseUntil = Date.now() + JOB_LEASE_MS; - return record; - }); - return updated; - }); -} - -function startJobLeaseHeartbeat( - store: DurableEvalStore, - key: string, - workerId: string, -) { - return startLeaseHeartbeat(async () => { - const updated = await updateJobRecord(store, key, (job) => { - if ( - job.workerId !== workerId || - (job.status !== "preparing" && - job.status !== "submitting" && - job.status !== "submitted") - ) { - return undefined; - } - job.leaseUntil = Date.now() + JOB_LEASE_MS; - return job; - }); - return updated?.value.workerId === workerId; - }); -} - -async function withJobLeaseHeartbeat({ - store, - key, - workerId, - signal, - operation, -}: { - store: DurableEvalStore; - key: string; - workerId: string; - signal: AbortSignal; - operation: () => Promise; -}) { - if (signal.aborted) throw new DurableEvalAbortError(); - const heartbeat = startJobLeaseHeartbeat(store, key, workerId); - try { - return await awaitWithSignal(Promise.resolve().then(operation), signal); - } finally { - await heartbeat.stop(); - } -} - -function startLeaseHeartbeat(renew: () => Promise) { - let stopped = false; - let timer: ReturnType | undefined; - let pending = Promise.resolve(); - const schedule = () => { - if (stopped) return; - timer = setTimeout(() => { - pending = renew() - .then((owned) => { - if (!owned) stopped = true; - }) - .finally(schedule); - }, LEASE_HEARTBEAT_MS); - }; - schedule(); - return { - async stop() { - stopped = true; - if (timer) clearTimeout(timer); - await pending; - }, - }; -} - -function awaitWithSignal( - operation: Promise, - signal: AbortSignal, -): Promise { - if (signal.aborted) { - return Promise.reject(new DurableEvalAbortError()); - } - return new Promise((resolve, reject) => { - const abort = () => reject(new DurableEvalAbortError()); - signal.addEventListener("abort", abort, { once: true }); - operation.then(resolve, reject).finally(() => { - signal.removeEventListener("abort", abort); - }); - }); +async function readJson(store: DurableEvalStore, key: string) { + const value = await store.read(key); + return value ? (JSON.parse(decoder.decode(value)) as T) : undefined; } -class DurableEvalAbortError extends Error { - constructor() { - super("Durable eval operation was aborted"); - this.name = "AbortError"; - } +async function writeJson(store: DurableEvalStore, key: string, value: unknown) { + await store.write(key, encoder.encode(stableStringify(value))); } -function stableStringify(value: unknown): string { +function stableStringify(value: unknown) { return JSON.stringify(value, (_key, nested) => { if (nested && typeof nested === "object" && !Array.isArray(nested)) { return Object.fromEntries( @@ -3952,41 +1478,46 @@ function stableStringify(value: unknown): string { function assertJsonValue(value: unknown, label: string): JsonValue { try { const serialized = JSON.stringify(value); - if (serialized === undefined) { + if (serialized === undefined) throw new Error("value serializes to undefined"); - } return JSON.parse(serialized) as JsonValue; } catch (error) { throw new Error(`${label} must be JSON serializable`, { cause: error }); } } -function serializeError(error: unknown): SerializedError { - if (error instanceof Error) { - return { name: error.name, message: error.message, stack: error.stack }; +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 { message: String(error) }; + return value.id; } -function isErrorCode(error: unknown, code: string) { +function isBatchTask( + value: unknown, +): value is DurableBatchTask { return ( - typeof error === "object" && - error !== null && - "code" in error && - error.code === code + typeof value === "object" && + value !== null && + "kind" in value && + value.kind === BATCH_TASK_KIND ); } -function validateShard(shard: { index: number; count: number }) { - if ( - !Number.isInteger(shard.index) || - !Number.isInteger(shard.count) || - shard.count < 1 || - shard.index < 0 || - shard.index >= shard.count - ) { - throw new Error(`Invalid DurableEval shard ${shard.index}/${shard.count}`); - } +function isBatchScorer( + value: unknown, +): value is DurableBatchScorer { + return ( + typeof value === "object" && + value !== null && + "kind" in value && + value.kind === BATCH_SCORER_KIND + ); } function isIterable(value: unknown): value is Iterable { @@ -4001,9 +1532,7 @@ function isAsyncIterable(value: unknown): value is AsyncIterable { ); } -async function* toAsyncIterable( - value: Iterable | AsyncIterable, -): AsyncGenerator { +async function* toAsyncIterable(value: Iterable | AsyncIterable) { if (isAsyncIterable(value)) { for await (const item of value) yield item; } else { @@ -4011,31 +1540,6 @@ async function* toAsyncIterable( } } -async function* toAbortableAsyncIterable( - value: Iterable | AsyncIterable, - signal: AbortSignal, -): AsyncGenerator { - if (isAsyncIterable(value)) { - const iterator = value[Symbol.asyncIterator](); - while (true) { - const next = await awaitWithSignal(iterator.next(), signal); - if (next.done) return; - yield next.value; - } - } else { - for (const item of value) { - if (signal.aborted) throw new DurableEvalAbortError(); - yield item; - } - } -} - -function delay(ms: number) { - return new Promise((resolve) => setTimeout(resolve, Math.max(ms, 0))); -} - -function timeRemaining(deadlineAt: number | undefined) { - return deadlineAt === undefined - ? undefined - : Math.max(deadlineAt - Date.now(), 0); +function asError(error: unknown) { + return error instanceof Error ? error : new Error(String(error)); } diff --git a/js/src/exports.ts b/js/src/exports.ts index 0e9415917..cfd8bfc38 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -264,12 +264,7 @@ export { export type { DurableEvalStore } from "./durable-eval"; -export { - BatchScorer, - BatchTask, - DurableEval, - DurableEvalNotSubmittedError, -} from "./durable-eval"; +export { BatchScorer, BatchTask, DurableEval } from "./durable-eval"; export { agentAssertionScorer } from "./agent-assertions"; diff --git a/js/src/framework.ts b/js/src/framework.ts index 8461799d7..8628eb0e9 100644 --- a/js/src/framework.ts +++ b/js/src/framework.ts @@ -456,14 +456,6 @@ export type EvaluatorFile = { reporter?: ReporterDef | string; }; }; - durableEvaluators?: Record< - string, - { - // Kept opaque here to avoid coupling the existing Eval framework to the - // additive durable evaluator's generic surface. - definition: unknown; - } - >; reporters: { [reporterName: string]: ReporterDef }; }; @@ -567,7 +559,6 @@ globalThis._evals = { prompts: [], parameters: [], evaluators: {}, - durableEvaluators: {}, reporters: {}, }; From 83f5fff4214ecd0b30cdeaf770cdb60f8c46bda4 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Tue, 4 Aug 2026 13:31:43 +0000 Subject: [PATCH 09/13] Refine durable eval runtime integration --- .../durable-eval-webhook/scenario.test.ts | 73 +- .../durable-eval-webhook/scenario.ts | 20 +- js/README.md | 19 +- js/src/durable-eval.test.ts | 89 +- js/src/durable-eval.ts | 879 +++++++++++------- js/src/framework.ts | 428 +++++---- js/src/logger.ts | 39 +- 7 files changed, 954 insertions(+), 593 deletions(-) diff --git a/e2e/scenarios/durable-eval-webhook/scenario.test.ts b/e2e/scenarios/durable-eval-webhook/scenario.test.ts index 2a5efd08f..ebe696f92 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.test.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.test.ts @@ -11,26 +11,57 @@ const scenarioDir = await prepareScenarioDir({ }); test("durable eval collects webhook sub-batches and logs completed rows", async () => { - await withScenarioHarness(async ({ runScenarioDir, testRunEvents }) => { - await runScenarioDir({ scenarioDir }); + await withScenarioHarness( + async ({ events, runScenarioDir, testRunEvents }) => { + await runScenarioDir({ scenarioDir }); - const evalSpans = findAllSpans(testRunEvents(), "eval"); - const webhookSpans = evalSpans.filter( - (event) => event.metadata?.kind === "webhook", - ); - expect(webhookSpans).toHaveLength(3); - expect(webhookSpans.map((event) => event.output).sort()).toEqual([2, 4, 6]); - expect( - webhookSpans - .map((event) => event.scores) - .sort((left, right) => - JSON.stringify(left).localeCompare(JSON.stringify(right)), - ), - ).toEqual([{ exact: 1 }, { exact: 1 }, { exact: 1 }]); - 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 evalSpans = findAllSpans(testRunEvents(), "eval"); + const webhookSpans = evalSpans.filter( + (event) => event.metadata?.kind === "webhook", + ); + expect(webhookSpans).toHaveLength(3); + expect(webhookSpans.map((event) => event.output).sort()).toEqual([ + 2, 4, 6, + ]); + expect( + webhookSpans + .map((event) => event.scores) + .sort((left, right) => + JSON.stringify(left).localeCompare(JSON.stringify(right)), + ), + ).toEqual([{ exact: 1 }, { exact: 1 }, { exact: 1 }]); + 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); + expect(taskSpans.map((event) => event.output).sort()).toEqual([2, 4, 6]); + + const scoreSpans = findAllSpans(events(), "exact"); + expect(scoreSpans).toHaveLength(3); + expect(scoreSpans.map((event) => event.scores)).toEqual([ + { exact: 1 }, + { exact: 1 }, + { exact: 1 }, + ]); + expect(scoreSpans.map((event) => event.metadata?.method)).toEqual([ + "shared-eval-runtime", + "shared-eval-runtime", + "shared-eval-runtime", + ]); + + const classifierSpans = findAllSpans(events(), "quality"); + expect(classifierSpans).toHaveLength(3); + expect(webhookSpans.map((event) => event.row.classifications)).toEqual([ + { quality: [{ id: "pass", label: "Pass" }] }, + { quality: [{ id: "pass", label: "Pass" }] }, + { quality: [{ id: "pass", label: "Pass" }] }, + ]); + }, + ); }); diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts index 37e3a8a4c..0020dc8f6 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -87,15 +87,26 @@ async function main() { task, scores: [ function exact({ output, expected }) { - return output === expected ? 1 : 0; + return { + name: "exact", + score: output === expected ? 1 : 0, + metadata: { method: "shared-eval-runtime" }, + }; + }, + ], + classifiers: [ + function quality({ output, expected }) { + return { + name: "quality", + id: output === expected ? "pass" : "fail", + label: output === expected ? "Pass" : "Fail", + }; }, ], }, ); - const waiting = await definition.start({ - runId: `durable-${testRunId}`, - }); + const waiting = await definition.start(); if (waiting.status !== "waiting" || jobs.size !== 2) { throw new Error("Durable eval did not pause with two webhook batches"); } @@ -107,6 +118,7 @@ async function main() { if (!externalId) throw new Error("Durable eval stopped before completion"); completedJobs.add(externalId); processed = await definition.processBatchResult({ + runId: waiting.runId, externalId, }); } diff --git a/js/README.md b/js/README.md index 610894ccf..13bf76b73 100644 --- a/js/README.md +++ b/js/README.md @@ -75,7 +75,10 @@ const supportEval = DurableEval("Support bot", { async submit(items, context) { const batch = await provider.submit({ idempotencyKey: context.batchId, - metadata: { durableBatchId: context.batchId }, + metadata: { + durableRunId: context.runId, + durableBatchId: context.batchId, + }, items, }); return { id: batch.id }; @@ -102,9 +105,8 @@ const supportEval = DurableEval("Support bot", { ], }); -const result = await supportEval.start({ - runId: "release-2026-07-27", -}); +const result = await supportEval.start(); +const { runId } = result; ``` `start()` initializes the run, submits every ready task sub-batch, and returns. @@ -136,7 +138,7 @@ results, submits newly ready work, and returns without sleeping: ```typescript const result = await supportEval.poll({ - runId: "release-2026-07-27", + runId, }); if (result.status === "waiting" && result.pending.poll > 0) { @@ -151,7 +153,7 @@ mode: ```typescript { status: "waiting", - runId: "release-2026-07-27", + runId, pending: { poll: 2, webhook: 1 }, } ``` @@ -161,7 +163,7 @@ collecting results, or advancing the workflow: ```typescript const status = await supportEval.status({ - runId: "release-2026-07-27", + runId, }); ``` @@ -210,8 +212,11 @@ store its results through `processBatchResult()`: 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. DurableEval saved it from submit()'s handle. externalId: batch.id, // The SDK-generated ID passed to submit(); include it in provider metadata diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts index 3c4f70f1d..387fe5850 100644 --- a/js/src/durable-eval.test.ts +++ b/js/src/durable-eval.test.ts @@ -38,10 +38,7 @@ describe("DurableEval", () => { return output === expected ? 1 : 0; }, ], - }).start({ - runId: "local-run", - noSendLogs: true, - }); + }).start({ noSendLogs: true }); expect(result).toMatchObject({ status: "completed", @@ -50,6 +47,20 @@ describe("DurableEval", () => { expect(task).toHaveBeenCalledTimes(2); }); + test("generates a new run id for every start", async () => { + const durable = DurableEval("generated-runs", { + store: new MemoryStore(), + 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("polls each existing task and scorer sub-batch once", async () => { const taskJobs = new Map< string, @@ -118,21 +129,18 @@ describe("DurableEval", () => { task, scores: [scorer], }); - const options = { - runId: "polling-run", - noSendLogs: true, - }; - - await expect(durable.start(options)).resolves.toEqual({ + const waiting = await durable.start({ noSendLogs: true }); + expect(waiting).toMatchObject({ status: "waiting", - runId: "polling-run", + 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: "polling-run", + runId: waiting.runId, pending: { poll: 2, webhook: 0 }, }); expect(taskJobs.size).toBe(2); @@ -140,7 +148,7 @@ describe("DurableEval", () => { await expect(durable.poll(options)).resolves.toEqual({ status: "waiting", - runId: "polling-run", + runId: waiting.runId, pending: { poll: 2, webhook: 0 }, }); expect(scoreJobs.size).toBe(2); @@ -231,30 +239,24 @@ describe("DurableEval", () => { scores: [scorer], }); - await expect( - durable.start({ runId: "webhook-run", noSendLogs: true }), - ).resolves.toEqual({ - status: "waiting", - runId: "webhook-run", - pending: { poll: 0, webhook: 2 }, - }); - expect(taskJobs.size).toBe(2); - - await expect( - durable.start({ runId: "webhook-run", noSendLogs: true }), - ).resolves.toEqual({ + const waiting = await durable.start({ noSendLogs: true }); + expect(waiting).toMatchObject({ status: "waiting", - runId: "webhook-run", + 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("DurableEval run missing-run is missing"); for (const [index, externalId] of taskIds.entries()) { - const result = await durable.processBatchResult( - { externalId }, - { noSendLogs: true }, - ); + const result = await durable.processBatchResult({ runId, externalId }); expect(result).toMatchObject({ status: "waiting", pending: { poll: 0, webhook: index === 0 ? 1 : 2 }, @@ -265,10 +267,7 @@ describe("DurableEval", () => { const scoreIds = [...scoreJobs.keys()]; let result; for (const externalId of scoreIds) { - result = await durable.processBatchResult( - { externalId }, - { noSendLogs: true }, - ); + result = await durable.processBatchResult({ runId, externalId }); } expect(result).toMatchObject({ status: "completed", @@ -402,14 +401,11 @@ describe("DurableEval", () => { task, scores: [scorer], }); - const options = { - runId: "workflow-run", - noSendLogs: true, - }; - - await expect(durable.start(options)).resolves.toMatchObject({ + const waiting = await durable.start({ noSendLogs: true }); + expect(waiting).toMatchObject({ status: "waiting", }); + const options = { runId: waiting.runId }; await expect(durable.poll(options)).resolves.toMatchObject({ status: "waiting", }); @@ -433,7 +429,18 @@ describe("DurableEval", () => { DurableEval("missing-ids", { store: new MemoryStore(), data: [{ input: "hello" }], - task: (input) => input, + task: BatchTask({ + async submit() { + return { id: "unused" }; + }, + completion: { + mode: "webhook", + externalId: (handle) => handle.id, + }, + async collect() { + return []; + }, + }), scores: [], }).start({ noSendLogs: true }), ).rejects.toThrow("requires id, upsert_id, or caseId"); diff --git a/js/src/durable-eval.ts b/js/src/durable-eval.ts index 160de1893..489698658 100644 --- a/js/src/durable-eval.ts +++ b/js/src/durable-eval.ts @@ -1,23 +1,40 @@ -import { type Score, SpanTypeAttribute } from "../util/index"; -import { type EvalParameters, type InferParameters } from "./eval-parameters"; +import { makeScorerPropagatedEvent, SpanTypeAttribute } from "../util/index"; import { - type EvalData, - type EvalHooks, + type EvalParameters, + type InferParameters, + validateParameters, +} from "./eval-parameters"; +import { + _internalInitEvaluatorExperiment, + _internalPrepareEvaluatorClassification, + _internalPrepareEvaluatorScore, + _internalResolveEvaluatorData, + _internalRunEvaluatorTask, + buildLocalSummary as buildEvaluatorLocalSummary, + callEvaluatorData, + classifierName, + type EvalClassifier, + type Evaluator, + type EvaluatorDef, + type EvalResult, type EvalScorer, type EvalScorerArgs, type EvalTask, type OneOrMoreScores, + runEvaluator, } from "./framework"; import iso from "./isomorph"; import { type BaseMetadata, - type BraintrustState, type DefaultMetadataType, type EvalCase, type Experiment, type ExperimentSummary, NOOP_SPAN, - init as initExperiment, + type Span, + _internalResumeSpan, + _internalStartSpanWithInitialMerge, + logError as logSpanError, newId, } from "./logger"; @@ -394,9 +411,11 @@ export type DurableEvaluator< Expected = void, Metadata extends BaseMetadata = DefaultMetadataType, Parameters extends EvalParameters = EvalParameters, -> = { +> = Omit< + Evaluator, + "task" | "scores" | "timeout" | "signal" | "maxConcurrency" | "update" +> & { store: DurableEvalStore; - data: EvalData; caseId?: ( datum: EvalCase, ) => string | Promise; @@ -410,26 +429,21 @@ export type DurableEvaluator< Parameters, JsonValue >; - scores: Array< + scores?: Array< | EvalScorer | DurableBatchScorer >; - parameters?: InferParameters; - experimentName?: string; - description?: string; - metadata?: Record; - tags?: string[]; - trialCount?: number; - projectId?: string; - state?: BraintrustState; }; -export interface DurableEvalRuntimeOptions { - runId?: string; +export interface DurableEvalStartOptions< + Parameters extends EvalParameters = EvalParameters, +> { + parameters?: InferParameters; noSendLogs?: boolean; } export type DurableBatchResult = { + runId: string; batchId?: string; externalId?: string; }; @@ -469,17 +483,12 @@ export interface DurableEvalDefinition< Metadata, Parameters >; - start(options?: DurableEvalRuntimeOptions): Promise; - status( - options: DurableEvalRuntimeOptions & { runId: string }, - ): Promise; - poll( - options: DurableEvalRuntimeOptions & { runId: string }, - ): Promise; - processBatchResult( - result: DurableBatchResult, - options?: Omit, + start( + options?: DurableEvalStartOptions, ): Promise; + status(options: { runId: string }): Promise; + poll(options: { runId: string }): Promise; + processBatchResult(result: DurableBatchResult): Promise; } type DurableCaseRecord = { @@ -490,10 +499,15 @@ type DurableCaseRecord = { metadata: JsonValue; tags?: string[]; taskComplete: boolean; + taskLogged: boolean; output?: JsonValue; + rootSpan?: string; taskNodeOutputs: Record; scores: Record; + loggedScores: Record; scoreNodeOutputs: Record>; + classifications: Record; + loggedClassifications: Record; }; type DurableBatchRecord = { @@ -513,17 +527,14 @@ type DurableRunState = { projectName: string; evalName: string; experimentName: string; + noSendLogs: boolean; + parameters: JsonValue; status: "running" | "completed"; summary?: ExperimentSummary; cases: DurableCaseRecord[]; batches: DurableBatchRecord[]; }; -type DurableBatchLocator = { - runKey: string; - batchId: string; -}; - class DurableEvalDefinitionImpl< Input, Output, @@ -552,27 +563,22 @@ class DurableEvalDefinitionImpl< this.evalName = evaluator.experimentName ?? projectName; } - start(options: DurableEvalRuntimeOptions = {}): Promise { + start( + options: DurableEvalStartOptions = {}, + ): Promise { return startDurableEval(this, options); } - status( - options: DurableEvalRuntimeOptions & { runId: string }, - ): Promise { + status(options: { runId: string }): Promise { return getDurableEvalStatus(this, options); } - poll( - options: DurableEvalRuntimeOptions & { runId: string }, - ): Promise { + poll(options: { runId: string }): Promise { return pollDurableEval(this, options); } - processBatchResult( - result: DurableBatchResult, - options: Omit = {}, - ): Promise { - return processDurableBatchResult(this, result, options); + processBatchResult(result: DurableBatchResult): Promise { + return processDurableBatchResult(this, result); } } @@ -603,28 +609,94 @@ async function startDurableEval< Metadata, Parameters >, - options: DurableEvalRuntimeOptions, + options: DurableEvalStartOptions, ): Promise { const store = definition.evaluator.store; - const runId = options.runId ?? newId(); + const runId = newId(); const key = runKey(definition.projectName, definition.evalName, runId); - let state = await readJson(store, key); - if (!state) { - state = { + const { data } = callEvaluatorData(definition.evaluator.data); + const parameters = await validateParameters( + options.parameters ?? {}, + definition.evaluator.parameters, + ); + const experimentName = + definition.evaluator.experimentName ?? `${definition.evalName}-${runId}`; + const experiment = await _internalInitEvaluatorExperiment( + definition.projectName, + { ...definition.evaluator, data } as unknown as Evaluator< + Input, + Output, + Expected, + Metadata, + Parameters + >, + data, + { + disabled: options.noSendLogs ?? false, + experimentName, + update: true, + }, + ); + if ( + !isBatchTask(definition.evaluator.task) && + !(definition.evaluator.scores ?? []).some(isBatchScorer) + ) { + const result = await runEvaluator( + experiment, + { + ...definition.evaluator, + projectName: definition.projectName, + evalName: definition.evalName, + data, + } as unknown as EvaluatorDef< + Input, + Output, + Expected, + Metadata, + Parameters + >, + { + start: () => undefined, + stop: () => undefined, + increment: () => undefined, + }, + [], + undefined, + parameters, + true, + true, + ); + const state: DurableRunState = { schemaVersion: CHECKPOINT_VERSION, runId, projectName: definition.projectName, evalName: definition.evalName, - experimentName: - definition.evaluator.experimentName ?? - `${definition.evalName}-${runId}`, - status: "running", - cases: await materializeCases(definition.evaluator), + experimentName, + noSendLogs: options.noSendLogs ?? false, + parameters: assertJsonValue(parameters, "eval parameters"), + status: "completed", + summary: result.summary, + cases: [], batches: [], }; + await experiment?.flush(); await writeJson(store, key, state); + return currentStatus(definition, state); } - return advanceDurableEval(definition, state, store, key, options); + const state: DurableRunState = { + schemaVersion: CHECKPOINT_VERSION, + runId, + projectName: definition.projectName, + evalName: definition.evalName, + experimentName, + noSendLogs: options.noSendLogs ?? false, + parameters: assertJsonValue(parameters, "eval parameters"), + status: "running", + cases: await materializeCases(definition, data, experiment), + batches: [], + }; + await writeJson(store, key, state); + return advanceDurableEval(definition, state, store, key, experiment); } async function getDurableEvalStatus< @@ -641,7 +713,7 @@ async function getDurableEvalStatus< Metadata, Parameters >, - options: DurableEvalRuntimeOptions & { runId: string }, + options: { runId: string }, ): Promise { const store = definition.evaluator.store; const state = await readJson( @@ -667,34 +739,33 @@ async function processDurableBatchResult< Parameters >, result: DurableBatchResult, - options: Omit, ): Promise { if (!result.batchId && !result.externalId) { throw new Error("Batch results require batchId or externalId"); } const store = definition.evaluator.store; - const locator = await locateBatch( - store, - definition.projectName, - definition.evalName, - result, - ); - if (!locator) { - throw new Error("No submitted batch matches this result"); + const key = runKey(definition.projectName, definition.evalName, result.runId); + const state = await readJson(store, key); + if (!state) throw new Error(`DurableEval run ${result.runId} is missing`); + const byBatch = result.batchId + ? state.batches.find((candidate) => candidate.id === result.batchId) + : undefined; + const byExternal = result.externalId + ? state.batches.find( + (candidate) => candidate.externalId === result.externalId, + ) + : undefined; + if (byBatch && byExternal && byBatch.id !== byExternal.id) { + throw new Error("batchId and externalId identify different batches"); } - const state = await readJson(store, locator.runKey); - if (!state) - throw new Error(`DurableEval run for batch ${locator.batchId} is missing`); - const batch = state.batches.find( - (candidate) => candidate.id === locator.batchId, - ); - if (!batch) throw new Error(`Batch ${locator.batchId} is missing`); + const batch = byBatch ?? byExternal; + if (!batch) throw new Error("No submitted batch matches this result"); if (batch.status !== "complete") { await collectBatch(definition, state, batch); batch.status = "complete"; - await writeJson(store, locator.runKey, state); + await writeJson(store, key, state); } - return advanceDurableEval(definition, state, store, locator.runKey, options); + return advanceDurableEval(definition, state, store, key); } async function pollDurableEval< @@ -711,7 +782,7 @@ async function pollDurableEval< Metadata, Parameters >, - options: DurableEvalRuntimeOptions & { runId: string }, + options: { runId: string }, ): Promise { const store = definition.evaluator.store; const key = runKey( @@ -747,7 +818,30 @@ async function pollDurableEval< batch.status = "complete"; await writeJson(store, key, state); } - return advanceDurableEval(definition, state, store, key, options); + return advanceDurableEval(definition, state, store, key); +} + +async function openDurableExperiment( + definition: DurableEvalDefinition, + state: DurableRunState, +) { + const data: EvalCase[] = []; + return await _internalInitEvaluatorExperiment( + definition.projectName, + { ...definition.evaluator, data } as unknown as Evaluator< + unknown, + unknown, + unknown, + BaseMetadata, + EvalParameters + >, + data, + { + disabled: state.noSendLogs, + experimentName: state.experimentName, + update: true, + }, + ); } async function advanceDurableEval< @@ -767,16 +861,21 @@ async function advanceDurableEval< state: DurableRunState, store: DurableEvalStore, key: string, - options: Omit, + existingExperiment?: Experiment | null, ): Promise { if (state.status === "completed") return currentStatus(definition, state); - await runTaskStage(definition, state, store, key); + const experiment = + existingExperiment === undefined + ? await openDurableExperiment(definition, state) + : existingExperiment; + await runTaskStage(definition, state, store, key, experiment); + await logCompletedTasks(definition, state, store, key, experiment); if (state.cases.some((record) => !record.taskComplete)) { return currentStatus(definition, state); } - await runScoreStages(definition, state, store, key); - const scorerNames = resolveScorers(definition.evaluator.scores).map( + await runScoreStages(definition, state, store, key, experiment); + const scorerNames = resolveScorers(definition.evaluator.scores ?? []).map( ({ name }) => name, ); if ( @@ -787,7 +886,7 @@ async function advanceDurableEval< return currentStatus(definition, state); } - state.summary = await finishExperiment(definition, state, options.noSendLogs); + state.summary = await finishExperiment(definition, state, experiment); state.status = "completed"; await writeJson(store, key, state); return currentStatus(definition, state); @@ -816,6 +915,112 @@ function currentStatus( return { status: "waiting", runId: state.runId, pending }; } +async function startCaseRoot( + definition: DurableEvalDefinition, + state: DurableRunState, + record: DurableCaseRecord, + experiment: Experiment | null, +): Promise { + if (!experiment) return NOOP_SPAN; + const datum = record.datum as EvalCase; + return _internalStartSpanWithInitialMerge({ + ...(definition.evaluator.state + ? { state: definition.evaluator.state } + : {}), + parent: await experiment.export(), + name: "eval", + spanId: deterministicId(`${state.runId}:${record.id}:span`), + spanAttributes: { type: SpanTypeAttribute.EVAL }, + event: { + id: deterministicId(`${state.runId}:${record.id}:row`), + input: datum.input, + expected: "expected" in datum ? datum.expected : undefined, + tags: datum.tags, + origin: datum.origin, + }, + }); +} + +async function logTaskResult( + definition: DurableEvalDefinition, + state: DurableRunState, + record: DurableCaseRecord, + experiment: Experiment | null, + task?: EvalTask, +) { + const datum = record.datum as EvalCase; + const root = await startCaseRoot(definition, state, record, experiment); + try { + if (task) { + const result = await root.traced( + (span) => + _internalRunEvaluatorTask( + task, + datum, + record.trialIndex, + state.parameters as Record, + span, + ), + { + name: "task", + spanId: deterministicId(`${state.runId}:${record.id}:task`), + spanAttributes: { type: SpanTypeAttribute.TASK }, + event: { input: datum.input }, + }, + ); + record.output = assertJsonValue( + result.output, + `task output for ${record.caseId}`, + ); + record.metadata = assertJsonValue(result.metadata, "task metadata"); + record.tags = result.tags; + record.taskComplete = true; + } else { + await root.traced((span) => span.log({ output: record.output }), { + name: "task", + spanId: deterministicId(`${state.runId}:${record.id}:task`), + spanAttributes: { type: SpanTypeAttribute.TASK }, + event: { input: datum.input }, + }); + } + root.log({ + output: record.output, + expected: "expected" in datum ? datum.expected : undefined, + metadata: { + ...(record.metadata as Record), + durable_eval: { + run_id: state.runId, + case_id: record.caseId, + trial_index: record.trialIndex, + }, + }, + tags: record.tags, + }); + record.rootSpan = await root.export(); + record.taskLogged = true; + } catch (error) { + logSpanError(root, error); + throw error; + } finally { + root.end(); + await experiment?.flush(); + } +} + +async function logCompletedTasks( + definition: DurableEvalDefinition, + state: DurableRunState, + store: DurableEvalStore, + key: string, + experiment: Experiment | null, +) { + for (const record of state.cases) { + if (!record.taskComplete || record.taskLogged) continue; + await logTaskResult(definition, state, record, experiment); + await writeJson(store, key, state); + } +} + async function runTaskStage< Input, Output, @@ -833,6 +1038,7 @@ async function runTaskStage< state: DurableRunState, store: DurableEvalStore, key: string, + experiment: Experiment | null, ) { if (isBatchTask(definition.evaluator.task)) { await ensureWorkflowBatches(definition, state, store, key, "task"); @@ -848,32 +1054,7 @@ async function runTaskStage< >; for (const record of state.cases) { if (record.taskComplete) continue; - const datum = record.datum as EvalCase; - const metadata = { ...(record.metadata as Record) }; - const hooks: EvalHooks = { - meta(value) { - Object.assign(metadata, value); - }, - metadata: metadata as EvalHooks< - Expected, - Metadata, - Parameters - >["metadata"], - expected: ("expected" in datum ? datum.expected : undefined) as Expected, - span: NOOP_SPAN, - parameters: (definition.evaluator.parameters ?? - {}) as InferParameters, - reportProgress: () => undefined, - trialIndex: record.trialIndex, - tags: record.tags, - }; - record.output = assertJsonValue( - await task(datum.input, hooks), - `task output for ${record.caseId}`, - ); - record.metadata = assertJsonValue(hooks.metadata, "task metadata"); - record.tags = hooks.tags; - record.taskComplete = true; + await logTaskResult(definition, state, record, experiment, task); await writeJson(store, key, state); } } @@ -895,35 +1076,182 @@ async function runScoreStages< state: DurableRunState, store: DurableEvalStore, key: string, + experiment: Experiment | null, ) { - const scorers = resolveScorers(definition.evaluator.scores); + const scorers = resolveScorers(definition.evaluator.scores ?? []); for (const { name, scorer } of scorers) { if (isBatchScorer(scorer)) { + for (const record of state.cases) { + if (name in record.scores && !record.loggedScores[name]) { + await evaluateAndLogScore( + definition, + state, + record, + name, + experiment, + ); + await writeJson(store, key, state); + } + } await ensureWorkflowBatches(definition, state, store, key, "score", name); continue; } for (const record of state.cases) { - if (name in record.scores) continue; - const datum = record.datum as EvalCase; - const value = await ( - scorer as EvalScorer - )({ - input: datum.input, - ...(datum.tags ? { tags: datum.tags } : {}), - ...(datum.id ? { id: datum.id } : {}), - ...(datum.upsert_id ? { upsert_id: datum.upsert_id } : {}), - ...(datum.trialCount ? { trialCount: datum.trialCount } : {}), - ...("expected" in datum ? { expected: datum.expected } : {}), - metadata: record.metadata as Metadata, - output: record.output as Output, - } as unknown as EvalScorerArgs); - record.scores[name] = assertJsonValue( - normalizeScores(value, name), - `scorer ${name} output`, + if (record.loggedScores[name]) continue; + await evaluateAndLogScore( + definition, + state, + record, + name, + experiment, + scorer, ); await writeJson(store, key, state); } } + + for (const [index, classifier] of ( + definition.evaluator.classifiers ?? [] + ).entries()) { + const name = classifierName(classifier, index); + for (const record of state.cases) { + if (record.loggedClassifications[name]) continue; + await evaluateAndLogClassification( + definition, + state, + record, + name, + classifier, + experiment, + ); + await writeJson(store, key, state); + } + } +} + +function scorerArgs(record: DurableCaseRecord) { + const datum = record.datum as EvalCase; + return { + ...datum, + metadata: record.metadata, + output: record.output, + } as EvalScorerArgs; +} + +function resumeCaseRoot( + definition: DurableEvalDefinition, + record: DurableCaseRecord, + experiment: Experiment | null, +) { + if (!experiment) return NOOP_SPAN; + if (!record.rootSpan) { + throw new Error(`Durable eval case ${record.caseId} has no root span`); + } + return _internalResumeSpan({ + exported: record.rootSpan, + state: definition.evaluator.state, + }); +} + +async function evaluateAndLogScore( + definition: DurableEvalDefinition, + state: DurableRunState, + record: DurableCaseRecord, + name: string, + experiment: Experiment | null, + scorer?: EvalScorer, +) { + const root = resumeCaseRoot(definition, record, experiment); + try { + const rootExport = await root.export(); + const prepared = await root.traced( + async (span) => { + const value = scorer + ? await scorer(scorerArgs(record)) + : (record.scores[name] as OneOrMoreScores); + if (scorer) { + record.scores[name] = assertJsonValue(value, `scorer ${name} output`); + } + const result = _internalPrepareEvaluatorScore(value, name); + if (result.results !== null) { + span.log({ + output: result.output, + metadata: result.metadata, + scores: result.scores, + }); + } + return result; + }, + { + name, + spanId: deterministicId(`${state.runId}:${record.id}:score:${name}`), + spanAttributes: { + type: SpanTypeAttribute.SCORE, + purpose: "scorer", + }, + propagatedEvent: makeScorerPropagatedEvent(rootExport || undefined), + event: { input: scorerArgs(record) }, + }, + ); + if (prepared.scores) root.log({ scores: prepared.scores }); + record.loggedScores[name] = true; + } catch (error) { + logSpanError(root, error); + throw error; + } finally { + root.end(); + await experiment?.flush(); + } +} + +async function evaluateAndLogClassification( + definition: DurableEvalDefinition, + state: DurableRunState, + record: DurableCaseRecord, + name: string, + classifier: EvalClassifier, + experiment: Experiment | null, +) { + const root = resumeCaseRoot(definition, record, experiment); + try { + const rootExport = await root.export(); + const prepared = await root.traced( + async (span) => { + const value = await classifier(scorerArgs(record)); + record.classifications[name] = assertJsonValue( + value, + `classifier ${name} output`, + ); + const result = _internalPrepareEvaluatorClassification(value, name); + if (result.results !== null) { + span.log({ output: result.output, metadata: result.metadata }); + } + return result; + }, + { + name, + spanId: deterministicId( + `${state.runId}:${record.id}:classification:${name}`, + ), + spanAttributes: { + type: SpanTypeAttribute.CLASSIFIER, + purpose: "scorer", + }, + propagatedEvent: makeScorerPropagatedEvent(rootExport || undefined), + event: { input: scorerArgs(record) }, + }, + ); + if (prepared.classifications) { + root.log({ classifications: prepared.classifications }); + } + record.loggedClassifications[name] = true; + } catch (error) { + logSpanError(root, error); + throw error; + } finally { + root.end(); + await experiment?.flush(); + } } async function ensureWorkflowBatches( @@ -965,7 +1293,7 @@ async function ensureWorkflowBatches( const batchId = newId(); const context = { runId: state.runId, batchId }; const items = records.map((record) => - itemForNode(definition, record, kind, scorerName, node), + itemForNode(state.parameters, record, kind, scorerName, node), ); const handle = assertJsonValue( await node.processor.submit(items, context), @@ -990,7 +1318,6 @@ async function ensureWorkflowBatches( }; state.batches.push(batch); await writeJson(store, key, state); - await indexBatch(store, definition, batch, key); } } } @@ -1047,7 +1374,7 @@ async function collectBatch( record.taskComplete = true; } else { record.scores[batch.scorerName!] = assertJsonValue( - normalizeScores(output as OneOrMoreScores, batch.scorerName!), + output, `score for ${id}`, ); } @@ -1090,7 +1417,7 @@ function workflowForStage( } stage = definition.evaluator.task; } else { - const scorer = resolveScorers(definition.evaluator.scores).find( + const scorer = resolveScorers(definition.evaluator.scores ?? []).find( ({ name }) => name === scorerName, )?.scorer; if (!isBatchScorer(scorer)) { @@ -1124,7 +1451,7 @@ function workflowForStage( } function itemForNode( - definition: DurableEvalDefinition, + parameters: JsonValue, record: DurableCaseRecord, kind: "task" | "score", scorerName: string | undefined, @@ -1132,7 +1459,7 @@ function itemForNode( ) { const rootItem = kind === "task" - ? taskBatchItem(record, definition.evaluator.parameters ?? {}) + ? taskBatchItem(record, parameters) : scorerBatchItem(record); const outputs = nodeOutputsFor(record, kind, scorerName); return node.item( @@ -1156,25 +1483,42 @@ function nodeOutputsFor( return (record.scoreNodeOutputs[scorerName!] ??= {}); } -async function materializeCases( - evaluator: DurableEvaluator, +async function materializeCases< + Input, + Output, + Expected, + Metadata extends BaseMetadata, + Parameters extends EvalParameters, +>( + definition: DurableEvalDefinition< + Input, + Output, + Expected, + Metadata, + Parameters + >, + data: Evaluator["data"], + experiment: Experiment | null, ): Promise { - const raw = - typeof evaluator.data === "function" ? evaluator.data() : evaluator.data; - const data = raw instanceof Promise ? await raw : raw; - if (typeof data === "object" && data !== null && "_type" in data) { - throw new Error("DurableEval does not support BaseExperiment data sources"); - } - if (!isIterable(data) && !isAsyncIterable(data)) { - throw new Error("DurableEval data must be iterable"); - } + const evaluator = definition.evaluator; + const iterable = await _internalResolveEvaluatorData( + { + data, + projectName: definition.projectName, + projectId: evaluator.projectId, + state: evaluator.state, + }, + experiment, + ); const records: DurableCaseRecord[] = []; const seen = new Set(); - for await (const datum of toAsyncIterable(data)) { + for await (const datum of iterable) { const caseId = datum.id ?? datum.upsert_id ?? - (evaluator.caseId ? await evaluator.caseId(datum) : undefined); + (evaluator.caseId + ? await evaluator.caseId(datum as EvalCase) + : undefined); if (!caseId) { throw new Error( "Every DurableEval case requires id, upsert_id, or caseId", @@ -1199,9 +1543,13 @@ async function materializeCases( ), tags: datum.tags, taskComplete: false, + taskLogged: false, taskNodeOutputs: {}, scores: {}, + loggedScores: {}, scoreNodeOutputs: {}, + classifications: {}, + loggedClassifications: {}, }); } } @@ -1237,107 +1585,64 @@ function scorerBatchItem(record: DurableCaseRecord) { async function finishExperiment( definition: DurableEvalDefinition, state: DurableRunState, - noSendLogs = false, + experiment: Experiment | null, ) { - const scorerNames = resolveScorers(definition.evaluator.scores).map( + const scorerNames = resolveScorers(definition.evaluator.scores ?? []).map( ({ name }) => name, ); - if (noSendLogs) { - return buildLocalSummary( - state, - definition.projectName, - state.experimentName, - scorerNames, + const results = state.cases.map((record) => { + const datum = record.datum as EvalCase; + const scores = Object.assign( + {}, + ...scorerNames.map( + (name) => + _internalPrepareEvaluatorScore( + record.scores[name] as OneOrMoreScores, + name, + ).scores ?? {}, + ), ); - } - const experiment = initExperiment({ - state: definition.evaluator.state, - ...(definition.evaluator.projectId - ? { projectId: definition.evaluator.projectId } - : { project: definition.projectName }), - experiment: state.experimentName, - update: true, - description: definition.evaluator.description, - metadata: definition.evaluator.metadata, - tags: definition.evaluator.tags, - setCurrent: false, - }); - for (const record of state.cases) { - logCase(experiment, state.runId, record, scorerNames); - } - await experiment.flush(); - return await experiment.summarize(); -} - -function logCase( - experiment: Experiment, - runId: string, - record: DurableCaseRecord, - scorerNames: string[], -) { - const datum = record.datum as EvalCase; - const scores = Object.assign( - {}, - ...scorerNames.map((name) => record.scores[name] ?? {}), - ) as Record; - const span = experiment.startSpan({ - name: "eval", - spanId: deterministicId(`${runId}:${record.id}:span`), - spanAttributes: { type: SpanTypeAttribute.EVAL }, - event: { - id: deterministicId(`${runId}:${record.id}:row`), - input: datum.input, - expected: "expected" in datum ? datum.expected : undefined, + const classifications = Object.assign( + {}, + ...Object.entries(record.classifications).map( + ([name, value]) => + _internalPrepareEvaluatorClassification(value as never, name) + .classifications ?? {}, + ), + ); + return { + ...datum, output: record.output, - scores, - metadata: { - ...(record.metadata as Record), - durable_eval: { - run_id: runId, - case_id: record.caseId, - trial_index: record.trialIndex, - }, - }, + metadata: record.metadata, tags: record.tags, - }, + scores, + error: undefined, + ...(Object.keys(classifications).length > 0 ? { classifications } : {}), + } as EvalResult; }); - span.end(); -} - -function buildLocalSummary( - state: DurableRunState, - projectName: string, - experimentName: string, - scorerNames: string[], -): ExperimentSummary { - const totals: Record = {}; - for (const record of state.cases) { - for (const scorerName of scorerNames) { - const scores = record.scores[scorerName] as Record; - for (const [name, score] of Object.entries(scores)) { - if (score === null) continue; - const total = totals[name] ?? { total: 0, count: 0 }; - total.total += score; - total.count++; - totals[name] = total; - } + if (!experiment) { + return buildEvaluatorLocalSummary( + { + ...definition.evaluator, + projectName: definition.projectName, + evalName: state.experimentName, + } as unknown as EvaluatorDef, + results, + ); + } + await experiment.flush(); + let comparisonExperimentId = definition.evaluator.baseExperimentId; + if (!comparisonExperimentId) { + try { + comparisonExperimentId = await experiment._getBaseExperimentId(); + } catch { + comparisonExperimentId = undefined; } } - return { - projectName, - experimentName, - scores: Object.fromEntries( - Object.entries(totals).map(([name, value]) => [ - name, - { - name, - score: value.total / value.count, - improvements: 0, - regressions: 0, - }, - ]), - ), - }; + return await experiment.summarize({ + summarizeScores: definition.evaluator.summarizeScores, + ...(comparisonExperimentId ? { comparisonExperimentId } : {}), + }); } function resolveScorers( @@ -1354,88 +1659,10 @@ function resolveScorers( })); } -function normalizeScores(value: OneOrMoreScores, defaultName: string) { - if (value === null) return { [defaultName]: null }; - if (typeof value === "number") return { [defaultName]: value }; - const values = Array.isArray(value) ? value : [value]; - return Object.fromEntries( - values.map((score: Score) => [score.name ?? defaultName, score.score]), - ); -} - -async function indexBatch( - store: DurableEvalStore, - definition: DurableEvalDefinition, - batch: DurableBatchRecord, - key: string, -) { - const locator = { runKey: key, batchId: batch.id }; - await writeJson( - store, - batchIndexKey( - definition.projectName, - definition.evalName, - "batch", - batch.id, - ), - locator, - ); - if (batch.externalId) { - await writeJson( - store, - batchIndexKey( - definition.projectName, - definition.evalName, - "external", - batch.externalId, - ), - locator, - ); - } -} - -async function locateBatch( - store: DurableEvalStore, - projectName: string, - evalName: string, - result: DurableBatchResult, -) { - const byBatch = result.batchId - ? await readJson( - store, - batchIndexKey(projectName, evalName, "batch", result.batchId), - ) - : undefined; - const byExternal = result.externalId - ? await readJson( - store, - batchIndexKey(projectName, evalName, "external", result.externalId), - ) - : undefined; - if ( - byBatch && - byExternal && - (byBatch.runKey !== byExternal.runKey || - byBatch.batchId !== byExternal.batchId) - ) { - throw new Error("batchId and externalId identify different batches"); - } - return byBatch ?? byExternal; -} - function runKey(projectName: string, evalName: string, runId: string) { return `durable-eval/v2/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; } -function batchIndexKey( - projectName: string, - evalName: string, - type: "batch" | "external", - id: string, -) { - return `durable-eval/v2/index/${contentVersion(encoder.encode(`${projectName}\0${evalName}`))}/${type}/${contentVersion(encoder.encode(id))}`; -} - function deterministicId(value: string) { const hex = contentVersion(encoder.encode(value)) .padEnd(32, "0") @@ -1520,26 +1747,6 @@ function isBatchScorer( ); } -function isIterable(value: unknown): value is Iterable { - return ( - typeof value === "object" && value !== null && Symbol.iterator in value - ); -} - -function isAsyncIterable(value: unknown): value is AsyncIterable { - return ( - typeof value === "object" && value !== null && Symbol.asyncIterator in value - ); -} - -async function* toAsyncIterable(value: Iterable | AsyncIterable) { - if (isAsyncIterable(value)) { - for await (const item of value) yield item; - } else { - for (const item of value) yield item; - } -} - function asError(error: unknown) { return error instanceof Error ? error : new Error(String(error)); } diff --git a/js/src/framework.ts b/js/src/framework.ts index 8628eb0e9..8dff8bafd 100644 --- a/js/src/framework.ts +++ b/js/src/framework.ts @@ -498,6 +498,38 @@ async function getExperimentParametersRef( }; } +export async function _internalInitEvaluatorExperiment( + projectName: string, + evaluator: Evaluator, + data: EvalData, + options: { + disabled?: boolean; + experimentName?: string; + update?: boolean; + } = {}, +): Promise { + if (options.disabled) return null; + const { baseExperiment } = callEvaluatorData(data); + const parameters = await getExperimentParametersRef(evaluator.parameters); + return initExperiment(evaluator.state, { + ...(evaluator.projectId + ? { projectId: evaluator.projectId } + : { project: projectName }), + experiment: options.experimentName ?? evaluator.experimentName, + description: evaluator.description, + metadata: evaluator.metadata, + tags: evaluator.tags, + isPublic: evaluator.isPublic, + update: options.update ?? evaluator.update, + baseExperiment: evaluator.baseExperimentName ?? baseExperiment, + baseExperimentId: evaluator.baseExperimentId, + gitMetadataSettings: evaluator.gitMetadataSettings, + repoInfo: evaluator.repoInfo, + dataset: Dataset.isDataset(data) ? data : undefined, + parameters, + }); +} + export function callEvaluatorData< Input, Expected, @@ -546,6 +578,65 @@ function isIterable(value: unknown): value is Iterable { ); } +export async function _internalResolveEvaluatorData( + evaluator: Pick< + EvaluatorDef, + "data" | "projectName" | "projectId" | "state" + >, + experiment: Experiment | null, +): Promise>> { + if (typeof evaluator.data === "string") { + throw new Error("Unimplemented: string data paths"); + } + let dataResult = + typeof evaluator.data === "function" ? evaluator.data() : evaluator.data; + + if ("_type" in dataResult) { + if (dataResult._type !== "BaseExperiment") { + throw new Error("Invalid _type"); + } + if (!experiment) { + throw new Error( + "Cannot use BaseExperiment() without connecting to Braintrust (you most likely set --no-send-logs)", + ); + } + let name = dataResult.name; + if (isEmpty(name)) { + const baseExperiment = await experiment.fetchBaseExperiment(); + if (!baseExperiment) { + throw new Error("BaseExperiment() failed to fetch base experiment"); + } + name = baseExperiment.name; + } + + dataResult = initExperiment(evaluator.state, { + ...(evaluator.projectId + ? { projectId: evaluator.projectId } + : { project: evaluator.projectName }), + experiment: name, + open: true, + }).asDataset(); + } + + const resolvedDataResult = + dataResult instanceof Promise ? await dataResult : dataResult; + if (isAsyncIterable>(resolvedDataResult)) { + return resolvedDataResult; + } + if ( + Array.isArray(resolvedDataResult) || + isIterable>(resolvedDataResult) + ) { + const iterable = resolvedDataResult as Iterable>; + return (async function* () { + for (const datum of iterable) yield datum; + })(); + } + throw new Error( + "Evaluator data must be an array, iterable, or async iterable", + ); +} + declare global { var _evals: EvaluatorFile; @@ -716,33 +807,13 @@ export async function Eval< const resolvedReporter = options.reporter || defaultReporter; try { - const { data, baseExperiment: defaultBaseExperiment } = callEvaluatorData( - evaluator.data, + const { data } = callEvaluatorData(evaluator.data); + const experiment = await _internalInitEvaluatorExperiment( + name, + evaluator, + data, + { disabled: Boolean(options.parent || options.noSendLogs) }, ); - const parameters = await getExperimentParametersRef(evaluator.parameters); - // NOTE: This code is duplicated with initExperiment in js/src/cli.ts. Make sure - // to update that if you change this. - const experiment = - options.parent || options.noSendLogs - ? null - : initExperiment(evaluator.state, { - ...(evaluator.projectId - ? { projectId: evaluator.projectId } - : { project: name }), - experiment: evaluator.experimentName, - description: evaluator.description, - metadata: evaluator.metadata, - tags: evaluator.tags, - isPublic: evaluator.isPublic, - update: evaluator.update, - baseExperiment: - evaluator.baseExperimentName ?? defaultBaseExperiment, - baseExperimentId: evaluator.baseExperimentId, - gitMetadataSettings: evaluator.gitMetadataSettings, - repoInfo: evaluator.repoInfo, - dataset: Dataset.isDataset(data) ? data : undefined, - parameters, - }); // Ensure experiment ID is resolved before tasks start for OTEL parent attribute support // The Experiment constructor starts resolution (fire-and-forget), but we await here to ensure completion @@ -905,6 +976,42 @@ export function classifierName( return classifier.name || `classifier_${classifier_idx}`; } +export async function _internalRunEvaluatorTask( + task: EvalTask, + datum: EvalCase, + trialIndex: number, + parameters: Record, + span: Span, + reportProgress: (event: TaskProgressEvent) => void = () => undefined, +): Promise<{ + output: unknown; + metadata: Record; + tags: string[]; +}> { + const metadata: Record = { + ...("metadata" in datum ? datum.metadata : {}), + }; + const hooks: EvalHooks, EvalParameters> = { + meta(value) { + Object.assign(metadata, value); + }, + metadata, + expected: "expected" in datum ? datum.expected : undefined, + span, + parameters, + reportProgress, + trialIndex, + tags: [...(datum.tags ?? [])], + }; + const output = await task(datum.input, hooks); + span.log({ output }); + return { + output, + metadata: hooks.metadata, + tags: hooks.tags ?? [], + }; +} + function buildSpanMetadata( results: Array<{ name: string; metadata?: Record }>, ) { @@ -930,6 +1037,50 @@ function buildSpanScores( return { resultMetadata: buildSpanMetadata(results), scoresRecord }; } +export function _internalPrepareEvaluatorScore( + scoreValue: OneOrMoreScores, + name: string, +): { + results: Score[] | null; + output?: unknown; + metadata?: Record; + scores?: Record; +} { + if (scoreValue === null) return { results: null }; + if (Array.isArray(scoreValue)) { + for (const score of scoreValue) { + if (!(typeof score === "object" && !isEmpty(score))) { + throw new Error( + `When returning an array of scores, each score must be a non-empty object. Got: ${JSON.stringify(score)}`, + ); + } + } + } + const results: Score[] = Array.isArray(scoreValue) + ? scoreValue + : typeof scoreValue === "object" && !isEmpty(scoreValue) + ? [scoreValue] + : [{ name, score: scoreValue }]; + const { resultMetadata, scoresRecord } = buildSpanScores(results); + const fields = (score: Score) => { + const { metadata: _metadata, name: _name, ...rest } = score; + return rest; + }; + return { + results, + output: + results.length === 1 + ? fields(results[0]) + : results.reduce( + (previous, score) => + mergeDicts(previous, { [score.name ?? name]: fields(score) }), + {}, + ), + metadata: resultMetadata, + scores: scoresRecord, + }; +} + async function runInScorerSpan( rootSpan: Span, spanName: string, @@ -996,6 +1147,40 @@ function toClassificationItem(c: Classification): ClassificationItem { }; } +export function _internalPrepareEvaluatorClassification( + value: OneOrMoreClassifications, + name: string, +): { + results: Classification[] | null; + output?: unknown; + metadata?: Record; + classifications?: Record; +} { + if (value === null) return { results: null }; + const results = (Array.isArray(value) ? value : [value]).map((result) => + validateClassificationResult(result, name), + ); + const classifications: Record = {}; + for (const result of results) { + (classifications[result.name] ??= []).push(toClassificationItem(result)); + } + return { + results, + output: + results.length === 1 + ? toClassificationItem(results[0]) + : results.reduce( + (previous, result) => + mergeDicts(previous, { + [result.name]: toClassificationItem(result), + }), + {}, + ), + metadata: buildSpanMetadata(results), + classifications, + }; +} + function logScoringFailures( kind: string, failures: { name: string; error: unknown }[], @@ -1076,69 +1261,14 @@ async function runEvaluatorInternal( (evaluator.state ?? _internalGetGlobalState())?.spanCache?.start(); } try { - if (typeof evaluator.data === "string") { - throw new Error("Unimplemented: string data paths"); - } - let dataResult = - typeof evaluator.data === "function" ? evaluator.data() : evaluator.data; - parameters = await validateParameters( parameters ?? {}, evaluator.parameters, ); - - if ("_type" in dataResult) { - if (dataResult._type !== "BaseExperiment") { - // For some reason, the typesystem won't let me check if dataResult._type === "BaseExperiment" - throw new Error("Invalid _type"); - } - if (!experiment) { - throw new Error( - "Cannot use BaseExperiment() without connecting to Braintrust (you most likely set --no-send-logs)", - ); - } - let name = dataResult.name; - if (isEmpty(name)) { - const baseExperiment = await experiment.fetchBaseExperiment(); - if (!baseExperiment) { - throw new Error("BaseExperiment() failed to fetch base experiment"); - } - name = baseExperiment.name; - } - - dataResult = initExperiment(evaluator.state, { - ...(evaluator.projectId - ? { projectId: evaluator.projectId } - : { project: evaluator.projectName }), - experiment: name, - open: true, - }).asDataset(); - } - - const resolvedDataResult = - dataResult instanceof Promise ? await dataResult : dataResult; - - const dataIterable: AsyncIterable> = (() => { - if (isAsyncIterable>(resolvedDataResult)) { - return resolvedDataResult; - } - if ( - Array.isArray(resolvedDataResult) || - isIterable>(resolvedDataResult) - ) { - const iterable = resolvedDataResult as Iterable< - EvalCase - >; - return (async function* () { - for (const datum of iterable) { - yield datum; - } - })(); - } - throw new Error( - "Evaluator data must be an array, iterable, or async iterable", - ); - })(); + const dataIterable = await _internalResolveEvaluatorData( + evaluator, + experiment, + ); progressReporter.start(evaluator.evalName, 0); @@ -1252,13 +1382,11 @@ async function runEvaluatorInternal( }) : undefined; - let metadata: Record = { - ...("metadata" in datum ? datum.metadata : {}), - }; + let metadata: Record = {}; const expected = "expected" in datum ? datum.expected : undefined; let output: unknown = undefined; let error: unknown | undefined = undefined; - let tags: string[] = [...(datum.tags ?? [])]; + let tags: string[] = []; const scores: Record = {}; const classifications: Record = {}; const scorerNames = (evaluator.scores ?? []).map(scorerName); @@ -1267,22 +1395,15 @@ async function runEvaluatorInternal( ); let unhandledScores: string[] | null = scorerNames; try { - const meta = (o: Record) => - (metadata = { ...metadata, ...o }); - - await rootSpan.traced( - async (span: Span) => { - const hooksForTask: EvalHooks< - unknown, - Record, - EvalParameters - > = { - meta, - metadata, - expected, + const taskResult = await rootSpan.traced( + (span: Span) => + _internalRunEvaluatorTask( + evaluator.task, + datum, + trialIndex, + parameters ?? {}, span, - parameters: parameters ?? {}, - reportProgress: (event: TaskProgressEvent) => { + (event) => { stream?.({ ...event, id: rootSpan.id, @@ -1291,27 +1412,16 @@ async function runEvaluatorInternal( object_type: "task", }); }, - trialIndex, - tags, - }; - - const outputResult = evaluator.task(datum.input, hooksForTask); - if (outputResult instanceof Promise) { - output = await outputResult; - } else { - output = outputResult; - } - - tags = hooksForTask.tags ?? []; - - span.log({ output }); - }, + ), { name: "task", spanAttributes: { type: SpanTypeAttribute.TASK }, event: { input: datum.input }, }, ); + output = taskResult.output; + metadata = taskResult.metadata; + tags = taskResult.tags; if (tags.length) { rootSpan.log({ output, metadata, expected, tags }); } else { @@ -1334,11 +1444,6 @@ async function runEvaluatorInternal( await rootSpan.export(), ); - const getOtherFields = (s: Score) => { - const { metadata: _metadata, name: _name, ...rest } = s; - return rest; - }; - const [scoreResults, classificationResults] = await Promise.all([ Promise.all( (evaluator.scores ?? []).map((score, score_idx) => @@ -1352,44 +1457,17 @@ async function runEvaluatorInternal( const scoreValue = await Promise.resolve( score(scoringArgs), ); - if (scoreValue === null) return null; - if (Array.isArray(scoreValue)) { - for (const s of scoreValue) { - if (!(typeof s === "object" && !isEmpty(s))) { - throw new Error( - `When returning an array of scores, each score must be a non-empty object. Got: ${JSON.stringify(s)}`, - ); - } - } - } - const results: Score[] = Array.isArray(scoreValue) - ? scoreValue - : typeof scoreValue === "object" && !isEmpty(scoreValue) - ? [scoreValue] - : [ - { - name: scorerNames[score_idx], - score: scoreValue, - }, - ]; - const { resultMetadata, scoresRecord } = - buildSpanScores(results); - const resultOutput = - results.length === 1 - ? getOtherFields(results[0]) - : results.reduce( - (prev, s) => - mergeDicts(prev, { - [s.name]: getOtherFields(s), - }), - {}, - ); + const prepared = _internalPrepareEvaluatorScore( + scoreValue, + scorerNames[score_idx], + ); + if (prepared.results === null) return null; span.log({ - output: resultOutput, - metadata: resultMetadata, - scores: scoresRecord, + output: prepared.output, + metadata: prepared.metadata, + scores: prepared.scores, }); - return results; + return prepared.results; }, ), ), @@ -1406,32 +1484,16 @@ async function runEvaluatorInternal( const classifierValue = await Promise.resolve( classifier(scoringArgs), ); - if (classifierValue === null) return null; - const rawResults = ( - Array.isArray(classifierValue) - ? classifierValue - : [classifierValue] - ).map((result) => - validateClassificationResult( - result, - classifierNames[idx], - ), + const prepared = _internalPrepareEvaluatorClassification( + classifierValue, + classifierNames[idx], ); - const resultOutput = - rawResults.length === 1 - ? toClassificationItem(rawResults[0]) - : rawResults.reduce( - (prev, r) => - mergeDicts(prev, { - [r.name]: toClassificationItem(r), - }), - {}, - ); + if (prepared.results === null) return null; span.log({ - output: resultOutput, - metadata: buildSpanMetadata(rawResults), + output: prepared.output, + metadata: prepared.metadata, }); - return rawResults; + return prepared.results; }, ), ), diff --git a/js/src/logger.ts b/js/src/logger.ts index fab3b3ead..76676c9e7 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -272,9 +272,13 @@ type StartSpanEventArgs = ExperimentLogPartialArgs & Partial; const INITIAL_SPAN_WRITE_AS_MERGE = Symbol( "braintrust.initial-span-write-as-merge", ); +const RESUME_SPAN_WITHOUT_INITIAL_WRITE = Symbol( + "braintrust.resume-span-without-initial-write", +); type InitialSpanWriteAsMergeArg = { readonly [INITIAL_SPAN_WRITE_AS_MERGE]?: true; + readonly [RESUME_SPAN_WITHOUT_INITIAL_WRITE]?: true; }; export type StartSpanArgs = { @@ -2204,6 +2208,37 @@ export function updateSpan({ }); } +/** @internal Rehydrate an exported root span so work can continue in another process. */ +export function _internalResumeSpan({ + exported, + state, +}: { + exported: string; + state?: BraintrustState; +}): Span { + const resolvedState = state ?? _globalState; + const components = SpanComponentsV4.fromStr(exported); + const { row_id, root_span_id, span_id } = components.data; + if (!row_id || !root_span_id || !span_id) { + throw new Error("Only exported root spans can be resumed"); + } + return new SpanImpl({ + state: resolvedState, + parentObjectType: components.data.object_type, + parentObjectId: new LazyValue( + spanComponentsToObjectIdLambda(resolvedState, components), + ), + parentComputeObjectMetadataArgs: undefined, + parentSpanIds: { parentSpanIds: [], rootSpanId: root_span_id }, + spanId: span_id, + event: { id: row_id }, + propagatedEvent: (components.data.propagated_event ?? undefined) as + | StartSpanEventArgs + | undefined, + [RESUME_SPAN_WITHOUT_INITIAL_WRITE]: true, + }); +} + /** * An opaque W3C trace-context, as returned by * {@link extractTraceContextFromHeaders}. @@ -7694,7 +7729,9 @@ export class SpanImpl implements Span { // Deterministic spans can be initialized concurrently by separate // workflow executions, so their first write must not replace later merges. this.isMerge = args[INITIAL_SPAN_WRITE_AS_MERGE] === true; - this.logInternal({ event, internalData }); + if (!args[RESUME_SPAN_WITHOUT_INITIAL_WRITE]) { + this.logInternal({ event, internalData }); + } this.isMerge = true; } From 855ea97ef02d61e9efd6464fa1840055e31792fc Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 10 Aug 2026 10:24:21 +0000 Subject: [PATCH 10/13] Update PR #2297 --- .changeset/durable-batches-evaluate.md | 2 +- .../durable-eval-webhook/scenario.test.ts | 29 +- .../durable-eval-webhook/scenario.ts | 82 +++- js/README.md | 189 --------- js/src/durable-eval.test.ts | 83 ++-- js/src/durable-eval.ts | 385 ++++++++++++++++-- js/src/exports.ts | 8 +- knip.jsonc | 3 - 8 files changed, 501 insertions(+), 280 deletions(-) diff --git a/.changeset/durable-batches-evaluate.md b/.changeset/durable-batches-evaluate.md index a0486161b..6067dc216 100644 --- a/.changeset/durable-batches-evaluate.md +++ b/.changeset/durable-batches-evaluate.md @@ -2,4 +2,4 @@ "braintrust": minor --- -feat: Add batch/durable evals api +feat: Add experimental batch/durable evals API diff --git a/e2e/scenarios/durable-eval-webhook/scenario.test.ts b/e2e/scenarios/durable-eval-webhook/scenario.test.ts index ebe696f92..4a0b3eae9 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.test.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.test.ts @@ -10,7 +10,7 @@ const scenarioDir = await prepareScenarioDir({ scenarioDir: resolveScenarioDir(import.meta.url), }); -test("durable eval collects webhook sub-batches and logs completed rows", async () => { +test("durable eval collects task and scorer webhook sub-batches", async () => { await withScenarioHarness( async ({ events, runScenarioDir, testRunEvents }) => { await runScenarioDir({ scenarioDir }); @@ -29,7 +29,11 @@ test("durable eval collects webhook sub-batches and logs completed rows", async .sort((left, right) => JSON.stringify(left).localeCompare(JSON.stringify(right)), ), - ).toEqual([{ exact: 1 }, { exact: 1 }, { exact: 1 }]); + ).toEqual([ + { batch_exact: 1, exact: 1 }, + { batch_exact: 1, exact: 1 }, + { batch_exact: 1, exact: 1 }, + ]); expect(webhookSpans.map((event) => event.metadata?.durable_eval)).toEqual( [ expect.objectContaining({ run_id: expect.any(String) }), @@ -42,19 +46,32 @@ test("durable eval collects webhook sub-batches and logs completed rows", async expect(taskSpans).toHaveLength(3); expect(taskSpans.map((event) => event.output).sort()).toEqual([2, 4, 6]); - const scoreSpans = findAllSpans(events(), "exact"); - expect(scoreSpans).toHaveLength(3); - expect(scoreSpans.map((event) => event.scores)).toEqual([ + const exactScoreSpans = findAllSpans(events(), "exact"); + expect(exactScoreSpans).toHaveLength(3); + expect(exactScoreSpans.map((event) => event.scores)).toEqual([ { exact: 1 }, { exact: 1 }, { exact: 1 }, ]); - expect(scoreSpans.map((event) => event.metadata?.method)).toEqual([ + expect(exactScoreSpans.map((event) => event.metadata?.method)).toEqual([ "shared-eval-runtime", "shared-eval-runtime", "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 classifierSpans = findAllSpans(events(), "quality"); expect(classifierSpans).toHaveLength(3); expect(webhookSpans.map((event) => event.row.classifications)).toEqual([ diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts index 0020dc8f6..ed6f370e0 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -1,26 +1,19 @@ -import { BatchTask, DurableEval, type DurableEvalStore } from "braintrust"; +import { + BatchScorer, + BatchTask, + defineDurableEval, + DurableEvalMemoryStore, +} from "braintrust"; import { getTestRunId, runMain, scopedName, } from "../../helpers/scenario-runtime"; -class MemoryStore implements DurableEvalStore { - private readonly values = new Map(); - - async read(key: string) { - return this.values.get(key)?.slice(); - } - - async write(key: string, value: Uint8Array) { - this.values.set(key, value.slice()); - } -} - async function main() { const testRunId = getTestRunId(); - const store = new MemoryStore(); - const jobs = new Map>(); + const store = new DurableEvalMemoryStore(); + const jobs = new Map(); const task = BatchTask< number, number, @@ -42,7 +35,11 @@ async function main() { externalId: (handle) => handle.id, }, async collect(handle) { - return (jobs.get(handle.id) ?? []).map((item) => ({ + const items = (jobs.get(handle.id) ?? []) as Array<{ + id: string; + input: number; + }>; + return items.map((item) => ({ id: item.id, output: item.input * 2, })); @@ -62,7 +59,11 @@ async function main() { externalId: (handle) => handle.id, }, async collect(handle) { - return (jobs.get(handle.id) ?? []).map((item) => ({ + const items = (jobs.get(handle.id) ?? []) as Array<{ + id: string; + input: number; + }>; + return items.map((item) => ({ id: item.id, output: item.input, })); @@ -70,7 +71,41 @@ async function main() { }); }, }); - const definition = DurableEval( + const scorer = 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: { + mode: "webhook", + externalId: (handle) => handle.id, + }, + async collect(handle) { + const items = (jobs.get(handle.id) ?? []) as Array<{ + id: string; + output: number; + expected: number; + }>; + return items.map((item) => ({ + id: item.id, + score: { + name: "batch_exact", + score: item.output === item.expected ? 1 : 0, + metadata: { method: "batch-provider" }, + }, + })); + }, + }); + const definition = defineDurableEval( scopedName("e2e-durable-eval-webhook-project", testRunId), { store, @@ -93,6 +128,7 @@ async function main() { metadata: { method: "shared-eval-runtime" }, }; }, + scorer, ], classifiers: [ function quality({ output, expected }) { @@ -111,16 +147,20 @@ async function main() { throw new Error("Durable eval did not pause with two webhook batches"); } - let processed = waiting; + let completed = false; const completedJobs = new Set(); - while (completedJobs.size < jobs.size || processed.status !== "completed") { + 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); - processed = await definition.processBatchResult({ + 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"); } } diff --git a/js/README.md b/js/README.md index 13bf76b73..da6da81a6 100644 --- a/js/README.md +++ b/js/README.md @@ -44,195 +44,6 @@ async function main() { main().catch(console.error); ``` -## Durable evaluations - -`DurableEval` 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 -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). - -```typescript -import { BatchTask, DurableEval, type DurableEvalStore } from "braintrust"; - -const store: DurableEvalStore = checkpointStore; -const supportEval = DurableEval("Support bot", { - store, - data: [ - { - id: "password-reset", - input: "How do I reset my password?", - expected: "Open account settings...", - }, - ], - task: BatchTask({ - // Each provider job contains at most 500 eval cases. - batchSize: 500, - - // Submit one sub-batch and return a JSON-serializable provider handle. - async submit(items, context) { - const batch = await provider.submit({ - idempotencyKey: context.batchId, - metadata: { - durableRunId: context.runId, - durableBatchId: context.batchId, - }, - items, - }); - return { id: batch.id }; - }, - - completion: { - // "webhook" waits for processBatchResult(). Use "poll" with a poll() - // callback when the provider does not send completion events. - mode: "webhook", - externalId: (handle) => handle.id, - }, - - async collect(handle) { - return (await provider.results(handle.id)).map((item) => ({ - id: item.id, - output: item.output, - })); - }, - }), - scores: [ - function exact({ output, expected }) { - return output === expected ? 1 : 0; - }, - ], -}); - -const result = await supportEval.start(); -const { runId } = result; -``` - -`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. - -### Polling - -Polling adapters report the provider's current status through `completion`: - -```typescript -completion: { - mode: "poll", - async poll(handle) { - const batch = await provider.getBatch(handle.id); - if (batch.status === "completed") return { status: "complete" }; - if (batch.status === "failed") { - return { status: "failed", error: batch.error }; - } - return { status: "pending" }; - }, -}, -``` - -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: - -```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 workflow: - -```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. - -### Multi-stage workflows - -The direct `BatchTask({ submit, completion, collect })` form remains the -one-batch shorthand. Use `workflow` when a task or scorer requires multiple -provider batch operations: - -```typescript -task: BatchTask({ - workflow(workflow) { - const draft = workflow.batch("draft", { - input: ({ input }) => ({ prompt: input }), - batchSize: 500, - submit: submitDraftBatch, - completion: draftCompletion, - collect: collectDraftBatch, - }); - - return workflow.batch("revise", { - needs: { draft }, - input: ({ input }, { draft }) => ({ original: input, draft }), - batchSize: 500, - submit: submitRevisionBatch, - completion: revisionCompletion, - collect: collectRevisionBatch, - }); - }, -}), -``` - -Every named batch is a persisted workflow node. `needs` can express sequential -operations, parallel branches, and joins. The returned node supplies the final -task output or scorer result. - -### 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. DurableEval saved it from submit()'s handle. - 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 handle. - 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 workflow node, whose `collect()` results are -stored before the eval advances. Webhook idempotency and provider failure -handling remain application responsibilities for now. - ## Auto-Instrumentation Braintrust can automatically instrument popular AI SDKs (OpenAI, Anthropic, Vercel AI SDK, and others) to log calls without manual wrapper code. diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts index 387fe5850..005209b95 100644 --- a/js/src/durable-eval.test.ts +++ b/js/src/durable-eval.test.ts @@ -3,31 +3,64 @@ import { configureNode } from "./node/config"; import { BatchScorer, BatchTask, - DurableEval, + defineDurableEval, + DurableEvalMemoryStore, + DurableEvalRedisStore, type DurableBatchScorerItem, type DurableBatchTaskItem, - type DurableEvalStore, } from "./durable-eval"; configureNode(); -class MemoryStore implements DurableEvalStore { - private readonly values = new Map(); +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]); - async read(key: string) { - return this.values.get(key)?.slice(); - } + await store.write("run", value); + value[0] = 9; - async write(key: string, value: Uint8Array) { - this.values.set(key, value.slice()); - } -} + 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(); + }); + + 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) => { + values.set(key, value); + return "OK"; + }), + }; + const store = new DurableEvalRedisStore(client, { + keyPrefix: "evals:", + }); + + await store.write("run", new Uint8Array([0, 255, 1])); + + expect(client.set).toHaveBeenCalledWith("evals:run", "AP8B"); + 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({ + get: async () => 42, + set: async () => "OK", + }).read("invalid"), + ).rejects.toThrow("expected GET to return a string"); + }); +}); -describe("DurableEval", () => { +describe("defineDurableEval", () => { test("runs ordinary tasks and scorers", async () => { const task = vi.fn((input: number) => input * 2); - const result = await DurableEval("local", { - store: new MemoryStore(), + const result = await defineDurableEval("local", { + store: new DurableEvalMemoryStore(), data: [ { id: "one", input: 1, expected: 2 }, { id: "two", input: 2, expected: 4 }, @@ -48,8 +81,8 @@ describe("DurableEval", () => { }); test("generates a new run id for every start", async () => { - const durable = DurableEval("generated-runs", { - store: new MemoryStore(), + const durable = defineDurableEval("generated-runs", { + store: new DurableEvalMemoryStore(), data: [{ input: 1 }], task: (input) => input, scores: [() => 1], @@ -118,8 +151,8 @@ describe("DurableEval", () => { }, }); - const store = new MemoryStore(); - const durable = DurableEval("polling-batches", { + const store = new DurableEvalMemoryStore(); + const durable = defineDurableEval("polling-batches", { store, data: [1, 2, 3].map((input) => ({ id: `case-${input}`, @@ -175,7 +208,7 @@ describe("DurableEval", () => { }); test("processes task and scorer webhook batches through one method", async () => { - const store = new MemoryStore(); + const store = new DurableEvalMemoryStore(); const taskJobs = new Map< string, DurableBatchTaskItem>[] @@ -228,7 +261,7 @@ describe("DurableEval", () => { })); }, }); - const durable = DurableEval("webhook-batches", { + const durable = defineDurableEval("webhook-batches", { store, data: [1, 2, 3].map((input) => ({ id: `case-${input}`, @@ -254,7 +287,7 @@ describe("DurableEval", () => { runId: "missing-run", externalId: taskIds[0], }), - ).rejects.toThrow("DurableEval run missing-run is missing"); + ).rejects.toThrow("Durable eval run missing-run is missing"); for (const [index, externalId] of taskIds.entries()) { const result = await durable.processBatchResult({ runId, externalId }); expect(result).toMatchObject({ @@ -390,8 +423,8 @@ describe("DurableEval", () => { }); }, }); - const store = new MemoryStore(); - const durable = DurableEval("workflow", { + const store = new DurableEvalMemoryStore(); + const durable = defineDurableEval("workflow", { store, data: [1, 2, 3].map((input) => ({ id: `case-${input}`, @@ -426,8 +459,8 @@ describe("DurableEval", () => { test("requires stable case ids", async () => { await expect( - DurableEval("missing-ids", { - store: new MemoryStore(), + defineDurableEval("missing-ids", { + store: new DurableEvalMemoryStore(), data: [{ input: "hello" }], task: BatchTask({ async submit() { diff --git a/js/src/durable-eval.ts b/js/src/durable-eval.ts index 489698658..a2c54c100 100644 --- a/js/src/durable-eval.ts +++ b/js/src/durable-eval.ts @@ -1,4 +1,9 @@ -import { makeScorerPropagatedEvent, SpanTypeAttribute } from "../util/index"; +import { + base64ToUint8Array, + makeScorerPropagatedEvent, + SpanTypeAttribute, + uint8ArrayToBase64, +} from "../util/index"; import { type EvalParameters, type InferParameters, @@ -46,31 +51,84 @@ const CHECKPOINT_VERSION = 2; const DEFAULT_BATCH_SIZE = 1_000; type JsonPrimitive = string | number | boolean | null; -export type JsonValue = - | JsonPrimitive - | JsonValue[] - | { [key: string]: JsonValue }; +type JsonValue = JsonPrimitive | JsonValue[] | { [key: string]: JsonValue }; /** * Minimal persistence used to reconnect provider webhooks with submitted - * batches. DurableEval does not require any Braintrust backend changes. + * batches. Durable 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 { read(key: string): Promise; write(key: string, value: Uint8Array): Promise; } -export interface DurableBatchContext { +/** + * Stores durable 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 { + private readonly values = new Map(); + + async read(key: string): Promise { + return this.values.get(key)?.slice(); + } + + async write(key: string, value: Uint8Array): Promise { + this.values.set(key, value.slice()); + } +} + +/** + * 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 + * `@upstash/redis` can be passed directly; other clients with compatible async + * `get` and `set` methods are also supported. + * + * @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 { + private readonly keyPrefix: string; + + constructor( + private readonly client: { + get(key: string): Promise; + set(key: string, value: string): Promise; + }, + options: { keyPrefix?: string } = {}, + ) { + this.keyPrefix = options.keyPrefix ?? "braintrust:"; + } + + async read(key: string): Promise { + 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"); + } + return base64ToUint8Array(value); + } + + async write(key: string, value: Uint8Array): Promise { + await this.client.set(`${this.keyPrefix}${key}`, uint8ArrayToBase64(value)); + } +} + +interface DurableBatchContext { runId: string; batchId: string; } -export type DurableBatchPoll = +type DurableBatchPoll = | { status: "pending" } | { status: "complete" } | { status: "failed"; error: unknown }; -export type DurableBatchCompletion = +type DurableBatchCompletion = | { mode: "poll"; poll( @@ -98,7 +156,7 @@ export interface DurableBatchTaskItem< trialIndex: number; } -export type DurableBatchTaskResult = +type DurableBatchTaskResult = | { id: string; output: Output; @@ -117,23 +175,23 @@ export type DurableBatchScorerItem< trialIndex: number; }; -export type DurableBatchScorerResult = +type DurableBatchScorerResult = | { id: string; score: OneOrMoreScores } | { id: string; error: unknown }; -export interface DurableBatchProcessor { +interface DurableBatchProcessor { batchSize?: number; submit(items: Item[], context: DurableBatchContext): Promise; completion: DurableBatchCompletion; collect(handle: Handle, context: DurableBatchContext): Promise; } -export interface DurableWorkflowBatchItem { +interface DurableWorkflowBatchItem { id: string; input: Input; } -export type DurableWorkflowBatchResult = +type DurableWorkflowBatchResult = | { id: string; output: Output; @@ -144,7 +202,7 @@ export type DurableWorkflowBatchResult = const WORKFLOW_NODE_OUTPUT: unique symbol = Symbol("DurableWorkflowNodeOutput"); -export interface DurableWorkflowNode { +interface DurableWorkflowNode { readonly [WORKFLOW_NODE_OUTPUT]: Output; } @@ -156,10 +214,7 @@ type DurableWorkflowNodeOutputs = { : never; }; -export interface DurableWorkflowBuilder< - RootItem, - Metadata extends BaseMetadata, -> { +interface DurableWorkflowBuilder { batch< Output, Needs extends DurableWorkflowNodeMap = Record, @@ -198,7 +253,7 @@ type DurableWorkflowDefinition = { outputNode: string; }; -export interface DurableBatchTask< +interface DurableBatchTask< Input, Output, Expected, @@ -215,7 +270,7 @@ export interface DurableBatchTask< readonly workflow?: DurableWorkflowDefinition; } -export interface DurableBatchScorer< +interface DurableBatchScorer< Input, Output, Expected, @@ -232,6 +287,11 @@ export interface DurableBatchScorer< readonly workflow?: DurableWorkflowDefinition; } +/** + * Defines a task that runs through asynchronous provider batch operations. + * + * @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 BatchTask< Input, Output, @@ -291,6 +351,11 @@ export function BatchTask( >; } +/** + * Defines a scorer that runs through asynchronous provider batch operations. + * + * @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 BatchScorer< Input, Output, @@ -405,7 +470,7 @@ function buildWorkflow( return { nodes, outputNode }; } -export type DurableEvaluator< +type DurableEvaluator< Input, Output, Expected = void, @@ -435,20 +500,20 @@ export type DurableEvaluator< >; }; -export interface DurableEvalStartOptions< +interface DurableEvalStartOptions< Parameters extends EvalParameters = EvalParameters, > { parameters?: InferParameters; noSendLogs?: boolean; } -export type DurableBatchResult = { +type DurableBatchResult = { runId: string; batchId?: string; externalId?: string; }; -export type DurableEvalResult = +type DurableEvalResult = | { status: "waiting"; runId: string; @@ -467,7 +532,7 @@ export type DurableEvalResult = summary: ExperimentSummary; }; -export interface DurableEvalDefinition< +interface DurableEvalDefinition< Input, Output, Expected = void, @@ -582,7 +647,259 @@ class DurableEvalDefinitionImpl< } } -export function DurableEval< +/* + * Internal usage notes. Keep these out of the public README while + * defineDurableEval() is experimental. + * + * ## Durable 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 + * 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 + * 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`: + * + * - [`redis`](https://github.com/redis/node-redis) (node-redis), including the + * lower-level `@redis/client` package + * - [`ioredis`](https://github.com/redis/ioredis) + * - [`@upstash/redis`](https://upstash.com/docs/redis/sdks/ts/deployment), + * including its platform-specific entrypoints + * + * Choose the example for your client: + * + * node-redis (`redis` or `@redis/client`): + * + * ```typescript + * import { createClient } from "redis"; + * import { DurableEvalRedisStore } from "braintrust"; + * + * const nodeRedis = await createClient({ url: process.env.REDIS_URL! }).connect(); + * const redisStore = new DurableEvalRedisStore(nodeRedis); + * ``` + * + * ioredis: + * + * ```typescript + * import Redis from "ioredis"; + * import { DurableEvalRedisStore } from "braintrust"; + * + * const ioRedis = new Redis(process.env.REDIS_URL!); + * const redisStore = new DurableEvalRedisStore(ioRedis); + * ``` + * + * Upstash: + * + * ```typescript + * import { Redis } from "@upstash/redis"; + * import { DurableEvalRedisStore } from "braintrust"; + * + * const upstashRedis = Redis.fromEnv(); + * const redisStore = new DurableEvalRedisStore(upstashRedis); + * ``` + * + * Other clients are compatible when `get(key)` resolves to a string, `null`, or + * `undefined`, and `set(key, value)` accepts a string value. For local testing, + * `new DurableEvalMemoryStore()` 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"; + * + * const supportEval = defineDurableEval("Support bot", { + * store: redisStore, + * data: [ + * { + * id: "password-reset", + * input: "How do I reset my password?", + * expected: "Open account settings...", + * }, + * ], + * task: BatchTask({ + * // Each provider job contains at most 500 eval cases. + * batchSize: 500, + * + * // Submit one sub-batch and return a JSON-serializable provider handle. + * async submit(items, context) { + * const batch = await provider.submit({ + * idempotencyKey: context.batchId, + * metadata: { + * durableRunId: context.runId, + * durableBatchId: context.batchId, + * }, + * items, + * }); + * return { id: batch.id }; + * }, + * + * completion: { + * // "webhook" waits for processBatchResult(). Use "poll" with a poll() + * // callback when the provider does not send completion events. + * mode: "webhook", + * externalId: (handle) => handle.id, + * }, + * + * async collect(handle) { + * return (await provider.results(handle.id)).map((item) => ({ + * id: item.id, + * output: item.output, + * })); + * }, + * }), + * scores: [ + * function exact({ output, expected }) { + * return output === expected ? 1 : 0; + * }, + * ], + * }); + * + * const result = await supportEval.start(); + * const { runId } = result; + * ``` + * + * `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. + * + * ### Polling + * + * Polling adapters report the provider's current status through `completion`: + * + * ```typescript + * completion: { + * mode: "poll", + * async poll(handle) { + * const batch = await provider.getBatch(handle.id); + * if (batch.status === "completed") return { status: "complete" }; + * if (batch.status === "failed") { + * return { status: "failed", error: batch.error }; + * } + * return { status: "pending" }; + * }, + * }, + * ``` + * + * 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: + * + * ```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 workflow: + * + * ```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. + * + * ### Multi-stage workflows + * + * The direct `BatchTask({ submit, completion, collect })` form remains the + * one-batch shorthand. Use `workflow` when a task or scorer requires multiple + * provider batch operations: + * + * ```typescript + * task: BatchTask({ + * workflow(workflow) { + * const draft = workflow.batch("draft", { + * input: ({ input }) => ({ prompt: input }), + * batchSize: 500, + * submit: submitDraftBatch, + * completion: draftCompletion, + * collect: collectDraftBatch, + * }); + * + * return workflow.batch("revise", { + * needs: { draft }, + * input: ({ input }, { draft }) => ({ original: input, draft }), + * batchSize: 500, + * submit: submitRevisionBatch, + * completion: revisionCompletion, + * collect: collectRevisionBatch, + * }); + * }, + * }), + * ``` + * + * Every named batch is a persisted workflow node. `needs` can express sequential + * operations, parallel branches, and joins. The returned node supplies the final + * task output or scorer result. + * + * ### 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 handle. + * 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 handle. + * 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 workflow node, whose `collect()` results are + * stored before the eval advances. Webhook idempotency and provider failure + * handling remain application responsibilities for now. + */ + +/** + * Defines a durable 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< Input, Output, Expected = void, @@ -720,7 +1037,7 @@ async function getDurableEvalStatus< store, runKey(definition.projectName, definition.evalName, options.runId), ); - if (!state) throw new Error(`DurableEval run ${options.runId} is missing`); + if (!state) throw new Error(`Durable eval run ${options.runId} is missing`); return currentStatus(definition, state); } @@ -746,7 +1063,7 @@ async function processDurableBatchResult< const store = definition.evaluator.store; const key = runKey(definition.projectName, definition.evalName, result.runId); const state = await readJson(store, key); - if (!state) throw new Error(`DurableEval run ${result.runId} is missing`); + if (!state) throw new Error(`Durable eval run ${result.runId} is missing`); const byBatch = result.batchId ? state.batches.find((candidate) => candidate.id === result.batchId) : undefined; @@ -791,7 +1108,7 @@ async function pollDurableEval< options.runId, ); const state = await readJson(store, key); - if (!state) throw new Error(`DurableEval run ${options.runId} is missing`); + if (!state) throw new Error(`Durable eval run ${options.runId} is missing`); const batches = state.batches.filter((batch) => { if (batch.status === "complete") return false; @@ -898,7 +1215,7 @@ function currentStatus( ): DurableEvalResult { if (state.status === "completed") { if (!state.summary) { - throw new Error(`DurableEval run ${state.runId} has no saved summary`); + throw new Error(`Durable eval run ${state.runId} has no saved summary`); } return { status: "completed", @@ -1521,15 +1838,15 @@ async function materializeCases< : undefined); if (!caseId) { throw new Error( - "Every DurableEval case requires id, upsert_id, or caseId", + "Every durable eval case requires id, upsert_id, or caseId", ); } if (seen.has(caseId)) - throw new Error(`Duplicate DurableEval case id: ${caseId}`); + throw new Error(`Duplicate durable 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 DurableEval case ${caseId}`); + throw new Error(`Invalid trialCount for durable eval case ${caseId}`); } for (let trialIndex = 0; trialIndex < trialCount; trialIndex++) { records.push({ diff --git a/js/src/exports.ts b/js/src/exports.ts index e2098a7f9..bccb0bd37 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -268,7 +268,13 @@ export { export type { DurableEvalStore } from "./durable-eval"; -export { BatchScorer, BatchTask, DurableEval } from "./durable-eval"; +export { + BatchScorer, + BatchTask, + defineDurableEval, + DurableEvalMemoryStore, + DurableEvalRedisStore, +} from "./durable-eval"; export { agentAssertionScorer } from "./agent-assertions"; diff --git a/knip.jsonc b/knip.jsonc index 7337fa278..9752cb9a2 100644 --- a/knip.jsonc +++ b/knip.jsonc @@ -10,9 +10,6 @@ ], "ignoreIssues": { "**/generated_types.ts": ["exports", "types"], - // These support the inferred signatures of the intentionally small public - // DurableEval API and must remain exported for declaration bundling. - "js/src/durable-eval.ts": ["types"], }, "workspaces": { "dev-packages/seinfeld": { From 51ac4c7496f178f63cacd79bde47e127fe58782f Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 10 Aug 2026 14:37:15 +0000 Subject: [PATCH 11/13] Update PR #2297 --- .gitignore | 1 - .../durable-eval-webhook/scenario.ts | 19 +- js/src/durable-eval.test.ts | 429 ++++++++++- js/src/durable-eval.ts | 665 ++++++++++++++++-- js/src/framework.ts | 51 +- js/src/isomorph.ts | 1 - js/src/node/config.ts | 1 - 7 files changed, 1047 insertions(+), 120 deletions(-) diff --git a/.gitignore b/.gitignore index d6531a16e..bf65941bd 100644 --- a/.gitignore +++ b/.gitignore @@ -12,7 +12,6 @@ dist !.aiderignore .pnpm-store **/.bt-tmp -**/.braintrust/evals docker-compose.override.yml Dockerfile.local diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts index ed6f370e0..7685843ee 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -14,6 +14,10 @@ async function main() { const testRunId = getTestRunId(); const store = new DurableEvalMemoryStore(); const jobs = new Map(); + const webhookCompletion = { + mode: "webhook" as const, + externalId: (handle: { id: string }) => handle.id, + }; const task = BatchTask< number, number, @@ -30,10 +34,7 @@ async function main() { jobs.set(id, items); return { id }; }, - completion: { - mode: "webhook", - externalId: (handle) => handle.id, - }, + completion: webhookCompletion, async collect(handle) { const items = (jobs.get(handle.id) ?? []) as Array<{ id: string; @@ -54,10 +55,7 @@ async function main() { jobs.set(id, items); return { id }; }, - completion: { - mode: "webhook", - externalId: (handle) => handle.id, - }, + completion: webhookCompletion, async collect(handle) { const items = (jobs.get(handle.id) ?? []) as Array<{ id: string; @@ -85,10 +83,7 @@ async function main() { jobs.set(id, items); return { id }; }, - completion: { - mode: "webhook", - externalId: (handle) => handle.id, - }, + completion: webhookCompletion, async collect(handle) { const items = (jobs.get(handle.id) ?? []) as Array<{ id: string; diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts index 005209b95..dae9e95f1 100644 --- a/js/src/durable-eval.test.ts +++ b/js/src/durable-eval.test.ts @@ -8,6 +8,7 @@ import { DurableEvalRedisStore, type DurableBatchScorerItem, type DurableBatchTaskItem, + type DurableEvalStore, } from "./durable-eval"; configureNode(); @@ -25,6 +26,14 @@ describe("durable eval stores", () => { 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 () => { @@ -54,6 +63,84 @@ describe("durable eval stores", () => { }).read("invalid"), ).rejects.toThrow("expected GET to return a string"); }); + + 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?: { NX?: boolean; GET?: boolean }, + ) => { + expect(options).toEqual({ 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); + + 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: string[]) => { + expect(options).toEqual(["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); + + 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?: { nx?: boolean; get?: boolean }, + ) => { + expect(options).toEqual({ 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); + + 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", () => { @@ -94,6 +181,83 @@ describe("defineDurableEval", () => { expect(first.runId).not.toBe(second.runId); }); + 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 = 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({ + schemaVersion: 4, + status: "running", + caseIds: ["one:trial:0", "two:trial:0"], + }); + 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, @@ -217,6 +381,11 @@ describe("defineDurableEval", () => { string, DurableBatchScorerItem[] >(); + let taskCollectCount = 0; + let releaseTaskCollect!: () => void; + const taskBatchesCollecting = new Promise((resolve) => { + releaseTaskCollect = resolve; + }); const task = BatchTask< number, number, @@ -236,6 +405,9 @@ describe("defineDurableEval", () => { externalId: (handle) => handle.id, }, async collect(handle) { + taskCollectCount++; + if (taskCollectCount === 2) releaseTaskCollect(); + await taskBatchesCollecting; return (taskJobs.get(handle.id) ?? []).map((item) => ({ id: item.id, output: item.input * 2, @@ -288,14 +460,16 @@ describe("defineDurableEval", () => { externalId: taskIds[0], }), ).rejects.toThrow("Durable eval run missing-run is missing"); - for (const [index, externalId] of taskIds.entries()) { - const result = await durable.processBatchResult({ runId, externalId }); - expect(result).toMatchObject({ - status: "waiting", - pending: { poll: 0, webhook: index === 0 ? 1 : 2 }, - }); - } + 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; @@ -309,6 +483,182 @@ describe("defineDurableEval", () => { }); }); + 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 = BatchTask< + number, + number, + number, + void, + Record, + { id: string } + >({ + async submit(items) { + taskItems = items; + return { id: "task-provider" }; + }, + completion: { + mode: "webhook", + externalId: (handle) => handle.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 = BatchScorer({ + name: "exact", + submit: scoreSubmit, + completion: { + mode: "webhook", + externalId: (handle) => handle.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("preserves parallel workflow outputs from concurrent webhooks", async () => { + const jobs = new Map>(); + let collectCount = 0; + let releaseCollect!: () => void; + const bothCollecting = new Promise((resolve) => { + releaseCollect = resolve; + }); + const joinSubmit = vi.fn( + async (items: Array<{ id: string; input: unknown }>) => { + jobs.set("join", items); + return { id: "join" }; + }, + ); + const completion = { + mode: "webhook" as const, + externalId: (handle: { id: string }) => handle.id, + }; + const task = BatchTask>({ + workflow(workflow) { + const doubled = workflow.batch("double", { + input: (item) => item.input, + async submit(items) { + jobs.set("double", items); + return { id: "double" }; + }, + completion, + async collect() { + collectCount++; + if (collectCount === 2) releaseCollect(); + await bothCollecting; + return (jobs.get("double") ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) * 2, + })); + }, + }); + const incremented = workflow.batch("increment", { + input: (item) => item.input, + async submit(items) { + jobs.set("increment", items); + return { id: "increment" }; + }, + completion, + async collect() { + collectCount++; + if (collectCount === 2) releaseCollect(); + await bothCollecting; + return (jobs.get("increment") ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) + 1, + })); + }, + }); + return workflow.batch("join", { + needs: { doubled, incremented }, + input: (_item, outputs) => outputs, + submit: joinSubmit, + completion, + async collect() { + return (jobs.get("join") ?? []).map((item) => { + const input = item.input as { + doubled: number; + incremented: number; + }; + return { + id: item.id, + output: input.doubled + input.incremented, + }; + }); + }, + }); + }, + }); + const durable = defineDurableEval("parallel-webhooks", { + store: new DurableEvalMemoryStore(), + data: [{ id: "one", input: 2 }], + task, + }); + const waiting = await durable.start({ noSendLogs: true }); + + await Promise.all( + ["double", "increment"].map((externalId) => + durable.processBatchResult({ + runId: waiting.runId, + externalId, + }), + ), + ); + + expect(joinSubmit).toHaveBeenCalledTimes(1); + await expect( + durable.processBatchResult({ + runId: waiting.runId, + externalId: "join", + }), + ).resolves.toMatchObject({ status: "completed" }); + }); + test("runs multi-stage task and scorer workflows", async () => { const jobs = new Map>(); const submit = async ( @@ -457,6 +807,71 @@ describe("defineDurableEval", () => { }); }); + test("supports workflow and 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 = BatchTask>( + { + workflow(workflow) { + return workflow.batch("constructor", { + input: (item) => item.input, + 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 = BatchScorer({ + 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", { diff --git a/js/src/durable-eval.ts b/js/src/durable-eval.ts index a2c54c100..553981016 100644 --- a/js/src/durable-eval.ts +++ b/js/src/durable-eval.ts @@ -47,7 +47,7 @@ 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 CHECKPOINT_VERSION = 2; +const CHECKPOINT_VERSION = 4; const DEFAULT_BATCH_SIZE = 1_000; type JsonPrimitive = string | number | boolean | null; @@ -55,13 +55,19 @@ type JsonValue = JsonPrimitive | JsonValue[] | { [key: string]: JsonValue }; /** * Minimal persistence used to reconnect provider webhooks with submitted - * batches. Durable evaluations do not require any Braintrust backend changes. + * batches. Each run, case, and batch is stored under its own key. Durable + * 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 { read(key: string): Promise; write(key: string, value: Uint8Array): Promise; + /** Atomically stores `value` when `key` is absent and returns its stored value. */ + getOrSet( + key: string, + value: Uint8Array, + ): Promise<{ value: Uint8Array; created: boolean }>; } /** @@ -80,14 +86,20 @@ export class DurableEvalMemoryStore implements DurableEvalStore { async write(key: string, value: Uint8Array): Promise { this.values.set(key, value.slice()); } + + async getOrSet(key: string, value: Uint8Array) { + const existing = this.values.get(key); + if (existing) return { value: existing.slice(), created: false }; + this.values.set(key, value.slice()); + return { value: value.slice(), created: true }; + } } /** * 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 - * `@upstash/redis` can be passed directly; other clients with compatible async - * `get` and `set` methods are also supported. + * `@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. */ @@ -116,6 +128,39 @@ export class DurableEvalRedisStore implements DurableEvalStore { async write(key: string, value: Uint8Array): Promise { await this.client.set(`${this.keyPrefix}${key}`, uint8ArrayToBase64(value)); } + + async getOrSet(key: string, value: Uint8Array) { + const redisKey = `${this.keyPrefix}${key}`; + const encoded = uint8ArrayToBase64(value); + const client = this.client as typeof this.client & { + createScript?: unknown; + defineCommand?: unknown; + sendCommand?: unknown; + }; + const set = client.set as unknown as ( + ...args: unknown[] + ) => Promise; + let setOptions: unknown[]; + if (typeof client.defineCommand === "function") { + setOptions = ["NX", "GET"]; + } else if (typeof client.sendCommand === "function") { + setOptions = [{ NX: true, GET: true }]; + } else if (typeof client.createScript === "function") { + setOptions = [{ nx: true, get: true }]; + } else { + throw new Error( + "DurableEvalRedisStore 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", + ); + } + return { value: base64ToUint8Array(existing), created: false }; + } } interface DurableBatchContext { @@ -575,6 +620,20 @@ type DurableCaseRecord = { loggedClassifications: Record; }; +type DurableCaseBaseRecord = Pick< + DurableCaseRecord, + "id" | "caseId" | "trialIndex" | "datum" | "metadata" | "tags" +>; + +type DurableTaskResultRecord = Pick< + DurableCaseRecord, + "output" | "metadata" | "tags" +> & { taskComplete: true }; + +type DurableTaskLogRecord = Pick & { + taskLogged: true; +}; + type DurableBatchRecord = { id: string; kind: "task" | "score"; @@ -600,6 +659,10 @@ type DurableRunState = { batches: DurableBatchRecord[]; }; +type DurableRunRecord = Omit & { + caseIds: string[]; +}; + class DurableEvalDefinitionImpl< Input, Output, @@ -997,7 +1060,7 @@ async function startDurableEval< batches: [], }; await experiment?.flush(); - await writeJson(store, key, state); + await writeRunRecord(store, key, state); return currentStatus(definition, state); } const state: DurableRunState = { @@ -1012,7 +1075,8 @@ async function startDurableEval< cases: await materializeCases(definition, data, experiment), batches: [], }; - await writeJson(store, key, state); + await writeCaseBaseRecords(store, key, state.cases); + await writeRunRecord(store, key, state); return advanceDurableEval(definition, state, store, key, experiment); } @@ -1033,7 +1097,8 @@ async function getDurableEvalStatus< options: { runId: string }, ): Promise { const store = definition.evaluator.store; - const state = await readJson( + const state = await readRunState( + definition, store, runKey(definition.projectName, definition.evalName, options.runId), ); @@ -1062,7 +1127,7 @@ async function processDurableBatchResult< } const store = definition.evaluator.store; const key = runKey(definition.projectName, definition.evalName, result.runId); - const state = await readJson(store, key); + 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) @@ -1078,11 +1143,17 @@ async function processDurableBatchResult< const batch = byBatch ?? byExternal; if (!batch) throw new Error("No submitted batch matches this result"); if (batch.status !== "complete") { - await collectBatch(definition, state, batch); + const records = await collectBatch(definition, state, batch); batch.status = "complete"; - await writeJson(store, key, state); + await writeCaseRecords(store, key, records); + await writeBatchRecords(store, key, [batch]); } - return advanceDurableEval(definition, state, store, key); + return advanceDurableEval( + definition, + (await readRunState(definition, store, key))!, + store, + key, + ); } async function pollDurableEval< @@ -1107,7 +1178,7 @@ async function pollDurableEval< definition.evalName, options.runId, ); - const state = await readJson(store, key); + const state = await readRunState(definition, store, key); if (!state) throw new Error(`Durable eval run ${options.runId} is missing`); const batches = state.batches.filter((batch) => { @@ -1128,14 +1199,26 @@ async function pollDurableEval< }), })), ); + const changedCases = new Map(); + const changedBatches: DurableBatchRecord[] = []; for (const { batch, result } of results) { if (result.status === "failed") throw asError(result.error); if (result.status !== "complete") continue; - await collectBatch(definition, state, batch); + for (const record of await collectBatch(definition, state, batch)) { + changedCases.set(record.id, record); + } batch.status = "complete"; - await writeJson(store, key, state); + changedBatches.push(batch); + } + if (changedBatches.length > 0) { + await writeCaseRecords(store, key, [...changedCases.values()]); + await writeBatchRecords(store, key, changedBatches); } - return advanceDurableEval(definition, state, store, key); + const currentState = + changedBatches.length > 0 + ? (await readRunState(definition, store, key))! + : state; + return advanceDurableEval(definition, currentState, store, key); } async function openDurableExperiment( @@ -1187,25 +1270,44 @@ async function advanceDurableEval< : existingExperiment; await runTaskStage(definition, state, store, key, experiment); await logCompletedTasks(definition, state, store, key, experiment); - if (state.cases.some((record) => !record.taskComplete)) { + state = (await readRunState(definition, store, key)) ?? state; + if ( + state.cases.some((record) => !record.taskComplete || !record.taskLogged) + ) { return currentStatus(definition, state); } await runScoreStages(definition, state, store, key, experiment); + state = (await readRunState(definition, store, key)) ?? state; const scorerNames = resolveScorers(definition.evaluator.scores ?? []).map( ({ name }) => name, ); + const classifierNames = (definition.evaluator.classifiers ?? []).map( + classifierName, + ); if ( - state.cases.some((record) => - scorerNames.some((name) => !(name in record.scores)), + state.cases.some( + (record) => + scorerNames.some( + (name) => + !Object.hasOwn(record.scores, name) || + !Object.hasOwn(record.loggedScores, name), + ) || + classifierNames.some( + (name) => !Object.hasOwn(record.loggedClassifications, name), + ), ) ) { return currentStatus(definition, state); } + if (!(await claimAction(store, key, "finish"))) { + const latest = await readRunState(definition, store, key); + return currentStatus(definition, latest ?? state); + } state.summary = await finishExperiment(definition, state, experiment); state.status = "completed"; - await writeJson(store, key, state); + await writeRunRecord(store, key, state); return currentStatus(definition, state); } @@ -1320,7 +1422,6 @@ async function logTaskResult( throw error; } finally { root.end(); - await experiment?.flush(); } } @@ -1331,10 +1432,16 @@ async function logCompletedTasks( key: string, experiment: Experiment | null, ) { + const changed: DurableCaseRecord[] = []; for (const record of state.cases) { if (!record.taskComplete || record.taskLogged) continue; + if (!(await claimAction(store, key, "task-log", record.id))) continue; await logTaskResult(definition, state, record, experiment); - await writeJson(store, key, state); + changed.push(record); + } + if (changed.length > 0) { + await experiment?.flush(); + await writeCaseRecords(store, key, changed); } } @@ -1369,10 +1476,16 @@ async function runTaskStage< Metadata, Parameters >; + const changed: DurableCaseRecord[] = []; for (const record of state.cases) { if (record.taskComplete) continue; + if (!(await claimAction(store, key, "task", record.id))) continue; await logTaskResult(definition, state, record, experiment, task); - await writeJson(store, key, state); + changed.push(record); + } + if (changed.length > 0) { + await experiment?.flush(); + await writeCaseRecords(store, key, changed); } } @@ -1396,10 +1509,23 @@ async function runScoreStages< experiment: Experiment | null, ) { const scorers = resolveScorers(definition.evaluator.scores ?? []); + const changed = new Map(); + const persistChangedCases = async () => { + if (changed.size === 0) return; + await experiment?.flush(); + await writeCaseRecords(store, key, [...changed.values()]); + changed.clear(); + }; for (const { name, scorer } of scorers) { if (isBatchScorer(scorer)) { for (const record of state.cases) { - if (name in record.scores && !record.loggedScores[name]) { + if ( + Object.hasOwn(record.scores, name) && + !Object.hasOwn(record.loggedScores, name) + ) { + if (!(await claimAction(store, key, "score-log", record.id, name))) { + continue; + } await evaluateAndLogScore( definition, state, @@ -1407,14 +1533,16 @@ async function runScoreStages< name, experiment, ); - await writeJson(store, key, state); + changed.set(record.id, record); } } + await persistChangedCases(); await ensureWorkflowBatches(definition, state, store, key, "score", name); continue; } for (const record of state.cases) { - if (record.loggedScores[name]) continue; + if (Object.hasOwn(record.loggedScores, name)) continue; + if (!(await claimAction(store, key, "score", record.id, name))) continue; await evaluateAndLogScore( definition, state, @@ -1423,7 +1551,7 @@ async function runScoreStages< experiment, scorer, ); - await writeJson(store, key, state); + changed.set(record.id, record); } } @@ -1432,7 +1560,10 @@ async function runScoreStages< ).entries()) { const name = classifierName(classifier, index); for (const record of state.cases) { - if (record.loggedClassifications[name]) continue; + if (Object.hasOwn(record.loggedClassifications, name)) continue; + if (!(await claimAction(store, key, "classification", record.id, name))) { + continue; + } await evaluateAndLogClassification( definition, state, @@ -1441,9 +1572,10 @@ async function runScoreStages< classifier, experiment, ); - await writeJson(store, key, state); + changed.set(record.id, record); } } + await persistChangedCases(); } function scorerArgs(record: DurableCaseRecord) { @@ -1517,7 +1649,6 @@ async function evaluateAndLogScore( throw error; } finally { root.end(); - await experiment?.flush(); } } @@ -1567,7 +1698,6 @@ async function evaluateAndLogClassification( throw error; } finally { root.end(); - await experiment?.flush(); } } @@ -1580,34 +1710,38 @@ async function ensureWorkflowBatches( scorerName?: string, ) { const workflow = workflowForStage(definition, kind, scorerName); + const plans = plannedBatches( + definition, + state.runId, + state.cases.map(({ id }) => id), + ); + const casesById = new Map(state.cases.map((record) => [record.id, record])); for (const node of workflow.nodes) { - const batchSize = node.processor.batchSize ?? DEFAULT_BATCH_SIZE; - if (!Number.isInteger(batchSize) || batchSize < 1) { - throw new Error( - `Invalid batchSize for ${scorerName ? `${scorerName}.${node.name}` : node.name}`, - ); - } - const assigned = new Set( - state.batches - .filter( - (batch) => - batch.kind === kind && - batch.scorerName === scorerName && - batch.nodeName === node.name, - ) - .flatMap((batch) => batch.itemIds), + const nodePlans = plans.filter( + (plan) => + plan.kind === kind && + plan.scorerName === scorerName && + plan.nodeName === node.name, ); - const eligible = state.cases.filter((record) => { - const outputs = nodeOutputsFor(record, kind, scorerName); - return ( - !assigned.has(record.id) && - !(node.name in outputs) && - Object.values(node.needs).every((dependency) => dependency in outputs) + for (const plan of nodePlans) { + if (state.batches.some(({ id }) => id === plan.id)) continue; + const records = plan.itemIds.map((id) => casesById.get(id)!); + const ready = records.every((record) => { + const outputs = nodeOutputsFor(record, kind, scorerName); + return ( + !Object.hasOwn(outputs, node.name) && + Object.values(node.needs).every((dependency) => + Object.hasOwn(outputs, dependency), + ) + ); + }); + if (!ready) continue; + const batchId = plan.id; + const claim = await store.getOrSet( + claimRecordKey(key, "batch", batchId), + encoder.encode(batchId), ); - }); - for (let offset = 0; offset < eligible.length; offset += batchSize) { - const records = eligible.slice(offset, offset + batchSize); - const batchId = newId(); + if (!claim.created) continue; const context = { runId: state.runId, batchId }; const items = records.map((record) => itemForNode(state.parameters, record, kind, scorerName, node), @@ -1634,7 +1768,7 @@ async function ensureWorkflowBatches( status: "submitted", }; state.batches.push(batch); - await writeJson(store, key, state); + await writeBatchRecords(store, key, [batch]); } } } @@ -1660,6 +1794,7 @@ async function collectBatch( } 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)) { @@ -1671,6 +1806,7 @@ async function collectBatch( seen.add(id); if ("error" in result) throw asError(result.error); const record = state.cases.find((candidate) => candidate.id === id)!; + records.push(record); const output = assertJsonValue( node.result(result), `output for ${batch.nodeName} item ${id}`, @@ -1702,6 +1838,7 @@ async function collectBatch( `Batch ${batch.id} did not return results for: ${missing.join(", ")}`, ); } + return records; } function processorForBatch( @@ -1797,7 +1934,7 @@ function nodeOutputsFor( scorerName?: string, ) { if (kind === "task") return record.taskNodeOutputs; - return (record.scoreNodeOutputs[scorerName!] ??= {}); + return (record.scoreNodeOutputs[scorerName!] ??= Object.create(null)); } async function materializeCases< @@ -1861,12 +1998,12 @@ async function materializeCases< tags: datum.tags, taskComplete: false, taskLogged: false, - taskNodeOutputs: {}, - scores: {}, - loggedScores: {}, - scoreNodeOutputs: {}, - classifications: {}, - loggedClassifications: {}, + taskNodeOutputs: Object.create(null), + scores: Object.create(null), + loggedScores: Object.create(null), + scoreNodeOutputs: Object.create(null), + classifications: Object.create(null), + loggedClassifications: Object.create(null), }); } } @@ -1909,22 +2046,22 @@ async function finishExperiment( ); const results = state.cases.map((record) => { const datum = record.datum as EvalCase; - const scores = Object.assign( - {}, - ...scorerNames.map( - (name) => + const scores = Object.fromEntries( + scorerNames.flatMap((name) => + Object.entries( _internalPrepareEvaluatorScore( record.scores[name] as OneOrMoreScores, name, ).scores ?? {}, + ), ), ); - const classifications = Object.assign( - {}, - ...Object.entries(record.classifications).map( - ([name, value]) => + const classifications = Object.fromEntries( + Object.entries(record.classifications).flatMap(([name, value]) => + Object.entries( _internalPrepareEvaluatorClassification(value as never, name) .classifications ?? {}, + ), ), ); return { @@ -1977,7 +2114,391 @@ function resolveScorers( } function runKey(projectName: string, evalName: string, runId: string) { - return `durable-eval/v2/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; + return `durable-eval/v4/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; +} + +function encodedKeyPart(value: string) { + return uint8ArrayToBase64(encoder.encode(value)) + .replaceAll("+", "-") + .replaceAll("/", "_") + .replace(/=+$/, ""); +} + +function caseRecordKey( + key: string, + caseId: string, + kind: + | "base" + | "task" + | "task-log" + | "task-node" + | "score" + | "score-log" + | "score-node" + | "classification" + | "classification-log", + ...names: string[] +) { + const suffix = names.map(encodedKeyPart).join("/"); + return `${key}/cases/${encodedKeyPart(caseId)}/${kind}${suffix ? `/${suffix}` : ""}`; +} + +function batchRecordKey(key: string, batchId: string) { + return `${key}/batches/${encodedKeyPart(batchId)}`; +} + +function claimRecordKey(key: string, kind: string, ...parts: string[]) { + const identity = stableStringify([kind, parts]); + return `${key}/claims/${contentVersion(encoder.encode(identity))}`; +} + +async function claimAction( + store: DurableEvalStore, + key: string, + kind: string, + ...parts: string[] +) { + return ( + await store.getOrSet( + claimRecordKey(key, kind, ...parts), + encoder.encode("claimed"), + ) + ).created; +} + +type DurableBatchPlan = Omit< + DurableBatchRecord, + "handle" | "externalId" | "status" +>; + +function plannedBatches( + definition: DurableEvalDefinition, + runId: string, + caseIds: string[], +) { + const stages: Array<{ kind: "task" | "score"; scorerName?: string }> = []; + if (isBatchTask(definition.evaluator.task)) stages.push({ kind: "task" }); + for (const { name, scorer } of resolveScorers( + definition.evaluator.scores ?? [], + )) { + if (isBatchScorer(scorer)) { + stages.push({ kind: "score", scorerName: name }); + } + } + const plans: DurableBatchPlan[] = []; + for (const { kind, scorerName } of stages) { + const workflow = workflowForStage(definition, kind, scorerName); + for (const node of workflow.nodes) { + const batchSize = node.processor.batchSize ?? DEFAULT_BATCH_SIZE; + if (!Number.isInteger(batchSize) || batchSize < 1) { + throw new Error( + `Invalid batchSize for ${scorerName ? `${scorerName}.${node.name}` : node.name}`, + ); + } + for (let offset = 0; offset < caseIds.length; offset += batchSize) { + const itemIds = caseIds.slice(offset, offset + batchSize); + plans.push({ + id: deterministicId( + stableStringify([runId, kind, scorerName, node.name, itemIds]), + ), + kind, + scorerName, + nodeName: node.name, + itemIds, + }); + } + } + } + return plans; +} + +async function readCaseRecord( + definition: DurableEvalDefinition, + store: DurableEvalStore, + key: string, + id: string, +) { + const scorers = resolveScorers(definition.evaluator.scores ?? []); + const classifiers = (definition.evaluator.classifiers ?? []).map( + classifierName, + ); + const taskNodes = isBatchTask(definition.evaluator.task) + ? workflowForStage(definition, "task").nodes + : []; + const scoreNodes = scorers.flatMap(({ name, scorer }) => + isBatchScorer(scorer) + ? workflowForStage(definition, "score", name).nodes.map((node) => ({ + scorerName: name, + nodeName: node.name, + })) + : [], + ); + const [ + base, + task, + taskLog, + taskNodeValues, + scoreValues, + scoreLogValues, + scoreNodeValues, + classificationValues, + classificationLogValues, + ] = await Promise.all([ + readJson(store, caseRecordKey(key, id, "base")), + readJson(store, caseRecordKey(key, id, "task")), + readJson(store, caseRecordKey(key, id, "task-log")), + Promise.all( + taskNodes.map(async ({ name }) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "task-node", name), + ), + })), + ), + Promise.all( + scorers.map(async ({ name }) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "score", name), + ), + })), + ), + Promise.all( + scorers.map(async ({ name }) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "score-log", name), + ), + })), + ), + Promise.all( + scoreNodes.map(async ({ scorerName, nodeName }) => ({ + scorerName, + nodeName, + value: await readJson( + store, + caseRecordKey(key, id, "score-node", scorerName, nodeName), + ), + })), + ), + Promise.all( + classifiers.map(async (name) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "classification", name), + ), + })), + ), + Promise.all( + classifiers.map(async (name) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "classification-log", name), + ), + })), + ), + ]); + if (!base) throw new Error(`Durable eval case ${id} is missing`); + const taskNodeOutputs: Record = Object.create(null); + for (const { name, value } of taskNodeValues) { + if (value !== undefined) taskNodeOutputs[name] = value; + } + const scores: Record = Object.create(null); + for (const { name, value } of scoreValues) { + if (value !== undefined) scores[name] = value; + } + const loggedScores: Record = Object.create(null); + for (const { name, value } of scoreLogValues) { + if (value) loggedScores[name] = true; + } + const scoreNodeOutputs: Record< + string, + Record + > = Object.create(null); + for (const { scorerName, nodeName, value } of scoreNodeValues) { + if (value === undefined) continue; + const outputs = (scoreNodeOutputs[scorerName] ??= Object.create(null)); + outputs[nodeName] = value; + } + const classifications: Record = Object.create(null); + for (const { name, value } of classificationValues) { + if (value !== undefined) classifications[name] = value; + } + const loggedClassifications: Record = Object.create(null); + for (const { name, value } of classificationLogValues) { + if (value) loggedClassifications[name] = true; + } + return { + ...base, + metadata: task?.metadata ?? base.metadata, + tags: task ? task.tags : base.tags, + taskComplete: task !== undefined, + taskLogged: taskLog !== undefined, + output: task?.output, + rootSpan: taskLog?.rootSpan, + taskNodeOutputs, + scores, + loggedScores, + scoreNodeOutputs, + classifications, + loggedClassifications, + } satisfies DurableCaseRecord; +} + +async function readRunState( + definition: DurableEvalDefinition, + store: DurableEvalStore, + key: string, +) { + const record = await readJson(store, key); + if (!record) return undefined; + const plans = plannedBatches(definition, record.runId, record.caseIds); + const [cases, batchRecords] = await Promise.all([ + Promise.all( + record.caseIds.map((id) => readCaseRecord(definition, store, key, id)), + ), + Promise.all( + plans.map(async ({ id }) => { + return readJson(store, batchRecordKey(key, id)); + }), + ), + ]); + const batches = batchRecords.filter( + (value): value is DurableBatchRecord => value !== undefined, + ); + const { caseIds: _caseIds, ...state } = record; + return { ...state, cases, batches }; +} + +async function writeRunRecord( + store: DurableEvalStore, + key: string, + state: DurableRunState, +) { + const { cases, batches: _batches, ...record } = state; + await writeJson(store, key, { + ...record, + caseIds: cases.map(({ id }) => id), + } satisfies DurableRunRecord); +} + +async function writeCaseBaseRecords( + store: DurableEvalStore, + key: string, + records: DurableCaseRecord[], +) { + await Promise.all( + records.map(({ id, caseId, trialIndex, datum, metadata, tags }) => + writeJson(store, caseRecordKey(key, id, "base"), { + id, + caseId, + trialIndex, + datum, + metadata, + tags, + } satisfies DurableCaseBaseRecord), + ), + ); +} + +async function writeCaseRecords( + store: DurableEvalStore, + key: string, + records: DurableCaseRecord[], +) { + const writes: Promise[] = []; + for (const record of records) { + if (record.taskComplete) { + writes.push( + writeJson(store, caseRecordKey(key, record.id, "task"), { + output: record.output, + metadata: record.metadata, + tags: record.tags, + taskComplete: true, + } satisfies DurableTaskResultRecord), + ); + } + if (record.taskLogged) { + writes.push( + writeJson(store, caseRecordKey(key, record.id, "task-log"), { + rootSpan: record.rootSpan, + taskLogged: true, + } satisfies DurableTaskLogRecord), + ); + } + for (const [name, value] of Object.entries(record.taskNodeOutputs)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "task-node", name), + value, + ), + ); + } + for (const [name, value] of Object.entries(record.scores)) { + writes.push( + writeJson(store, caseRecordKey(key, record.id, "score", name), value), + ); + } + for (const name of Object.keys(record.loggedScores)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "score-log", name), + true, + ), + ); + } + for (const [scorerName, outputs] of Object.entries( + record.scoreNodeOutputs, + )) { + for (const [nodeName, value] of Object.entries(outputs)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "score-node", scorerName, nodeName), + value, + ), + ); + } + } + for (const [name, value] of Object.entries(record.classifications)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "classification", name), + value, + ), + ); + } + for (const name of Object.keys(record.loggedClassifications)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "classification-log", name), + true, + ), + ); + } + } + await Promise.all(writes); +} + +async function writeBatchRecords( + store: DurableEvalStore, + key: string, + records: DurableBatchRecord[], +) { + await Promise.all( + records.map((record) => + writeJson(store, batchRecordKey(key, record.id), record), + ), + ); } function deterministicId(value: string) { diff --git a/js/src/framework.ts b/js/src/framework.ts index 8dff8bafd..6371afa9c 100644 --- a/js/src/framework.ts +++ b/js/src/framework.ts @@ -1,6 +1,5 @@ import { makeScorerPropagatedEvent, - mergeDicts, Classification, ClassificationItem, Score, @@ -1017,9 +1016,8 @@ function buildSpanMetadata( ) { return results.length === 1 ? results[0].metadata - : results.reduce( - (prev, s) => mergeDicts(prev, { [s.name]: s.metadata }), - {}, + : Object.fromEntries( + results.map((result) => [result.name, result.metadata]), ); } @@ -1030,9 +1028,8 @@ function buildSpanScores( metadata?: Record; }>, ) { - const scoresRecord = results.reduce( - (prev, s) => mergeDicts(prev, { [s.name]: s.score }), - {}, + const scoresRecord = Object.fromEntries( + results.map((result) => [result.name, result.score]), ); return { resultMetadata: buildSpanMetadata(results), scoresRecord }; } @@ -1056,11 +1053,14 @@ export function _internalPrepareEvaluatorScore( } } } - const results: Score[] = Array.isArray(scoreValue) - ? scoreValue - : typeof scoreValue === "object" && !isEmpty(scoreValue) - ? [scoreValue] - : [{ name, score: scoreValue }]; + let results: Score[]; + if (Array.isArray(scoreValue)) { + results = scoreValue; + } else if (typeof scoreValue === "object" && !isEmpty(scoreValue)) { + results = [scoreValue]; + } else { + results = [{ name, score: scoreValue }]; + } const { resultMetadata, scoresRecord } = buildSpanScores(results); const fields = (score: Score) => { const { metadata: _metadata, name: _name, ...rest } = score; @@ -1071,10 +1071,8 @@ export function _internalPrepareEvaluatorScore( output: results.length === 1 ? fields(results[0]) - : results.reduce( - (previous, score) => - mergeDicts(previous, { [score.name ?? name]: fields(score) }), - {}, + : Object.fromEntries( + results.map((score) => [score.name ?? name, fields(score)]), ), metadata: resultMetadata, scores: scoresRecord, @@ -1160,7 +1158,8 @@ export function _internalPrepareEvaluatorClassification( const results = (Array.isArray(value) ? value : [value]).map((result) => validateClassificationResult(result, name), ); - const classifications: Record = {}; + const classifications: Record = + Object.create(null); for (const result of results) { (classifications[result.name] ??= []).push(toClassificationItem(result)); } @@ -1169,12 +1168,11 @@ export function _internalPrepareEvaluatorClassification( output: results.length === 1 ? toClassificationItem(results[0]) - : results.reduce( - (previous, result) => - mergeDicts(previous, { - [result.name]: toClassificationItem(result), - }), - {}, + : Object.fromEntries( + results.map((result) => [ + result.name, + toClassificationItem(result), + ]), ), metadata: buildSpanMetadata(results), classifications, @@ -1387,8 +1385,9 @@ async function runEvaluatorInternal( let output: unknown = undefined; let error: unknown | undefined = undefined; let tags: string[] = []; - const scores: Record = {}; - const classifications: Record = {}; + const scores: Record = Object.create(null); + const classifications: Record = + Object.create(null); const scorerNames = (evaluator.scores ?? []).map(scorerName); const classifierNames = (evaluator.classifiers ?? []).map( classifierName, @@ -1795,7 +1794,7 @@ function ensureScoreAccumulator( // eslint-disable-next-line @typescript-eslint/no-explicit-any results: EvalResult[], ) { - const accumulator: ScoreAccumulator = {}; + const accumulator: ScoreAccumulator = Object.create(null); for (const result of results) { accumulateScores(accumulator, result.scores); } diff --git a/js/src/isomorph.ts b/js/src/isomorph.ts index a4076bb6a..d23ddcc7e 100644 --- a/js/src/isomorph.ts +++ b/js/src/isomorph.ts @@ -79,7 +79,6 @@ interface Common { writeFile?: (filename: string, data: string | Uint8Array) => Promise; readFile?: (filename: string) => Promise; readdir?: (path: string) => Promise; - rename?: (oldPath: string, newPath: string) => Promise; utimes?: (path: string, atime: Date, mtime: Date) => Promise; unlink?: (path: string) => Promise; // eslint-disable-next-line @typescript-eslint/no-explicit-any diff --git a/js/src/node/config.ts b/js/src/node/config.ts index 5cfa5076b..6b3f37802 100644 --- a/js/src/node/config.ts +++ b/js/src/node/config.ts @@ -120,7 +120,6 @@ export function configureNode() { iso.writeFile = fs.writeFile; iso.readFile = fs.readFile; iso.readdir = fs.readdir; - iso.rename = fs.rename; iso.stat = fs.stat; iso.statSync = fsSync.statSync; iso.utimes = fs.utimes; From 1513c80059534a90c7563ea1f2665e1fd98065c5 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:28:53 +0000 Subject: [PATCH 12/13] Remove workflow API from durable evals --- .../durable-eval-webhook/scenario.ts | 58 +- js/src/durable-eval.test.ts | 275 +------- js/src/durable-eval.ts | 617 +++--------------- 3 files changed, 135 insertions(+), 815 deletions(-) diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts index 7685843ee..a7645a678 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -25,48 +25,22 @@ async function main() { { testRunId: string; kind: string }, Record >({ - workflow(workflow) { - const generated = workflow.batch("generate", { - batchSize: 2, - input: (item) => item.input, - async submit(items) { - const id = `generate-${jobs.size + 1}`; - jobs.set(id, items); - return { id }; - }, - completion: webhookCompletion, - async collect(handle) { - const items = (jobs.get(handle.id) ?? []) as Array<{ - id: string; - input: number; - }>; - return items.map((item) => ({ - id: item.id, - output: item.input * 2, - })); - }, - }); - return workflow.batch("finalize", { - needs: { generated }, - input: (_item, { generated }) => generated, - batchSize: 2, - async submit(items) { - const id = `finalize-${jobs.size + 1}`; - jobs.set(id, items); - return { id }; - }, - completion: webhookCompletion, - async collect(handle) { - const items = (jobs.get(handle.id) ?? []) as Array<{ - id: string; - input: number; - }>; - return items.map((item) => ({ - id: item.id, - output: item.input, - })); - }, - }); + batchSize: 2, + async submit(items) { + const id = `task-${jobs.size + 1}`; + jobs.set(id, items); + return { id }; + }, + completion: webhookCompletion, + async collect(handle) { + const items = (jobs.get(handle.id) ?? []) as Array<{ + id: string; + input: number; + }>; + return items.map((item) => ({ + id: item.id, + output: item.input * 2, + })); }, }); const scorer = BatchScorer< diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts index dae9e95f1..b13d52c09 100644 --- a/js/src/durable-eval.test.ts +++ b/js/src/durable-eval.test.ts @@ -242,7 +242,7 @@ describe("defineDurableEval", () => { ([key]) => !key.includes("/cases/") && !key.includes("/batches/"), )?.[1]; expect(run).toMatchObject({ - schemaVersion: 4, + schemaVersion: 1, status: "running", caseIds: ["one:trial:0", "two:trial:0"], }); @@ -561,253 +561,7 @@ describe("defineDurableEval", () => { }); }); - test("preserves parallel workflow outputs from concurrent webhooks", async () => { - const jobs = new Map>(); - let collectCount = 0; - let releaseCollect!: () => void; - const bothCollecting = new Promise((resolve) => { - releaseCollect = resolve; - }); - const joinSubmit = vi.fn( - async (items: Array<{ id: string; input: unknown }>) => { - jobs.set("join", items); - return { id: "join" }; - }, - ); - const completion = { - mode: "webhook" as const, - externalId: (handle: { id: string }) => handle.id, - }; - const task = BatchTask>({ - workflow(workflow) { - const doubled = workflow.batch("double", { - input: (item) => item.input, - async submit(items) { - jobs.set("double", items); - return { id: "double" }; - }, - completion, - async collect() { - collectCount++; - if (collectCount === 2) releaseCollect(); - await bothCollecting; - return (jobs.get("double") ?? []).map((item) => ({ - id: item.id, - output: (item.input as number) * 2, - })); - }, - }); - const incremented = workflow.batch("increment", { - input: (item) => item.input, - async submit(items) { - jobs.set("increment", items); - return { id: "increment" }; - }, - completion, - async collect() { - collectCount++; - if (collectCount === 2) releaseCollect(); - await bothCollecting; - return (jobs.get("increment") ?? []).map((item) => ({ - id: item.id, - output: (item.input as number) + 1, - })); - }, - }); - return workflow.batch("join", { - needs: { doubled, incremented }, - input: (_item, outputs) => outputs, - submit: joinSubmit, - completion, - async collect() { - return (jobs.get("join") ?? []).map((item) => { - const input = item.input as { - doubled: number; - incremented: number; - }; - return { - id: item.id, - output: input.doubled + input.incremented, - }; - }); - }, - }); - }, - }); - const durable = defineDurableEval("parallel-webhooks", { - store: new DurableEvalMemoryStore(), - data: [{ id: "one", input: 2 }], - task, - }); - const waiting = await durable.start({ noSendLogs: true }); - - await Promise.all( - ["double", "increment"].map((externalId) => - durable.processBatchResult({ - runId: waiting.runId, - externalId, - }), - ), - ); - - expect(joinSubmit).toHaveBeenCalledTimes(1); - await expect( - durable.processBatchResult({ - runId: waiting.runId, - externalId: "join", - }), - ).resolves.toMatchObject({ status: "completed" }); - }); - - test("runs multi-stage task and scorer workflows", async () => { - const jobs = new Map>(); - const submit = async ( - prefix: string, - items: Array<{ id: string; input: unknown }>, - ) => { - const id = `${prefix}-${jobs.size + 1}`; - jobs.set(id, items); - return { id }; - }; - const completion = { - mode: "poll" as const, - async poll() { - return { status: "complete" as const }; - }, - }; - - const task = BatchTask>( - { - workflow(w) { - const doubled = w.batch("double", { - batchSize: 2, - input: (item) => item.input, - submit: (items) => submit("double", items), - completion, - async collect(handle) { - return (jobs.get(handle.id) ?? []).map((item) => ({ - id: item.id, - output: (item.input as number) * 2, - })); - }, - }); - const incremented = w.batch("increment", { - needs: { doubled }, - input: (_item, outputs) => outputs.doubled, - submit: (items) => submit("increment", items), - completion, - async collect(handle) { - return (jobs.get(handle.id) ?? []).map((item) => ({ - id: item.id, - output: (item.input as number) + 1, - })); - }, - }); - const decremented = w.batch("decrement", { - needs: { doubled }, - input: (_item, outputs) => outputs.doubled, - submit: (items) => submit("decrement", items), - completion, - async collect(handle) { - return (jobs.get(handle.id) ?? []).map((item) => ({ - id: item.id, - output: (item.input as number) - 1, - })); - }, - }); - return w.batch("combine", { - needs: { incremented, decremented }, - input: (_item, outputs) => outputs, - submit: (items) => submit("combine", items), - completion, - async collect(handle) { - return (jobs.get(handle.id) ?? []).map((item) => { - const value = item.input as { - incremented: number; - decremented: number; - }; - return { - id: item.id, - output: (value.incremented + value.decremented) / 2, - }; - }); - }, - }); - }, - }, - ); - const scorer = BatchScorer({ - name: "exact", - workflow(w) { - const comparison = w.batch("compare", { - input: (item) => ({ - output: item.output, - expected: item.expected, - }), - submit: (items) => submit("compare", items), - completion, - async collect(handle) { - return (jobs.get(handle.id) ?? []).map((item) => { - const value = item.input as { - output: number; - expected: number; - }; - return { - id: item.id, - output: value.output === value.expected, - }; - }); - }, - }); - return w.batch("score", { - needs: { comparison }, - input: (_item, outputs) => outputs.comparison, - submit: (items) => submit("score", items), - completion, - async collect(handle) { - return (jobs.get(handle.id) ?? []).map((item) => ({ - id: item.id, - output: item.input ? 1 : 0, - })); - }, - }); - }, - }); - const store = new DurableEvalMemoryStore(); - const durable = defineDurableEval("workflow", { - 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", - }); - const options = { runId: waiting.runId }; - await expect(durable.poll(options)).resolves.toMatchObject({ - status: "waiting", - }); - await expect(durable.poll(options)).resolves.toMatchObject({ - status: "waiting", - }); - await expect(durable.poll(options)).resolves.toMatchObject({ - status: "waiting", - }); - await expect(durable.poll(options)).resolves.toMatchObject({ - status: "waiting", - }); - await expect(durable.poll(options)).resolves.toMatchObject({ - status: "completed", - summary: { scores: { exact: { score: 1 } } }, - }); - }); - - test("supports workflow and scorer names inherited from Object.prototype", async () => { + test("supports scorer names inherited from Object.prototype", async () => { const jobs = new Map< string, Array<{ id: string; input: number; expected?: number; output?: number }> @@ -820,21 +574,16 @@ describe("defineDurableEval", () => { }; const task = BatchTask>( { - workflow(workflow) { - return workflow.batch("constructor", { - input: (item) => item.input, - 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, - })); - }, - }); + 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, + })); }, }, ); diff --git a/js/src/durable-eval.ts b/js/src/durable-eval.ts index 553981016..380f385d6 100644 --- a/js/src/durable-eval.ts +++ b/js/src/durable-eval.ts @@ -47,7 +47,7 @@ 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 CHECKPOINT_VERSION = 4; +const CHECKPOINT_VERSION = 1; const DEFAULT_BATCH_SIZE = 1_000; type JsonPrimitive = string | number | boolean | null; @@ -231,73 +231,6 @@ interface DurableBatchProcessor { collect(handle: Handle, context: DurableBatchContext): Promise; } -interface DurableWorkflowBatchItem { - id: string; - input: Input; -} - -type DurableWorkflowBatchResult = - | { - id: string; - output: Output; - metadata?: Metadata; - tags?: string[]; - } - | { id: string; error: unknown }; - -const WORKFLOW_NODE_OUTPUT: unique symbol = Symbol("DurableWorkflowNodeOutput"); - -interface DurableWorkflowNode { - readonly [WORKFLOW_NODE_OUTPUT]: Output; -} - -type DurableWorkflowNodeMap = Record>; - -type DurableWorkflowNodeOutputs = { - [Name in keyof Nodes]: Nodes[Name] extends DurableWorkflowNode - ? Output - : never; -}; - -interface DurableWorkflowBuilder { - batch< - Output, - Needs extends DurableWorkflowNodeMap = Record, - Input = RootItem, - Handle extends JsonValue = JsonValue, - >( - name: string, - processor: DurableBatchProcessor< - DurableWorkflowBatchItem, - DurableWorkflowBatchResult, - Handle - > & { - needs?: Needs; - input?: ( - item: RootItem, - outputs: DurableWorkflowNodeOutputs, - ) => Input; - }, - ): DurableWorkflowNode; -} - -type DurableWorkflowNodeDefinition = { - name: string; - needs: Record; - item: ( - rootItem: unknown, - outputs: Record, - id: string, - ) => unknown; - processor: DurableBatchProcessor; - result: (result: any) => unknown; -}; - -type DurableWorkflowDefinition = { - nodes: DurableWorkflowNodeDefinition[]; - outputNode: string; -}; - interface DurableBatchTask< Input, Output, @@ -307,12 +240,11 @@ interface DurableBatchTask< Handle extends JsonValue, > { readonly kind: typeof BATCH_TASK_KIND; - readonly processor?: DurableBatchProcessor< + readonly processor: DurableBatchProcessor< DurableBatchTaskItem, DurableBatchTaskResult, Handle >; - readonly workflow?: DurableWorkflowDefinition; } interface DurableBatchScorer< @@ -324,12 +256,11 @@ interface DurableBatchScorer< > { readonly kind: typeof BATCH_SCORER_KIND; name: string; - readonly processor?: DurableBatchProcessor< + readonly processor: DurableBatchProcessor< DurableBatchScorerItem, DurableBatchScorerResult, Handle >; - readonly workflow?: DurableWorkflowDefinition; } /** @@ -350,50 +281,8 @@ export function BatchTask< DurableBatchTaskResult, Handle >, -): DurableBatchTask; -export function BatchTask< - Input, - Output, - Expected = void, - Metadata extends BaseMetadata = DefaultMetadataType, - Parameters extends EvalParameters = EvalParameters, ->(config: { - workflow( - builder: DurableWorkflowBuilder< - DurableBatchTaskItem, - Metadata - >, - ): DurableWorkflowNode; -}): DurableBatchTask; -export function BatchTask( - config: - | DurableBatchProcessor - | { - workflow( - builder: DurableWorkflowBuilder, - ): DurableWorkflowNode; - }, -): DurableBatchTask< - unknown, - unknown, - unknown, - BaseMetadata, - EvalParameters, - JsonValue -> { - return { - kind: BATCH_TASK_KIND, - ...("workflow" in config - ? { workflow: buildWorkflow(config.workflow) } - : { processor: config }), - } as DurableBatchTask< - unknown, - unknown, - unknown, - BaseMetadata, - EvalParameters, - JsonValue - >; +): DurableBatchTask { + return { kind: BATCH_TASK_KIND, processor }; } /** @@ -413,106 +302,12 @@ export function BatchScorer< DurableBatchScorerResult, Handle > & { name: string }, -): DurableBatchScorer; -export function BatchScorer< - Input, - Output, - Expected = void, - Metadata extends BaseMetadata = DefaultMetadataType, ->(config: { - name: string; - workflow( - builder: DurableWorkflowBuilder< - DurableBatchScorerItem, - Metadata - >, - ): DurableWorkflowNode; -}): DurableBatchScorer; -export function BatchScorer( - config: - | (DurableBatchProcessor & { name: string }) - | { - name: string; - workflow( - builder: DurableWorkflowBuilder, - ): DurableWorkflowNode; - }, -): DurableBatchScorer { +): DurableBatchScorer { return { kind: BATCH_SCORER_KIND, - name: config.name, - ...("workflow" in config - ? { workflow: buildWorkflow(config.workflow) } - : { processor: config }), - } as DurableBatchScorer; -} - -function buildWorkflow( - define: ( - builder: DurableWorkflowBuilder, - ) => DurableWorkflowNode, -): DurableWorkflowDefinition { - const nodes: DurableWorkflowNodeDefinition[] = []; - const nodeNames = new Map, string>(); - const builder: DurableWorkflowBuilder = { - batch(name, config) { - if (!name.trim()) throw new Error("Workflow batch names cannot be empty"); - if (nodes.some((node) => node.name === name)) { - throw new Error(`Duplicate workflow batch name: ${name}`); - } - const needs = Object.fromEntries( - Object.entries(config.needs ?? {}).map(([alias, dependency]) => { - const dependencyName = nodeNames.get(dependency); - if (!dependencyName) { - throw new Error( - `Workflow batch ${name} depends on an unknown or later batch`, - ); - } - return [alias, dependencyName]; - }), - ); - const { input, needs: _needs, ...processor } = config; - const handle = {} as DurableWorkflowNode; - nodeNames.set(handle, name); - nodes.push({ - name, - needs, - item(rootItem, outputs, id) { - return { - id, - input: input - ? input(rootItem as RootItem, outputs as never) - : rootItem, - }; - }, - processor: processor as DurableBatchProcessor, - result: (result) => result.output, - }); - return handle as never; - }, - }; - const output = define(builder); - const outputNode = nodeNames.get(output); - if (!outputNode) { - throw new Error("A batch workflow must return one of its batch nodes"); - } - const reachable = new Set(); - const visit = (name: string) => { - if (reachable.has(name)) return; - reachable.add(name); - const node = nodes.find((candidate) => candidate.name === name)!; - Object.values(node.needs).forEach(visit); + name: processor.name, + processor, }; - visit(outputNode); - const unused = nodes.filter((node) => !reachable.has(node.name)); - if (unused.length > 0) { - throw new Error( - `Batch workflow contains nodes that do not contribute to its output: ${unused - .map((node) => node.name) - .join(", ")}`, - ); - } - return { nodes, outputNode }; } type DurableEvaluator< @@ -612,10 +407,8 @@ type DurableCaseRecord = { taskLogged: boolean; output?: JsonValue; rootSpan?: string; - taskNodeOutputs: Record; scores: Record; loggedScores: Record; - scoreNodeOutputs: Record>; classifications: Record; loggedClassifications: Record; }; @@ -638,7 +431,6 @@ type DurableBatchRecord = { id: string; kind: "task" | "score"; scorerName?: string; - nodeName: string; itemIds: string[]; handle: JsonValue; externalId?: string; @@ -882,7 +674,7 @@ class DurableEvalDefinitionImpl< * ``` * * Use `status()` to read the same information without polling providers, - * collecting results, or advancing the workflow: + * collecting results, or advancing the evaluation: * * ```typescript * const status = await supportEval.status({ @@ -893,39 +685,6 @@ class DurableEvalDefinitionImpl< * Completed statuses have zero pending batches and include the saved experiment * summary. They can be read repeatedly without logging the eval again. * - * ### Multi-stage workflows - * - * The direct `BatchTask({ submit, completion, collect })` form remains the - * one-batch shorthand. Use `workflow` when a task or scorer requires multiple - * provider batch operations: - * - * ```typescript - * task: BatchTask({ - * workflow(workflow) { - * const draft = workflow.batch("draft", { - * input: ({ input }) => ({ prompt: input }), - * batchSize: 500, - * submit: submitDraftBatch, - * completion: draftCompletion, - * collect: collectDraftBatch, - * }); - * - * return workflow.batch("revise", { - * needs: { draft }, - * input: ({ input }, { draft }) => ({ original: input, draft }), - * batchSize: 500, - * submit: submitRevisionBatch, - * completion: revisionCompletion, - * collect: collectRevisionBatch, - * }); - * }, - * }), - * ``` - * - * Every named batch is a persisted workflow node. `needs` can express sequential - * operations, parallel branches, and joins. The returned node supplies the final - * task output or scorer result. - * * ### Webhook processing * * When the provider reports that any task or scorer batch completed, fetch and @@ -952,9 +711,9 @@ class DurableEvalDefinitionImpl< * ``` * * The method accepts either `externalId` or `batchId`. The stored batch locator - * identifies the task or scorer workflow node, whose `collect()` results are - * stored before the eval advances. Webhook idempotency and provider failure - * handling remain application responsibilities for now. + * 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. */ /** @@ -1183,13 +942,17 @@ async function pollDurableEval< const batches = state.batches.filter((batch) => { if (batch.status === "complete") return false; - return processorForBatch(definition, batch).completion.mode === "poll"; + return ( + processorForStage(definition, batch.kind, batch.scorerName).completion + .mode === "poll" + ); }); const results = await Promise.all( batches.map(async (batch) => ({ batch, result: await ( - processorForBatch(definition, batch).completion as Extract< + processorForStage(definition, batch.kind, batch.scorerName) + .completion as Extract< DurableBatchCompletion, { mode: "poll" } > @@ -1329,7 +1092,10 @@ function currentStatus( const pending = { poll: 0, webhook: 0 }; for (const batch of state.batches) { if (batch.status === "complete") continue; - pending[processorForBatch(definition, batch).completion.mode]++; + pending[ + processorForStage(definition, batch.kind, batch.scorerName).completion + .mode + ]++; } return { status: "waiting", runId: state.runId, pending }; } @@ -1465,7 +1231,7 @@ async function runTaskStage< experiment: Experiment | null, ) { if (isBatchTask(definition.evaluator.task)) { - await ensureWorkflowBatches(definition, state, store, key, "task"); + await ensureBatches(definition, state, store, key, "task"); return; } @@ -1537,7 +1303,7 @@ async function runScoreStages< } } await persistChangedCases(); - await ensureWorkflowBatches(definition, state, store, key, "score", name); + await ensureBatches(definition, state, store, key, "score", name); continue; } for (const record of state.cases) { @@ -1701,7 +1467,7 @@ async function evaluateAndLogClassification( } } -async function ensureWorkflowBatches( +async function ensureBatches( definition: DurableEvalDefinition, state: DurableRunState, store: DurableEvalStore, @@ -1709,67 +1475,57 @@ async function ensureWorkflowBatches( kind: "task" | "score", scorerName?: string, ) { - const workflow = workflowForStage(definition, kind, scorerName); + const processor = processorForStage(definition, kind, scorerName); const plans = plannedBatches( definition, state.runId, state.cases.map(({ id }) => id), ); const casesById = new Map(state.cases.map((record) => [record.id, record])); - for (const node of workflow.nodes) { - const nodePlans = plans.filter( - (plan) => - plan.kind === kind && - plan.scorerName === scorerName && - plan.nodeName === node.name, + 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) => + kind === "task" + ? !record.taskComplete + : !Object.hasOwn(record.scores, scorerName!), ); - for (const plan of nodePlans) { - if (state.batches.some(({ id }) => id === plan.id)) continue; - const records = plan.itemIds.map((id) => casesById.get(id)!); - const ready = records.every((record) => { - const outputs = nodeOutputsFor(record, kind, scorerName); - return ( - !Object.hasOwn(outputs, node.name) && - Object.values(node.needs).every((dependency) => - Object.hasOwn(outputs, dependency), - ) - ); - }); - if (!ready) continue; - const batchId = plan.id; - const claim = await store.getOrSet( - claimRecordKey(key, "batch", batchId), - encoder.encode(batchId), - ); - if (!claim.created) continue; - const context = { runId: state.runId, batchId }; - const items = records.map((record) => - itemForNode(state.parameters, record, kind, scorerName, node), - ); - const handle = assertJsonValue( - await node.processor.submit(items, context), - `handle for batch ${batchId}`, - ); - const externalId = - node.processor.completion.mode === "webhook" - ? node.processor.completion.externalId(handle, context) - : undefined; - if (externalId !== undefined && !externalId.trim()) { - throw new Error(`Batch ${batchId} produced an empty externalId`); - } - const batch: DurableBatchRecord = { - id: batchId, - kind, - scorerName, - nodeName: node.name, - itemIds: records.map((record) => record.id), - handle, - externalId, - status: "submitted", - }; - state.batches.push(batch); - await writeBatchRecords(store, key, [batch]); + if (!ready) continue; + const batchId = plan.id; + const claim = await store.getOrSet( + claimRecordKey(key, "batch", batchId), + encoder.encode(batchId), + ); + if (!claim.created) continue; + const context = { runId: state.runId, batchId }; + const items = records.map((record) => + kind === "task" + ? taskBatchItem(record, state.parameters) + : scorerBatchItem(record), + ); + const handle = assertJsonValue( + await processor.submit(items, context), + `handle for batch ${batchId}`, + ); + const externalId = + processor.completion.mode === "webhook" + ? processor.completion.externalId(handle, context) + : undefined; + if (externalId !== undefined && !externalId.trim()) { + throw new Error(`Batch ${batchId} produced an empty externalId`); } + const batch: DurableBatchRecord = { + id: batchId, + kind, + scorerName, + itemIds: records.map((record) => record.id), + handle, + externalId, + status: "submitted", + }; + state.batches.push(batch); + await writeBatchRecords(store, key, [batch]); } } @@ -1778,15 +1534,7 @@ async function collectBatch( state: DurableRunState, batch: DurableBatchRecord, ) { - const workflow = workflowForStage(definition, batch.kind, batch.scorerName); - const node = workflow.nodes.find( - (candidate) => candidate.name === batch.nodeName, - ); - if (!node) - throw new Error( - `Definition no longer contains batch node ${batch.nodeName}`, - ); - const processor = node.processor; + const processor = processorForStage(definition, batch.kind, batch.scorerName); const context = { runId: state.runId, batchId: batch.id }; const results = await processor.collect(batch.handle, context); if (!Array.isArray(results)) { @@ -1807,15 +1555,11 @@ async function collectBatch( if ("error" in result) throw asError(result.error); const record = state.cases.find((candidate) => candidate.id === id)!; records.push(record); - const output = assertJsonValue( - node.result(result), - `output for ${batch.nodeName} item ${id}`, - ); - nodeOutputsFor(record, batch.kind, batch.scorerName)[batch.nodeName] = - output; - if (batch.nodeName !== workflow.outputNode) continue; if (batch.kind === "task") { - record.output = output; + record.output = assertJsonValue( + result.output, + `task output for item ${id}`, + ); if ("metadata" in result && result.metadata !== undefined) { record.metadata = assertJsonValue( result.metadata, @@ -1827,8 +1571,8 @@ async function collectBatch( record.taskComplete = true; } else { record.scores[batch.scorerName!] = assertJsonValue( - output, - `score for ${id}`, + result.score, + `score output for item ${id}`, ); } } @@ -1841,100 +1585,28 @@ async function collectBatch( return records; } -function processorForBatch( - definition: DurableEvalDefinition, - batch: DurableBatchRecord, -): DurableBatchProcessor { - const node = workflowForStage( - definition, - batch.kind, - batch.scorerName, - ).nodes.find((candidate) => candidate.name === batch.nodeName); - if (!node) - throw new Error( - `Definition no longer contains batch node ${batch.nodeName}`, - ); - return node.processor; -} - -function workflowForStage( +function processorForStage( definition: DurableEvalDefinition, kind: "task" | "score", scorerName?: string, -): DurableWorkflowDefinition { - let stage: - | DurableBatchTask - | DurableBatchScorer; +): DurableBatchProcessor { if (kind === "task") { if (!isBatchTask(definition.evaluator.task)) { throw new Error("Definition no longer contains the batch task"); } - stage = definition.evaluator.task; - } else { - const scorer = resolveScorers(definition.evaluator.scores ?? []).find( - ({ name }) => name === scorerName, - )?.scorer; - if (!isBatchScorer(scorer)) { - throw new Error(`Definition no longer contains scorer ${scorerName}`); - } - stage = scorer; + return definition.evaluator.task.processor as DurableBatchProcessor< + any, + any, + JsonValue + >; } - if (stage.workflow) return stage.workflow; - if (!stage.processor) { - throw new Error("Batch definition has neither a processor nor a workflow"); + const scorer = resolveScorers(definition.evaluator.scores ?? []).find( + ({ name }) => name === scorerName, + )?.scorer; + if (!isBatchScorer(scorer)) { + throw new Error(`Definition no longer contains scorer ${scorerName}`); } - return { - outputNode: "$batch", - nodes: [ - { - name: "$batch", - needs: {}, - item: (rootItem) => rootItem, - processor: stage.processor as DurableBatchProcessor< - any, - any, - JsonValue - >, - result: - kind === "task" - ? (result) => result.output - : (result) => result.score, - }, - ], - }; -} - -function itemForNode( - parameters: JsonValue, - record: DurableCaseRecord, - kind: "task" | "score", - scorerName: string | undefined, - node: DurableWorkflowNodeDefinition, -) { - const rootItem = - kind === "task" - ? taskBatchItem(record, parameters) - : scorerBatchItem(record); - const outputs = nodeOutputsFor(record, kind, scorerName); - return node.item( - rootItem, - Object.fromEntries( - Object.entries(node.needs).map(([alias, dependency]) => [ - alias, - outputs[dependency], - ]), - ), - record.id, - ); -} - -function nodeOutputsFor( - record: DurableCaseRecord, - kind: "task" | "score", - scorerName?: string, -) { - if (kind === "task") return record.taskNodeOutputs; - return (record.scoreNodeOutputs[scorerName!] ??= Object.create(null)); + return scorer.processor as DurableBatchProcessor; } async function materializeCases< @@ -1998,10 +1670,8 @@ async function materializeCases< tags: datum.tags, taskComplete: false, taskLogged: false, - taskNodeOutputs: Object.create(null), scores: Object.create(null), loggedScores: Object.create(null), - scoreNodeOutputs: Object.create(null), classifications: Object.create(null), loggedClassifications: Object.create(null), }); @@ -2114,7 +1784,7 @@ function resolveScorers( } function runKey(projectName: string, evalName: string, runId: string) { - return `durable-eval/v4/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; + return `durable-eval/v1/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; } function encodedKeyPart(value: string) { @@ -2131,10 +1801,8 @@ function caseRecordKey( | "base" | "task" | "task-log" - | "task-node" | "score" | "score-log" - | "score-node" | "classification" | "classification-log", ...names: string[] @@ -2187,26 +1855,24 @@ function plannedBatches( } const plans: DurableBatchPlan[] = []; for (const { kind, scorerName } of stages) { - const workflow = workflowForStage(definition, kind, scorerName); - for (const node of workflow.nodes) { - const batchSize = node.processor.batchSize ?? DEFAULT_BATCH_SIZE; - if (!Number.isInteger(batchSize) || batchSize < 1) { - throw new Error( - `Invalid batchSize for ${scorerName ? `${scorerName}.${node.name}` : node.name}`, - ); - } - for (let offset = 0; offset < caseIds.length; offset += batchSize) { - const itemIds = caseIds.slice(offset, offset + batchSize); - plans.push({ - id: deterministicId( - stableStringify([runId, kind, scorerName, node.name, itemIds]), - ), - kind, - scorerName, - nodeName: node.name, - itemIds, - }); - } + 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); + plans.push({ + id: deterministicId( + stableStringify([runId, kind, scorerName, itemIds]), + ), + kind, + scorerName, + itemIds, + }); } } return plans; @@ -2222,40 +1888,18 @@ async function readCaseRecord( const classifiers = (definition.evaluator.classifiers ?? []).map( classifierName, ); - const taskNodes = isBatchTask(definition.evaluator.task) - ? workflowForStage(definition, "task").nodes - : []; - const scoreNodes = scorers.flatMap(({ name, scorer }) => - isBatchScorer(scorer) - ? workflowForStage(definition, "score", name).nodes.map((node) => ({ - scorerName: name, - nodeName: node.name, - })) - : [], - ); const [ base, task, taskLog, - taskNodeValues, scoreValues, scoreLogValues, - scoreNodeValues, classificationValues, classificationLogValues, ] = await Promise.all([ readJson(store, caseRecordKey(key, id, "base")), readJson(store, caseRecordKey(key, id, "task")), readJson(store, caseRecordKey(key, id, "task-log")), - Promise.all( - taskNodes.map(async ({ name }) => ({ - name, - value: await readJson( - store, - caseRecordKey(key, id, "task-node", name), - ), - })), - ), Promise.all( scorers.map(async ({ name }) => ({ name, @@ -2274,16 +1918,6 @@ async function readCaseRecord( ), })), ), - Promise.all( - scoreNodes.map(async ({ scorerName, nodeName }) => ({ - scorerName, - nodeName, - value: await readJson( - store, - caseRecordKey(key, id, "score-node", scorerName, nodeName), - ), - })), - ), Promise.all( classifiers.map(async (name) => ({ name, @@ -2304,10 +1938,6 @@ async function readCaseRecord( ), ]); if (!base) throw new Error(`Durable eval case ${id} is missing`); - const taskNodeOutputs: Record = Object.create(null); - for (const { name, value } of taskNodeValues) { - if (value !== undefined) taskNodeOutputs[name] = value; - } const scores: Record = Object.create(null); for (const { name, value } of scoreValues) { if (value !== undefined) scores[name] = value; @@ -2316,15 +1946,6 @@ async function readCaseRecord( for (const { name, value } of scoreLogValues) { if (value) loggedScores[name] = true; } - const scoreNodeOutputs: Record< - string, - Record - > = Object.create(null); - for (const { scorerName, nodeName, value } of scoreNodeValues) { - if (value === undefined) continue; - const outputs = (scoreNodeOutputs[scorerName] ??= Object.create(null)); - outputs[nodeName] = value; - } const classifications: Record = Object.create(null); for (const { name, value } of classificationValues) { if (value !== undefined) classifications[name] = value; @@ -2341,10 +1962,8 @@ async function readCaseRecord( taskLogged: taskLog !== undefined, output: task?.output, rootSpan: taskLog?.rootSpan, - taskNodeOutputs, scores, loggedScores, - scoreNodeOutputs, classifications, loggedClassifications, } satisfies DurableCaseRecord; @@ -2431,15 +2050,6 @@ async function writeCaseRecords( } satisfies DurableTaskLogRecord), ); } - for (const [name, value] of Object.entries(record.taskNodeOutputs)) { - writes.push( - writeJson( - store, - caseRecordKey(key, record.id, "task-node", name), - value, - ), - ); - } for (const [name, value] of Object.entries(record.scores)) { writes.push( writeJson(store, caseRecordKey(key, record.id, "score", name), value), @@ -2454,19 +2064,6 @@ async function writeCaseRecords( ), ); } - for (const [scorerName, outputs] of Object.entries( - record.scoreNodeOutputs, - )) { - for (const [nodeName, value] of Object.entries(outputs)) { - writes.push( - writeJson( - store, - caseRecordKey(key, record.id, "score-node", scorerName, nodeName), - value, - ), - ); - } - } for (const [name, value] of Object.entries(record.classifications)) { writes.push( writeJson( From a42ecbff74c05c417fe49b70ee5d047593f61897 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:32:16 +0000 Subject: [PATCH 13/13] Add durable eval workflows --- .../durable-eval-webhook/scenario.ts | 58 +- js/src/durable-eval.test.ts | 275 +++++++- js/src/durable-eval.ts | 617 +++++++++++++++--- 3 files changed, 815 insertions(+), 135 deletions(-) diff --git a/e2e/scenarios/durable-eval-webhook/scenario.ts b/e2e/scenarios/durable-eval-webhook/scenario.ts index a7645a678..7685843ee 100644 --- a/e2e/scenarios/durable-eval-webhook/scenario.ts +++ b/e2e/scenarios/durable-eval-webhook/scenario.ts @@ -25,22 +25,48 @@ async function main() { { 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(handle) { - const items = (jobs.get(handle.id) ?? []) as Array<{ - id: string; - input: number; - }>; - return items.map((item) => ({ - id: item.id, - output: item.input * 2, - })); + workflow(workflow) { + const generated = workflow.batch("generate", { + batchSize: 2, + input: (item) => item.input, + async submit(items) { + const id = `generate-${jobs.size + 1}`; + jobs.set(id, items); + return { id }; + }, + completion: webhookCompletion, + async collect(handle) { + const items = (jobs.get(handle.id) ?? []) as Array<{ + id: string; + input: number; + }>; + return items.map((item) => ({ + id: item.id, + output: item.input * 2, + })); + }, + }); + return workflow.batch("finalize", { + needs: { generated }, + input: (_item, { generated }) => generated, + batchSize: 2, + async submit(items) { + const id = `finalize-${jobs.size + 1}`; + jobs.set(id, items); + return { id }; + }, + completion: webhookCompletion, + async collect(handle) { + const items = (jobs.get(handle.id) ?? []) as Array<{ + id: string; + input: number; + }>; + return items.map((item) => ({ + id: item.id, + output: item.input, + })); + }, + }); }, }); const scorer = BatchScorer< diff --git a/js/src/durable-eval.test.ts b/js/src/durable-eval.test.ts index b13d52c09..17f53bf14 100644 --- a/js/src/durable-eval.test.ts +++ b/js/src/durable-eval.test.ts @@ -242,7 +242,7 @@ describe("defineDurableEval", () => { ([key]) => !key.includes("/cases/") && !key.includes("/batches/"), )?.[1]; expect(run).toMatchObject({ - schemaVersion: 1, + schemaVersion: 2, status: "running", caseIds: ["one:trial:0", "two:trial:0"], }); @@ -561,7 +561,253 @@ describe("defineDurableEval", () => { }); }); - test("supports scorer names inherited from Object.prototype", async () => { + test("preserves parallel workflow outputs from concurrent webhooks", async () => { + const jobs = new Map>(); + let collectCount = 0; + let releaseCollect!: () => void; + const bothCollecting = new Promise((resolve) => { + releaseCollect = resolve; + }); + const joinSubmit = vi.fn( + async (items: Array<{ id: string; input: unknown }>) => { + jobs.set("join", items); + return { id: "join" }; + }, + ); + const completion = { + mode: "webhook" as const, + externalId: (handle: { id: string }) => handle.id, + }; + const task = BatchTask>({ + workflow(workflow) { + const doubled = workflow.batch("double", { + input: (item) => item.input, + async submit(items) { + jobs.set("double", items); + return { id: "double" }; + }, + completion, + async collect() { + collectCount++; + if (collectCount === 2) releaseCollect(); + await bothCollecting; + return (jobs.get("double") ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) * 2, + })); + }, + }); + const incremented = workflow.batch("increment", { + input: (item) => item.input, + async submit(items) { + jobs.set("increment", items); + return { id: "increment" }; + }, + completion, + async collect() { + collectCount++; + if (collectCount === 2) releaseCollect(); + await bothCollecting; + return (jobs.get("increment") ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) + 1, + })); + }, + }); + return workflow.batch("join", { + needs: { doubled, incremented }, + input: (_item, outputs) => outputs, + submit: joinSubmit, + completion, + async collect() { + return (jobs.get("join") ?? []).map((item) => { + const input = item.input as { + doubled: number; + incremented: number; + }; + return { + id: item.id, + output: input.doubled + input.incremented, + }; + }); + }, + }); + }, + }); + const durable = defineDurableEval("parallel-webhooks", { + store: new DurableEvalMemoryStore(), + data: [{ id: "one", input: 2 }], + task, + }); + const waiting = await durable.start({ noSendLogs: true }); + + await Promise.all( + ["double", "increment"].map((externalId) => + durable.processBatchResult({ + runId: waiting.runId, + externalId, + }), + ), + ); + + expect(joinSubmit).toHaveBeenCalledTimes(1); + await expect( + durable.processBatchResult({ + runId: waiting.runId, + externalId: "join", + }), + ).resolves.toMatchObject({ status: "completed" }); + }); + + test("runs multi-stage task and scorer workflows", async () => { + const jobs = new Map>(); + const submit = async ( + prefix: string, + items: Array<{ id: string; input: unknown }>, + ) => { + const id = `${prefix}-${jobs.size + 1}`; + jobs.set(id, items); + return { id }; + }; + const completion = { + mode: "poll" as const, + async poll() { + return { status: "complete" as const }; + }, + }; + + const task = BatchTask>( + { + workflow(w) { + const doubled = w.batch("double", { + batchSize: 2, + input: (item) => item.input, + submit: (items) => submit("double", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) * 2, + })); + }, + }); + const incremented = w.batch("increment", { + needs: { doubled }, + input: (_item, outputs) => outputs.doubled, + submit: (items) => submit("increment", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) + 1, + })); + }, + }); + const decremented = w.batch("decrement", { + needs: { doubled }, + input: (_item, outputs) => outputs.doubled, + submit: (items) => submit("decrement", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: (item.input as number) - 1, + })); + }, + }); + return w.batch("combine", { + needs: { incremented, decremented }, + input: (_item, outputs) => outputs, + submit: (items) => submit("combine", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => { + const value = item.input as { + incremented: number; + decremented: number; + }; + return { + id: item.id, + output: (value.incremented + value.decremented) / 2, + }; + }); + }, + }); + }, + }, + ); + const scorer = BatchScorer({ + name: "exact", + workflow(w) { + const comparison = w.batch("compare", { + input: (item) => ({ + output: item.output, + expected: item.expected, + }), + submit: (items) => submit("compare", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => { + const value = item.input as { + output: number; + expected: number; + }; + return { + id: item.id, + output: value.output === value.expected, + }; + }); + }, + }); + return w.batch("score", { + needs: { comparison }, + input: (_item, outputs) => outputs.comparison, + submit: (items) => submit("score", items), + completion, + async collect(handle) { + return (jobs.get(handle.id) ?? []).map((item) => ({ + id: item.id, + output: item.input ? 1 : 0, + })); + }, + }); + }, + }); + const store = new DurableEvalMemoryStore(); + const durable = defineDurableEval("workflow", { + 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", + }); + const options = { runId: waiting.runId }; + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "waiting", + }); + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "waiting", + }); + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "waiting", + }); + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "waiting", + }); + await expect(durable.poll(options)).resolves.toMatchObject({ + status: "completed", + summary: { scores: { exact: { score: 1 } } }, + }); + }); + + test("supports workflow and scorer names inherited from Object.prototype", async () => { const jobs = new Map< string, Array<{ id: string; input: number; expected?: number; output?: number }> @@ -574,16 +820,21 @@ describe("defineDurableEval", () => { }; const task = BatchTask>( { - 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, - })); + workflow(workflow) { + return workflow.batch("constructor", { + input: (item) => item.input, + 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, + })); + }, + }); }, }, ); diff --git a/js/src/durable-eval.ts b/js/src/durable-eval.ts index 380f385d6..113b920be 100644 --- a/js/src/durable-eval.ts +++ b/js/src/durable-eval.ts @@ -47,7 +47,7 @@ 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 CHECKPOINT_VERSION = 1; +const CHECKPOINT_VERSION = 2; const DEFAULT_BATCH_SIZE = 1_000; type JsonPrimitive = string | number | boolean | null; @@ -231,6 +231,73 @@ interface DurableBatchProcessor { collect(handle: Handle, context: DurableBatchContext): Promise; } +interface DurableWorkflowBatchItem { + id: string; + input: Input; +} + +type DurableWorkflowBatchResult = + | { + id: string; + output: Output; + metadata?: Metadata; + tags?: string[]; + } + | { id: string; error: unknown }; + +const WORKFLOW_NODE_OUTPUT: unique symbol = Symbol("DurableWorkflowNodeOutput"); + +interface DurableWorkflowNode { + readonly [WORKFLOW_NODE_OUTPUT]: Output; +} + +type DurableWorkflowNodeMap = Record>; + +type DurableWorkflowNodeOutputs = { + [Name in keyof Nodes]: Nodes[Name] extends DurableWorkflowNode + ? Output + : never; +}; + +interface DurableWorkflowBuilder { + batch< + Output, + Needs extends DurableWorkflowNodeMap = Record, + Input = RootItem, + Handle extends JsonValue = JsonValue, + >( + name: string, + processor: DurableBatchProcessor< + DurableWorkflowBatchItem, + DurableWorkflowBatchResult, + Handle + > & { + needs?: Needs; + input?: ( + item: RootItem, + outputs: DurableWorkflowNodeOutputs, + ) => Input; + }, + ): DurableWorkflowNode; +} + +type DurableWorkflowNodeDefinition = { + name: string; + needs: Record; + item: ( + rootItem: unknown, + outputs: Record, + id: string, + ) => unknown; + processor: DurableBatchProcessor; + result: (result: any) => unknown; +}; + +type DurableWorkflowDefinition = { + nodes: DurableWorkflowNodeDefinition[]; + outputNode: string; +}; + interface DurableBatchTask< Input, Output, @@ -240,11 +307,12 @@ interface DurableBatchTask< Handle extends JsonValue, > { readonly kind: typeof BATCH_TASK_KIND; - readonly processor: DurableBatchProcessor< + readonly processor?: DurableBatchProcessor< DurableBatchTaskItem, DurableBatchTaskResult, Handle >; + readonly workflow?: DurableWorkflowDefinition; } interface DurableBatchScorer< @@ -256,11 +324,12 @@ interface DurableBatchScorer< > { readonly kind: typeof BATCH_SCORER_KIND; name: string; - readonly processor: DurableBatchProcessor< + readonly processor?: DurableBatchProcessor< DurableBatchScorerItem, DurableBatchScorerResult, Handle >; + readonly workflow?: DurableWorkflowDefinition; } /** @@ -281,8 +350,50 @@ export function BatchTask< DurableBatchTaskResult, Handle >, -): DurableBatchTask { - return { kind: BATCH_TASK_KIND, processor }; +): DurableBatchTask; +export function BatchTask< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, + Parameters extends EvalParameters = EvalParameters, +>(config: { + workflow( + builder: DurableWorkflowBuilder< + DurableBatchTaskItem, + Metadata + >, + ): DurableWorkflowNode; +}): DurableBatchTask; +export function BatchTask( + config: + | DurableBatchProcessor + | { + workflow( + builder: DurableWorkflowBuilder, + ): DurableWorkflowNode; + }, +): DurableBatchTask< + unknown, + unknown, + unknown, + BaseMetadata, + EvalParameters, + JsonValue +> { + return { + kind: BATCH_TASK_KIND, + ...("workflow" in config + ? { workflow: buildWorkflow(config.workflow) } + : { processor: config }), + } as DurableBatchTask< + unknown, + unknown, + unknown, + BaseMetadata, + EvalParameters, + JsonValue + >; } /** @@ -302,12 +413,106 @@ export function BatchScorer< DurableBatchScorerResult, Handle > & { name: string }, -): DurableBatchScorer { +): DurableBatchScorer; +export function BatchScorer< + Input, + Output, + Expected = void, + Metadata extends BaseMetadata = DefaultMetadataType, +>(config: { + name: string; + workflow( + builder: DurableWorkflowBuilder< + DurableBatchScorerItem, + Metadata + >, + ): DurableWorkflowNode; +}): DurableBatchScorer; +export function BatchScorer( + config: + | (DurableBatchProcessor & { name: string }) + | { + name: string; + workflow( + builder: DurableWorkflowBuilder, + ): DurableWorkflowNode; + }, +): DurableBatchScorer { return { kind: BATCH_SCORER_KIND, - name: processor.name, - processor, + name: config.name, + ...("workflow" in config + ? { workflow: buildWorkflow(config.workflow) } + : { processor: config }), + } as DurableBatchScorer; +} + +function buildWorkflow( + define: ( + builder: DurableWorkflowBuilder, + ) => DurableWorkflowNode, +): DurableWorkflowDefinition { + const nodes: DurableWorkflowNodeDefinition[] = []; + const nodeNames = new Map, string>(); + const builder: DurableWorkflowBuilder = { + batch(name, config) { + if (!name.trim()) throw new Error("Workflow batch names cannot be empty"); + if (nodes.some((node) => node.name === name)) { + throw new Error(`Duplicate workflow batch name: ${name}`); + } + const needs = Object.fromEntries( + Object.entries(config.needs ?? {}).map(([alias, dependency]) => { + const dependencyName = nodeNames.get(dependency); + if (!dependencyName) { + throw new Error( + `Workflow batch ${name} depends on an unknown or later batch`, + ); + } + return [alias, dependencyName]; + }), + ); + const { input, needs: _needs, ...processor } = config; + const handle = {} as DurableWorkflowNode; + nodeNames.set(handle, name); + nodes.push({ + name, + needs, + item(rootItem, outputs, id) { + return { + id, + input: input + ? input(rootItem as RootItem, outputs as never) + : rootItem, + }; + }, + processor: processor as DurableBatchProcessor, + result: (result) => result.output, + }); + return handle as never; + }, + }; + const output = define(builder); + const outputNode = nodeNames.get(output); + if (!outputNode) { + throw new Error("A batch workflow must return one of its batch nodes"); + } + const reachable = new Set(); + const visit = (name: string) => { + if (reachable.has(name)) return; + reachable.add(name); + const node = nodes.find((candidate) => candidate.name === name)!; + Object.values(node.needs).forEach(visit); }; + visit(outputNode); + const unused = nodes.filter((node) => !reachable.has(node.name)); + if (unused.length > 0) { + throw new Error( + `Batch workflow contains nodes that do not contribute to its output: ${unused + .map((node) => node.name) + .join(", ")}`, + ); + } + return { nodes, outputNode }; } type DurableEvaluator< @@ -407,8 +612,10 @@ type DurableCaseRecord = { taskLogged: boolean; output?: JsonValue; rootSpan?: string; + taskNodeOutputs: Record; scores: Record; loggedScores: Record; + scoreNodeOutputs: Record>; classifications: Record; loggedClassifications: Record; }; @@ -431,6 +638,7 @@ type DurableBatchRecord = { id: string; kind: "task" | "score"; scorerName?: string; + nodeName: string; itemIds: string[]; handle: JsonValue; externalId?: string; @@ -674,7 +882,7 @@ class DurableEvalDefinitionImpl< * ``` * * Use `status()` to read the same information without polling providers, - * collecting results, or advancing the evaluation: + * collecting results, or advancing the workflow: * * ```typescript * const status = await supportEval.status({ @@ -685,6 +893,39 @@ class DurableEvalDefinitionImpl< * Completed statuses have zero pending batches and include the saved experiment * summary. They can be read repeatedly without logging the eval again. * + * ### Multi-stage workflows + * + * The direct `BatchTask({ submit, completion, collect })` form remains the + * one-batch shorthand. Use `workflow` when a task or scorer requires multiple + * provider batch operations: + * + * ```typescript + * task: BatchTask({ + * workflow(workflow) { + * const draft = workflow.batch("draft", { + * input: ({ input }) => ({ prompt: input }), + * batchSize: 500, + * submit: submitDraftBatch, + * completion: draftCompletion, + * collect: collectDraftBatch, + * }); + * + * return workflow.batch("revise", { + * needs: { draft }, + * input: ({ input }, { draft }) => ({ original: input, draft }), + * batchSize: 500, + * submit: submitRevisionBatch, + * completion: revisionCompletion, + * collect: collectRevisionBatch, + * }); + * }, + * }), + * ``` + * + * Every named batch is a persisted workflow node. `needs` can express sequential + * operations, parallel branches, and joins. The returned node supplies the final + * task output or scorer result. + * * ### Webhook processing * * When the provider reports that any task or scorer batch completed, fetch and @@ -711,9 +952,9 @@ class DurableEvalDefinitionImpl< * ``` * * 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. + * identifies the task or scorer workflow node, whose `collect()` results are + * stored before the eval advances. Webhook idempotency and provider failure + * handling remain application responsibilities for now. */ /** @@ -942,17 +1183,13 @@ async function pollDurableEval< const batches = state.batches.filter((batch) => { if (batch.status === "complete") return false; - return ( - processorForStage(definition, batch.kind, batch.scorerName).completion - .mode === "poll" - ); + return processorForBatch(definition, batch).completion.mode === "poll"; }); const results = await Promise.all( batches.map(async (batch) => ({ batch, result: await ( - processorForStage(definition, batch.kind, batch.scorerName) - .completion as Extract< + processorForBatch(definition, batch).completion as Extract< DurableBatchCompletion, { mode: "poll" } > @@ -1092,10 +1329,7 @@ function currentStatus( 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 - ]++; + pending[processorForBatch(definition, batch).completion.mode]++; } return { status: "waiting", runId: state.runId, pending }; } @@ -1231,7 +1465,7 @@ async function runTaskStage< experiment: Experiment | null, ) { if (isBatchTask(definition.evaluator.task)) { - await ensureBatches(definition, state, store, key, "task"); + await ensureWorkflowBatches(definition, state, store, key, "task"); return; } @@ -1303,7 +1537,7 @@ async function runScoreStages< } } await persistChangedCases(); - await ensureBatches(definition, state, store, key, "score", name); + await ensureWorkflowBatches(definition, state, store, key, "score", name); continue; } for (const record of state.cases) { @@ -1467,7 +1701,7 @@ async function evaluateAndLogClassification( } } -async function ensureBatches( +async function ensureWorkflowBatches( definition: DurableEvalDefinition, state: DurableRunState, store: DurableEvalStore, @@ -1475,57 +1709,67 @@ async function ensureBatches( kind: "task" | "score", scorerName?: string, ) { - const processor = processorForStage(definition, kind, scorerName); + const workflow = workflowForStage(definition, kind, scorerName); const plans = plannedBatches( 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) => - kind === "task" - ? !record.taskComplete - : !Object.hasOwn(record.scores, scorerName!), + for (const node of workflow.nodes) { + const nodePlans = plans.filter( + (plan) => + plan.kind === kind && + plan.scorerName === scorerName && + plan.nodeName === node.name, ); - if (!ready) continue; - const batchId = plan.id; - const claim = await store.getOrSet( - claimRecordKey(key, "batch", batchId), - encoder.encode(batchId), - ); - if (!claim.created) continue; - const context = { runId: state.runId, batchId }; - const items = records.map((record) => - kind === "task" - ? taskBatchItem(record, state.parameters) - : scorerBatchItem(record), - ); - const handle = assertJsonValue( - await processor.submit(items, context), - `handle for batch ${batchId}`, - ); - const externalId = - processor.completion.mode === "webhook" - ? processor.completion.externalId(handle, context) - : undefined; - if (externalId !== undefined && !externalId.trim()) { - throw new Error(`Batch ${batchId} produced an empty externalId`); + for (const plan of nodePlans) { + if (state.batches.some(({ id }) => id === plan.id)) continue; + const records = plan.itemIds.map((id) => casesById.get(id)!); + const ready = records.every((record) => { + const outputs = nodeOutputsFor(record, kind, scorerName); + return ( + !Object.hasOwn(outputs, node.name) && + Object.values(node.needs).every((dependency) => + Object.hasOwn(outputs, dependency), + ) + ); + }); + if (!ready) continue; + const batchId = plan.id; + const claim = await store.getOrSet( + claimRecordKey(key, "batch", batchId), + encoder.encode(batchId), + ); + if (!claim.created) continue; + const context = { runId: state.runId, batchId }; + const items = records.map((record) => + itemForNode(state.parameters, record, kind, scorerName, node), + ); + const handle = assertJsonValue( + await node.processor.submit(items, context), + `handle for batch ${batchId}`, + ); + const externalId = + node.processor.completion.mode === "webhook" + ? node.processor.completion.externalId(handle, context) + : undefined; + if (externalId !== undefined && !externalId.trim()) { + throw new Error(`Batch ${batchId} produced an empty externalId`); + } + const batch: DurableBatchRecord = { + id: batchId, + kind, + scorerName, + nodeName: node.name, + itemIds: records.map((record) => record.id), + handle, + externalId, + status: "submitted", + }; + state.batches.push(batch); + await writeBatchRecords(store, key, [batch]); } - const batch: DurableBatchRecord = { - id: batchId, - kind, - scorerName, - itemIds: records.map((record) => record.id), - handle, - externalId, - status: "submitted", - }; - state.batches.push(batch); - await writeBatchRecords(store, key, [batch]); } } @@ -1534,7 +1778,15 @@ async function collectBatch( state: DurableRunState, batch: DurableBatchRecord, ) { - const processor = processorForStage(definition, batch.kind, batch.scorerName); + const workflow = workflowForStage(definition, batch.kind, batch.scorerName); + const node = workflow.nodes.find( + (candidate) => candidate.name === batch.nodeName, + ); + if (!node) + throw new Error( + `Definition no longer contains batch node ${batch.nodeName}`, + ); + const processor = node.processor; const context = { runId: state.runId, batchId: batch.id }; const results = await processor.collect(batch.handle, context); if (!Array.isArray(results)) { @@ -1555,11 +1807,15 @@ async function collectBatch( if ("error" in result) throw asError(result.error); const record = state.cases.find((candidate) => candidate.id === id)!; records.push(record); + const output = assertJsonValue( + node.result(result), + `output for ${batch.nodeName} item ${id}`, + ); + nodeOutputsFor(record, batch.kind, batch.scorerName)[batch.nodeName] = + output; + if (batch.nodeName !== workflow.outputNode) continue; if (batch.kind === "task") { - record.output = assertJsonValue( - result.output, - `task output for item ${id}`, - ); + record.output = output; if ("metadata" in result && result.metadata !== undefined) { record.metadata = assertJsonValue( result.metadata, @@ -1571,8 +1827,8 @@ async function collectBatch( record.taskComplete = true; } else { record.scores[batch.scorerName!] = assertJsonValue( - result.score, - `score output for item ${id}`, + output, + `score for ${id}`, ); } } @@ -1585,28 +1841,100 @@ async function collectBatch( return records; } -function processorForStage( +function processorForBatch( + definition: DurableEvalDefinition, + batch: DurableBatchRecord, +): DurableBatchProcessor { + const node = workflowForStage( + definition, + batch.kind, + batch.scorerName, + ).nodes.find((candidate) => candidate.name === batch.nodeName); + if (!node) + throw new Error( + `Definition no longer contains batch node ${batch.nodeName}`, + ); + return node.processor; +} + +function workflowForStage( definition: DurableEvalDefinition, kind: "task" | "score", scorerName?: string, -): DurableBatchProcessor { +): DurableWorkflowDefinition { + let stage: + | DurableBatchTask + | DurableBatchScorer; if (kind === "task") { if (!isBatchTask(definition.evaluator.task)) { throw new Error("Definition no longer contains the batch task"); } - return definition.evaluator.task.processor as DurableBatchProcessor< - any, - any, - JsonValue - >; + stage = definition.evaluator.task; + } else { + const scorer = resolveScorers(definition.evaluator.scores ?? []).find( + ({ name }) => name === scorerName, + )?.scorer; + if (!isBatchScorer(scorer)) { + throw new Error(`Definition no longer contains scorer ${scorerName}`); + } + stage = scorer; } - const scorer = resolveScorers(definition.evaluator.scores ?? []).find( - ({ name }) => name === scorerName, - )?.scorer; - if (!isBatchScorer(scorer)) { - throw new Error(`Definition no longer contains scorer ${scorerName}`); + if (stage.workflow) return stage.workflow; + if (!stage.processor) { + throw new Error("Batch definition has neither a processor nor a workflow"); } - return scorer.processor as DurableBatchProcessor; + return { + outputNode: "$batch", + nodes: [ + { + name: "$batch", + needs: {}, + item: (rootItem) => rootItem, + processor: stage.processor as DurableBatchProcessor< + any, + any, + JsonValue + >, + result: + kind === "task" + ? (result) => result.output + : (result) => result.score, + }, + ], + }; +} + +function itemForNode( + parameters: JsonValue, + record: DurableCaseRecord, + kind: "task" | "score", + scorerName: string | undefined, + node: DurableWorkflowNodeDefinition, +) { + const rootItem = + kind === "task" + ? taskBatchItem(record, parameters) + : scorerBatchItem(record); + const outputs = nodeOutputsFor(record, kind, scorerName); + return node.item( + rootItem, + Object.fromEntries( + Object.entries(node.needs).map(([alias, dependency]) => [ + alias, + outputs[dependency], + ]), + ), + record.id, + ); +} + +function nodeOutputsFor( + record: DurableCaseRecord, + kind: "task" | "score", + scorerName?: string, +) { + if (kind === "task") return record.taskNodeOutputs; + return (record.scoreNodeOutputs[scorerName!] ??= Object.create(null)); } async function materializeCases< @@ -1670,8 +1998,10 @@ async function materializeCases< tags: datum.tags, taskComplete: false, taskLogged: false, + taskNodeOutputs: Object.create(null), scores: Object.create(null), loggedScores: Object.create(null), + scoreNodeOutputs: Object.create(null), classifications: Object.create(null), loggedClassifications: Object.create(null), }); @@ -1784,7 +2114,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 `durable-eval/v2/runs/${contentVersion(encoder.encode(`${projectName}\0${evalName}\0${runId}`))}`; } function encodedKeyPart(value: string) { @@ -1801,8 +2131,10 @@ function caseRecordKey( | "base" | "task" | "task-log" + | "task-node" | "score" | "score-log" + | "score-node" | "classification" | "classification-log", ...names: string[] @@ -1855,24 +2187,26 @@ function plannedBatches( } const plans: DurableBatchPlan[] = []; 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); - plans.push({ - id: deterministicId( - stableStringify([runId, kind, scorerName, itemIds]), - ), - kind, - scorerName, - itemIds, - }); + const workflow = workflowForStage(definition, kind, scorerName); + for (const node of workflow.nodes) { + const batchSize = node.processor.batchSize ?? DEFAULT_BATCH_SIZE; + if (!Number.isInteger(batchSize) || batchSize < 1) { + throw new Error( + `Invalid batchSize for ${scorerName ? `${scorerName}.${node.name}` : node.name}`, + ); + } + for (let offset = 0; offset < caseIds.length; offset += batchSize) { + const itemIds = caseIds.slice(offset, offset + batchSize); + plans.push({ + id: deterministicId( + stableStringify([runId, kind, scorerName, node.name, itemIds]), + ), + kind, + scorerName, + nodeName: node.name, + itemIds, + }); + } } } return plans; @@ -1888,18 +2222,40 @@ async function readCaseRecord( const classifiers = (definition.evaluator.classifiers ?? []).map( classifierName, ); + const taskNodes = isBatchTask(definition.evaluator.task) + ? workflowForStage(definition, "task").nodes + : []; + const scoreNodes = scorers.flatMap(({ name, scorer }) => + isBatchScorer(scorer) + ? workflowForStage(definition, "score", name).nodes.map((node) => ({ + scorerName: name, + nodeName: node.name, + })) + : [], + ); const [ base, task, taskLog, + taskNodeValues, scoreValues, scoreLogValues, + scoreNodeValues, classificationValues, classificationLogValues, ] = await Promise.all([ readJson(store, caseRecordKey(key, id, "base")), readJson(store, caseRecordKey(key, id, "task")), readJson(store, caseRecordKey(key, id, "task-log")), + Promise.all( + taskNodes.map(async ({ name }) => ({ + name, + value: await readJson( + store, + caseRecordKey(key, id, "task-node", name), + ), + })), + ), Promise.all( scorers.map(async ({ name }) => ({ name, @@ -1918,6 +2274,16 @@ async function readCaseRecord( ), })), ), + Promise.all( + scoreNodes.map(async ({ scorerName, nodeName }) => ({ + scorerName, + nodeName, + value: await readJson( + store, + caseRecordKey(key, id, "score-node", scorerName, nodeName), + ), + })), + ), Promise.all( classifiers.map(async (name) => ({ name, @@ -1938,6 +2304,10 @@ async function readCaseRecord( ), ]); if (!base) throw new Error(`Durable eval case ${id} is missing`); + const taskNodeOutputs: Record = Object.create(null); + for (const { name, value } of taskNodeValues) { + if (value !== undefined) taskNodeOutputs[name] = value; + } const scores: Record = Object.create(null); for (const { name, value } of scoreValues) { if (value !== undefined) scores[name] = value; @@ -1946,6 +2316,15 @@ async function readCaseRecord( for (const { name, value } of scoreLogValues) { if (value) loggedScores[name] = true; } + const scoreNodeOutputs: Record< + string, + Record + > = Object.create(null); + for (const { scorerName, nodeName, value } of scoreNodeValues) { + if (value === undefined) continue; + const outputs = (scoreNodeOutputs[scorerName] ??= Object.create(null)); + outputs[nodeName] = value; + } const classifications: Record = Object.create(null); for (const { name, value } of classificationValues) { if (value !== undefined) classifications[name] = value; @@ -1962,8 +2341,10 @@ async function readCaseRecord( taskLogged: taskLog !== undefined, output: task?.output, rootSpan: taskLog?.rootSpan, + taskNodeOutputs, scores, loggedScores, + scoreNodeOutputs, classifications, loggedClassifications, } satisfies DurableCaseRecord; @@ -2050,6 +2431,15 @@ async function writeCaseRecords( } satisfies DurableTaskLogRecord), ); } + for (const [name, value] of Object.entries(record.taskNodeOutputs)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "task-node", name), + value, + ), + ); + } for (const [name, value] of Object.entries(record.scores)) { writes.push( writeJson(store, caseRecordKey(key, record.id, "score", name), value), @@ -2064,6 +2454,19 @@ async function writeCaseRecords( ), ); } + for (const [scorerName, outputs] of Object.entries( + record.scoreNodeOutputs, + )) { + for (const [nodeName, value] of Object.entries(outputs)) { + writes.push( + writeJson( + store, + caseRecordKey(key, record.id, "score-node", scorerName, nodeName), + value, + ), + ); + } + } for (const [name, value] of Object.entries(record.classifications)) { writes.push( writeJson(