diff --git a/package.json b/package.json index 2e74d4b..6c39183 100644 --- a/package.json +++ b/package.json @@ -30,6 +30,7 @@ "./scorer": "./src/harness/scorer.ts", "./solver": "./src/harness/solver.ts", "./parquet": "./src/results/parquet.ts", + "./partial-outcome-store": "./src/results/partial-outcome-store.ts", "./parquet-schema": "./src/results/parquet-schema.ts", "./progress": "./src/harness/progress.ts", "./result-store": "./src/results/result-store.ts", diff --git a/src/harness/run.test.ts b/src/harness/run.test.ts index 803a334..e41f45e 100644 --- a/src/harness/run.test.ts +++ b/src/harness/run.test.ts @@ -26,7 +26,12 @@ import { Dataset } from "./dataset"; import type { ModelService } from "./model"; import { Model } from "./model"; import type { CheckpointStore, ProgressReporter } from "./progress"; -import { runBenchmark } from "./run"; +import { + makeProgressReporter, + ProgressReporter as ProgressReporterTag, +} from "./progress"; +import type { SampleOutcome } from "./run"; +import { aggregateOutcomes, runBenchmark, sampleEpochKey } from "./run"; import { Scorer } from "./scorer"; import type { SolverService } from "./solver"; import { systemMessage, chain, generate, Solver } from "./solver"; @@ -397,4 +402,87 @@ describe("runBenchmark", () => { expect(result.usage.generationTimeMs).toBe(200); expect(result.usage.outputTokens).toBe(0); }); + it("evaluates only the sample-epochs not listed in skipSampleEpochs", async () => { + const model = fakeModel(() => "Answer: B"); + const solver = chain( + systemMessage("You are a helpful assistant."), + generate(model.service, { temperature: 0.5 }) + ); + const result = await runPromise( + runBenchmark({ + epochs: 2, + maxConcurrency: 2, + skipSampleEpochs: new Set([ + sampleEpochKey("s-correct", 0), + sampleEpochKey("s-wrong", 1), + ]), + }).pipe(provide(makeLayers(model, solver))) + ); + const evaluated = result.sampleScores.map((score) => + sampleEpochKey(score.sampleId, score.epoch) + ); + expect(evaluated.toSorted()).toEqual([ + sampleEpochKey("s-correct", 1), + sampleEpochKey("s-wrong", 0), + ]); + }); + it("offsets progress counts by the skipped sample-epochs when resuming", async () => { + const model = fakeModel(() => "Answer: B"); + const solver = chain( + systemMessage("You are a helpful assistant."), + generate(model.service, { temperature: 0.5 }) + ); + const reported: number[] = []; + const reporterLayer = layerSucceed( + ProgressReporterTag, + makeProgressReporter({ + onSampleComplete: (count) => { + reported.push(count); + }, + }) + ); + const layers = mergeAll( + fakeDatasetLayer(SAMPLES), + layerSucceed(Solver, Solver.of(solver)), + layerSucceed(Scorer, Scorer.of(mcqScorer)), + model.layer, + reporterLayer, + noopCheckpointLayer + ); + await runPromise( + runBenchmark({ + epochs: 2, + maxConcurrency: 1, + skipSampleEpochs: new Set([ + sampleEpochKey("s-correct", 0), + sampleEpochKey("s-wrong", 1), + ]), + }).pipe(provide(layers)) + ); + expect(reported.toSorted()).toEqual([3, 4]); + }); + it("reports every completed outcome through onOutcome as it folds", async () => { + const model = fakeModel((input) => + input.includes("Q1") ? "Answer: B" : "Answer: A" + ); + const solver = chain( + systemMessage("You are a helpful assistant."), + generate(model.service, { temperature: 0.5 }) + ); + const collected: SampleOutcome[] = []; + const result = await runPromise( + runBenchmark({ + epochs: 1, + maxConcurrency: 2, + onOutcome: (outcome) => { + collected.push(outcome); + }, + }).pipe(provide(makeLayers(model, solver))) + ); + const collectedIds = collected + .map((outcome) => outcome.sampleScore.sampleId) + .toSorted(); + expect(collectedIds).toEqual(["s-correct", "s-wrong"]); + expect(aggregateOutcomes(collected)).toEqual(result); + }); }); diff --git a/src/harness/run.ts b/src/harness/run.ts index 9135416..b1bfc00 100644 --- a/src/harness/run.ts +++ b/src/harness/run.ts @@ -11,6 +11,7 @@ import { } from "effect/Effect"; import type { Stream } from "effect/Stream"; import { + filter as streamFilter, flatMap as streamFlatMap, fromIterable as streamFromIterable, mapEffect as streamMapEffect, @@ -54,6 +55,8 @@ export interface RunConfig { }; readonly degradeSolverErrors?: boolean; readonly logAnnotations?: Readonly>; + readonly skipSampleEpochs?: ReadonlySet; + readonly onOutcome?: (outcome: SampleOutcome) => void; } export interface RunResult { @@ -73,12 +76,29 @@ interface FoldAccumulator { usage: UsageTotals; } -type EvalOutcome = { +export type SampleOutcome = { sampleScore: SampleScore; usage?: ModelUsage; generationTimeMs?: number; }; +export function sampleEpochKey(sampleId: string, epoch: number): string { + return `${sampleId}:${epoch}`; +} + +export function aggregateOutcomes( + outcomes: readonly SampleOutcome[] +): RunResult { + let acc: FoldAccumulator = { + scores: [], + usage: { ...ZERO_USAGE }, + }; + for (const outcome of outcomes) { + acc = accumulateOutcome(acc, outcome); + } + return finalizeRun(acc); +} + function sampleEpochStream( dataset: DatasetService, epochs: number, @@ -107,12 +127,12 @@ function sampleEpochStream( function evalWithProgress( sampleEpoch: SampleEpoch, evaluate: Effect< - EvalOutcome, + SampleOutcome, ModelError | SolverError, Solver | Scorer | ProgressReporter | CheckpointStore > ): Effect< - EvalOutcome, + SampleOutcome, ModelError | SolverError, Solver | Scorer | ProgressReporter | CheckpointStore > { @@ -139,7 +159,7 @@ function evalWithProgress( function accumulateOutcome( acc: FoldAccumulator, - item: EvalOutcome + item: SampleOutcome ): FoldAccumulator { acc.scores.push(item.sampleScore); const u = item.usage; @@ -163,9 +183,9 @@ function finalizeRun(acc: FoldAccumulator): RunResult { } function applyReplayedUsage( - outcome: EvalOutcome, + outcome: SampleOutcome, replayed: ReplayedUsage | undefined -): EvalOutcome { +): SampleOutcome { if (replayed === undefined) { return outcome; } @@ -193,7 +213,7 @@ interface EvaluateOneOpts { function evaluateOne( opts: EvaluateOneOpts ): Effect< - EvalOutcome, + SampleOutcome, ModelError | SolverError, Solver | Scorer | ProgressReporter | CheckpointStore > { @@ -296,7 +316,7 @@ interface ErrorOutcomeOpts { readonly explanation: string; } -function errorOutcome(opts: ErrorOutcomeOpts): EvalOutcome { +function errorOutcome(opts: ErrorOutcomeOpts): SampleOutcome { const { sample, epoch, value, explanation } = opts; const score: Score = { value, @@ -334,11 +354,19 @@ export function runBenchmark( > { return Dataset.pipe( effectFlatMap((dataset) => { + const skipKeys = config.skipSampleEpochs; const sampleEpochs = sampleEpochStream( dataset, config.epochs, config.range + ).pipe( + streamFilter( + (se) => + skipKeys === undefined || + !skipKeys.has(sampleEpochKey(se.sample.id, se.epoch)) + ) ); + const completedOffset = skipKeys?.size ?? 0; const initialAcc: FoldAccumulator = { scores: [], usage: { ...ZERO_USAGE }, @@ -364,8 +392,11 @@ export function runBenchmark( streamRunFoldEffect(initialAcc, (acc, item) => effectGen(function* () { const updated = accumulateOutcome(acc, item); + config.onOutcome?.(item); const reporter = yield* ProgressReporter; - yield* reporter.onSampleComplete(updated.scores.length); + yield* reporter.onSampleComplete( + completedOffset + updated.scores.length + ); return updated; }) ), diff --git a/src/results/partial-outcome-store.test.ts b/src/results/partial-outcome-store.test.ts new file mode 100644 index 0000000..ad03307 --- /dev/null +++ b/src/results/partial-outcome-store.test.ts @@ -0,0 +1,144 @@ +import { describe, expect, it } from "bun:test"; + +import { MessageRole, ScoreValue } from "../harness/core"; +import type { SampleOutcome } from "../harness/run"; +import { assertLeft, assertRight } from "../internal/testing"; +import { parseSchema } from "../internal/zod"; +import type { PartialOutcomesPayload } from "./partial-outcome-store"; +import { + isSameRunScope, + PartialOutcomesPayloadSchema, +} from "./partial-outcome-store"; + +const OUTCOME: SampleOutcome = { + sampleScore: { + sampleId: "s-1", + epoch: 0, + score: { + value: ScoreValue.Correct, + answer: "B", + explanation: "matched target", + trajectory: { kind: "verifier_log", log: "verified B" }, + }, + messages: [{ role: MessageRole.Assistant, content: "Answer: B" }], + responseItems: [{ type: "message", id: "resp-1" }], + requestBody: { model: "test-model", temperature: 0 }, + generationIds: ["gen-1"], + metadata: { category: "math" }, + input: "Q1 target B", + target: "B", + }, + usage: { + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + reasoningTokens: 2, + totalCost: 0.001, + serverToolUse: { + webSearchRequests: 1, + toolCallsRequested: 2, + toolCallsExecuted: 2, + }, + }, + generationTimeMs: 100, +}; + +const PAYLOAD: PartialOutcomesPayload = { + scope: { epochs: 2, range: { start: 0, end: 10 } }, + outcomes: [OUTCOME], +}; + +describe("PartialOutcomesPayloadSchema", () => { + it("round-trips a JSON-serialized payload back to an equal value", () => { + const wire: unknown = JSON.parse(JSON.stringify(PAYLOAD)); + + const parsed = parseSchema(PartialOutcomesPayloadSchema, wire); + + assertRight(parsed); + expect(parsed.right).toEqual(PAYLOAD); + }); + + it("round-trips a payload with only required fields", () => { + const minimal: PartialOutcomesPayload = { + scope: { epochs: 1 }, + outcomes: [ + { + sampleScore: { + sampleId: "s-2", + epoch: 1, + score: { + value: ScoreValue.Skipped, + answer: null, + explanation: "interrupted", + }, + }, + }, + ], + }; + const wire: unknown = JSON.parse(JSON.stringify(minimal)); + + const parsed = parseSchema(PartialOutcomesPayloadSchema, wire); + + assertRight(parsed); + expect(parsed.right).toEqual(minimal); + }); + + it("rejects an outcome with an unknown score value", () => { + const wire: unknown = JSON.parse( + JSON.stringify({ + ...PAYLOAD, + outcomes: [ + { + ...OUTCOME, + sampleScore: { + ...OUTCOME.sampleScore, + score: { ...OUTCOME.sampleScore.score, value: "X" }, + }, + }, + ], + }) + ); + + const parsed = parseSchema(PartialOutcomesPayloadSchema, wire); + + assertLeft(parsed); + }); + + it("rejects a payload without a run scope", () => { + const wire: unknown = JSON.parse( + JSON.stringify({ outcomes: PAYLOAD.outcomes }) + ); + + const parsed = parseSchema(PartialOutcomesPayloadSchema, wire); + + assertLeft(parsed); + }); +}); + +describe("isSameRunScope", () => { + it("matches identical scopes with and without ranges", () => { + expect(isSameRunScope({ epochs: 2 }, { epochs: 2 })).toBe(true); + expect( + isSameRunScope( + { epochs: 2, range: { start: 0, end: 10 } }, + { epochs: 2, range: { start: 0, end: 10 } } + ) + ).toBe(true); + }); + + it("rejects scopes with different epochs", () => { + expect(isSameRunScope({ epochs: 2 }, { epochs: 3 })).toBe(false); + }); + + it("rejects scopes with different or missing ranges", () => { + expect( + isSameRunScope( + { epochs: 2, range: { start: 0, end: 10 } }, + { epochs: 2, range: { start: 0, end: 5 } } + ) + ).toBe(false); + expect( + isSameRunScope({ epochs: 2, range: { start: 0, end: 10 } }, { epochs: 2 }) + ).toBe(false); + }); +}); diff --git a/src/results/partial-outcome-store.ts b/src/results/partial-outcome-store.ts new file mode 100644 index 0000000..60e2909 --- /dev/null +++ b/src/results/partial-outcome-store.ts @@ -0,0 +1,126 @@ +import { ChatMessageSchema, ScoreValue } from "../harness/core"; +import type { + ModelUsage, + ResponseItem, + Score, + ScorerTrajectory, + ServerToolUseCounts, +} from "../harness/core"; +import type { SampleScore } from "../harness/metric"; +import type { SampleOutcome } from "../harness/run"; +import type { ZodShape } from "../internal/zod"; +import { z } from "../internal/zod"; + +export interface PartialOutcomeRunScope { + readonly epochs: number; + readonly range?: { + readonly start?: number; + readonly end?: number; + }; +} + +export interface PartialOutcomesPayload { + readonly scope: PartialOutcomeRunScope; + readonly outcomes: readonly SampleOutcome[]; +} + +export interface PartialOutcomeStoreService { + readonly read: () => Promise; + readonly write: (payload: PartialOutcomesPayload) => Promise; + readonly remove: () => Promise; +} + +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: z.number(), + 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: z.number().optional(), + toolCallsRequested: z.number().optional(), + toolCallsExecuted: z.number().optional(), +} satisfies ZodShape); + +const ModelUsageSchema = z.object({ + inputTokens: z.number().optional(), + outputTokens: z.number().optional(), + totalTokens: z.number().optional(), + reasoningTokens: z.number().optional(), + totalCost: z.number().optional(), + serverToolUse: ServerToolUseCountsSchema.optional(), +} satisfies ZodShape); + +export const SampleOutcomeSchema = z.object({ + sampleScore: SampleScoreSchema, + usage: ModelUsageSchema.optional(), + generationTimeMs: z.number().optional(), +} satisfies ZodShape); + +const PartialOutcomeRunScopeRangeSchema = z.object({ + start: z.number().optional(), + end: z.number().optional(), +} satisfies ZodShape>); + +export const PartialOutcomeRunScopeSchema = z.object({ + epochs: z.number(), + range: PartialOutcomeRunScopeRangeSchema.optional(), +} satisfies ZodShape); + +export const PartialOutcomesPayloadSchema = z.object({ + scope: PartialOutcomeRunScopeSchema, + outcomes: z.array(SampleOutcomeSchema).readonly(), +} satisfies ZodShape); + +export function isSameRunScope( + a: PartialOutcomeRunScope, + b: PartialOutcomeRunScope +): boolean { + return ( + a.epochs === b.epochs && + a.range?.start === b.range?.start && + a.range?.end === b.range?.end + ); +} diff --git a/src/runner/run-by-id.test.ts b/src/runner/run-by-id.test.ts index 93a92a6..a0ac97a 100644 --- a/src/runner/run-by-id.test.ts +++ b/src/runner/run-by-id.test.ts @@ -1,13 +1,26 @@ import { describe, expect, it } from "bun:test"; -import { succeed } from "effect/Effect"; -import { fail as layerFail, succeed as layerSucceed } from "effect/Layer"; +import { dieMessage, succeed } from "effect/Effect"; +import { + fail as layerFail, + mergeAll, + succeed as layerSucceed, +} from "effect/Layer"; import { fromIterable } from "effect/Stream"; import type { HostBenchmarkRunConfig } from "../benchmarks/benchmark-config"; +import { mcqScorer } from "../benchmarks/scorers/mcq/scorer"; import type { Benchmark } from "../benchmarks/types"; +import { MessageRole, ScoreValue } from "../harness/core"; import { Dataset } from "../harness/dataset"; +import type { SampleOutcome } from "../harness/run"; +import { Scorer } from "../harness/scorer"; +import { generate, Solver } from "../harness/solver"; import { assertLeft, assertRight } from "../internal/testing"; +import type { + PartialOutcomesPayload, + PartialOutcomeStoreService, +} from "../results/partial-outcome-store"; import { datasetSizeById, runBenchmarkById } from "./run-by-id"; const HOST_BENCHMARK: Benchmark = { @@ -103,3 +116,193 @@ describe("benchmark runner by id", () => { expect(sizeResult.left).toBe(runResult.left); }); }); + +const RUNNABLE_SAMPLES = [ + { id: "s-1", input: "Q1 target B", target: { text: "B" } }, + { id: "s-2", input: "Q2 target B", target: { text: "B" } }, +] as const; + +function makeRunnableHostBenchmark( + solverCalls: string[] +): Benchmark { + const datasetLayer = layerSucceed(Dataset, { + stream: () => fromIterable(RUNNABLE_SAMPLES), + size: succeed(RUNNABLE_SAMPLES.length), + }); + const solver = generate( + { + generate: (messages) => { + const userMsg = + messages.find((m) => m.role === MessageRole.User)?.content ?? ""; + solverCalls.push(userMsg); + return succeed({ + completion: "Answer: B", + message: { role: MessageRole.Assistant, content: "Answer: B" }, + generationTimeMs: 100, + }); + }, + }, + { temperature: 0 } + ); + return { + ...HOST_BENCHMARK, + makeLayer: () => + mergeAll( + datasetLayer, + layerSucceed(Solver, Solver.of(solver)), + layerSucceed(Scorer, Scorer.of(mcqScorer)) + ), + }; +} + +function priorOutcome(sampleId: string, epoch: number): SampleOutcome { + return { + sampleScore: { + sampleId, + epoch, + score: { + value: ScoreValue.Correct, + answer: "B", + explanation: "persisted", + }, + }, + }; +} + +function makePartialStore(payload: PartialOutcomesPayload | null): { + store: PartialOutcomeStoreService; + removals: number[]; +} { + const removals: number[] = []; + return { + store: { + read: () => Promise.resolve(payload), + write: () => Promise.resolve(), + remove: () => { + removals.push(1); + return Promise.resolve(); + }, + }, + removals, + }; +} + +describe("runBenchmarkById partial outcome resume", () => { + it("skips sample-epochs persisted under the same run scope", async () => { + const solverCalls: string[] = []; + const { store, removals } = makePartialStore({ + scope: { epochs: 1 }, + outcomes: [priorOutcome("s-1", 0)], + }); + + const result = await runBenchmarkById({ + benchmarkId: HOST_BENCHMARK.id, + hostBenchmark: makeRunnableHostBenchmark(solverCalls), + apiKey: "unused", + benchmarkConfig: HOST_CONFIG, + epochs: 1, + maxConcurrency: 1, + sessionId: "test", + partialOutcomeStore: store, + }); + + assertRight(result); + expect(solverCalls).toEqual(["Q2 target B"]); + expect(result.right.result.sampleScores).toHaveLength(2); + expect(removals).toHaveLength(1); + }); + + it("discards persisted outcomes from a mismatched run scope", async () => { + const solverCalls: string[] = []; + const { store } = makePartialStore({ + scope: { epochs: 2 }, + outcomes: [priorOutcome("s-1", 0)], + }); + + const result = await runBenchmarkById({ + benchmarkId: HOST_BENCHMARK.id, + hostBenchmark: makeRunnableHostBenchmark(solverCalls), + apiKey: "unused", + benchmarkConfig: HOST_CONFIG, + epochs: 1, + maxConcurrency: 1, + sessionId: "test", + partialOutcomeStore: store, + }); + + assertRight(result); + expect(solverCalls.toSorted()).toEqual(["Q1 target B", "Q2 target B"]); + expect(result.right.result.sampleScores).toHaveLength(2); + }); + + it("starts fresh when the partial store read fails", async () => { + const solverCalls: string[] = []; + const store: PartialOutcomeStoreService = { + read: () => Promise.reject(new Error("read failed")), + write: () => Promise.resolve(), + remove: () => Promise.resolve(), + }; + + const result = await runBenchmarkById({ + benchmarkId: HOST_BENCHMARK.id, + hostBenchmark: makeRunnableHostBenchmark(solverCalls), + apiKey: "unused", + benchmarkConfig: HOST_CONFIG, + epochs: 1, + maxConcurrency: 1, + sessionId: "test", + partialOutcomeStore: store, + }); + + assertRight(result); + expect(solverCalls).toHaveLength(2); + }); + + it("keeps the partial store when result persistence fails", async () => { + const solverCalls: string[] = []; + const { store, removals } = makePartialStore({ + scope: { epochs: 1 }, + outcomes: [priorOutcome("s-1", 0)], + }); + + const result = await runBenchmarkById({ + benchmarkId: HOST_BENCHMARK.id, + hostBenchmark: makeRunnableHostBenchmark(solverCalls), + apiKey: "unused", + benchmarkConfig: HOST_CONFIG, + epochs: 1, + maxConcurrency: 1, + sessionId: "test", + partialOutcomeStore: store, + resultStore: { write: () => dieMessage("results store down") }, + }); + + assertRight(result); + expect(result.right.resultsPath).toBeNull(); + expect(removals).toHaveLength(0); + }); + + it("removes the partial store after result persistence succeeds", async () => { + const solverCalls: string[] = []; + const { store, removals } = makePartialStore({ + scope: { epochs: 1 }, + outcomes: [priorOutcome("s-1", 0)], + }); + + const result = await runBenchmarkById({ + benchmarkId: HOST_BENCHMARK.id, + hostBenchmark: makeRunnableHostBenchmark(solverCalls), + apiKey: "unused", + benchmarkConfig: HOST_CONFIG, + epochs: 1, + maxConcurrency: 1, + sessionId: "test", + partialOutcomeStore: store, + resultStore: { write: () => succeed("/tmp/results.parquet") }, + }); + + assertRight(result); + expect(result.right.resultsPath).toBe("/tmp/results.parquet"); + expect(removals).toHaveLength(1); + }); +}); diff --git a/src/runner/run-by-id.ts b/src/runner/run-by-id.ts index 59308ac..1a046ed 100644 --- a/src/runner/run-by-id.ts +++ b/src/runner/run-by-id.ts @@ -31,12 +31,22 @@ import { NOOP_PROGRESS_REPORTER, ProgressReporter, } from "../harness/progress"; -import type { RunResult, RunConfig } from "../harness/run"; -import { runBenchmark } from "../harness/run"; +import type { RunResult, RunConfig, SampleOutcome } from "../harness/run"; +import { + aggregateOutcomes, + runBenchmark, + sampleEpochKey, +} from "../harness/run"; import { runHarnessPromise } from "../internal/effect-logger"; import type { AsyncEither } from "../internal/either"; import { Either } from "../internal/either"; import { wLog } from "../internal/log"; +import type { + PartialOutcomeRunScope, + PartialOutcomesPayload, + PartialOutcomeStoreService, +} from "../results/partial-outcome-store"; +import { isSameRunScope } from "../results/partial-outcome-store"; import type { ResultStoreService } from "../results/result-store"; import { GenerationResolver, @@ -64,6 +74,7 @@ export interface RunBenchmarkInput { readonly checkpointStore?: CheckpointStoreService; readonly abortSignal?: AbortSignal; readonly resultStore?: ResultStoreService; + readonly partialOutcomeStore?: PartialOutcomeStoreService; readonly maxOutputTokensCeiling?: number; } @@ -72,14 +83,24 @@ export interface RunBenchmarkOutput { readonly resultsPath: string | null; } -export function runBenchmarkById( +export async function runBenchmarkById( input: RunBenchmarkInput ): AsyncEither { const benchmarkResult = resolveRunBenchmark(input); if (Either.isLeft(benchmarkResult)) { - return Promise.resolve(Either.left(benchmarkResult.left)); + return Either.left(benchmarkResult.left); } const { benchmark, benchmarkLayer } = benchmarkResult.right; + const partialStore = input.partialOutcomeStore; + const runScope: PartialOutcomeRunScope = { + epochs: input.epochs, + ...(input.range !== undefined && { range: input.range }), + }; + const priorOutcomes = + partialStore === undefined + ? [] + : await readPartialOutcomes(partialStore, runScope); + const collectedOutcomes: SampleOutcome[] = []; const progressLayer = layerSucceed( ProgressReporter, input.progressReporter ?? NOOP_PROGRESS_REPORTER @@ -104,6 +125,21 @@ export function runBenchmarkById( run_attempt: `${input.runAttempt}`, }), }, + ...(partialStore !== undefined && { + onOutcome: (outcome: SampleOutcome) => { + collectedOutcomes.push(outcome); + }, + }), + ...(priorOutcomes.length > 0 && { + skipSampleEpochs: new Set( + priorOutcomes.map((outcome) => + sampleEpochKey( + outcome.sampleScore.sampleId, + outcome.sampleScore.epoch + ) + ) + ), + }), }; const fullBenchmarkLayer = benchmarkLayer.pipe( layerProvide(FetchHttpClient.layer) @@ -124,34 +160,117 @@ export function runBenchmarkById( const runOpts = input.abortSignal !== undefined ? { signal: input.abortSignal } : undefined; const program = runBenchmark(runConfig).pipe(provide(layers)); - return runHarnessPromise( - input.runAttempt === undefined - ? program - : withRunAttempt(input.runAttempt, program), - runOpts - ) - .then((result) => { - if (input.resultStore !== undefined) { - return runHarnessPromise( - input.resultStore.write({ - result, - benchmark, - benchmarkConfig: input.benchmarkConfig, - epochs: input.epochs, - sessionId: input.sessionId, - }) - ) - .then((resultsPath) => Either.right({ result, resultsPath })) - .catch((storeErr) => { - wLog("Failed to persist benchmark results", { - error: String(storeErr), - }); - return Either.right({ result, resultsPath: null }); - }); - } - return Either.right({ result, resultsPath: null }); - }) - .catch((error) => Either.left(String(error))); + let harnessResult: RunResult; + try { + harnessResult = await runHarnessPromise( + input.runAttempt === undefined + ? program + : withRunAttempt(input.runAttempt, program), + runOpts + ); + } catch (error) { + await flushPartialOutcomes({ + partialStore, + abortSignal: input.abortSignal, + runScope, + outcomes: [...priorOutcomes, ...collectedOutcomes], + newOutcomeCount: collectedOutcomes.length, + }); + return Either.left(String(error)); + } + const result = + priorOutcomes.length > 0 + ? aggregateOutcomes([...priorOutcomes, ...collectedOutcomes]) + : harnessResult; + if (input.resultStore === undefined) { + await removePartialOutcomes(partialStore); + return Either.right({ result, resultsPath: null }); + } + try { + const resultsPath = await runHarnessPromise( + input.resultStore.write({ + result, + benchmark, + benchmarkConfig: input.benchmarkConfig, + epochs: input.epochs, + sessionId: input.sessionId, + }) + ); + await removePartialOutcomes(partialStore); + return Either.right({ result, resultsPath }); + } catch (storeErr) { + wLog("Failed to persist benchmark results", { + error: String(storeErr), + }); + return Either.right({ result, resultsPath: null }); + } +} + +async function readPartialOutcomes( + partialStore: PartialOutcomeStoreService, + runScope: PartialOutcomeRunScope +): Promise { + let payload: PartialOutcomesPayload | null; + try { + payload = await partialStore.read(); + } catch (error) { + wLog("Failed to read partial benchmark outcomes; starting fresh", { + error: String(error), + }); + return []; + } + if (payload === null) { + return []; + } + if (!isSameRunScope(payload.scope, runScope)) { + wLog("Discarding partial benchmark outcomes from a mismatched run scope", { + persisted_epochs: payload.scope.epochs, + current_epochs: runScope.epochs, + persisted_range: JSON.stringify(payload.scope.range ?? null), + current_range: JSON.stringify(runScope.range ?? null), + discarded_outcome_count: payload.outcomes.length, + }); + return []; + } + return payload.outcomes; +} + +async function removePartialOutcomes( + partialStore: PartialOutcomeStoreService | undefined +): Promise { + if (partialStore === undefined) { + return; + } + await partialStore.remove().catch(() => {}); +} + +async function flushPartialOutcomes(opts: { + readonly partialStore: PartialOutcomeStoreService | undefined; + readonly abortSignal: AbortSignal | undefined; + readonly runScope: PartialOutcomeRunScope; + readonly outcomes: readonly SampleOutcome[]; + readonly newOutcomeCount: number; +}): Promise { + const { partialStore, abortSignal, runScope, outcomes, newOutcomeCount } = + opts; + if ( + partialStore === undefined || + abortSignal?.aborted !== true || + newOutcomeCount === 0 + ) { + return; + } + try { + await partialStore.write({ scope: runScope, outcomes }); + wLog("Flushed partial benchmark outcomes after abort", { + outcome_count: outcomes.length, + new_outcome_count: newOutcomeCount, + }); + } catch (error) { + wLog("Failed to flush partial benchmark outcomes after abort", { + error: String(error), + }); + } } export function datasetSizeById(