diff --git a/src/cli/commands/compare.ts b/src/cli/commands/compare.ts index fefae08..5bf4fbc 100644 --- a/src/cli/commands/compare.ts +++ b/src/cli/commands/compare.ts @@ -17,6 +17,7 @@ interface CompareArgs { compareId?: string sample?: number sampleType?: SampleType + seed?: number limit?: number force?: boolean } @@ -50,6 +51,13 @@ export function parseCompareArgs(args: string[]): CompareArgs | null { logger.error(`Invalid sample type: ${type}. Valid types: consecutive, random`) return null } + } else if (arg === "--seed") { + const seed = Number(args[++i]) + if (!Number.isInteger(seed) || seed < 0) { + logger.error(`Invalid seed: ${seed}. Expected a non-negative integer.`) + return null + } + parsed.seed = seed } else if (arg === "-l" || arg === "--limit") { parsed.limit = parseInt(args[++i], 10) } else if (arg === "--force") { @@ -83,6 +91,9 @@ export async function compareCommand(args: string[]): Promise { console.log(` -m, --answering-model Answering model (default: ${DEFAULT_ANSWERING_MODEL})`) console.log(" -s, --sample Sample N questions per category") console.log(" --sample-type Sample type: consecutive (default), random") + console.log( + " --seed N Seed for --sample-type random (logged and stored when omitted)" + ) console.log(" -l, --limit Limit total number of questions") console.log(" --compare-id Compare ID (for resuming)") console.log(" --force Clear existing comparison and start fresh") @@ -122,6 +133,7 @@ export async function compareCommand(args: string[]): Promise { mode: "sample", sampleType: parsed.sampleType || "consecutive", perCategory: parsed.sample, + seed: parsed.seed, } } else if (parsed.limit) { sampling = { diff --git a/src/cli/commands/run.ts b/src/cli/commands/run.ts index 46d1f14..1807b92 100644 --- a/src/cli/commands/run.ts +++ b/src/cli/commands/run.ts @@ -20,6 +20,7 @@ interface RunArgs { limit?: number sample?: number sampleType?: SampleType + seed?: number force?: boolean fromPhase?: PhaseId concurrency?: ConcurrencyConfig @@ -60,6 +61,13 @@ export function parseRunArgs(args: string[]): RunArgs | null { logger.error(`Invalid sample type: ${type}. Valid types: consecutive, random`) return null } + } else if (arg === "--seed") { + const seed = Number(args[++i]) + if (!Number.isInteger(seed) || seed < 0) { + logger.error(`Invalid seed: ${seed}. Expected a non-negative integer.`) + return null + } + parsed.seed = seed } else if (arg === "-f" || arg === "--from-phase") { const phase = args[++i] as PhaseId if (PHASE_ORDER.includes(phase)) { @@ -121,6 +129,9 @@ export async function runCommand(args: string[]): Promise { console.log(` -m, --answering-model Answering model (default: ${DEFAULT_ANSWERING_MODEL})`) console.log(" -s, --sample Sample N questions per category") console.log(" --sample-type Sample type: consecutive (default), random") + console.log( + " --seed N Seed for --sample-type random (logged and stored when omitted)" + ) console.log(" -l, --limit Limit total number of questions to process") console.log(` -f, --from-phase Start from phase: ${PHASE_ORDER.join(", ")}`) console.log(" --concurrency N Default concurrency for all phases") @@ -191,6 +202,7 @@ export async function runCommand(args: string[]): Promise { mode: "sample", sampleType: parsed.sampleType || "consecutive", perCategory: parsed.sample, + seed: parsed.seed, } } else if (parsed.limit) { sampling = { diff --git a/src/orchestrator/batch.ts b/src/orchestrator/batch.ts index 239c6ae..4204e31 100644 --- a/src/orchestrator/batch.ts +++ b/src/orchestrator/batch.ts @@ -3,6 +3,7 @@ import type { BenchmarkName } from "../types/benchmark" import type { SamplingConfig } from "../types/checkpoint" import type { BenchmarkResult } from "../types/unified" import { orchestrator, CheckpointManager } from "./index" +import { selectQuestionsBySampling } from "./sampling" import { createBenchmark } from "../benchmarks" import { logger } from "../utils/logger" import { existsSync, mkdirSync, readFileSync, writeFileSync, rmSync } from "fs" @@ -52,36 +53,6 @@ function generateCompareId(): string { return `compare-${date}-${time}` } -function selectQuestionsBySampling( - allQuestions: { questionId: string; questionType: string }[], - sampling: SamplingConfig -): string[] { - if (sampling.mode === "full") { - return allQuestions.map((q) => q.questionId) - } - if (sampling.mode === "limit" && sampling.limit) { - return allQuestions.slice(0, sampling.limit).map((q) => q.questionId) - } - if (sampling.mode === "sample" && sampling.perCategory) { - const byType: Record = {} - for (const q of allQuestions) { - if (!byType[q.questionType]) byType[q.questionType] = [] - byType[q.questionType].push(q) - } - const selected: string[] = [] - for (const questions of Object.values(byType)) { - if (sampling.sampleType === "random") { - const shuffled = [...questions].sort(() => Math.random() - 0.5) - selected.push(...shuffled.slice(0, sampling.perCategory).map((q) => q.questionId)) - } else { - selected.push(...questions.slice(0, sampling.perCategory).map((q) => q.questionId)) - } - } - return selected - } - return allQuestions.map((q) => q.questionId) -} - export class BatchManager { private getComparePath(compareId: string): string { return join(COMPARE_DIR, compareId) @@ -155,8 +126,17 @@ export class BatchManager { const allQuestions = benchmarkInstance.getQuestions() let targetQuestionIds: string[] + let resolvedSampling = sampling if (sampling) { - targetQuestionIds = selectQuestionsBySampling(allQuestions, sampling) + const selection = selectQuestionsBySampling(allQuestions, sampling) + targetQuestionIds = selection.questionIds + // Persist the seed that was actually used so the comparison can be rebuilt + // on the same subset later. + resolvedSampling = + selection.seed === undefined ? sampling : { ...sampling, seed: selection.seed } + if (selection.seed !== undefined) { + logger.info(`Random sampling seed: ${selection.seed} (reuse with --seed ${selection.seed})`) + } } else { targetQuestionIds = allQuestions.map((q) => q.questionId) } @@ -168,7 +148,7 @@ export class BatchManager { benchmark, judge: judgeModel, answeringModel, - sampling, + sampling: resolvedSampling, targetQuestionIds, runs: providers.map((provider) => ({ provider, diff --git a/src/orchestrator/index.ts b/src/orchestrator/index.ts index 0db1244..5031eaa 100644 --- a/src/orchestrator/index.ts +++ b/src/orchestrator/index.ts @@ -7,6 +7,7 @@ import { createProvider } from "../providers" import { createBenchmark } from "../benchmarks" import { createJudge } from "../judges" import { CheckpointManager } from "./checkpoint" +import { selectQuestionsBySampling } from "./sampling" import { getProviderConfig, getJudgeConfig } from "../utils/config" import { resolveModel } from "../utils/models" import { logger } from "../utils/logger" @@ -31,40 +32,6 @@ export interface OrchestratorOptions { phases?: ("ingest" | "indexing" | "search" | "answer" | "evaluate" | "report")[] } -function selectQuestionsBySampling( - allQuestions: { questionId: string; questionType: string }[], - sampling: SamplingConfig -): string[] { - if (sampling.mode === "full") { - return allQuestions.map((q) => q.questionId) - } - - if (sampling.mode === "limit" && sampling.limit) { - return allQuestions.slice(0, sampling.limit).map((q) => q.questionId) - } - - if (sampling.mode === "sample" && sampling.perCategory) { - const byType: Record = {} - for (const q of allQuestions) { - if (!byType[q.questionType]) byType[q.questionType] = [] - byType[q.questionType].push(q) - } - - const selected: string[] = [] - for (const questions of Object.values(byType)) { - if (sampling.sampleType === "random") { - const shuffled = [...questions].sort(() => Math.random() - 0.5) - selected.push(...shuffled.slice(0, sampling.perCategory).map((q) => q.questionId)) - } else { - selected.push(...questions.slice(0, sampling.perCategory).map((q) => q.questionId)) - } - } - return selected - } - - return allQuestions.map((q) => q.questionId) -} - export class Orchestrator { private checkpointManager: CheckpointManager @@ -217,11 +184,20 @@ export class Orchestrator { targetQuestionIds = questionIds } else if (sampling) { logger.info(`Using sampling mode: ${sampling.mode}`) - targetQuestionIds = selectQuestionsBySampling(allQuestions, sampling) - checkpoint.sampling = sampling + const selection = selectQuestionsBySampling(allQuestions, sampling) + targetQuestionIds = selection.questionIds + // Record the resolved seed, not the requested config, so the checkpoint + // is enough to reproduce this exact subset. + checkpoint.sampling = + selection.seed === undefined ? sampling : { ...sampling, seed: selection.seed } logger.info( `Sampling selected ${targetQuestionIds.length} questions from ${allQuestions.length} total` ) + if (selection.seed !== undefined) { + logger.info( + `Random sampling seed: ${selection.seed} (reuse with --seed ${selection.seed})` + ) + } } else if (effectiveLimit) { logger.info(`Using limit: ${effectiveLimit}`) targetQuestionIds = allQuestions.slice(0, effectiveLimit).map((q) => q.questionId) diff --git a/src/orchestrator/sampling.test.ts b/src/orchestrator/sampling.test.ts new file mode 100644 index 0000000..c3a6749 --- /dev/null +++ b/src/orchestrator/sampling.test.ts @@ -0,0 +1,218 @@ +import { describe, expect, test } from "bun:test" +import type { SamplingConfig } from "../types/checkpoint" +import { + createSeededRandom, + selectQuestionsBySampling, + shuffle, + type SelectableQuestion, +} from "./sampling" + +const TRIALS = 12000 +const PERMUTATION_SIZE = 6 +const EXPECTED_PER_CELL = TRIALS / PERMUTATION_SIZE // 2000 + +// ±150 is ~3.7 standard deviations for this many trials (σ ≈ 41), and the seeded +// runs below are deterministic, so this cannot flake. +const TOLERANCE = 150 + +const BASE = [0, 1, 2, 3, 4, 5] + +/** + * counts[element][position] over many trials. A uniform shuffle puts every + * element in every position equally often. + */ +function positionCounts(permute: (trial: number) => number[]): number[][] { + const counts = Array.from({ length: PERMUTATION_SIZE }, () => new Array(PERMUTATION_SIZE).fill(0)) + for (let trial = 0; trial < TRIALS; trial++) { + const permuted = permute(trial) + for (let position = 0; position < PERMUTATION_SIZE; position++) { + counts[permuted[position]][position]++ + } + } + return counts +} + +function maxDeviation(counts: number[][]): number { + return Math.max(...counts.flat().map((c) => Math.abs(c - EXPECTED_PER_CELL))) +} + +function questions(spec: Record): SelectableQuestion[] { + const out: SelectableQuestion[] = [] + for (const [questionType, count] of Object.entries(spec)) { + for (let i = 0; i < count; i++) { + out.push({ questionId: `${questionType}-${i}`, questionType }) + } + } + return out +} + +const sample = (overrides: Partial = {}): SamplingConfig => ({ + mode: "sample", + sampleType: "random", + perCategory: 3, + ...overrides, +}) + +describe("shuffle", () => { + test("returns a permutation and leaves the input untouched", () => { + const input = [...BASE] + const result = shuffle(input, createSeededRandom(7)) + + expect([...result].sort()).toEqual(BASE) + expect(input).toEqual(BASE) + expect(result).not.toBe(input) + }) + + test("places every element in every position equally often", () => { + const counts = positionCounts((trial) => shuffle(BASE, createSeededRandom(trial + 1))) + + expect(maxDeviation(counts)).toBeLessThan(TOLERANCE) + }) + + test("the `sort(() => Math.random() - 0.5)` it replaces is measurably biased", () => { + // The defect this module exists to fix, kept as a characterisation test: the + // comparator ignores its arguments, so elements stay near their original + // positions far more often than chance. Deviations here run past 1000 — + // 25x the standard deviation of a genuine shuffle — while a uniform shuffle + // clears TOLERANCE above. + const counts = positionCounts(() => [...BASE].sort(() => Math.random() - 0.5)) + + expect(maxDeviation(counts)).toBeGreaterThan(300) + }) + + test("handles empty and single-element inputs", () => { + expect(shuffle([], createSeededRandom(1))).toEqual([]) + expect(shuffle(["only"], createSeededRandom(1))).toEqual(["only"]) + }) +}) + +describe("createSeededRandom", () => { + test("is reproducible for a given seed and differs across seeds", () => { + const first = Array.from({ length: 5 }, createSeededRandom(99)) + const again = Array.from({ length: 5 }, createSeededRandom(99)) + const other = Array.from({ length: 5 }, createSeededRandom(100)) + + expect(again).toEqual(first) + expect(other).not.toEqual(first) + }) + + test("stays within [0, 1)", () => { + const rand = createSeededRandom(12345) + for (let i = 0; i < 5000; i++) { + const value = rand() + expect(value).toBeGreaterThanOrEqual(0) + expect(value).toBeLessThan(1) + } + }) +}) + +describe("selectQuestionsBySampling", () => { + const pool = questions({ "single-session": 8, "multi-session": 8, temporal: 8 }) + + test("full mode selects every question in order", () => { + const selection = selectQuestionsBySampling(pool, { mode: "full" }) + + expect(selection.questionIds).toEqual(pool.map((q) => q.questionId)) + expect(selection.seed).toBeUndefined() + }) + + test("limit mode takes the first N overall", () => { + const selection = selectQuestionsBySampling(pool, { mode: "limit", limit: 4 }) + + expect(selection.questionIds).toEqual([ + "single-session-0", + "single-session-1", + "single-session-2", + "single-session-3", + ]) + }) + + test("consecutive sampling takes the first N of each category and needs no seed", () => { + const selection = selectQuestionsBySampling( + pool, + sample({ sampleType: "consecutive", perCategory: 2 }) + ) + + expect(selection.questionIds).toEqual([ + "single-session-0", + "single-session-1", + "multi-session-0", + "multi-session-1", + "temporal-0", + "temporal-1", + ]) + expect(selection.seed).toBeUndefined() + }) + + test("random sampling takes N per category without duplicates", () => { + const selection = selectQuestionsBySampling(pool, sample({ seed: 2024 })) + + expect(selection.questionIds).toHaveLength(9) + expect(new Set(selection.questionIds).size).toBe(9) + for (const type of ["single-session", "multi-session", "temporal"]) { + expect(selection.questionIds.filter((id) => id.startsWith(type))).toHaveLength(3) + } + }) + + test("the same seed reproduces the same subset in the same order", () => { + const first = selectQuestionsBySampling(pool, sample({ seed: 2024 })) + const again = selectQuestionsBySampling(pool, sample({ seed: 2024 })) + + expect(again.questionIds).toEqual(first.questionIds) + expect(again.seed).toBe(2024) + }) + + test("a different seed selects a different subset", () => { + const first = selectQuestionsBySampling(pool, sample({ seed: 1 })) + const other = selectQuestionsBySampling(pool, sample({ seed: 2 })) + + expect(other.questionIds).not.toEqual(first.questionIds) + }) + + test("generates and reports a seed when none is supplied", () => { + const selection = selectQuestionsBySampling(pool, sample()) + + expect(selection.seed).toBeDefined() + expect(Number.isInteger(selection.seed)).toBe(true) + + // The reported seed is enough to reproduce the run. + const replayed = selectQuestionsBySampling(pool, sample({ seed: selection.seed })) + expect(replayed.questionIds).toEqual(selection.questionIds) + }) + + test("random sampling reaches questions beyond the first N of a category", () => { + // The bias this fixes showed up as the front of each category being selected + // far more often than the tail. Across seeds, every question should be + // reachable. + const seen = new Set() + for (let seed = 0; seed < 200; seed++) { + for (const id of selectQuestionsBySampling(pool, sample({ seed })).questionIds) { + seen.add(id) + } + } + + expect(seen.size).toBe(pool.length) + }) + + test("takes the whole category when it holds fewer than perCategory questions", () => { + const small = questions({ rare: 2 }) + const selection = selectQuestionsBySampling(small, sample({ perCategory: 5, seed: 3 })) + + expect(selection.questionIds.sort()).toEqual(["rare-0", "rare-1"]) + }) + + test("keeps category order stable for integer-like question types", () => { + // A plain object would hand back "1" before "10"; insertion order is what + // callers see for consecutive sampling. + const numeric = [ + { questionId: "b", questionType: "10" }, + { questionId: "a", questionType: "1" }, + ] + const selection = selectQuestionsBySampling( + numeric, + sample({ sampleType: "consecutive", perCategory: 1 }) + ) + + expect(selection.questionIds).toEqual(["b", "a"]) + }) +}) diff --git a/src/orchestrator/sampling.ts b/src/orchestrator/sampling.ts new file mode 100644 index 0000000..fb66b7a --- /dev/null +++ b/src/orchestrator/sampling.ts @@ -0,0 +1,99 @@ +import type { SamplingConfig } from "../types/checkpoint" + +/** The only question fields sampling needs. */ +export interface SelectableQuestion { + questionId: string + questionType: string +} + +export interface SamplingSelection { + /** Question IDs to run, in selection order. */ + questionIds: string[] + /** The seed actually used, present only when the selection was randomised. */ + seed?: number +} + +const UINT32 = 0x100000000 + +/** + * mulberry32: a small PRNG with 32 bits of state, deterministic for a given + * seed. Determinism is the point — a sampled run records its seed, so the same + * subset can be re-selected later or handed to another provider. + */ +export function createSeededRandom(seed: number): () => number { + let state = seed >>> 0 + return () => { + state = (state + 0x6d2b79f5) >>> 0 + let t = state + t = Math.imul(t ^ (t >>> 15), t | 1) + t ^= t + Math.imul(t ^ (t >>> 7), t | 61) + return ((t ^ (t >>> 14)) >>> 0) / UINT32 + } +} + +/** A fresh seed for a run that did not ask for a specific one. */ +export function randomSeed(): number { + return Math.floor(Math.random() * UINT32) +} + +/** + * Fisher–Yates: every permutation equally likely. + * + * `[...arr].sort(() => Math.random() - 0.5)` is not a shuffle. The comparator + * ignores its arguments, which violates the consistency `sort` requires, so the + * permutation depends on the engine's sort algorithm and leaves elements near + * their original positions far more often than chance. For a sampler that then + * takes the first N, that bias lands squarely on the selected slice. + */ +export function shuffle(items: readonly T[], rand: () => number = Math.random): T[] { + const out = [...items] + for (let i = out.length - 1; i > 0; i--) { + const j = Math.floor(rand() * (i + 1)) + ;[out[i], out[j]] = [out[j], out[i]] + } + return out +} + +/** + * Resolve a sampling config into the concrete set of questions to run. + * + * Shared by the single-run orchestrator and the batch comparison path so both + * select identically — the two had drifting copies of this logic. + */ +export function selectQuestionsBySampling( + allQuestions: SelectableQuestion[], + sampling: SamplingConfig +): SamplingSelection { + if (sampling.mode === "full") { + return { questionIds: allQuestions.map((q) => q.questionId) } + } + + if (sampling.mode === "limit" && sampling.limit) { + return { questionIds: allQuestions.slice(0, sampling.limit).map((q) => q.questionId) } + } + + if (sampling.mode === "sample" && sampling.perCategory) { + // A Map keeps insertion order for every key; a plain object would reorder + // integer-like question types ahead of the rest. + const byType = new Map() + for (const q of allQuestions) { + const bucket = byType.get(q.questionType) + if (bucket) bucket.push(q) + else byType.set(q.questionType, [q]) + } + + // One generator across all categories, so the seed reproduces the selection + // as a whole rather than each category in isolation. + const seed = sampling.sampleType === "random" ? (sampling.seed ?? randomSeed()) : undefined + const rand = seed === undefined ? undefined : createSeededRandom(seed) + + const questionIds: string[] = [] + for (const questions of byType.values()) { + const ordered = rand ? shuffle(questions, rand) : questions + questionIds.push(...ordered.slice(0, sampling.perCategory).map((q) => q.questionId)) + } + return { questionIds, seed } + } + + return { questionIds: allQuestions.map((q) => q.questionId) } +} diff --git a/src/types/checkpoint.ts b/src/types/checkpoint.ts index f8f1180..0b45ed6 100644 --- a/src/types/checkpoint.ts +++ b/src/types/checkpoint.ts @@ -109,6 +109,12 @@ export interface SamplingConfig { sampleType?: SampleType perCategory?: number limit?: number + /** + * Seed for `sampleType: "random"`. Supply one to reproduce a previous + * selection; when omitted, a seed is generated and recorded here so the run + * stays auditable. + */ + seed?: number } export interface RunCheckpoint {