From 07c34a6a683fae1023a6f8f2a0d3a5edb924888c Mon Sep 17 00:00:00 2001 From: Luke Parke <5702154+LukasParke@users.noreply.github.com> Date: Wed, 5 Aug 2026 14:17:20 -0500 Subject: [PATCH] feat(harness): durable per-(sample, epoch) result store + fiber-scoped request context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the monorepo evals-product stack's harness plumbing onto the standalone layout (openrouter-web#31077 review rounds included): - SampleResultStore: durable per-(sample, epoch) records with a pinned wire schema; completed entries seed the accumulator on retry so a re-run chunk reproduces the same aggregate. Degraded (error-synthesized) scores are written only at end of run and ignored by the resume skip-list, so retries re-attempt them for free. Failed writes fail the run after bounded retries — a successful run guarantees durable records. - request-context: fiber-scoped (sampleId, epoch) identity, stamped into each request's x_bench extension by the OpenRouter model layer, so a budget-gateway can key request coalescing per epoch (multi-epoch runs must not replay one epoch's answer). OpenRouter ignores the field on direct calls. - wandr primary score/reward mean exclude Skipped (infrastructure- degraded) samples, matching aggregateScores' accuracy denominator. - runBenchmarkById accepts an optional sampleResultStore (noop default). --- package.json | 4 +- src/benchmarks/wandr/scorer.test.ts | 63 +++- src/benchmarks/wandr/scorer.ts | 19 +- src/harness/request-context.ts | 29 ++ src/harness/run-sample-resume.test.ts | 444 ++++++++++++++++++++++++ src/harness/run.test.ts | 39 ++- src/harness/run.ts | 323 +++++++++++++---- src/harness/sample-result-store.test.ts | 179 ++++++++++ src/harness/sample-result-store.ts | 268 ++++++++++++++ src/providers/openrouter-model.ts | 12 + src/results/parquet-schema.ts | 5 +- src/runner/run-by-id.ts | 16 +- test/helpers/noop-progress-layer.ts | 8 + 13 files changed, 1325 insertions(+), 84 deletions(-) create mode 100644 src/harness/request-context.ts create mode 100644 src/harness/run-sample-resume.test.ts create mode 100644 src/harness/sample-result-store.test.ts create mode 100644 src/harness/sample-result-store.ts diff --git a/package.json b/package.json index 00c92f4..0fa607e 100644 --- a/package.json +++ b/package.json @@ -23,7 +23,9 @@ "./progress": "./src/harness/progress.ts", "./result-store": "./src/results/result-store.ts", "./internal/log": "./src/internal/log.ts", - "./internal/effect-logger": "./src/internal/effect-logger.ts" + "./internal/effect-logger": "./src/internal/effect-logger.ts", + "./sample-result-store": "./src/harness/sample-result-store.ts", + "./request-context": "./src/harness/request-context.ts" }, "scripts": { "bench": "bun src/cli/index.ts", diff --git a/src/benchmarks/wandr/scorer.test.ts b/src/benchmarks/wandr/scorer.test.ts index d03cd3f..c63b66f 100644 --- a/src/benchmarks/wandr/scorer.test.ts +++ b/src/benchmarks/wandr/scorer.test.ts @@ -65,19 +65,51 @@ describe("WANDR score aggregation", () => { expect(wandrRunLevelScores(result)).toEqual([]); expect(wandrPrimaryScore(result)).toBeUndefined(); }); - it("includes degraded samples as zero rewards in the aggregate denominator", () => { + + it("excludes skipped (infrastructure-failure) samples from the mean and weight", () => { + const result = resultWith([1, 0]); + const validSample = result.sampleScores[0]!; + const skippedSample = result.sampleScores[1]!; + const partialFailure = { + ...result, + sampleScores: [ + validSample, + { + ...skippedSample, + score: { + ...skippedSample.score, + value: ScoreValue.Skipped, + explanation: "Solver error (skipped): timeout", + }, + }, + ], + }; + + expect(wandrRunLevelScores(partialFailure)).toEqual([ + { + name: "wandr", + metrics: Object.fromEntries( + WANDR_REWARD_NAMES.map((name) => [name, { value: 1 }]) + ), + }, + ]); + expect(wandrPrimaryScore(partialFailure)).toEqual({ value: 1, weight: 1 }); + }); + + it("counts evaluated samples without reward metadata as zero rewards", () => { const result = resultWith([1, 0]); const validSample = result.sampleScores[0]!; - const degradedSample = result.sampleScores[1]!; + const unparseableSample = result.sampleScores[1]!; const partialFailure = { ...result, sampleScores: [ validSample, { - ...degradedSample, + ...unparseableSample, score: { - ...degradedSample.score, - explanation: "Solver error: timeout", + ...unparseableSample.score, + value: ScoreValue.Incorrect, + explanation: "not reward json", }, }, ], @@ -95,4 +127,25 @@ describe("WANDR score aggregation", () => { weight: 2, }); }); + + it("returns no aggregate when every sample was skipped", () => { + const result = resultWith([1]); + const skippedSample = result.sampleScores[0]!; + const allSkipped = { + ...result, + sampleScores: [ + { + ...skippedSample, + score: { + ...skippedSample.score, + value: ScoreValue.Skipped, + explanation: "Solver error (skipped): sandbox down", + }, + }, + ], + }; + + expect(wandrRunLevelScores(allSkipped)).toEqual([]); + expect(wandrPrimaryScore(allSkipped)).toBeUndefined(); + }); }); diff --git a/src/benchmarks/wandr/scorer.ts b/src/benchmarks/wandr/scorer.ts index 3d64545..e73734c 100644 --- a/src/benchmarks/wandr/scorer.ts +++ b/src/benchmarks/wandr/scorer.ts @@ -53,6 +53,18 @@ export const wandrScorer: ScorerService = ( return succeed(score); }; +/** + * Samples that were actually evaluated. Skipped scores are infrastructure + * failures (degraded solver/model errors), excluded from the accuracy + * denominator by `aggregateScores` — the reward mean and primary-score weight + * must exclude them the same way. + */ +function scoredWandrSamples(result: RunResult): RunResult["sampleScores"] { + return result.sampleScores.filter( + (sample) => sample.score.value !== ScoreValue.Skipped + ); +} + export function wandrRunLevelScores(result: RunResult): readonly { name: string; metrics: Readonly< @@ -64,7 +76,7 @@ export function wandrRunLevelScores(result: RunResult): readonly { > >; }[] { - const rewards = result.sampleScores.flatMap((sample) => { + const rewards = scoredWandrSamples(result).flatMap((sample) => { const parsed = Either.try((): unknown => JSON.parse(sample.score.explanation) ); @@ -97,12 +109,13 @@ export function wandrRunLevelScores(result: RunResult): readonly { export function wandrPrimaryScore( result: RunResult ): BenchmarkPrimaryScore | undefined { - if (result.sampleScores.length === 0) { + const weight = scoredWandrSamples(result).length; + if (weight === 0) { return undefined; } const metrics = wandrRunLevelScores(result)[0]?.metrics; return { value: metrics?.["soft_f1_full"]?.value ?? 0, - weight: result.sampleScores.length, + weight, }; } diff --git a/src/harness/request-context.ts b/src/harness/request-context.ts new file mode 100644 index 0000000..9bfcf5b --- /dev/null +++ b/src/harness/request-context.ts @@ -0,0 +1,29 @@ +import type { Effect } from "effect/Effect"; +import { get, set, unsafeMake } from "effect/FiberRef"; + +/** + * Fiber-scoped (sample, epoch) identity of the evaluation currently issuing + * model calls. The model layer stamps it into each request's `x_bench` + * extension so the bench-gateway's request-coalescing hash can tell two + * epochs of the same sample apart — without it, identical bodies from + * different epochs coalesce and epoch 2 replays epoch 1's answer. Fiber-scoped + * for the same reason as `generationIdCollector`: each (sample, epoch) runs + * in its own child fiber under streamMapEffect concurrency. + */ +export interface BenchRequestContext { + readonly sampleId: string; + readonly epoch: number; +} + +export const benchRequestContext = unsafeMake( + undefined +); + +export function setBenchRequestContext( + context: BenchRequestContext +): Effect { + return set(benchRequestContext, context); +} + +export const getBenchRequestContext: Effect = + get(benchRequestContext); diff --git a/src/harness/run-sample-resume.test.ts b/src/harness/run-sample-resume.test.ts new file mode 100644 index 0000000..c433b50 --- /dev/null +++ b/src/harness/run-sample-resume.test.ts @@ -0,0 +1,444 @@ +import { describe, expect, it } from "bun:test"; + +import { fromIterable } from "effect/Chunk"; +import { + fail as effectFail, + orDie, + promise as effectPromise, + provide, + runPromise, + succeed as effectSucceed, +} from "effect/Effect"; +import type { Layer } from "effect/Layer"; +import { mergeAll, succeed as layerSucceed } from "effect/Layer"; +import { fromChunk } from "effect/Stream"; + +import { + noopProgressLayer, + noopCheckpointLayer, +} from "../../test/helpers/noop-progress-layer"; +import { mcqScorer } from "../benchmarks/scorers/mcq/scorer"; +import type { ModelOutput, Sample } from "./core"; +import { MessageRole, ModelError, ScoreValue, SolverError } from "./core"; +import { Dataset } from "./dataset"; +import type { ModelService } from "./model"; +import { Model } from "./model"; +import type { CheckpointStore, ProgressReporter } from "./progress"; +import { runBenchmark } from "./run"; +import type { + CompletedSampleEntry, + SampleResultStoreService, +} from "./sample-result-store"; +import { SampleResultStore } from "./sample-result-store"; +import { Scorer } from "./scorer"; +import { generate, Solver } from "./solver"; + +const SAMPLES: readonly Sample[] = [ + { id: "s-0", input: "Q0 target B", target: { text: "B" } }, + { id: "s-1", input: "Q1 target B", target: { text: "B" } }, + { id: "s-2", input: "Q2 target B", target: { text: "B" } }, +]; + +const MODEL_RESPONSE: ModelOutput = { + completion: "Answer: B", + message: { role: MessageRole.Assistant, content: "Answer: B" }, + usage: { + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + reasoningTokens: 0, + totalCost: 0.001, + }, + generationTimeMs: 100, +}; + +const ANSWERING_MODEL: ModelService = { + generate: () => effectSucceed(MODEL_RESPONSE), +}; + +/** Records every write and replays a seeded set of prior records. */ +function recordingStore(seed: readonly CompletedSampleEntry[]): { + readonly service: SampleResultStoreService; + readonly writes: CompletedSampleEntry[]; +} { + const writes: CompletedSampleEntry[] = []; + return { + writes, + service: { + write: async (entry) => { + writes.push(entry); + }, + list: async (range) => + seed.filter( + (entry) => + (range.start === undefined || entry.sampleIndex >= range.start) && + (range.end === undefined || entry.sampleIndex < range.end) + ), + }, + }; +} + +function completedEntry( + sampleIndex: number, + epoch: number +): CompletedSampleEntry { + return { + sampleIndex, + epoch, + sampleScore: { + sampleId: `s-${sampleIndex}`, + epoch, + score: { value: ScoreValue.Correct, answer: "B", explanation: "stored" }, + }, + usage: { + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + reasoningTokens: 0, + totalCost: 0.001, + }, + generationTimeMs: 100, + }; +} + +function makeLayers( + store: SampleResultStoreService, + model: ModelService = ANSWERING_MODEL +): Layer< + | Dataset + | Solver + | Scorer + | Model + | SampleResultStore + | ProgressReporter + | CheckpointStore +> { + return mergeAll( + layerSucceed( + Dataset, + Dataset.of({ + stream: (opts) => + fromChunk( + fromIterable( + SAMPLES.slice(opts?.start ?? 0, opts?.end ?? SAMPLES.length) + ) + ), + size: effectSucceed(SAMPLES.length), + }) + ), + layerSucceed(Solver, Solver.of(generate(model, { temperature: 0 }))), + layerSucceed(Scorer, Scorer.of(mcqScorer)), + layerSucceed(Model, Model.of(model)), + layerSucceed(SampleResultStore, store), + noopProgressLayer, + noopCheckpointLayer + ); +} + +/** Resolve once `condition` holds, so a test can observe a write mid-run. */ +async function waitFor(condition: () => boolean): Promise { + while (!condition()) { + // oxlint-disable-next-line eslint/no-await-in-loop -- polling an in-flight run + await new Promise((resolve) => { + setTimeout(resolve, 5); + }); + } +} + +describe("runBenchmark sample-result resume", () => { + it("persists one record per completed (sample, epoch)", async () => { + const store = recordingStore([]); + + const result = await runPromise( + runBenchmark({ epochs: 2, maxConcurrency: 2 }).pipe( + provide(makeLayers(store.service)) + ) + ); + + expect(result.sampleScores).toHaveLength(6); + expect( + store.writes + .map((w) => `${w.sampleIndex}:${w.epoch}`) + .toSorted((a, b) => a.localeCompare(b)) + ).toEqual(["0:0", "0:1", "1:0", "1:1", "2:0", "2:1"]); + expect(store.writes[0]?.usage?.totalCost).toBe(0.001); + expect(store.writes[0]?.generationTimeMs).toBe(100); + }); + + it("skips stored (sample, epoch) pairs and reproduces the uninterrupted aggregate", async () => { + const fresh = await runPromise( + runBenchmark({ epochs: 1, maxConcurrency: 2 }).pipe( + provide(makeLayers(recordingStore([]).service)) + ) + ); + + const resumed = recordingStore([ + completedEntry(0, 0), + completedEntry(1, 0), + ]); + const result = await runPromise( + runBenchmark({ epochs: 1, maxConcurrency: 2 }).pipe( + provide(makeLayers(resumed.service)) + ) + ); + + expect(resumed.writes.map((w) => w.sampleIndex)).toEqual([2]); + expect(result.metrics).toEqual(fresh.metrics); + expect(result.usage).toEqual(fresh.usage); + expect(result.sampleScores.map((s) => s.sampleId).toSorted()).toEqual([ + "s-0", + "s-1", + "s-2", + ]); + }); + + it("keys stored records on the absolute sample index when running a chunk range", async () => { + const resumed = recordingStore([completedEntry(1, 0)]); + + const result = await runPromise( + runBenchmark({ + epochs: 1, + maxConcurrency: 1, + range: { start: 1, end: 3 }, + }).pipe(provide(makeLayers(resumed.service))) + ); + + expect(resumed.writes.map((w) => w.sampleIndex)).toEqual([2]); + expect(result.metrics.totalQuestions).toBe(2); + }); + + it("ignores stored records whose epoch is outside the configured epoch count", async () => { + const resumed = recordingStore([completedEntry(0, 5)]); + + const result = await runPromise( + runBenchmark({ epochs: 1, maxConcurrency: 2 }).pipe( + provide(makeLayers(resumed.service)) + ) + ); + + expect(resumed.writes.map((w) => w.sampleIndex)).toEqual([0, 1, 2]); + expect(result.sampleScores).toHaveLength(3); + }); + + it("records scores synthesized from an exhausted model error as degraded, at the end of the run", async () => { + const store = recordingStore([]); + const rateLimitedModel: ModelService = { + generate: () => + effectFail( + new ModelError({ + message: "OpenRouter HTTP 429: rate-limited", + status: 429, + }) + ), + }; + + const result = await runPromise( + runBenchmark({ epochs: 1, maxConcurrency: 1 }).pipe( + provide(makeLayers(store.service, rateLimitedModel)) + ) + ); + + expect(result.metrics.skippedQuestions).toBe(3); + expect(store.writes).toHaveLength(3); + expect(store.writes.every((w) => w.degraded === true)).toBe(true); + }); + + it("writes a degraded record after every non-degraded write of the run", async () => { + const store = recordingStore([]); + const flakyFirstSampleModel: ModelService = { + generate: (messages) => { + const userMsg = + messages.find((m) => m.role === MessageRole.User)?.content ?? ""; + if (typeof userMsg === "string" && userMsg.includes("Q0")) { + return effectFail( + new ModelError({ + message: "OpenRouter HTTP 429: rate-limited", + status: 429, + }) + ); + } + return effectSucceed(MODEL_RESPONSE); + }, + }; + + await runPromise( + runBenchmark({ epochs: 1, maxConcurrency: 1 }).pipe( + provide(makeLayers(store.service, flakyFirstSampleModel)) + ) + ); + + expect(store.writes.map((w) => w.degraded === true)).toEqual([ + false, + false, + true, + ]); + expect(store.writes.at(-1)?.sampleIndex).toBe(0); + }); + + it("re-runs samples whose stored record is degraded", async () => { + const degradedSeed: CompletedSampleEntry = { + ...completedEntry(0, 0), + degraded: true, + sampleScore: { + sampleId: "s-0", + epoch: 0, + score: { + value: ScoreValue.Skipped, + answer: null, + explanation: "Model error (skipped)", + }, + }, + }; + const resumed = recordingStore([degradedSeed, completedEntry(1, 0)]); + + const result = await runPromise( + runBenchmark({ epochs: 1, maxConcurrency: 2 }).pipe( + provide(makeLayers(resumed.service)) + ) + ); + + expect(resumed.writes.map((w) => w.sampleIndex).toSorted()).toEqual([0, 2]); + expect(result.metrics.skippedQuestions).toBe(0); + expect(result.metrics.totalQuestions).toBe(3); + }); + + it("scores a degraded solver error as Skipped, out of the accuracy denominator", async () => { + const store = recordingStore([]); + const layers = mergeAll( + layerSucceed( + Dataset, + Dataset.of({ + stream: () => fromChunk(fromIterable(SAMPLES.slice(0, 1))), + size: effectSucceed(1), + }) + ), + layerSucceed( + Solver, + Solver.of(() => + effectFail(new SolverError({ message: "sandbox exec failed" })) + ) + ), + layerSucceed(Scorer, Scorer.of(mcqScorer)), + layerSucceed(Model, Model.of(ANSWERING_MODEL)), + layerSucceed(SampleResultStore, store.service), + noopProgressLayer, + noopCheckpointLayer + ); + + const result = await runPromise( + runBenchmark({ + epochs: 1, + maxConcurrency: 1, + degradeSolverErrors: true, + }).pipe(provide(layers)) + ); + + expect(result.metrics.skippedQuestions).toBe(1); + expect(result.sampleScores[0]?.score.value).toBe(ScoreValue.Skipped); + expect(store.writes).toHaveLength(1); + expect(store.writes[0]?.degraded).toBe(true); + }); + + it("records a finished result without waiting for an earlier slow sample", async () => { + const store = recordingStore([]); + let releaseSlowSample = (): void => {}; + const slowSampleStarted = new Promise((resolve) => { + const markStarted = resolve; + const gate = new Promise((_resolve) => { + releaseSlowSample = _resolve; + }); + const slowModel: ModelService = { + generate: (messages) => { + const userMsg = + messages.find((m) => m.role === MessageRole.User)?.content ?? ""; + if (typeof userMsg !== "string" || !userMsg.includes("Q0")) { + return effectSucceed(MODEL_RESPONSE); + } + return effectPromise(async () => { + markStarted(); + await gate; + return MODEL_RESPONSE; + }).pipe(orDie); + }, + }; + // oxlint-disable-next-line eslint/no-void -- fire-and-forget: the test observes the run via the store + void runPromise( + runBenchmark({ epochs: 1, maxConcurrency: 3 }).pipe( + provide(makeLayers(store.service, slowModel)) + ) + ); + }); + + await slowSampleStarted; + await waitFor(() => store.writes.length === 2); + + expect(store.writes.map((w) => w.sampleIndex).toSorted()).toEqual([1, 2]); + releaseSlowSample(); + }); + + it("re-evaluates everything when listing stored records fails", async () => { + const failingStore: SampleResultStoreService = { + write: async () => {}, + list: async () => { + throw new Error("gcs unavailable"); + }, + }; + + const result = await runPromise( + runBenchmark({ epochs: 1, maxConcurrency: 2 }).pipe( + provide(makeLayers(failingStore)) + ) + ); + + expect(result.sampleScores).toHaveLength(3); + }); + + it("retries a failed record write within the sample and completes once it succeeds", async () => { + let attempts = 0; + const writes: CompletedSampleEntry[] = []; + const flakyStore: SampleResultStoreService = { + write: async (entry) => { + attempts += 1; + if (attempts <= 2) { + throw new Error("gcs unavailable"); + } + writes.push(entry); + }, + list: async () => [], + }; + + const result = await runPromise( + runBenchmark({ + epochs: 1, + maxConcurrency: 1, + range: { start: 0, end: 1 }, + }).pipe(provide(makeLayers(flakyStore))) + ); + + expect(attempts).toBe(3); + expect(writes.map((w) => w.sampleIndex)).toEqual([0]); + expect(result.metrics.totalQuestions).toBe(1); + }); + + it("fails the run when persisting a record still fails after bounded retries", async () => { + let attempts = 0; + const failingStore: SampleResultStoreService = { + write: async () => { + attempts += 1; + throw new Error("gcs unavailable"); + }, + list: async () => [], + }; + + await expect( + runPromise( + runBenchmark({ + epochs: 1, + maxConcurrency: 1, + range: { start: 0, end: 1 }, + }).pipe(provide(makeLayers(failingStore))) + ) + ).rejects.toThrow("Failed to persist sample result"); + expect(attempts).toBe(4); + }); +}); diff --git a/src/harness/run.test.ts b/src/harness/run.test.ts index 7216f9f..20fd9db 100644 --- a/src/harness/run.test.ts +++ b/src/harness/run.test.ts @@ -16,6 +16,7 @@ import { fromChunk } from "effect/Stream"; import { noopProgressLayer, noopCheckpointLayer, + noopSampleResultLayer, } from "../../test/helpers/noop-progress-layer"; import { mcqScorer } from "../benchmarks/scorers/mcq/scorer"; import { runHarnessPromise } from "../internal/effect-logger"; @@ -27,13 +28,12 @@ import type { ModelService } from "./model"; import { Model } from "./model"; import type { CheckpointStore, ProgressReporter } from "./progress"; import { runBenchmark } from "./run"; +import type { SampleResultStore } from "./sample-result-store"; import { Scorer } from "./scorer"; import type { SolverService } from "./solver"; import { systemMessage, chain, generate, Solver } from "./solver"; -const infoSpies: { - mockRestore: () => void; -}[] = []; +const infoSpies: { mockRestore: () => void }[] = []; afterEach(() => { for (const info of infoSpies.splice(0)) { info.mockRestore(); @@ -94,14 +94,22 @@ function makeLayers( layer: Layer; }, solverService: ReturnType -): Layer { +): Layer< + | Dataset + | Solver + | Scorer + | ProgressReporter + | CheckpointStore + | SampleResultStore +> { return mergeAll( fakeDatasetLayer(SAMPLES), layerSucceed(Solver, Solver.of(solverService)), layerSucceed(Scorer, Scorer.of(mcqScorer)), model.layer, noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); } describe("runBenchmark", () => { @@ -130,7 +138,8 @@ describe("runBenchmark", () => { layerSucceed(Solver, Solver.of(solver)), layerSucceed(Scorer, Scorer.of(mcqScorer)), noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); await runHarnessPromise( runBenchmark({ @@ -184,7 +193,8 @@ describe("runBenchmark", () => { layerSucceed(Scorer, Scorer.of(mcqScorer)), model.layer, noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); const result = await runPromise( runBenchmark({ epochs: 1, maxConcurrency: 1 }).pipe(provide(layers)) @@ -225,7 +235,8 @@ describe("runBenchmark", () => { layerSucceed(Solver, Solver.of(solver)), layerSucceed(Scorer, Scorer.of(mcqScorer)), noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); const result = await runPromise( runBenchmark({ epochs: 1, maxConcurrency: 1 }).pipe(provide(layers)) @@ -241,7 +252,8 @@ describe("runBenchmark", () => { layerSucceed(Scorer, Scorer.of(mcqScorer)), model.layer, noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); const result = await runPromise( runBenchmark({ @@ -280,7 +292,8 @@ describe("runBenchmark", () => { layerSucceed(Scorer, Scorer.of(mcqScorer)), model.layer, noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); const result = await runPromise( runBenchmark({ epochs: 1, maxConcurrency: 2 }).pipe(provide(layers)) @@ -319,7 +332,8 @@ describe("runBenchmark", () => { layerSucceed(Scorer, Scorer.of(mcqScorer)), model.layer, noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); const result = await runPromise( runBenchmark({ epochs: 1, maxConcurrency: 2 }).pipe(provide(layers)) @@ -362,7 +376,8 @@ describe("runBenchmark", () => { layerSucceed(Scorer, Scorer.of(mcqScorer)), model.layer, noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); const result = await runPromise( runBenchmark({ epochs: 1, maxConcurrency: 2 }).pipe(provide(layers)) diff --git a/src/harness/run.ts b/src/harness/run.ts index aa7ba00..1230c0a 100644 --- a/src/harness/run.ts +++ b/src/harness/run.ts @@ -1,16 +1,24 @@ import type { Effect } from "effect/Effect"; import { + catchAll as effectCatchAll, catchTags, fail as effectFail, flatMap as effectFlatMap, gen as effectGen, + logWarning, map as effectMap, annotateLogs, + retry as effectRetry, + tap as effectTap, + tapError as effectTapError, + tryPromise as effectTryPromise, withLogSpan, succeed as effectSucceed, } from "effect/Effect"; +import { exponential, intersect, jittered, recurs } from "effect/Schedule"; import type { Stream } from "effect/Stream"; import { + filter as streamFilter, flatMap as streamFlatMap, fromIterable as streamFromIterable, mapEffect as streamMapEffect, @@ -42,6 +50,12 @@ import { Dataset } from "./dataset"; import type { AggregateMetrics, SampleScore } from "./metric"; import { aggregateScores } from "./metric"; import { CheckpointStore, ProgressReporter } from "./progress"; +import { setBenchRequestContext } from "./request-context"; +import type { CompletedSampleEntry } from "./sample-result-store"; +import { + SampleResultStore, + SampleResultStoreError, +} from "./sample-result-store"; import { Scorer } from "./scorer"; import { Solver } from "./solver"; @@ -68,6 +82,9 @@ interface SampleEpoch { readonly sampleIndex: number; } +/** Services every per-sample evaluation reaches for. */ +type EvalContext = Solver | Scorer | ProgressReporter | CheckpointStore; + interface FoldAccumulator { scores: SampleScore[]; usage: UsageTotals; @@ -77,8 +94,27 @@ type EvalOutcome = { sampleScore: SampleScore; usage?: ModelUsage; generationTimeMs?: number; + /** + * The score was synthesized from a model/solver failure rather than an actual + * evaluation. Such outcomes are persisted only at the end of the run, marked + * `degraded` in the record, and ignored by the retry skip-list — so a later + * activity attempt re-runs the sample instead of freezing a capacity failure + * into the results, while the finalization fold can still count them. + */ + degraded?: true; }; +/** Persist the run's degraded outcomes, deferred to after every sample has settled. */ +function persistDegradedOutcomes( + outcomes: Iterable +): Effect { + return effectGen(function* () { + for (const [sampleEpoch, outcome] of outcomes) { + yield* writeEntry(sampleEpoch, outcome); + } + }); +} + function sampleEpochStream( dataset: DatasetService, epochs: number, @@ -106,16 +142,8 @@ function sampleEpochStream( function evalWithProgress( sampleEpoch: SampleEpoch, - evaluate: Effect< - EvalOutcome, - ModelError | SolverError, - Solver | Scorer | ProgressReporter | CheckpointStore - > -): Effect< - EvalOutcome, - ModelError | SolverError, - Solver | Scorer | ProgressReporter | CheckpointStore -> { + evaluate: Effect +): Effect { const { sample, epoch, sampleIndex } = sampleEpoch; return effectGen(function* () { const reporter = yield* ProgressReporter; @@ -137,6 +165,137 @@ function evalWithProgress( }); } +function sampleEpochKey(sampleIndex: number, epoch: number): string { + return `${sampleIndex}:${epoch}`; +} + +/** + * Completed (sample, epoch) records already durable for this range, keyed by + * `sampleIndex:epoch`. Degraded records are ignored — a retry re-runs those + * samples instead of freezing an infrastructure failure into the results. A + * store failure is never fatal: the run falls back to re-evaluating + * everything. + */ +function loadCompletedEntries( + config: RunConfig +): Effect, never, SampleResultStore> { + return effectGen(function* () { + const store = yield* SampleResultStore; + const noEntries: readonly CompletedSampleEntry[] = []; + const entries = yield* effectTryPromise({ + try: () => store.list(config.range ?? {}), + catch: (e: unknown) => + new SampleResultStoreError({ + message: `Failed to list sample results: ${String(e)}`, + }), + }).pipe( + effectCatchAll((error) => + logWarning("Failed to list existing sample results", { + error: error.message, + }).pipe(effectMap(() => noEntries)) + ) + ); + const inRange = entries + .filter( + (entry) => + entry.degraded !== true && + entry.epoch >= 0 && + entry.epoch < config.epochs + ) + .toSorted((a, b) => a.sampleIndex - b.sampleIndex || a.epoch - b.epoch); + return new Map( + inRange.map((entry) => [ + sampleEpochKey(entry.sampleIndex, entry.epoch), + entry, + ]) + ); + }); +} + +/** Fold the already-completed records in so the final aggregate matches an uninterrupted run. */ +function seedAccumulator( + entries: Iterable +): FoldAccumulator { + const acc: FoldAccumulator = { scores: [], usage: { ...ZERO_USAGE } }; + for (const entry of entries) { + accumulateOutcome(acc, { + sampleScore: entry.sampleScore, + ...(entry.usage !== undefined && { usage: entry.usage }), + ...(entry.generationTimeMs !== undefined && { + generationTimeMs: entry.generationTimeMs, + }), + }); + } + return acc; +} + +/** Max additional write attempts before a record write is declared failed. */ +const PERSIST_MAX_RETRIES = 3; +const PERSIST_BASE_DELAY = "200 millis"; + +/** + * Write one (sample, epoch) record with bounded retries. The store is the + * source of truth for completed work, so a write that still fails after the + * retries fails the run: completing successfully must mean the record is + * durable, letting the activity retry re-run only what is genuinely missing. + */ +function writeEntry( + sampleEpoch: SampleEpoch, + outcome: EvalOutcome +): Effect { + return effectGen(function* () { + const store = yield* SampleResultStore; + yield* effectTryPromise({ + try: () => + store.write({ + sampleIndex: sampleEpoch.sampleIndex, + epoch: sampleEpoch.epoch, + sampleScore: outcome.sampleScore, + ...(outcome.usage !== undefined && { usage: outcome.usage }), + ...(outcome.generationTimeMs !== undefined && { + generationTimeMs: outcome.generationTimeMs, + }), + ...(outcome.degraded === true && { degraded: true }), + }), + catch: (e: unknown) => + new SampleResultStoreError({ + message: `Failed to persist sample result: ${String(e)}`, + }), + }).pipe( + effectRetry( + exponential(PERSIST_BASE_DELAY).pipe( + jittered, + intersect(recurs(PERSIST_MAX_RETRIES)) + ) + ), + effectTapError((error) => + logWarning("Failed to persist sample result; failing the run", { + sample_id: outcome.sampleScore.sampleId, + sample_index: sampleEpoch.sampleIndex, + epoch: sampleEpoch.epoch, + degraded: outcome.degraded === true, + error: error.message, + }) + ) + ); + }); +} + +/** + * Persist one successful result as soon as it completes. Degraded outcomes + * are deferred to the end of the run (see {@link runBenchmark}) so a mid-run + * degraded record never shadows a later successful retry of the same + * (sample, epoch) within this attempt. + */ +function persistOutcome( + sampleEpoch: SampleEpoch, + outcome: EvalOutcome +): Effect { + return outcome.degraded === true + ? effectSucceed(undefined) + : writeEntry(sampleEpoch, outcome); +} + function accumulateOutcome( acc: FoldAccumulator, item: EvalOutcome @@ -169,16 +328,16 @@ interface EvaluateOneOpts { function evaluateOne( opts: EvaluateOneOpts -): Effect< - EvalOutcome, - ModelError | SolverError, - Solver | Scorer | ProgressReporter | CheckpointStore -> { +): Effect { const { sampleEpoch } = opts; const { sample, epoch } = sampleEpoch; const evaluation = effectGen(function* () { const solver = yield* Solver; const scorer = yield* Scorer; + /* Stamp the (sample, epoch) identity into this fiber so every model call + it issues carries it — the bench-gateway keys request coalescing on it, + otherwise identical bodies across epochs replay one epoch's answer. */ + yield* setBenchRequestContext({ sampleId: String(sample.id), epoch }); const state = yield* solver(initialTaskState(sample, epoch)); const score = yield* scorer(state, sample.target); return { @@ -222,14 +381,17 @@ function evaluateOne( }) ); }, + /* Degraded solver errors are harness/sandbox infrastructure failures, + not the model's fault, so they score Skipped — out of the accuracy + denominator, like exhausted retryable model errors. */ SolverError: (solverErr) => opts.degradeSolverErrors ? effectSucceed( errorOutcome({ sample, epoch, - value: ScoreValue.Incorrect, - explanation: `Solver error: ${solverErr.message}`, + value: ScoreValue.Skipped, + explanation: `Solver error (skipped): ${solverErr.message}`, }) ) : effectFail(solverErr), @@ -282,6 +444,7 @@ function errorOutcome(opts: ErrorOutcomeOpts): EvalOutcome { input: sample.input, target: sample.target.text, }, + degraded: true, }; } @@ -294,53 +457,91 @@ const ZERO_USAGE: UsageTotals = { generationTimeMs: 0, }; +/** + * Stream the dataset, fan out across (sample, epoch) pairs with bounded + * concurrency, solve + score each, and fold into metrics + usage. + * + * Non-degraded (sample, epoch) pairs already recorded in the + * `SampleResultStore` are skipped and their stored results seed the + * accumulator, so a retried chunk produces the same aggregate as an + * uninterrupted run. Scores synthesized from a model/solver failure are + * written only at the end of the run, marked `degraded` — durable for the + * finalization fold to count, but ignored by the retry skip-list so a retry + * still re-runs them. Persisting the records is part of completing the run: a + * write that fails after bounded retries fails the run, so a successful run + * guarantees every evaluated (sample, epoch) record is durable. + * + * `Dataset | Solver | Scorer | ProgressReporter | CheckpointStore | + * SampleResultStore` are yielded from the environment — provided by the + * benchmark layer + entry point. + */ export function runBenchmark( config: RunConfig ): Effect< RunResult, - ModelError | SolverError | DatasetError, - Dataset | Solver | Scorer | ProgressReporter | CheckpointStore + ModelError | SolverError | DatasetError | SampleResultStoreError, + | Dataset + | Solver + | Scorer + | ProgressReporter + | CheckpointStore + | SampleResultStore > { - return Dataset.pipe( - effectFlatMap((dataset) => { - const sampleEpochs = sampleEpochStream( - dataset, - config.epochs, - config.range - ); - const initialAcc: FoldAccumulator = { - scores: [], - usage: { ...ZERO_USAGE }, - }; - return sampleEpochs.pipe( - streamMapEffect( - (se) => - evalWithProgress( - se, - evaluateOne({ - sampleEpoch: se, - degradeSolverErrors: config.degradeSolverErrors ?? false, - }).pipe( - annotateLogs({ - sample_id: se.sample.id, - epoch: se.epoch, - }), - withLogSpan("sample") - ) - ), - { concurrency: config.maxConcurrency } - ), - streamRunFoldEffect(initialAcc, (acc, item) => - effectGen(function* () { - const updated = accumulateOutcome(acc, item); - const reporter = yield* ProgressReporter; - yield* reporter.onSampleComplete(updated.scores.length); - return updated; - }) - ), - effectMap(finalizeRun) - ); - }), - annotateLogs(config.logAnnotations ?? {}) - ); + return effectGen(function* () { + const dataset = yield* Dataset; + const completed = yield* loadCompletedEntries(config); + const sampleEpochs = sampleEpochStream( + dataset, + config.epochs, + config.range + ).pipe( + streamFilter( + (se) => !completed.has(sampleEpochKey(se.sampleIndex, se.epoch)) + ) + ); + const initialAcc = seedAccumulator(completed.values()); + const degradedOutcomes: (readonly [SampleEpoch, EvalOutcome])[] = []; + + const folded = yield* sampleEpochs.pipe( + streamMapEffect( + (se) => + evalWithProgress( + se, + evaluateOne({ + sampleEpoch: se, + degradeSolverErrors: config.degradeSolverErrors ?? false, + }).pipe( + annotateLogs({ + sample_id: se.sample.id, + epoch: se.epoch, + }), + withLogSpan("sample") + ) + ).pipe( + /* Persist inside the concurrent stage: the downstream fold sees + elements in input order, so a slow sample would otherwise hold + every finished result out of durable storage. Degraded outcomes + are collected instead and written after the stream settles. */ + effectTap((outcome) => { + if (outcome.degraded === true) { + degradedOutcomes.push([se, outcome]); + } + return persistOutcome(se, outcome); + }) + ), + { concurrency: config.maxConcurrency } + ), + streamRunFoldEffect(initialAcc, (acc, outcome) => + effectGen(function* () { + const updated = accumulateOutcome(acc, outcome); + const reporter = yield* ProgressReporter; + yield* reporter.onSampleComplete(updated.scores.length); + return updated; + }) + ) + ); + + yield* persistDegradedOutcomes(degradedOutcomes); + return finalizeRun(folded); + }).pipe(annotateLogs(config.logAnnotations ?? {})); } diff --git a/src/harness/sample-result-store.test.ts b/src/harness/sample-result-store.test.ts new file mode 100644 index 0000000..8294ec0 --- /dev/null +++ b/src/harness/sample-result-store.test.ts @@ -0,0 +1,179 @@ +import { describe, expect, it } from "bun:test"; + +import { assertLeft, assertRight } from "../internal/testing"; +import { parseSchema } from "../internal/zod"; +import { MessageRole, ScoreValue } from "./core"; +import type { + CompletedSampleEntry, + SampleResultEnvelope, +} from "./sample-result-store"; +import { + decodeSampleResultEntry, + encodeSampleResultRecord, + NOOP_SAMPLE_RESULT_STORE, + SAMPLE_RESULT_FORMAT_VERSION, + SampleResultRecordSchema, +} from "./sample-result-store"; + +const ENVELOPE: SampleResultEnvelope = { + parentWorkflowId: "parent-1", + childWorkflowId: "child-1", + chunkIndex: 3, + benchmarkId: "gpqa_diamond", + model: "openai/gpt-5", +}; + +const ENTRY: CompletedSampleEntry = { + sampleIndex: 42, + epoch: 1, + sampleScore: { + sampleId: "s-42", + epoch: 1, + score: { + value: ScoreValue.Correct, + answer: "B", + explanation: "matched target", + }, + messages: [{ role: MessageRole.Assistant, content: "Answer: B" }], + responseItems: [{ type: "message", id: "item-1" }], + generationIds: ["gen-1"], + metadata: { difficulty: "hard" }, + input: "Q42", + target: "B", + }, + usage: { + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + reasoningTokens: 2, + totalCost: 0.001, + }, + generationTimeMs: 1234, +}; + +describe("encodeSampleResultRecord", () => { + it("stamps the pinned envelope fields and absolute sample index", () => { + const record = encodeSampleResultRecord(ENVELOPE, ENTRY); + + expect(record.format_version).toBe(SAMPLE_RESULT_FORMAT_VERSION); + expect(record.parent_workflow_id).toBe("parent-1"); + expect(record.child_workflow_id).toBe("child-1"); + expect(record.chunk_index).toBe(3); + expect(record.benchmark_id).toBe("gpqa_diamond"); + expect(record.model).toBe("openai/gpt-5"); + expect(record.sample_id).toBe("s-42"); + expect(record.sample_index).toBe(42); + expect(record.epoch).toBe(1); + expect(record.usage).toEqual({ + input_tokens: 10, + output_tokens: 5, + total_tokens: 15, + reasoning_tokens: 2, + total_cost: 0.001, + }); + expect(record.generation_time_ms).toBe(1234); + expect(Date.parse(record.created_at)).not.toBeNaN(); + }); + + it("nulls usage and generation time when the outcome carried neither", () => { + const record = encodeSampleResultRecord(ENVELOPE, { + sampleIndex: 0, + epoch: 0, + sampleScore: { + sampleId: "s-0", + epoch: 0, + score: { + value: ScoreValue.Skipped, + answer: null, + explanation: "rate limited", + }, + }, + }); + + expect(record.usage).toBeNull(); + expect(record.generation_time_ms).toBeNull(); + }); + + it("defaults missing per-call usage fields to zero", () => { + const record = encodeSampleResultRecord(ENVELOPE, { + sampleIndex: 1, + epoch: 0, + sampleScore: { + sampleId: "s-1", + epoch: 0, + score: { + value: ScoreValue.Incorrect, + answer: "A", + explanation: "wrong", + }, + }, + usage: { totalCost: 0.5 }, + }); + + expect(record.usage).toEqual({ + input_tokens: 0, + output_tokens: 0, + total_tokens: 0, + reasoning_tokens: 0, + total_cost: 0.5, + }); + }); +}); + +describe("sample result record round-trip", () => { + it("restores the entry through JSON and Zod validation", () => { + const record = encodeSampleResultRecord(ENVELOPE, ENTRY); + const parsed = parseSchema( + SampleResultRecordSchema, + JSON.parse(JSON.stringify(record)) + ); + assertRight(parsed); + + expect(decodeSampleResultEntry(parsed.right)).toEqual(ENTRY); + }); + + it("round-trips the degraded marker and omits it on genuine evaluations", () => { + const degraded = encodeSampleResultRecord(ENVELOPE, { + ...ENTRY, + degraded: true, + }); + expect(degraded.degraded).toBe(true); + const parsed = parseSchema( + SampleResultRecordSchema, + JSON.parse(JSON.stringify(degraded)) + ); + assertRight(parsed); + expect(decodeSampleResultEntry(parsed.right)).toEqual({ + ...ENTRY, + degraded: true, + }); + + expect("degraded" in encodeSampleResultRecord(ENVELOPE, ENTRY)).toBe(false); + }); + + it("rejects a record whose format version is not the pinned one", () => { + const record = encodeSampleResultRecord(ENVELOPE, ENTRY); + const parsed = parseSchema(SampleResultRecordSchema, { + ...record, + format_version: 2, + }); + + assertLeft(parsed); + }); + + it("rejects a record missing the sample index", () => { + const { sample_index: _sampleIndex, ...withoutIndex } = + encodeSampleResultRecord(ENVELOPE, ENTRY); + const parsed = parseSchema(SampleResultRecordSchema, withoutIndex); + + assertLeft(parsed); + }); +}); + +describe("NOOP_SAMPLE_RESULT_STORE", () => { + it("accepts writes and reports nothing completed", async () => { + await NOOP_SAMPLE_RESULT_STORE.write(ENTRY); + + expect(await NOOP_SAMPLE_RESULT_STORE.list({})).toEqual([]); + }); +}); diff --git a/src/harness/sample-result-store.ts b/src/harness/sample-result-store.ts new file mode 100644 index 0000000..5f40fe3 --- /dev/null +++ b/src/harness/sample-result-store.ts @@ -0,0 +1,268 @@ +import { Tag } from "effect/Context"; +import { TaggedError } from "effect/Data"; + +import { z } from "../internal/zod"; +import { ScorerTrajectorySchema } from "../results/parquet-schema"; +import type { ModelUsage, Score } from "./core"; +import { ChatMessageSchema, ScoreValue } from "./core"; +import type { SampleScore } from "./metric"; + +/** + * Durable per-(sample, epoch) result records, persisted after each successful + * eval so a retried chunk can skip already-completed work. Distinct from + * `CheckpointStore` (in-flight sandbox state) and `ResultStore` (whole-chunk + * parquet): a record here is a finished outcome that survives activity + * retries. Records are idempotent — rewriting the same (sample, epoch) key + * produces the same content, so concurrent duplicate writes are harmless. + */ + +export const SAMPLE_RESULT_FORMAT_VERSION = 1; + +/** A sample-result store read or write failed. */ +// oxlint-disable-next-line unicorn/throw-new-error -- 1.74 false positive on Effect TaggedError class declaration +export class SampleResultStoreError extends TaggedError( + "SampleResultStoreError" +)<{ + readonly message: string; +}> {} + +const ScoreSchema = z.object({ + value: z.enum([ScoreValue.Correct, ScoreValue.Incorrect, ScoreValue.Skipped]), + answer: z.string().nullable(), + explanation: z.string(), + trajectory: ScorerTrajectorySchema.optional(), +}); + +/** + * The sample-score payload of a record: everything beyond `sample_id`/`epoch` + * (stored top-level) needed to reconstruct a {@link SampleScore}. Nested + * payloads (messages, response items) keep their canonical in-memory shapes. + */ +const SampleScorePayloadSchema = z.object({ + score: ScoreSchema, + messages: z.array(ChatMessageSchema).readonly().optional(), + response_items: z + .array(z.record(z.string(), z.unknown()).readonly()) + .readonly() + .optional(), + generation_ids: z.array(z.string()).readonly().optional(), + metadata: z.record(z.string(), z.unknown()).optional(), + input: z.string().optional(), + target: z.string().optional(), +}); + +const SampleUsageSchema = z.object({ + input_tokens: z.number(), + output_tokens: z.number(), + total_tokens: z.number(), + reasoning_tokens: z.number(), + total_cost: z.number(), +}); + +/** + * Pinned wire contract for one completed (sample, epoch) record. Downstream + * consumers (aggregate derivation, ingestion pipelines) are specced against + * this shape — do not change field names or the object path scheme + * (`samples////-.json`) + * without flagging it. + * + * `degraded: true` marks a score synthesized from a model/solver + * infrastructure failure, written only at the end of a chunk's run. Degraded + * records must be ignored by the retry skip-list (the sample is re-run for + * free) but counted by the finalization fold per their score value — + * Skipped into skipped_questions, Incorrect into the accuracy denominator. + */ +export const SampleResultRecordSchema = z.object({ + format_version: z.literal(SAMPLE_RESULT_FORMAT_VERSION), + parent_workflow_id: z.string(), + child_workflow_id: z.string(), + chunk_index: z.number().nullable(), + benchmark_id: z.string(), + model: z.string(), + sample_id: z.string(), + /** Absolute dataset index (not chunk-relative), unique per run regardless of chunking. */ + sample_index: z.number(), + epoch: z.number(), + sample_score: SampleScorePayloadSchema, + usage: SampleUsageSchema.nullable(), + generation_time_ms: z.number().nullable(), + /** Score synthesized from an infrastructure failure; absent on genuine evaluations. */ + degraded: z.literal(true).optional(), + created_at: z.string(), +}); + +export type SampleResultRecord = z.infer; + +/** One completed (sample, epoch) outcome, in harness in-memory shape. */ +export interface CompletedSampleEntry { + /** Absolute dataset index (not chunk-relative). */ + readonly sampleIndex: number; + readonly epoch: number; + readonly sampleScore: SampleScore; + readonly usage?: ModelUsage; + readonly generationTimeMs?: number; + /** Score synthesized from an infrastructure failure; never skip-seeded on retry. */ + readonly degraded?: true; +} + +export interface SampleResultStoreService { + /** Persist one completed (sample, epoch) result. Idempotent per key. */ + readonly write: (entry: CompletedSampleEntry) => Promise; + /** Existing completed entries whose sampleIndex falls in the half-open `[start, end)` range. */ + readonly list: (range: { + readonly start?: number; + readonly end?: number; + }) => Promise; +} + +export class SampleResultStore extends Tag( + "@openrouter/bench-harness/sample-result-store" +)() {} + +export const NOOP_SAMPLE_RESULT_STORE: SampleResultStoreService = { + write: async () => {}, + list: async () => [], +}; + +export interface SampleResultEnvelope { + readonly parentWorkflowId: string; + readonly childWorkflowId: string; + readonly chunkIndex: number | null; + readonly benchmarkId: string; + readonly model: string; +} + +export function encodeSampleResultRecord( + envelope: SampleResultEnvelope, + entry: CompletedSampleEntry +): SampleResultRecord { + const { sampleScore, usage } = entry; + return { + format_version: SAMPLE_RESULT_FORMAT_VERSION, + parent_workflow_id: envelope.parentWorkflowId, + child_workflow_id: envelope.childWorkflowId, + chunk_index: envelope.chunkIndex, + benchmark_id: envelope.benchmarkId, + model: envelope.model, + sample_id: sampleScore.sampleId, + sample_index: entry.sampleIndex, + epoch: entry.epoch, + sample_score: encodeSampleScorePayload(sampleScore), + usage: usage === undefined ? null : encodeUsage(usage), + generation_time_ms: entry.generationTimeMs ?? null, + ...(entry.degraded === true && { degraded: true }), + created_at: new Date().toISOString(), + }; +} + +/** + * Same rest-guard as {@link encodeSampleScorePayload}, for `ModelUsage` + * fields. `serverToolUse` is deliberately not part of the pinned v1 record + * contract — the resume aggregate only folds token/cost totals — so it is + * consumed here explicitly rather than silently dropped. + */ +function encodeUsage(usage: ModelUsage): SampleResultRecord["usage"] { + const { + inputTokens, + outputTokens, + totalTokens, + reasoningTokens, + totalCost, + serverToolUse: _serverToolUse, + ...rest + } = usage; + rest satisfies Record; + return { + input_tokens: inputTokens ?? 0, + output_tokens: outputTokens ?? 0, + total_tokens: totalTokens ?? 0, + reasoning_tokens: reasoningTokens ?? 0, + total_cost: totalCost ?? 0, + }; +} + +/** + * Rename the `SampleScore` fields a record carries beyond `sample_id` / + * `epoch`. Destructuring the whole score keeps the payload honest: adding a + * field to `SampleScore` without persisting it fails to compile on `rest` + * instead of silently dropping data. + */ +function encodeSampleScorePayload( + sampleScore: SampleScore +): SampleResultRecord["sample_score"] { + const { + sampleId: _sampleId, + epoch: _epoch, + score, + messages, + responseItems, + generationIds, + metadata, + input, + target, + ...rest + } = sampleScore; + rest satisfies Record; + return { + score: encodeScore(score), + ...(messages !== undefined && { messages }), + ...(responseItems !== undefined && { response_items: responseItems }), + ...(generationIds !== undefined && { generation_ids: generationIds }), + ...(metadata !== undefined && { metadata }), + ...(input !== undefined && { input }), + ...(target !== undefined && { target }), + }; +} + +/** Same rest-guard as {@link encodeSampleScorePayload}, for `Score` fields. */ +function encodeScore( + score: Score +): SampleResultRecord["sample_score"]["score"] { + const { value, answer, explanation, trajectory, ...rest } = score; + rest satisfies Record; + return { + value, + answer, + explanation, + ...(trajectory !== undefined && { trajectory }), + }; +} + +export function decodeSampleResultEntry( + record: SampleResultRecord +): CompletedSampleEntry { + const payload = record.sample_score; + const sampleScore: SampleScore = { + sampleId: record.sample_id, + epoch: record.epoch, + score: payload.score, + ...(payload.messages !== undefined && { messages: payload.messages }), + ...(payload.response_items !== undefined && { + responseItems: payload.response_items, + }), + ...(payload.generation_ids !== undefined && { + generationIds: payload.generation_ids, + }), + ...(payload.metadata !== undefined && { metadata: payload.metadata }), + ...(payload.input !== undefined && { input: payload.input }), + ...(payload.target !== undefined && { target: payload.target }), + }; + return { + sampleIndex: record.sample_index, + epoch: record.epoch, + sampleScore, + ...(record.usage !== null && { + usage: { + inputTokens: record.usage.input_tokens, + outputTokens: record.usage.output_tokens, + totalTokens: record.usage.total_tokens, + reasoningTokens: record.usage.reasoning_tokens, + totalCost: record.usage.total_cost, + }, + }), + ...(record.generation_time_ms !== null && { + generationTimeMs: record.generation_time_ms, + }), + ...(record.degraded === true && { degraded: true }), + }; +} diff --git a/src/providers/openrouter-model.ts b/src/providers/openrouter-model.ts index ce20fd7..a10c0c3 100644 --- a/src/providers/openrouter-model.ts +++ b/src/providers/openrouter-model.ts @@ -29,6 +29,7 @@ import type { GenerateConfig } from "../harness/model"; import { Model, stripVariantSuffix } from "../harness/model"; import type { ReasoningDetails } from "../harness/reasoning-details"; import { hasReasoningDetails } from "../harness/reasoning-details"; +import { getBenchRequestContext } from "../harness/request-context"; import { Either } from "../internal/either"; import { unknownErrorToString } from "../internal/errors"; import { isDefinedAndNotNull, isRecord } from "../internal/guards"; @@ -158,6 +159,7 @@ export function generate( : undefined; return gen(function* () { const startedAt = performance.now(); + const requestContext = yield* getBenchRequestContext; const body = { model, messages: messages.map(toApiMessage), @@ -177,6 +179,16 @@ export function generate( ...(sendSort && { provider: { sort: genConfig.sort } }), ...(autoRouterPlugin !== undefined && { plugins: [autoRouterPlugin] }), ...genConfig.extraBody, + /* Bench-gateway extension: identifies which (sample, epoch) issued the + call so the gateway's coalescing hash separates epochs. The gateway + strips `x_bench` before forwarding; OpenRouter ignores it if the + call goes direct. */ + ...(requestContext !== undefined && { + x_bench: { + sample_id: requestContext.sampleId, + sample_epoch: requestContext.epoch, + }, + }), }; const request = HttpClientRequest.post( `${opts.baseUrl}/chat/completions` diff --git a/src/results/parquet-schema.ts b/src/results/parquet-schema.ts index 4409ed0..20cde9f 100644 --- a/src/results/parquet-schema.ts +++ b/src/results/parquet-schema.ts @@ -6,7 +6,10 @@ export const RESULT_WRITER = "openrouter-bench" as const; export const ScorerTrajectorySchema = z.discriminatedUnion("kind", [ z.object({ kind: z.literal("verifier_log"), log: z.string() }), - z.object({ kind: z.literal("judge_runs"), runs: z.array(z.unknown()) }), + z.object({ + kind: z.literal("judge_runs"), + runs: z.array(z.unknown()).readonly(), + }), ]); export const BenchmarkResultRowSchema = z.object({ diff --git a/src/runner/run-by-id.ts b/src/runner/run-by-id.ts index b61a0f1..6cebe7e 100644 --- a/src/runner/run-by-id.ts +++ b/src/runner/run-by-id.ts @@ -22,6 +22,11 @@ import { } from "../harness/progress"; import type { RunResult, RunConfig } from "../harness/run"; import { runBenchmark } from "../harness/run"; +import type { SampleResultStoreService } from "../harness/sample-result-store"; +import { + NOOP_SAMPLE_RESULT_STORE, + SampleResultStore, +} from "../harness/sample-result-store"; import { runHarnessPromise } from "../internal/effect-logger"; import type { AsyncEither } from "../internal/either"; import { Either } from "../internal/either"; @@ -46,6 +51,8 @@ export interface RunBenchmarkInput { readonly checkpointStore?: CheckpointStoreService; readonly abortSignal?: AbortSignal; readonly resultStore?: ResultStoreService; + /** Durable per-(sample, epoch) results; enables skip/resume across activity retries. */ + readonly sampleResultStore?: SampleResultStoreService; } export interface RunBenchmarkOutput { @@ -81,6 +88,12 @@ export function runBenchmarkById( CheckpointStore, input.checkpointStore ?? NOOP_CHECKPOINT_STORE ); + + const sampleResultLayer = layerSucceed( + SampleResultStore, + input.sampleResultStore ?? NOOP_SAMPLE_RESULT_STORE + ); + const model = modelFromConfig(input.benchmarkConfig); const runConfig: RunConfig = { epochs: input.epochs, @@ -101,7 +114,8 @@ export function runBenchmarkById( const layers = layerMergeAll( fullBenchmarkLayer, progressLayer, - checkpointLayer + checkpointLayer, + sampleResultLayer ); const runOpts = input.abortSignal !== undefined ? { signal: input.abortSignal } : undefined; diff --git a/test/helpers/noop-progress-layer.ts b/test/helpers/noop-progress-layer.ts index 04b819b..682d8bd 100644 --- a/test/helpers/noop-progress-layer.ts +++ b/test/helpers/noop-progress-layer.ts @@ -6,6 +6,10 @@ import { NOOP_PROGRESS_REPORTER, ProgressReporter, } from "../../src/harness/progress"; +import { + NOOP_SAMPLE_RESULT_STORE, + SampleResultStore, +} from "../../src/harness/sample-result-store"; export const noopProgressLayer = layerSucceed( ProgressReporter, @@ -16,3 +20,7 @@ export const noopCheckpointLayer = layerSucceed( CheckpointStore, NOOP_CHECKPOINT_STORE ); +export const noopSampleResultLayer = layerSucceed( + SampleResultStore, + NOOP_SAMPLE_RESULT_STORE +);