feat(runner): persist and resume partial sample outcomes across aborted runs

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",

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);
+  });
 });

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<Record<string, string>>;
+  readonly skipSampleEpochs?: ReadonlySet<string>;
+  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;
           })
         ),

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);
+  });
+});