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 1/9] 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 2/9] 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 3/9] 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 4/9] 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 5/9] 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 6/9] 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 7/9] 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 8/9] 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 9/9] 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; }