diff --git a/package.json b/package.json index 2e74d4b..9557fa3 100644 --- a/package.json +++ b/package.json @@ -32,6 +32,7 @@ "./parquet": "./src/results/parquet.ts", "./parquet-schema": "./src/results/parquet-schema.ts", "./progress": "./src/harness/progress.ts", + "./sample-result-store": "./src/harness/sample-result-store.ts", "./result-store": "./src/results/result-store.ts", "./internal/log": "./src/internal/log.ts", "./internal/effect-logger": "./src/internal/effect-logger.ts" diff --git a/src/harness/run.test.ts b/src/harness/run.test.ts index 803a334..efcce8a 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,6 +28,7 @@ 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"; @@ -94,14 +96,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 +140,8 @@ describe("runBenchmark", () => { layerSucceed(Solver, Solver.of(solver)), layerSucceed(Scorer, Scorer.of(mcqScorer)), noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); await runHarnessPromise( runBenchmark({ @@ -184,7 +195,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 +237,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)) @@ -252,7 +265,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)) @@ -268,7 +282,8 @@ describe("runBenchmark", () => { layerSucceed(Scorer, Scorer.of(mcqScorer)), model.layer, noopProgressLayer, - noopCheckpointLayer + noopCheckpointLayer, + noopSampleResultLayer ); const result = await runPromise( runBenchmark({ @@ -307,7 +322,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)) @@ -346,7 +362,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)) @@ -389,7 +406,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 9135416..ab822d7 100644 --- a/src/harness/run.ts +++ b/src/harness/run.ts @@ -1,11 +1,13 @@ import type { Effect } from "effect/Effect"; import { + catchAll, catchTags, fail as effectFail, flatMap as effectFlatMap, gen as effectGen, map as effectMap, annotateLogs, + tryPromise, withLogSpan, succeed as effectSucceed, } from "effect/Effect"; @@ -18,6 +20,10 @@ import { zipWithIndex as streamZipWithIndex, } from "effect/Stream"; +import { Either } from "../internal/either"; +import { unknownErrorToString } from "../internal/errors"; +import { wLog } from "../internal/log"; +import { firstZodIssueMessage, parseSchema } from "../internal/zod"; import { resetGenerationIds } from "../runtime/generation-ids"; import type { ReplayedUsage } from "../runtime/generation-resolver"; import { resolveCollectedGenerations } from "../runtime/generation-resolver"; @@ -42,6 +48,11 @@ import { Dataset } from "./dataset"; import type { AggregateMetrics, SampleScore } from "./metric"; import { aggregateScores } from "./metric"; import { CheckpointStore, ProgressReporter } from "./progress"; +import { + PersistedSampleOutcomeSchema, + SampleResultStore, + sampleResultKey, +} from "./sample-result-store"; import { Scorer } from "./scorer"; import { Solver } from "./solver"; @@ -77,6 +88,7 @@ type EvalOutcome = { sampleScore: SampleScore; usage?: ModelUsage; generationTimeMs?: number; + isDegraded?: boolean; }; function sampleEpochStream( @@ -109,12 +121,12 @@ function evalWithProgress( evaluate: Effect< EvalOutcome, ModelError | SolverError, - Solver | Scorer | ProgressReporter | CheckpointStore + Solver | Scorer | ProgressReporter | CheckpointStore | SampleResultStore > ): Effect< EvalOutcome, ModelError | SolverError, - Solver | Scorer | ProgressReporter | CheckpointStore + Solver | Scorer | ProgressReporter | CheckpointStore | SampleResultStore > { const { sample, epoch, sampleIndex } = sampleEpoch; return effectGen(function* () { @@ -190,6 +202,68 @@ interface EvaluateOneOpts { readonly degradeSolverErrors: boolean; } +function evaluateOneResumable( + opts: EvaluateOneOpts +): Effect< + EvalOutcome, + ModelError | SolverError, + Solver | Scorer | ProgressReporter | CheckpointStore | SampleResultStore +> { + const { sample, epoch } = opts.sampleEpoch; + const key = sampleResultKey(sample.id, epoch); + return effectGen(function* () { + const store = yield* SampleResultStore; + const raw = yield* tryPromise({ + try: () => store.read(key), + catch: unknownErrorToString, + }).pipe( + catchAll((error) => { + wLog("Failed to read persisted sample outcome; re-evaluating", { + sample_result_key: key, + error, + }); + return effectSucceed(null); + }) + ); + if (raw !== null && raw !== undefined) { + const persisted = parseSchema(PersistedSampleOutcomeSchema, raw); + if (Either.isRight(persisted)) { + const { sampleId, epoch: persistedEpoch } = persisted.right.sampleScore; + if (sampleId === sample.id && persistedEpoch === epoch) { + return persisted.right; + } + wLog("Persisted sample outcome identity mismatch; re-evaluating", { + sample_result_key: key, + persisted_sample_id: sampleId, + persisted_epoch: persistedEpoch, + }); + } else { + wLog("Persisted sample outcome failed validation; re-evaluating", { + sample_result_key: key, + error: firstZodIssueMessage(persisted.left), + }); + } + } + const outcome = yield* evaluateOne(opts); + if (outcome.isDegraded === true) { + return outcome; + } + yield* tryPromise({ + try: () => store.write(key, outcome), + catch: unknownErrorToString, + }).pipe( + catchAll((error) => { + wLog("Failed to persist sample outcome; a resumed run will re-run it", { + sample_result_key: key, + error, + }); + return effectSucceed(undefined); + }) + ); + return outcome; + }); +} + function evaluateOne( opts: EvaluateOneOpts ): Effect< @@ -313,6 +387,7 @@ function errorOutcome(opts: ErrorOutcomeOpts): EvalOutcome { input: sample.input, target: sample.target.text, }, + isDegraded: true, }; } @@ -330,7 +405,12 @@ export function runBenchmark( ): Effect< RunResult, ModelError | SolverError | DatasetError, - Dataset | Solver | Scorer | ProgressReporter | CheckpointStore + | Dataset + | Solver + | Scorer + | ProgressReporter + | CheckpointStore + | SampleResultStore > { return Dataset.pipe( effectFlatMap((dataset) => { @@ -348,7 +428,7 @@ export function runBenchmark( (se) => evalWithProgress( se, - evaluateOne({ + evaluateOneResumable({ sampleEpoch: se, degradeSolverErrors: config.degradeSolverErrors ?? false, }).pipe( diff --git a/src/harness/sample-result-resume.test.ts b/src/harness/sample-result-resume.test.ts new file mode 100644 index 0000000..ea8edaf --- /dev/null +++ b/src/harness/sample-result-resume.test.ts @@ -0,0 +1,382 @@ +import { describe, expect, it, spyOn } from "bun:test"; + +import { fromIterable } from "effect/Chunk"; +import { + fail as effectFail, + flatMap as effectFlatMap, + runPromise, + succeed as effectSucceed, + provide, +} 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 { recordGenerationId } from "../runtime/generation-ids"; +import type { Sample } from "./core"; +import { MessageRole, ModelError, ScoreValue } from "./core"; +import { Dataset } from "./dataset"; +import type { ModelService } from "./model"; +import type { RunResult } from "./run"; +import { runBenchmark } from "./run"; +import type { SampleResultStoreService } from "./sample-result-store"; +import { + PersistedSampleOutcomeSchema, + SampleResultStore, +} from "./sample-result-store"; +import { Scorer } from "./scorer"; +import type { ScorerService } from "./scorer"; +import { systemMessage, chain, generate, Solver } from "./solver"; + +const SAMPLES: readonly Sample[] = [ + { id: "s1", input: "Q1 target B", target: { text: "B" } }, + { id: "s2", input: "Q2 target B", target: { text: "B" } }, + { id: "s3", input: "Q3 target B", target: { text: "B" } }, +]; +const EPOCHS = 2; +const ALL_KEYS = ["s1/0", "s1/1", "s2/0", "s2/1", "s3/0", "s3/1"]; + +function fakeDatasetLayer(samples: readonly Sample[]): Layer { + return layerSucceed( + Dataset, + Dataset.of({ + stream: (opts) => { + const start = opts?.start ?? 0; + const end = opts?.end ?? samples.length; + return fromChunk(fromIterable(samples.slice(start, end))); + }, + size: effectSucceed(samples.length), + }) + ); +} + +function fakeModel(): ModelService { + return { + generate: (messages) => { + const userMsg = + messages.find((m) => m.role === MessageRole.User)?.content ?? ""; + const completion = userMsg.includes("Q2") ? "Answer: A" : "Answer: B"; + return recordGenerationId(`fake-${userMsg}`).pipe( + effectFlatMap(() => + effectSucceed({ + completion, + message: { role: MessageRole.Assistant, content: completion }, + usage: { + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + totalCost: 0.001, + }, + generationTimeMs: 100, + }) + ) + ); + }, + }; +} + +interface Counters { + solver: number; + scorer: number; +} + +interface InMemoryStore extends SampleResultStoreService { + readonly map: Map; + readonly writtenKeys: string[]; +} + +function makeInMemoryStore(seed?: Map): InMemoryStore { + const map = new Map(seed ?? []); + const writtenKeys: string[] = []; + return { + map, + writtenKeys, + read: async (key) => { + const raw = map.get(key); + return raw === undefined ? null : JSON.parse(raw); + }, + write: async (key, data) => { + writtenKeys.push(key); + map.set(key, JSON.stringify(data)); + }, + }; +} + +function throwingStore(): SampleResultStoreService { + return { + read: async () => { + throw new Error("read exploded"); + }, + write: async () => { + throw new Error("write exploded"); + }, + }; +} + +async function runWithStore( + store: SampleResultStoreService, + overrides?: { + readonly model?: ModelService; + readonly scorer?: ScorerService; + } +): Promise<{ result: RunResult; counters: Counters }> { + const counters: Counters = { solver: 0, scorer: 0 }; + const model = overrides?.model ?? fakeModel(); + const innerSolver = chain( + systemMessage("You are a helpful assistant."), + generate(model, { temperature: 0.5 }) + ); + const countingSolver = Solver.of((state) => { + counters.solver += 1; + return innerSolver(state); + }); + const baseScorer = overrides?.scorer ?? mcqScorer; + const countingScorer: ScorerService = (state, target) => { + counters.scorer += 1; + return baseScorer(state, target); + }; + const layers = mergeAll( + fakeDatasetLayer(SAMPLES), + layerSucceed(Solver, countingSolver), + layerSucceed(Scorer, Scorer.of(countingScorer)), + noopProgressLayer, + noopCheckpointLayer, + layerSucceed(SampleResultStore, store) + ); + const result = await runPromise( + runBenchmark({ epochs: EPOCHS, maxConcurrency: 2 }).pipe(provide(layers)) + ); + return { result, counters }; +} + +function summarize(result: RunResult): { key: string; value: string }[] { + return result.sampleScores + .map((s) => ({ key: `${s.sampleId}/${s.epoch}`, value: s.score.value })) + .sort((a, b) => a.key.localeCompare(b.key)); +} + +let run1Store: InMemoryStore; +let run1Summary: { key: string; value: string }[]; +let run1Metrics: { accuracy: number; total: number; correct: number }; + +describe("sample-result resume", () => { + it("T1: fresh run evaluates all 6 and persists all 6 keys", async () => { + run1Store = makeInMemoryStore(); + const { result, counters } = await runWithStore(run1Store); + expect(result.sampleScores.length).toBe(6); + expect(counters.solver).toBe(6); + expect(counters.scorer).toBe(6); + expect([...run1Store.map.keys()].sort()).toEqual(ALL_KEYS); + for (const raw of run1Store.map.values()) { + const parsed = PersistedSampleOutcomeSchema.safeParse(JSON.parse(raw)); + expect(parsed.success).toBe(true); + } + run1Summary = summarize(result); + run1Metrics = { + accuracy: result.metrics.accuracy, + total: result.metrics.totalQuestions, + correct: result.metrics.correctAnswers, + }; + expect(run1Metrics.total).toBe(3); + expect(run1Metrics.accuracy).toBeCloseTo(2 / 3, 5); + }); + + it("T2: full resume — solver/scorer never invoked, identical outcome", async () => { + const store = makeInMemoryStore(run1Store.map); + const { result, counters } = await runWithStore(store); + expect(counters.solver).toBe(0); + expect(counters.scorer).toBe(0); + expect(summarize(result)).toEqual(run1Summary); + expect(result.metrics.accuracy).toBeCloseTo(run1Metrics.accuracy, 5); + expect(result.metrics.totalQuestions).toBe(run1Metrics.total); + expect(result.metrics.correctAnswers).toBe(run1Metrics.correct); + expect(store.writtenKeys).toEqual([]); + }); + + it("T3: partial resume — only the 4 missing keys are evaluated", async () => { + const seeded = new Map( + [...run1Store.map.entries()].filter(([k]) => ["s1/0", "s2/1"].includes(k)) + ); + const store = makeInMemoryStore(seeded); + const { result, counters } = await runWithStore(store); + expect(counters.solver).toBe(4); + expect(counters.scorer).toBe(4); + expect([...store.writtenKeys].sort()).toEqual([ + "s1/1", + "s2/0", + "s3/0", + "s3/1", + ]); + expect(summarize(result)).toEqual(run1Summary); + expect(result.metrics.accuracy).toBeCloseTo(run1Metrics.accuracy, 5); + }); + + it("T4: store whose read/write throw — run completes, all evaluated, failures logged", async () => { + const warn = spyOn(console, "warn").mockImplementation(() => {}); + try { + const { result, counters } = await runWithStore(throwingStore()); + expect(counters.solver).toBe(6); + expect(counters.scorer).toBe(6); + expect(result.sampleScores.length).toBe(6); + expect(summarize(result)).toEqual(run1Summary); + expect(result.metrics.accuracy).toBeCloseTo(run1Metrics.accuracy, 5); + const warned = warn.mock.calls.map((call) => + call.map((arg) => JSON.stringify(arg)).join(" ") + ); + expect( + warned.some((line) => + line.includes("Failed to read persisted sample outcome") + ) + ).toBe(true); + expect( + warned.some((line) => line.includes("Failed to persist sample outcome")) + ).toBe(true); + expect(warned.some((line) => line.includes("read exploded"))).toBe(true); + expect(warned.some((line) => line.includes("write exploded"))).toBe(true); + } finally { + warn.mockRestore(); + } + }); + + it("T5: degraded error outcomes are not persisted and are re-evaluated on resume", async () => { + const failingModel: ModelService = { + generate: (messages) => { + const userMsg = + messages.find((m) => m.role === MessageRole.User)?.content ?? ""; + if (userMsg.includes("Q1")) { + return effectFail( + new ModelError({ message: "OpenRouter HTTP 429", status: 429 }) + ); + } + return effectSucceed({ + completion: "Answer: B", + message: { role: MessageRole.Assistant, content: "Answer: B" }, + generationTimeMs: 100, + }); + }, + }; + const store = makeInMemoryStore(); + const run1 = await runWithStore(store, { model: failingModel }); + const skippedRun1 = run1.result.sampleScores.filter( + (s) => s.sampleId === "s1" + ); + expect(skippedRun1.every((s) => s.score.value === ScoreValue.Skipped)).toBe( + true + ); + expect([...store.map.keys()].sort()).toEqual([ + "s2/0", + "s2/1", + "s3/0", + "s3/1", + ]); + + const run2 = await runWithStore(store); + expect(run2.counters.solver).toBe(2); + expect(run2.counters.scorer).toBe(2); + const s1Run2 = run2.result.sampleScores.filter((s) => s.sampleId === "s1"); + expect(s1Run2.every((s) => s.score.value === ScoreValue.Correct)).toBe( + true + ); + expect([...store.map.keys()].sort()).toEqual(ALL_KEYS); + }); + + it("T7: unparseable persisted record is logged and re-evaluated", async () => { + const seeded = new Map(run1Store.map); + const corruptRaw = seeded.get("s1/0"); + if (corruptRaw === undefined) { + throw new Error("expected s1/0 in run1 store"); + } + const corrupt: unknown = JSON.parse(corruptRaw); + if (typeof corrupt !== "object" || corrupt === null) { + throw new Error("expected persisted record to be an object"); + } + seeded.set( + "s1/0", + JSON.stringify({ ...corrupt, usage: { inputTokens: 10.5 } }) + ); + const store = makeInMemoryStore(seeded); + const warn = spyOn(console, "warn").mockImplementation(() => {}); + try { + const { result, counters } = await runWithStore(store); + expect(counters.solver).toBe(1); + expect(counters.scorer).toBe(1); + expect(summarize(result)).toEqual(run1Summary); + expect(store.writtenKeys).toEqual(["s1/0"]); + const warned = warn.mock.calls.map((call) => + call.map((arg) => JSON.stringify(arg)).join(" ") + ); + expect( + warned.some( + (line) => + line.includes("Persisted sample outcome failed validation") && + line.includes("s1/0") && + line.includes("inputTokens") + ) + ).toBe(true); + } finally { + warn.mockRestore(); + } + }); + + it("T8: persisted record with mismatched identity is logged and re-evaluated", async () => { + const backing = makeInMemoryStore(run1Store.map); + const misMappingStore: SampleResultStoreService = { + read: (key) => backing.read(`s1/${key.split("/")[1]}`), + write: backing.write, + }; + const warn = spyOn(console, "warn").mockImplementation(() => {}); + try { + const { result, counters } = await runWithStore(misMappingStore); + expect(counters.solver).toBe(4); + expect(counters.scorer).toBe(4); + expect(summarize(result)).toEqual(run1Summary); + expect(result.metrics.totalQuestions).toBe(run1Metrics.total); + expect(result.metrics.accuracy).toBeCloseTo(run1Metrics.accuracy, 5); + const warned = warn.mock.calls.map((call) => + call.map((arg) => JSON.stringify(arg)).join(" ") + ); + expect( + warned.some( + (line) => + line.includes("Persisted sample outcome identity mismatch") && + line.includes("s2/0") && + line.includes("s1") + ) + ).toBe(true); + } finally { + warn.mockRestore(); + } + }); + + it("T6: scorer trajectory survives persistence and resume", async () => { + const trajectoryScorer: ScorerService = (state, target) => + mcqScorer(state, target).pipe( + effectFlatMap((score) => + effectSucceed({ + ...score, + trajectory: { + kind: "verifier_log" as const, + log: "pytest: 12 passed", + }, + }) + ) + ); + const store = makeInMemoryStore(); + await runWithStore(store, { scorer: trajectoryScorer }); + + const resumed = await runWithStore(store); + expect(resumed.counters.solver).toBe(0); + expect( + resumed.result.sampleScores.every( + (s) => + s.score.trajectory?.kind === "verifier_log" && + s.score.trajectory.log === "pytest: 12 passed" + ) + ).toBe(true); + }); +}); diff --git a/src/harness/sample-result-store.test.ts b/src/harness/sample-result-store.test.ts new file mode 100644 index 0000000..b9f1ea2 --- /dev/null +++ b/src/harness/sample-result-store.test.ts @@ -0,0 +1,40 @@ +import { describe, expect, it } from "bun:test"; + +import { ScoreValue } from "./core"; +import type { + PersistedSampleOutcome, + SampleResultStoreService, +} from "./sample-result-store"; +import { + namespacedSampleResultStore, + sampleResultKey, +} from "./sample-result-store"; + +const OUTCOME: PersistedSampleOutcome = { + sampleScore: { + sampleId: "s1", + epoch: 0, + score: { value: ScoreValue.Correct, answer: "B", explanation: "" }, + }, +}; + +describe("namespacedSampleResultStore", () => { + it("prefixes reads and writes so different sessions do not collide", async () => { + const map = new Map(); + const backing: SampleResultStoreService = { + read: async (key) => map.get(key) ?? null, + write: async (key, data) => { + map.set(key, data); + }, + }; + const sessionA = namespacedSampleResultStore("session-a", backing); + const sessionB = namespacedSampleResultStore("session-b", backing); + const key = sampleResultKey("s1", 0); + + await sessionA.write(key, OUTCOME); + + expect([...map.keys()]).toEqual(["session-a/s1/0"]); + expect(await sessionA.read(key)).toEqual(OUTCOME); + expect(await sessionB.read(key)).toBeNull(); + }); +}); diff --git a/src/harness/sample-result-store.ts b/src/harness/sample-result-store.ts new file mode 100644 index 0000000..39f6c89 --- /dev/null +++ b/src/harness/sample-result-store.ts @@ -0,0 +1,120 @@ +import { Tag } from "effect/Context"; + +import type { ZodShape } from "../internal/zod"; +import { z, zInt } from "../internal/zod"; +import type { + ModelUsage, + ResponseItem, + Score, + ScorerTrajectory, + ServerToolUseCounts, +} from "./core"; +import { ChatMessageSchema, ScoreValue } from "./core"; +import type { SampleScore } from "./metric"; + +export interface PersistedSampleOutcome { + readonly sampleScore: SampleScore; + readonly usage?: ModelUsage; + readonly generationTimeMs?: number; +} + +const ScoreValueSchema = z.enum([ + ScoreValue.Correct, + ScoreValue.Incorrect, + ScoreValue.Skipped, +]); + +type VerifierLogTrajectory = Extract< + ScorerTrajectory, + { kind: "verifier_log" } +>; +type JudgeRunsTrajectory = Extract; + +const ScorerTrajectorySchema: z.ZodType = + z.discriminatedUnion("kind", [ + z.object({ + kind: z.literal("verifier_log"), + log: z.string(), + } satisfies ZodShape), + z.object({ + kind: z.literal("judge_runs"), + runs: z.array(z.unknown()).readonly(), + } satisfies ZodShape), + ]); + +const ScoreSchema = z.object({ + value: ScoreValueSchema, + answer: z.string().nullable(), + explanation: z.string(), + trajectory: ScorerTrajectorySchema.optional(), +} satisfies ZodShape); + +const ResponseItemSchema: z.ZodType = z + .record(z.string(), z.unknown()) + .readonly(); + +const SampleScoreSchema = z.object({ + sampleId: z.string(), + epoch: zInt(), + score: ScoreSchema, + messages: z.array(ChatMessageSchema).readonly().optional(), + responseItems: z.array(ResponseItemSchema).readonly().optional(), + requestBody: z.record(z.string(), z.unknown()).readonly().optional(), + generationIds: z.array(z.string()).readonly().optional(), + metadata: z.record(z.string(), z.unknown()).readonly().optional(), + input: z.string().optional(), + target: z.string().optional(), +} satisfies ZodShape); + +const ServerToolUseCountsSchema = z.object({ + webSearchRequests: zInt().optional(), + toolCallsRequested: zInt().optional(), + toolCallsExecuted: zInt().optional(), +} satisfies ZodShape); + +const ModelUsageSchema = z.object({ + inputTokens: zInt().optional(), + outputTokens: zInt().optional(), + totalTokens: zInt().optional(), + reasoningTokens: zInt().optional(), + totalCost: z.number().optional(), + serverToolUse: ServerToolUseCountsSchema.optional(), +} satisfies ZodShape); + +export const PersistedSampleOutcomeSchema = z.object({ + sampleScore: SampleScoreSchema, + usage: ModelUsageSchema.optional(), + generationTimeMs: z.number().optional(), +} satisfies ZodShape); + +export function sampleResultKey(sampleId: string, epoch: number): string { + return `${sampleId}/${epoch}`; +} + +export interface SampleResultStoreService { + readonly read: (runRelativeKey: string) => Promise; + readonly write: ( + runRelativeKey: string, + data: PersistedSampleOutcome + ) => Promise; +} + +export class SampleResultStore extends Tag( + "@openrouter/bench-harness/sample-result-store" +)() {} + +export const NOOP_SAMPLE_RESULT_STORE: SampleResultStoreService = { + read: async () => null, + write: async () => {}, +}; + +export function namespacedSampleResultStore( + sessionId: string, + store: SampleResultStoreService +): SampleResultStoreService { + return { + read: (runRelativeKey) => store.read(`${sessionId}/${runRelativeKey}`), + write: (runRelativeKey, data) => + store.write(`${sessionId}/${runRelativeKey}`, data), + }; +} diff --git a/src/internal/zod.ts b/src/internal/zod.ts index 772cc0f..b53c8d2 100644 --- a/src/internal/zod.ts +++ b/src/internal/zod.ts @@ -5,7 +5,7 @@ import { Either } from "./either"; export { z }; export type ZodShape = { - [key in keyof T]: ZodType; + [key in keyof Required]: ZodType; }; export function parseSchema( diff --git a/src/runner/run-by-id.ts b/src/runner/run-by-id.ts index 59308ac..1863cc6 100644 --- a/src/runner/run-by-id.ts +++ b/src/runner/run-by-id.ts @@ -33,6 +33,12 @@ 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 { + namespacedSampleResultStore, + 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"; @@ -62,6 +68,7 @@ export interface RunBenchmarkInput { readonly datasetRetry?: RetryConfig; readonly progressReporter?: ProgressReporterService; readonly checkpointStore?: CheckpointStoreService; + readonly sampleResultStore?: SampleResultStoreService; readonly abortSignal?: AbortSignal; readonly resultStore?: ResultStoreService; readonly maxOutputTokensCeiling?: number; @@ -88,6 +95,12 @@ export function runBenchmarkById( CheckpointStore, input.checkpointStore ?? NOOP_CHECKPOINT_STORE ); + const sampleResultLayer = layerSucceed( + SampleResultStore, + input.sampleResultStore === undefined + ? NOOP_SAMPLE_RESULT_STORE + : namespacedSampleResultStore(input.sessionId, input.sampleResultStore) + ); const model = modelFromConfig(input.benchmarkConfig); const runConfig: RunConfig = { epochs: input.epochs, @@ -119,6 +132,7 @@ export function runBenchmarkById( fullBenchmarkLayer, progressLayer, checkpointLayer, + sampleResultLayer, resolverLayer ); const runOpts = diff --git a/test/helpers/noop-progress-layer.ts b/test/helpers/noop-progress-layer.ts index 04b819b..f98fde9 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,8 @@ export const noopCheckpointLayer = layerSucceed( CheckpointStore, NOOP_CHECKPOINT_STORE ); + +export const noopSampleResultLayer = layerSucceed( + SampleResultStore, + NOOP_SAMPLE_RESULT_STORE +);