diff --git a/src/harness/run.test.ts b/src/harness/run.test.ts index 422a4aa6..e0a520f1 100644 --- a/src/harness/run.test.ts +++ b/src/harness/run.test.ts @@ -172,6 +172,19 @@ describe("runBenchmark", () => { expect(result.usage.generationTimeMs).toBe(600); expect(result.sampleScores[0]?.generationIds).toEqual(["fake-Q1 target B"]); }); + it("runs only the requested epoch", async () => { + const model = fakeModel(() => "Answer: B"); + const solver = generate(model.service, { + temperature: 0.5, + reasoningEffort: "high", + }); + const result = await runPromise( + runBenchmark({ epochs: 3, epoch: 2, maxConcurrency: 4 }).pipe( + provide(makeLayers(model, solver)) + ) + ); + expect(result.sampleScores.map((score) => score.epoch)).toEqual([2, 2]); + }); it("captures per-sample message trajectories", async () => { const model = fakeModel(() => "Answer: B"); const solver = chain( diff --git a/src/harness/run.ts b/src/harness/run.ts index 7f004f5f..fc6e6833 100644 --- a/src/harness/run.ts +++ b/src/harness/run.ts @@ -49,6 +49,7 @@ import { Solver } from "./solver"; export interface RunConfig { readonly epochs: number; + readonly epoch?: number; readonly maxConcurrency: number; readonly range?: { readonly start?: number; @@ -83,7 +84,7 @@ type EvalOutcome = { function sampleEpochStream( dataset: DatasetService, - epochs: number, + epochs: readonly number[], range: | { readonly start?: number; @@ -96,7 +97,7 @@ function sampleEpochStream( streamZipWithIndex, streamFlatMap(([sample, i]) => streamFromIterable( - Array.from({ length: epochs }, (_, epoch) => ({ + epochs.map((epoch) => ({ sample, epoch, sampleIndex: baseIndex + i, @@ -337,7 +338,9 @@ export function runBenchmark( effectFlatMap((dataset) => { const sampleEpochs = sampleEpochStream( dataset, - config.epochs, + config.epoch === undefined + ? Array.from({ length: config.epochs }, (_, epoch) => epoch) + : [config.epoch], config.range ); const initialAcc: FoldAccumulator = { diff --git a/src/runner/run-by-id.ts b/src/runner/run-by-id.ts index e7e5c56e..aefe430d 100644 --- a/src/runner/run-by-id.ts +++ b/src/runner/run-by-id.ts @@ -56,6 +56,7 @@ export interface RunBenchmarkInput { readonly baseUrl?: string; readonly benchmarkConfig: BenchmarkRunConfig; readonly epochs: number; + readonly epoch?: number; readonly maxConcurrency: number; readonly range?: { readonly start?: number; @@ -96,6 +97,7 @@ export function runBenchmarkById( const model = modelFromConfig(input.benchmarkConfig); const runConfig: RunConfig = definedValues({ epochs: input.epochs, + epoch: input.epoch, maxConcurrency: input.maxConcurrency, range: input.range, degradeSolverErrors: benchmark.degradeSolverErrors,