From 69abad8075d68c7a6d0979f929e4bad0157eaee0 Mon Sep 17 00:00:00 2001 From: Lucas Burmeister Date: Tue, 4 Aug 2026 06:13:07 +0200 Subject: [PATCH 1/8] perf: parallelize benchmark candidate search --- CONTEXT.md | 13 + benchmarks/retrieval/fusion-core.d.mts | 15 + benchmarks/retrieval/fusion-core.mjs | 107 +++ benchmarks/retrieval/fusion-worker.mjs | 30 + benchmarks/retrieval/fusion.ts | 93 +++ benchmarks/retrieval/runner.ts | 231 ++++-- benchmarks/retrieval/weight-search.ts | 1002 ++++++++++++++++++------ benchmarks/retrieval/worker-pool.ts | 432 ++++++++++ benchmarks/tests/channels.test.ts | 119 ++- benchmarks/tests/worker-pool.test.ts | 202 +++++ 10 files changed, 1908 insertions(+), 336 deletions(-) create mode 100644 benchmarks/retrieval/fusion-core.d.mts create mode 100644 benchmarks/retrieval/fusion-core.mjs create mode 100644 benchmarks/retrieval/fusion-worker.mjs create mode 100644 benchmarks/retrieval/fusion.ts create mode 100644 benchmarks/retrieval/worker-pool.ts create mode 100644 benchmarks/tests/worker-pool.test.ts diff --git a/CONTEXT.md b/CONTEXT.md index 72b91c3..5441b6f 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -245,6 +245,19 @@ current beam elites, so a later coordinate cannot regress the best development c development folds and evaluated unchanged against static weights on excluded intent folds and repositories; authored query-form labels remain informed reference strata and are not router inputs. This router remains benchmark-only until holdouts justify a production change. +Serial benchmark weight and router searches use a benchmark-only prepared evaluator: each sample and +fusion method materializes per-chunk normalized or RRF contributions once, then candidate weights reuse +that data. Public `fuseRankings` semantics and its existing ranking, normalization, and typed-array caches +remain unchanged. The evaluator still materializes and sorts the full candidate union for every +configuration; partial top-K or metric-specialized evaluation remains future work. The explicit parallel +benchmark path uses a fixed native `node:worker_threads` pool (default `max(1, availableParallelism() - 1)`), +sends the prepared snapshot once per worker, batches candidate weight vectors, and keeps beam/cache/archive/ +selection state on the main thread. The search APIs are async so serial and worker execution share the +same selection algorithm; `PIX_BENCH_WORKERS=0` selects serial evaluation, and the runner's +`PIX_BENCH_SEARCH_MODE=serial` selects that serial pool mode. `PIX_BENCH_WORKER_BATCH_SIZE` bounds each +worker message. +Unavailable worker startup falls back to serial evaluation; task failures are surfaced and all workers are +terminated before the benchmark fails. Repository checkouts live under ignored `benchmarks/.cache/repos/`; generated artifacts live under ignored `benchmarks/results/`. Dense and Sparse vectors are held only in the production in-memory SQLite adapter during a benchmark run. See `benchmarks/README.md` and `benchmarks/BASELINE.md`. diff --git a/benchmarks/retrieval/fusion-core.d.mts b/benchmarks/retrieval/fusion-core.d.mts new file mode 100644 index 0000000..367be11 --- /dev/null +++ b/benchmarks/retrieval/fusion-core.d.mts @@ -0,0 +1,15 @@ +import type { QualitySummary } from "./types.js" +import type { EvaluationCandidate, EvaluationSnapshot } from "./worker-pool.js" + +export function evaluatePreparedContributions( + matrix: unknown, + weights: unknown, +): { + readonly chunkIndex: number + readonly score: number +}[] + +export function evaluateCandidate( + snapshot: EvaluationSnapshot, + candidate: EvaluationCandidate, +): QualitySummary diff --git a/benchmarks/retrieval/fusion-core.mjs b/benchmarks/retrieval/fusion-core.mjs new file mode 100644 index 0000000..d8a9f64 --- /dev/null +++ b/benchmarks/retrieval/fusion-core.mjs @@ -0,0 +1,107 @@ +export const evaluatePreparedContributions = (matrix, weights) => { + const activeMask = + Number(weights.identity > 0) | + (Number(weights.camelcase > 0) << 1) | + (Number(weights.bm25 > 0) << 2) | + (Number(weights.dense > 0) << 3) | + (Number(weights.sparse > 0) << 4) + if (activeMask === 0) return [] + + const entries = [] + for (let position = 0; position < matrix.chunkIndices.length; position++) { + const available = matrix.presence[position] & activeMask + if (available === 0) continue + let score = 0 + if ((available & 1) !== 0) score += weights.identity * matrix.values.identity[position] + if ((available & 2) !== 0) score += weights.camelcase * matrix.values.camelcase[position] + if ((available & 4) !== 0) score += weights.bm25 * matrix.values.bm25[position] + if ((available & 8) !== 0) score += weights.dense * matrix.values.dense[position] + if ((available & 16) !== 0) score += weights.sparse * matrix.values.sparse[position] + entries.push([matrix.chunkIndices[position], score]) + } + return entries + .sort( + ([leftChunkIndex, leftScore], [rightChunkIndex, rightScore]) => + rightScore - leftScore || leftChunkIndex - rightChunkIndex, + ) + .map(([chunkIndex, score]) => ({ chunkIndex, score })) +} + +const weightsForSample = (candidate, sampleIndex) => + Array.isArray(candidate.weights) ? candidate.weights[sampleIndex] : candidate.weights + +const recallAt = (ranked, targets, k) => { + if (targets.length === 0) return 1 + const returned = new Set(ranked.slice(0, k)) + const found = targets.filter((target) => target.some((index) => returned.has(index))).length + return found / targets.length +} + +const reciprocalRank = (ranked, targets) => { + const relevant = new Set(targets.flatMap((target) => target)) + const rank = ranked.findIndex((index) => relevant.has(index)) + return rank < 0 ? 0 : 1 / (rank + 1) +} + +const contextRecallAtBudget = (ranked, sample, budget) => { + let consumed = 0 + let rankedPrefix = 0 + for (const chunkIndex of ranked) { + const tokens = sample.contextTokens[chunkIndex] + if (tokens === undefined) { + rankedPrefix++ + continue + } + if (consumed + tokens > budget) break + consumed += tokens + rankedPrefix++ + } + return recallAt(ranked, sample.targets, rankedPrefix) +} + +export const evaluateCandidate = (snapshot, candidate) => { + let recall5 = 0 + let recall10 = 0 + let recall20 = 0 + let recall50 = 0 + let contextRecall = 0 + let meanReciprocalRank = 0 + let totalWeight = 0 + + for (let sampleIndex = 0; sampleIndex < snapshot.samples.length; sampleIndex++) { + const sample = snapshot.samples[sampleIndex] + const weight = sample.sampleWeight + if (weight <= 0) continue + const weights = weightsForSample(candidate, sampleIndex) + if (weights === undefined) throw new Error(`Missing weights for sample ${sampleIndex}`) + const ranked = evaluatePreparedContributions(sample.fusion, weights).map( + (entry) => entry.chunkIndex, + ) + recall5 += weight * recallAt(ranked, sample.targets, 5) + recall10 += weight * recallAt(ranked, sample.targets, 10) + recall20 += weight * recallAt(ranked, sample.targets, 20) + recall50 += weight * recallAt(ranked, sample.targets, 50) + contextRecall += weight * contextRecallAtBudget(ranked, sample, 4096) + meanReciprocalRank += weight * reciprocalRank(ranked, sample.targets) + totalWeight += weight + } + + if (totalWeight === 0) + return { + recallAt5: 0, + recallAt10: 0, + recallAt20: 0, + recallAt50: 0, + contextRecallAt4096: 0, + meanReciprocalRank: 0, + } + + return { + recallAt5: recall5 / totalWeight, + recallAt10: recall10 / totalWeight, + recallAt20: recall20 / totalWeight, + recallAt50: recall50 / totalWeight, + contextRecallAt4096: contextRecall / totalWeight, + meanReciprocalRank: meanReciprocalRank / totalWeight, + } +} diff --git a/benchmarks/retrieval/fusion-worker.mjs b/benchmarks/retrieval/fusion-worker.mjs new file mode 100644 index 0000000..65bc473 --- /dev/null +++ b/benchmarks/retrieval/fusion-worker.mjs @@ -0,0 +1,30 @@ +import { parentPort } from "node:worker_threads" + +import { evaluateCandidate } from "./fusion-core.mjs" + +if (parentPort === null) throw new Error("Fusion worker requires a parent port") + +let snapshot + +parentPort.on("message", (message) => { + try { + if (message.type === "init") { + snapshot = message.snapshot + parentPort.postMessage({ type: "ready" }) + return + } + if (message.type !== "evaluate" || snapshot === undefined) + throw new Error("Fusion worker received an invalid task") + parentPort.postMessage({ + type: "result", + taskId: message.taskId, + results: message.candidates.map((candidate) => evaluateCandidate(snapshot, candidate)), + }) + } catch (error) { + parentPort.postMessage({ + type: "error", + taskId: message.taskId, + message: error instanceof Error ? error.message : String(error), + }) + } +}) diff --git a/benchmarks/retrieval/fusion.ts b/benchmarks/retrieval/fusion.ts new file mode 100644 index 0000000..1f6e12b --- /dev/null +++ b/benchmarks/retrieval/fusion.ts @@ -0,0 +1,93 @@ +import type { RankedChunk } from "../../src/domain/ports.js" +import { + CHANNEL_NAMES, + type ChannelName, + type ChannelRankings, + type ChannelWeights, + type FusionMethod, +} from "../../src/domain/retrieval.js" +import { fuseRankings as productionFuseRankings } from "../../src/lib/retrieval/fusion.js" +import { evaluatePreparedContributions } from "./fusion-core.mjs" + +const DEFAULT_CANDIDATE_DEPTH = Number.POSITIVE_INFINITY + +/** Structured-cloneable contribution data shared with benchmark worker threads. */ +export interface PreparedFusionSnapshot { + readonly chunkIndices: readonly number[] + readonly presence: Uint8Array + readonly values: Readonly> +} + +/** Reusable benchmark-only evaluator for one prepared ranking set and fusion method. */ +export interface PreparedFusionEvaluator { + readonly evaluate: (weights: ChannelWeights) => RankedChunk[] + readonly snapshot: PreparedFusionSnapshot +} + +const preparedEvaluators = new WeakMap>() + +const prepareContributionMatrix = ( + method: FusionMethod, + rankings: ChannelRankings, + candidateDepth: number, +): PreparedFusionSnapshot => { + const positions = new Map() + const chunkIndices: number[] = [] + for (const channel of CHANNEL_NAMES) { + for (const entry of rankings[channel].slice(0, candidateDepth)) { + if (!positions.has(entry.chunkIndex)) { + positions.set(entry.chunkIndex, chunkIndices.length) + chunkIndices.push(entry.chunkIndex) + } + } + } + + const presence = new Uint8Array(chunkIndices.length) + const values: Record = { + identity: new Float64Array(chunkIndices.length), + camelcase: new Float64Array(chunkIndices.length), + bm25: new Float64Array(chunkIndices.length), + dense: new Float64Array(chunkIndices.length), + sparse: new Float64Array(chunkIndices.length), + } + + for (let channelIndex = 0; channelIndex < CHANNEL_NAMES.length; channelIndex++) { + const channel = CHANNEL_NAMES[channelIndex] + const channelWeights: ChannelWeights = { + identity: 0, + camelcase: 0, + bm25: 0, + dense: 0, + sparse: 0, + [channel]: 1, + } + for (const entry of productionFuseRankings(method, rankings, channelWeights, candidateDepth)) { + const position = positions.get(entry.chunkIndex) + if (position === undefined) continue + presence[position] |= 1 << channelIndex + values[channel][position] = entry.score + } + } + return { chunkIndices, presence, values } +} + +/** Prepare benchmark-only per-chunk contributions for repeated serial evaluation. */ +export const prepareFusion = ( + method: FusionMethod, + rankings: ChannelRankings, + candidateDepth = DEFAULT_CANDIDATE_DEPTH, +): PreparedFusionEvaluator => { + const key = `${method}:${candidateDepth}` + const cachedByKey = preparedEvaluators.get(rankings) + const cached = cachedByKey?.get(key) + if (cached !== undefined) return cached + + const matrix = prepareContributionMatrix(method, rankings, candidateDepth) + const evaluator: PreparedFusionEvaluator = { + evaluate: (weights) => evaluatePreparedContributions(matrix, weights), + snapshot: matrix, + } + if (cachedByKey === undefined) preparedEvaluators.set(rankings, new Map([[key, evaluator]])) + else cachedByKey.set(key, evaluator) + return evaluator +} diff --git a/benchmarks/retrieval/runner.ts b/benchmarks/retrieval/runner.ts index 2709eba..c1e04b1 100644 --- a/benchmarks/retrieval/runner.ts +++ b/benchmarks/retrieval/runner.ts @@ -40,12 +40,12 @@ import { type ValidationStrategy, } from "./types.js" import { - fitRecommendedEvidenceRouter, - fitRecommendedFusionWeights, - fitRecommendedWeights, - optimizeEvidenceRouter, - optimizeFusionWeights, - optimizeWeights, + fitRecommendedEvidenceRouterParallel, + fitRecommendedFusionWeightsParallel, + fitRecommendedWeightsParallel, + optimizeEvidenceRouterParallel, + optimizeFusionWeightsParallel, + optimizeWeightsParallel, evaluateProductionRouter, type WeightSearchSample, } from "./weight-search.js" @@ -115,6 +115,23 @@ const reportProgress = (message: string): void => { process.stderr.write(`[retrieval benchmark] ${message}\n`) } +const runParallelSearch = (operation: () => Promise): Effect.Effect => + Effect.tryPromise({ + try: operation, + catch: (cause) => (cause instanceof Error ? cause : new Error(String(cause))), + }) + +const splitSamples = ( + samples: readonly WeightSearchSample[], + isValidation: (sample: WeightSearchSample) => boolean, +): { + readonly development: readonly WeightSearchSample[] + readonly validation: readonly WeightSearchSample[] +} => ({ + development: samples.filter((sample) => !isValidation(sample)), + validation: samples.filter(isValidation), +}) + const isLongInput = (text: string): boolean => Buffer.byteLength(text, "utf8") / 4 > SINGLE_ITEM_ESTIMATED_TOKENS @@ -271,6 +288,8 @@ export const runRetrievalBenchmark = ( Effect.gen(function* () { const benchmarkStartedAt = performance.now() const config = profileConfig(profile) + const serialSearch = process.env.PIX_BENCH_SEARCH_MODE === "serial" + const searchOptions = serialSearch ? { workerCount: 0 } : {} const optimizationProfile = yield* selectOptimizationProfile() const groupedStrategy: ValidationStrategy = config.groupedFolds === 3 ? "grouped-3-fold" : "grouped-5-fold" @@ -494,42 +513,58 @@ export const runRetrievalBenchmark = ( } const weightSearchStartedAt = performance.now() - const weightSearch = config.legacyDiagnostics - ? [...sampleGroups.values()].flatMap((group) => { - const groupedFolds = Array.from({ length: config.groupedFolds }, (_, fold) => - optimizeWeights( + const weightSearch: BenchmarkArtifact["weightSearch"][number][] = [] + const recommendedWeights: BenchmarkArtifact["recommendedWeights"][number][] = [] + if (config.legacyDiagnostics) { + for (const group of sampleGroups.values()) { + for (let fold = 0; fold < config.groupedFolds; fold++) { + weightSearch.push( + yield* runParallelSearch(() => + optimizeWeightsParallel( + group.model, + group.queryKind, + groupedStrategy, + String(fold + 1), + group.samples.filter((sample) => sample.groupedFold !== fold), + group.samples.filter((sample) => sample.groupedFold === fold), + optimizationProfile, + searchOptions, + ), + ), + ) + } + const repositories = [...new Set(group.samples.map((sample) => sample.repository))] + if (config.repositoryHoldouts && repositories.length > 1) { + for (const repository of repositories) { + weightSearch.push( + yield* runParallelSearch(() => + optimizeWeightsParallel( + group.model, + group.queryKind, + "leave-one-repository-out", + repository, + group.samples.filter((sample) => sample.repository !== repository), + group.samples.filter((sample) => sample.repository === repository), + optimizationProfile, + searchOptions, + ), + ), + ) + } + } + recommendedWeights.push( + yield* runParallelSearch(() => + fitRecommendedWeightsParallel( group.model, group.queryKind, - groupedStrategy, - String(fold + 1), - group.samples.filter((sample) => sample.groupedFold !== fold), - group.samples.filter((sample) => sample.groupedFold === fold), + group.samples, optimizationProfile, + searchOptions, ), - ) - const repositories = [...new Set(group.samples.map((sample) => sample.repository))] - const repositoryFolds = - config.repositoryHoldouts && repositories.length > 1 - ? repositories.map((repository) => - optimizeWeights( - group.model, - group.queryKind, - "leave-one-repository-out", - repository, - group.samples.filter((sample) => sample.repository !== repository), - group.samples.filter((sample) => sample.repository === repository), - optimizationProfile, - ), - ) - : [] - return [...groupedFolds, ...repositoryFolds] - }) - : [] - const recommendedWeights = config.legacyDiagnostics - ? [...sampleGroups.values()].map((group) => - fitRecommendedWeights(group.model, group.queryKind, group.samples, optimizationProfile), + ), ) - : [] + } + } const weightSearchDurationMs = performance.now() - weightSearchStartedAt const fusionSearchStartedAt = performance.now() @@ -569,40 +604,57 @@ export const runRetrievalBenchmark = ( for (const fusion of config.fusionMethods) { reportProgress(`${model}: selecting static ${fusion} fusion weights`) for (let fold = 0; fold < config.groupedFolds; fold++) { + const split = splitSamples(samples, (sample) => sample.groupedFold === fold) fusionSearch.push( - optimizeFusionWeights( - model, - fusion, - groupedStrategy, - String(fold + 1), - samples.filter((sample) => sample.groupedFold !== fold), - samples.filter((sample) => sample.groupedFold === fold), - optimizationProfile, + yield* runParallelSearch(() => + optimizeFusionWeightsParallel( + model, + fusion, + groupedStrategy, + String(fold + 1), + split.development, + split.validation, + optimizationProfile, + searchOptions, + ), ), ) } if (config.repositoryHoldouts && repositories.length > 1) { for (const repository of repositories) { + const split = splitSamples(samples, (sample) => sample.repository === repository) fusionSearch.push( - optimizeFusionWeights( - model, - fusion, - "leave-one-repository-out", - repository, - samples.filter((sample) => sample.repository !== repository), - samples.filter((sample) => sample.repository === repository), - optimizationProfile, + yield* runParallelSearch(() => + optimizeFusionWeightsParallel( + model, + fusion, + "leave-one-repository-out", + repository, + split.development, + split.validation, + optimizationProfile, + searchOptions, + ), ), ) } } } } - const recommendedFusionWeights = [...samplesByModel].flatMap(([model, samples]) => - config.fusionMethods.map((fusion) => - fitRecommendedFusionWeights(model, fusion, samples, optimizationProfile), - ), - ) + const recommendedFusionWeights: BenchmarkArtifact["recommendedFusionWeights"][number][] = [] + for (const [model, samples] of samplesByModel) + for (const fusion of config.fusionMethods) + recommendedFusionWeights.push( + yield* runParallelSearch(() => + fitRecommendedFusionWeightsParallel( + model, + fusion, + samples, + optimizationProfile, + searchOptions, + ), + ), + ) const fusionSearchDurationMs = performance.now() - fusionSearchStartedAt const evidenceRouterSearchStartedAt = performance.now() @@ -610,48 +662,65 @@ export const runRetrievalBenchmark = ( for (const [model, samples] of samplesByModel) { for (const fusion of config.routerFusionMethods) { for (let fold = 0; fold < config.groupedFolds; fold++) { + const split = splitSamples(samples, (sample) => sample.groupedFold === fold) reportProgress( `${model}: selecting ${fusion} evidence router for grouped fold ${fold + 1}/${config.groupedFolds}`, ) evidenceRouterSearch.push( - ...optimizeEvidenceRouter( - model, - fusion, - groupedStrategy, - String(fold + 1), - samples.filter((sample) => sample.groupedFold !== fold), - samples.filter((sample) => sample.groupedFold === fold), - optimizationProfile, - ), + ...(yield* runParallelSearch(() => + optimizeEvidenceRouterParallel( + model, + fusion, + groupedStrategy, + String(fold + 1), + split.development, + split.validation, + optimizationProfile, + searchOptions, + ), + )), ) } const repositories = [...new Set(samples.map((sample) => sample.repository))] if (config.repositoryHoldouts && repositories.length > 1) { for (const repository of repositories) { + const split = splitSamples(samples, (sample) => sample.repository === repository) reportProgress( `${model}: selecting ${fusion} evidence router with ${repository} held out`, ) evidenceRouterSearch.push( - ...optimizeEvidenceRouter( - model, - fusion, - "leave-one-repository-out", - repository, - samples.filter((sample) => sample.repository !== repository), - samples.filter((sample) => sample.repository === repository), - optimizationProfile, - ), + ...(yield* runParallelSearch(() => + optimizeEvidenceRouterParallel( + model, + fusion, + "leave-one-repository-out", + repository, + split.development, + split.validation, + optimizationProfile, + searchOptions, + ), + )), ) } } } } reportProgress("fitting final evidence router on all samples") - const recommendedEvidenceRouters = [...samplesByModel].flatMap(([model, samples]) => - config.routerFusionMethods.flatMap((fusion) => - fitRecommendedEvidenceRouter(model, fusion, samples, optimizationProfile), - ), - ) + const recommendedEvidenceRouters: BenchmarkArtifact["recommendedEvidenceRouters"][number][] = [] + for (const [model, samples] of samplesByModel) + for (const fusion of config.routerFusionMethods) + recommendedEvidenceRouters.push( + ...(yield* runParallelSearch(() => + fitRecommendedEvidenceRouterParallel( + model, + fusion, + samples, + optimizationProfile, + searchOptions, + ), + )), + ) const evidenceRouterSearchDurationMs = performance.now() - evidenceRouterSearchStartedAt const embeddingDurationMs = embeddingRuns.reduce( diff --git a/benchmarks/retrieval/weight-search.ts b/benchmarks/retrieval/weight-search.ts index 6d9d500..7a7cf4a 100644 --- a/benchmarks/retrieval/weight-search.ts +++ b/benchmarks/retrieval/weight-search.ts @@ -19,7 +19,7 @@ import { type QueryTermCoverage, type RoutingEvidence, } from "../../src/lib/retrieval/evidence-router.js" -import { fuseRankings } from "../../src/lib/retrieval/fusion.js" +import { prepareFusion, type PreparedFusionEvaluator } from "./fusion.js" import { contextRecallAtBudget, recallAt, reciprocalRank } from "./metrics.js" import { SEARCH_PRIORITY_PROFILE, type OptimizationProfile } from "./optimization-profiles.js" import { @@ -40,6 +40,13 @@ import { type SearchBaselineComparison, type WeightSearchResult, } from "./types.js" +import { + createCandidateEvaluationPool, + createEvaluationSnapshot, + type CandidateEvaluationPool, + type CandidateEvaluationPoolOptions, + type EvaluationCandidate, +} from "./worker-pool.js" const CHANNELS: readonly ChannelName[] = CHANNEL_NAMES const WEIGHT_LEVELS = [0, 0.5, 1, 2] as const @@ -79,11 +86,61 @@ interface EvidenceSearchSample { readonly evidence: RoutingEvidence } +/** Native worker-pool controls for the explicit asynchronous benchmark search path. */ +export type ParallelSearchOptions = CandidateEvaluationPoolOptions + +const preparedFusionCache = new WeakMap< + WeightSearchSample, + Map +>() + +const preparedFusionEvaluatorFor = ( + sample: WeightSearchSample, + fusion: FusionMethod = "rrf", +): PreparedFusionEvaluator => { + let evaluators = preparedFusionCache.get(sample) + if (evaluators === undefined) { + evaluators = new Map() + preparedFusionCache.set(sample, evaluators) + } + let evaluator = evaluators.get(fusion) + if (evaluator === undefined) { + evaluator = prepareFusion(fusion, sample.rankings, SEARCH_CANDIDATE_DEPTH) + evaluators.set(fusion, evaluator) + } + return evaluator +} + const fuseWithWeights = ( - rankings: ChannelRankings, + sample: WeightSearchSample, weights: ChannelWeights, fusion: FusionMethod = "rrf", -): readonly RankedChunk[] => fuseRankings(fusion, rankings, weights, SEARCH_CANDIDATE_DEPTH) +): readonly RankedChunk[] => preparedFusionEvaluatorFor(sample, fusion).evaluate(weights) + +const createEvaluationPoolForSamples = async ( + samples: readonly WeightSearchSample[], + fusion: FusionMethod, + profile: OptimizationProfile, + options: CandidateEvaluationPoolOptions, +): Promise => + createCandidateEvaluationPool( + createEvaluationSnapshot( + samples.map((sample) => ({ + evaluator: preparedFusionEvaluatorFor(sample, fusion), + targets: sample.targets, + chunks: sample.chunks, + sampleWeight: profile.queryFormWeights[sample.queryKind], + })), + ), + options, + ) + +const routerEvaluationCandidate = ( + samples: readonly EvidenceSearchSample[], + config: EvidenceRouterConfig, +): EvaluationCandidate => ({ + weights: samples.map(({ evidence }) => routeWithEvidence(evidence, config)), +}) const summarizeRanked = ( samples: readonly T[], @@ -150,7 +207,7 @@ const summarize = ( ): QualitySummary => summarizeRanked( samples, - (sample) => fuseWithWeights(sample.rankings, weights, fusion), + (sample) => fuseWithWeights(sample, weights, fusion), (sample) => sample.targets, (sample) => sample.chunks, (sample) => profile.queryFormWeights[sample.queryKind], @@ -164,8 +221,7 @@ const summarizeEvidenceRouter = ( ): QualitySummary => summarizeRanked( samples, - ({ sample, evidence }) => - fuseWithWeights(sample.rankings, routeWithEvidence(evidence, config), fusion), + ({ sample, evidence }) => fuseWithWeights(sample, routeWithEvidence(evidence, config), fusion), ({ sample }) => sample.targets, ({ sample }) => sample.chunks, ({ sample }) => profile.queryFormWeights[sample.queryKind], @@ -180,7 +236,7 @@ const summarizeProductionRouter = ( samples, (sample) => fuseWithWeights( - sample.rankings, + sample, routeWithEvidence( buildRoutingEvidence(sample.query, sample.rankings, sample.termCoverage), PRODUCTION_COMPATIBILITY_CONFIG, @@ -514,27 +570,29 @@ const normalizeWeights = (weights: ChannelWeights): ChannelWeights => { const weightsKey = (weights: ChannelWeights): string => CHANNELS.map((channel) => weights[channel].toFixed(4)).join(":") -const rankWeightCandidates = ( - samples: readonly WeightSearchSample[], +const uniqueWeightCandidates = ( candidates: readonly ChannelWeights[], - limit: number, - qualityCache: Map, - fusion: FusionMethod, - objective: RouterObjective = "reranker-top20", - baseline?: QualitySummary, - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): readonly WeightCandidate[] => { +): readonly (readonly [string, ChannelWeights])[] => { const unique = new Map() for (const candidate of candidates) { const normalized = normalizeWeights(candidate) unique.set(weightsKey(normalized), normalized) } return [...unique] +} + +const sortWeightCandidates = ( + entries: readonly (readonly [string, ChannelWeights])[], + qualityCache: Map, + limit: number, + objective: RouterObjective, + baseline: QualitySummary | undefined, + profile: OptimizationProfile, +): readonly WeightCandidate[] => + entries .map(([key, weights]) => { - const cached = qualityCache.get(key) - if (cached !== undefined) return { weights, quality: cached } - const quality = summarize(samples, weights, fusion, profile) - qualityCache.set(key, quality) + const quality = qualityCache.get(key) + if (quality === undefined) throw new Error(`Missing cached quality for ${key}`) return { weights, quality } }) .sort( @@ -543,6 +601,42 @@ const rankWeightCandidates = ( activeChannelsKey(left.weights).localeCompare(activeChannelsKey(right.weights)), ) .slice(0, limit) + +const storeQualityResults = ( + pending: readonly (readonly [string, unknown])[], + qualities: readonly QualitySummary[], + qualityCache: Map, + errorMessage: string, +): void => { + for (let index = 0; index < pending.length; index++) { + const entry = pending[index] + const quality = qualities[index] + if (entry === undefined || quality === undefined) throw new Error(errorMessage) + qualityCache.set(entry[0], quality) + } +} + +const rankWeightCandidates = async ( + samples: readonly WeightSearchSample[], + candidates: readonly ChannelWeights[], + limit: number, + qualityCache: Map, + fusion: FusionMethod, + pool: CandidateEvaluationPool, + objective: RouterObjective = "reranker-top20", + baseline?: QualitySummary, + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, +): Promise => { + const unique = uniqueWeightCandidates(candidates) + const pending = unique.filter(([key]) => qualityCache.get(key) === undefined) + const qualities = await pool.evaluate(pending.map(([, weights]) => ({ weights }))) + storeQualityResults( + pending, + qualities, + qualityCache, + "Parallel weight evaluation returned an incomplete result", + ) + return sortWeightCandidates(unique, qualityCache, limit, objective, baseline, profile) } const withWeight = (weights: ChannelWeights, channel: ChannelName, value: number): ChannelWeights => @@ -551,38 +645,46 @@ const withWeight = (weights: ChannelWeights, channel: ChannelName, value: number [channel]: value, }) -const selectBestWeights = ( +const coordinateWeightCandidates = ( + beam: readonly WeightCandidate[], + channel: ChannelName, +): readonly ChannelWeights[] => [ + // Retain current elites so a later coordinate cannot regress the best development fit. + ...beam.map((candidate) => candidate.weights), + ...beam.flatMap((candidate) => + STATIC_FINE_WEIGHT_LEVELS.map((level) => withWeight(candidate.weights, channel, level)), + ), +] + +const selectBestWeights = async ( samples: readonly WeightSearchSample[], fusion: FusionMethod = "rrf", + pool: CandidateEvaluationPool, objective: RouterObjective = "reranker-top20", baseline?: QualitySummary, profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): { readonly weights: ChannelWeights; readonly quality: QualitySummary } => { +): Promise<{ readonly weights: ChannelWeights; readonly quality: QualitySummary }> => { const qualityCache = new Map() - let beam = rankWeightCandidates( + let beam = await rankWeightCandidates( samples, weightCandidates(), SEARCH_BEAM_WIDTH, qualityCache, fusion, + pool, objective, baseline, profile, ) for (let pass = 0; pass < SEARCH_PASSES; pass++) { for (const channel of CHANNELS) { - beam = rankWeightCandidates( + beam = await rankWeightCandidates( samples, - [ - // Retain current elites so a later coordinate cannot regress the best development fit. - ...beam.map((candidate) => candidate.weights), - ...beam.flatMap((candidate) => - STATIC_FINE_WEIGHT_LEVELS.map((level) => withWeight(candidate.weights, channel, level)), - ), - ], + coordinateWeightCandidates(beam, channel), SEARCH_BEAM_WIDTH, qualityCache, fusion, + pool, objective, baseline, profile, @@ -595,18 +697,21 @@ const selectBestWeights = ( const activeChannelsKey = (weights: ChannelWeights): string => CHANNELS.filter((channel) => weights[channel] > 0).join("+") -const selectBestWeightsPerSubset = ( - samples: readonly WeightSearchSample[], - fusion: FusionMethod, - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, +const selectWeightSubsetCandidates = ( + candidates: readonly ChannelWeights[], + qualities: readonly QualitySummary[], + profile: OptimizationProfile, ): readonly ChannelWeights[] => { const selected = new Map< string, { readonly weights: ChannelWeights; readonly quality: QualitySummary } >() - for (const weights of weightCandidates()) { + for (let index = 0; index < candidates.length; index++) { + const weights = candidates[index] + const quality = qualities[index] + if (weights === undefined || quality === undefined) + throw new Error("Weight subset search returned an incomplete result") const key = activeChannelsKey(weights) - const quality = summarize(samples, weights, fusion, profile) const current = selected.get(key) if ( current === undefined || @@ -617,6 +722,15 @@ const selectBestWeightsPerSubset = ( return [...selected.values()].map((entry) => entry.weights) } +const selectBestWeightsPerSubset = async ( + pool: CandidateEvaluationPool, + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, +): Promise => { + const candidates = weightCandidates() + const qualities = await pool.evaluate(candidates.map((weights) => ({ weights }))) + return selectWeightSubsetCandidates(candidates, qualities, profile) +} + /** Config field containing one coefficient per retrieval channel. */ type InfluenceName = | "scoreInfluence" @@ -873,6 +987,8 @@ interface RouterSearchContext { readonly fusion: FusionMethod readonly profile: OptimizationProfile readonly stats: SearchEvaluationStats + readonly fullPool: CandidateEvaluationPool + readonly proxyPool: CandidateEvaluationPool } const compareRouterCandidates = ( @@ -885,21 +1001,31 @@ const compareRouterCandidates = ( compareObjectiveQuality(left.quality, right.quality, objective, baseline, profile) || routerComplexity(left.config) - routerComplexity(right.config) -const selectRandomRouter = ( +const selectRandomRouter = async ( samples: readonly EvidenceSearchSample[], baseSeed: EvidenceRouterConfig, parameters: readonly RouterParameter[], - fusion: FusionMethod, + pool: CandidateEvaluationPool, baseline: QualitySummary, profile: OptimizationProfile, -): { readonly candidate: RouterCandidate; readonly candidates: number } => { - const candidates = buildRandomRouterSeeds(baseSeed, parameters).map((config) => ({ - config, - quality: summarizeEvidenceRouter(samples, config, fusion, profile), - })) +): Promise<{ readonly candidate: RouterCandidate; readonly candidates: number }> => { + const configs = buildRandomRouterSeeds(baseSeed, parameters) + const qualities = await pool.evaluate( + configs.map((config) => routerEvaluationCandidate(samples, config)), + ) + const candidates: RouterCandidate[] = [] + for (let index = 0; index < configs.length; index++) { + const config = configs[index] + const quality = qualities[index] + if (config === undefined || quality === undefined) + throw new Error("Parallel random router search returned an incomplete result") + candidates.push({ config, quality }) + } const candidate = [...candidates].sort((left, right) => compareRouterCandidates(left, right, "reranker-top20", baseline, profile), - )[0] ?? { config: baseSeed, quality: summarizeEvidenceRouter(samples, baseSeed, fusion, profile) } + )[0] + if (candidate === undefined) + throw new Error("Parallel random router search produced no candidate") return { candidate, candidates: candidates.length } } @@ -929,13 +1055,79 @@ const selectObjectiveCandidates = ( return [...selected.values()].slice(0, limit) } -const rankRouterCandidates = ( +interface RouterEvaluationResult { + readonly candidates: readonly RouterCandidate[] + readonly cacheHits: number + readonly evaluations: number +} + +const finalizeRouterCandidates = ( + context: RouterSearchContext, + rankedFull: readonly RouterCandidate[], + proxyKeys: readonly string[], + limit: number, +): readonly RouterCandidate[] => { + if (proxyKeys.length > 0) { + const fullRanks = new Map( + rankedFull.map((candidate, rank) => [routerKey(candidate.config), rank]), + ) + for (let left = 0; left < proxyKeys.length; left++) { + const leftRank = fullRanks.get(proxyKeys[left]) + if (leftRank === undefined) continue + for (let right = left + 1; right < proxyKeys.length; right++) { + const rightRank = fullRanks.get(proxyKeys[right]) + if (rightRank === undefined) continue + context.stats.proxyAgreementComparisons++ + if (leftRank < rightRank) context.stats.proxyAgreementMatches++ + } + } + } + for (const candidate of rankedFull) context.archive.set(routerKey(candidate.config), candidate) + const rankedElites = selectObjectiveCandidates( + [...context.elites.values(), ...rankedFull], + SEARCH_BEAM_WIDTH, + context.baseline, + context.profile, + ) + context.elites.clear() + for (const candidate of rankedElites) context.elites.set(routerKey(candidate.config), candidate) + return selectObjectiveCandidates(rankedFull, limit, context.baseline, context.profile) +} + +const evaluateRouterConfigs = async ( + entries: readonly (readonly [string, EvidenceRouterConfig])[], + samples: readonly EvidenceSearchSample[], + pool: CandidateEvaluationPool, + qualityCache: Map, +): Promise => { + const pending = entries.filter(([key]) => qualityCache.get(key) === undefined) + const qualities = await pool.evaluate( + pending.map(([, config]) => routerEvaluationCandidate(samples, config)), + ) + storeQualityResults( + pending, + qualities, + qualityCache, + "Parallel router evaluation returned an incomplete result", + ) + return { + candidates: entries.map(([key, config]) => { + const quality = qualityCache.get(key) + if (quality === undefined) throw new Error(`Missing cached router quality for ${key}`) + return { config, quality } + }), + cacheHits: entries.length - pending.length, + evaluations: pending.length, + } +} + +const rankRouterCandidates = async ( context: RouterSearchContext, configs: readonly EvidenceRouterConfig[], limit: number, protectedConfigs: readonly EvidenceRouterConfig[] = [], useProxy = true, -): readonly RouterCandidate[] => { +): Promise => { const unique = new Map() for (const config of configs) unique.set(routerKey(config), config) context.stats.rawCandidates += configs.length @@ -944,24 +1136,16 @@ const rankRouterCandidates = ( let fullCandidates = [...unique] let proxyKeys: readonly string[] = [] if (useProxy && context.proxySamples.length < context.samples.length) { - const rankedProxy = fullCandidates.map(([key, config]) => { - const cached = context.proxyQualityCache.get(key) - if (cached !== undefined) { - context.stats.proxyCacheHits++ - return { key, config, quality: cached } - } - const quality = summarizeEvidenceRouter( - context.proxySamples, - config, - context.fusion, - context.profile, - ) - context.proxyQualityCache.set(key, quality) - context.stats.proxyEvaluations++ - return { key, config, quality } - }) + const rankedProxy = await evaluateRouterConfigs( + fullCandidates, + context.proxySamples, + context.proxyPool, + context.proxyQualityCache, + ) + context.stats.proxyCacheHits += rankedProxy.cacheHits + context.stats.proxyEvaluations += rankedProxy.evaluations const selectedProxy = selectObjectiveCandidates( - rankedProxy.map((candidate) => ({ config: candidate.config, quality: candidate.quality })), + rankedProxy.candidates, limit * SEARCH_PROXY_PROMOTION_FACTOR, context.proxyBaseline, context.profile, @@ -977,47 +1161,15 @@ const rankRouterCandidates = ( } fullCandidates = [...selected] } - const rankedFull = fullCandidates.map(([key, config]) => { - const cached = context.qualityCache.get(key) - if (cached !== undefined) { - context.stats.fullCacheHits++ - return { config, quality: cached } - } - const quality = summarizeEvidenceRouter( - context.samples, - config, - context.fusion, - context.profile, - ) - context.qualityCache.set(key, quality) - context.stats.fullEvaluations++ - return { config, quality } - }) - if (proxyKeys.length > 0) { - const fullRanks = new Map( - rankedFull.map((candidate, rank) => [routerKey(candidate.config), rank]), - ) - for (let left = 0; left < proxyKeys.length; left++) { - const leftRank = fullRanks.get(proxyKeys[left]) - if (leftRank === undefined) continue - for (let right = left + 1; right < proxyKeys.length; right++) { - const rightRank = fullRanks.get(proxyKeys[right]) - if (rightRank === undefined) continue - context.stats.proxyAgreementComparisons++ - if (leftRank < rightRank) context.stats.proxyAgreementMatches++ - } - } - } - for (const candidate of rankedFull) context.archive.set(routerKey(candidate.config), candidate) - const rankedElites = selectObjectiveCandidates( - [...context.elites.values(), ...rankedFull], - SEARCH_BEAM_WIDTH, - context.baseline, - context.profile, + const rankedFull = await evaluateRouterConfigs( + fullCandidates, + context.samples, + context.fullPool, + context.qualityCache, ) - context.elites.clear() - for (const candidate of rankedElites) context.elites.set(routerKey(candidate.config), candidate) - return selectObjectiveCandidates(rankedFull, limit, context.baseline, context.profile) + context.stats.fullCacheHits += rankedFull.cacheHits + context.stats.fullEvaluations += rankedFull.evaluations + return finalizeRouterCandidates(context, rankedFull.candidates, proxyKeys, limit) } interface EvidenceRouterSelection { @@ -1040,39 +1192,81 @@ export const selectEligibleCandidate = ( } } -const selectBestEvidenceRouter = ( - samples: readonly WeightSearchSample[], - fusion: FusionMethod, - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): { - readonly selections: readonly EvidenceRouterSelection[] +interface RouterSearchPreparation { + readonly evidenceSamples: readonly EvidenceSearchSample[] + readonly proxySamples: readonly EvidenceSearchSample[] + readonly guardrailBaselines: GuardrailBaselines readonly productionQuality: QualitySummary - readonly proxyEvaluations: number - readonly fullEvaluations: number - readonly searchDiagnostics: RouterSearchDiagnostics - readonly randomCandidate: RouterCandidate - readonly randomCandidates: number -} => { + readonly proxyBaseline: QualitySummary + readonly stats: SearchEvaluationStats +} + +const prepareRouterSearch = ( + samples: readonly WeightSearchSample[], + profile: OptimizationProfile, +): RouterSearchPreparation => { const evidenceSamples = prepareEvidenceSamples(samples) const proxySamples = buildProxySamples(evidenceSamples) const guardrailBaselines = buildGuardrailBaselines(samples, profile) - const productionQuality = guardrailBaselines.overall const proxyBaselines = buildGuardrailBaselines( proxySamples.map(({ sample }) => sample), profile, ) - const stats: SearchEvaluationStats = { - rawCandidates: 0, - uniqueCandidates: 0, - proxyEvaluations: 0, - fullEvaluations: 0, - proxyCacheHits: 0, - fullCacheHits: 0, - proxyPromotions: 0, - proxyAgreementMatches: 0, - proxyAgreementComparisons: 0, - protectedEliteCount: 0, + return { + evidenceSamples, + proxySamples, + guardrailBaselines, + productionQuality: guardrailBaselines.overall, + proxyBaseline: proxyBaselines.overall, + stats: { + rawCandidates: 0, + uniqueCandidates: 0, + proxyEvaluations: 0, + fullEvaluations: 0, + proxyCacheHits: 0, + fullCacheHits: 0, + proxyPromotions: 0, + proxyAgreementMatches: 0, + proxyAgreementComparisons: 0, + protectedEliteCount: 0, + }, } +} + +const profileSeedFor = (profile: OptimizationProfile): EvidenceRouterConfig => ({ + baseWeights: profile.fusionConfig.baseWeights, + scoreInfluence: profile.fusionConfig.scoreInfluence, + geometryInfluence: profile.fusionConfig.geometryInfluence, + termCoverageInfluence: profile.fusionConfig.termCoverageInfluence, + pairwiseAgreementInfluence: profile.fusionConfig.pairwiseAgreementInfluence, + denseConfidenceInfluence: profile.fusionConfig.denseConfidenceInfluence, + identifierInfluence: profile.fusionConfig.identifierInfluence, + queryLengthInfluence: profile.fusionConfig.queryLengthInfluence, +}) + +const selectBestEvidenceRouter = async ( + samples: readonly WeightSearchSample[], + fusion: FusionMethod, + profile: OptimizationProfile, + fullPool: CandidateEvaluationPool, + proxyPool: CandidateEvaluationPool, +): Promise<{ + readonly selections: readonly EvidenceRouterSelection[] + readonly productionQuality: QualitySummary + readonly proxyEvaluations: number + readonly fullEvaluations: number + readonly searchDiagnostics: RouterSearchDiagnostics + readonly randomCandidate: RouterCandidate + readonly randomCandidates: number +}> => { + const { + evidenceSamples, + proxySamples, + guardrailBaselines, + productionQuality, + proxyBaseline, + stats, + } = prepareRouterSearch(samples, profile) const searchContext: RouterSearchContext = { samples: evidenceSamples, proxySamples, @@ -1081,38 +1275,37 @@ const selectBestEvidenceRouter = ( elites: new Map(), archive: new Map(), baseline: productionQuality, - proxyBaseline: proxyBaselines.overall, + proxyBaseline, fusion, profile, stats, + fullPool, + proxyPool, } - const profileSeed: EvidenceRouterConfig = { - baseWeights: profile.fusionConfig.baseWeights, - scoreInfluence: profile.fusionConfig.scoreInfluence, - geometryInfluence: profile.fusionConfig.geometryInfluence, - termCoverageInfluence: profile.fusionConfig.termCoverageInfluence, - pairwiseAgreementInfluence: profile.fusionConfig.pairwiseAgreementInfluence, - denseConfidenceInfluence: profile.fusionConfig.denseConfidenceInfluence, - identifierInfluence: profile.fusionConfig.identifierInfluence, - queryLengthInfluence: profile.fusionConfig.queryLengthInfluence, - } - const baseSeeds = [ - selectBestWeights(samples, fusion, "reranker-top20", productionQuality, profile).weights, - ...selectBestWeightsPerSubset(samples, fusion, profile), - ].map(emptyRouterConfig) + const profileSeed = profileSeedFor(profile) + const baseWeights = await selectBestWeights( + samples, + fusion, + fullPool, + "reranker-top20", + productionQuality, + profile, + ) + const subsetWeights = await selectBestWeightsPerSubset(fullPool, profile) + const baseSeeds = [baseWeights.weights, ...subsetWeights].map(emptyRouterConfig) baseSeeds.unshift(profileSeed) const parameters = routerParameters() const randomBaseSeed = baseSeeds[0] if (randomBaseSeed === undefined) throw new Error("Evidence router search has no base seed") - const randomSearch = selectRandomRouter( + const randomSearch = await selectRandomRouter( evidenceSamples, randomBaseSeed, parameters, - fusion, + fullPool, productionQuality, profile, ) - let beam = rankRouterCandidates( + let beam = await rankRouterCandidates( searchContext, [...baseSeeds, ...buildGlobalRouterSeeds(baseSeeds, parameters)], SEARCH_BEAM_WIDTH, @@ -1121,7 +1314,7 @@ const selectBestEvidenceRouter = ( for (let pass = 0; pass < SEARCH_PASSES; pass++) { const orderedParameters = pass % 2 === 0 ? parameters : [...parameters].reverse() for (const parameter of orderedParameters) { - beam = rankRouterCandidates( + beam = await rankRouterCandidates( searchContext, [ // Retain the current beam; the search context also protects full-quality elites. @@ -1167,8 +1360,23 @@ const selectBestEvidenceRouter = ( } } +const withCandidatePool = async ( + samples: readonly WeightSearchSample[], + fusion: FusionMethod, + profile: OptimizationProfile, + options: CandidateEvaluationPoolOptions, + operation: (pool: CandidateEvaluationPool) => Promise, +): Promise => { + const pool = await createEvaluationPoolForSamples(samples, fusion, profile, options) + try { + return await operation(pool) + } finally { + await pool.close() + } +} + /** Select weights on development samples, then evaluate unchanged on one validation fold. */ -export const optimizeWeights = ( +const optimizeWeightsWithPool = async ( model: string, queryKind: QueryKind, strategy: WeightSearchResult["strategy"], @@ -1176,8 +1384,16 @@ export const optimizeWeights = ( development: readonly WeightSearchSample[], validation: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): WeightSearchResult => { - const selected = selectBestWeights(development, "rrf", "reranker-top20", undefined, profile) + pool: CandidateEvaluationPool, +): Promise => { + const selected = await selectBestWeights( + development, + "rrf", + pool, + "reranker-top20", + undefined, + profile, + ) return { model, queryKind, @@ -1192,8 +1408,65 @@ export const optimizeWeights = ( } } +const optimizeWeightsWithOptions = ( + model: string, + queryKind: QueryKind, + strategy: WeightSearchResult["strategy"], + fold: string, + development: readonly WeightSearchSample[], + validation: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + options: ParallelSearchOptions, +): Promise => + withCandidatePool(development, "rrf", profile, options, (pool) => + optimizeWeightsWithPool( + model, + queryKind, + strategy, + fold, + development, + validation, + profile, + pool, + ), + ) + +export const optimizeWeights = ( + model: string, + queryKind: QueryKind, + strategy: WeightSearchResult["strategy"], + fold: string, + development: readonly WeightSearchSample[], + validation: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, +): Promise => + optimizeWeightsWithOptions(model, queryKind, strategy, fold, development, validation, profile, { + workerCount: 0, + }) + +export const optimizeWeightsParallel = ( + model: string, + queryKind: QueryKind, + strategy: WeightSearchResult["strategy"], + fold: string, + development: readonly WeightSearchSample[], + validation: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + options: ParallelSearchOptions = {}, +): Promise => + optimizeWeightsWithOptions( + model, + queryKind, + strategy, + fold, + development, + validation, + profile, + options, + ) + /** Select static weights for one fusion method, then evaluate them unchanged on a holdout. */ -export const optimizeFusionWeights = ( +const optimizeFusionWeightsWithPool = async ( model: string, fusion: FusionMethod, strategy: FusionSearchResult["strategy"], @@ -1201,8 +1474,16 @@ export const optimizeFusionWeights = ( development: readonly WeightSearchSample[], validationSamples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): FusionSearchResult => { - const selected = selectBestWeights(development, fusion, "reranker-top20", undefined, profile) + pool: CandidateEvaluationPool, +): Promise => { + const selected = await selectBestWeights( + development, + fusion, + pool, + "reranker-top20", + undefined, + profile, + ) const guardrailBaselines = buildGuardrailBaselines(development, profile) const holdoutProfile = unweightedProfile(profile) const validationQuality = summarize(validationSamples, selected.weights, fusion, profile) @@ -1233,6 +1514,70 @@ export const optimizeFusionWeights = ( } } +const optimizeFusionWeightsWithOptions = ( + model: string, + fusion: FusionMethod, + strategy: FusionSearchResult["strategy"], + fold: string, + development: readonly WeightSearchSample[], + validationSamples: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + options: ParallelSearchOptions, +): Promise => + withCandidatePool(development, fusion, profile, options, (pool) => + optimizeFusionWeightsWithPool( + model, + fusion, + strategy, + fold, + development, + validationSamples, + profile, + pool, + ), + ) + +export const optimizeFusionWeights = ( + model: string, + fusion: FusionMethod, + strategy: FusionSearchResult["strategy"], + fold: string, + development: readonly WeightSearchSample[], + validationSamples: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, +): Promise => + optimizeFusionWeightsWithOptions( + model, + fusion, + strategy, + fold, + development, + validationSamples, + profile, + { workerCount: 0 }, + ) + +export const optimizeFusionWeightsParallel = ( + model: string, + fusion: FusionMethod, + strategy: FusionSearchResult["strategy"], + fold: string, + development: readonly WeightSearchSample[], + validationSamples: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + options: ParallelSearchOptions = {}, +): Promise => + optimizeFusionWeightsWithOptions( + model, + fusion, + strategy, + fold, + development, + validationSamples, + profile, + options, + ) + /** Evaluate the current production router on one development/validation split. */ export const evaluateProductionRouter = ( model: string, @@ -1251,8 +1596,44 @@ export const evaluateProductionRouter = ( validation: summarizeProductionRouter(validation, profile), }) -/** Select one evidence-based router on development queries and evaluate it unchanged on a holdout. */ -export const optimizeEvidenceRouter = ( +const withParallelEvidencePools = async ( + samples: readonly WeightSearchSample[], + fusion: FusionMethod, + profile: OptimizationProfile, + options: ParallelSearchOptions, + operation: ( + selection: Awaited>, + fullPool: CandidateEvaluationPool, + ) => Promise, +): Promise => { + const evidenceSamples = prepareEvidenceSamples(samples) + const proxySamples = buildProxySamples(evidenceSamples) + const fullPool = await createEvaluationPoolForSamples( + evidenceSamples.map(({ sample }) => sample), + fusion, + profile, + options, + ) + let proxyPool: CandidateEvaluationPool | undefined + try { + proxyPool = + proxySamples === evidenceSamples + ? fullPool + : await createEvaluationPoolForSamples( + proxySamples.map(({ sample }) => sample), + fusion, + profile, + options, + ) + const selection = await selectBestEvidenceRouter(samples, fusion, profile, fullPool, proxyPool) + return await operation(selection, fullPool) + } finally { + if (proxyPool !== undefined && proxyPool !== fullPool) await proxyPool.close() + await fullPool.close() + } +} + +const optimizeEvidenceRouterWithPools = async ( model: string, fusion: FusionMethod, strategy: EvidenceRouterSearchResult["strategy"], @@ -1260,76 +1641,114 @@ export const optimizeEvidenceRouter = ( development: readonly WeightSearchSample[], validation: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): readonly EvidenceRouterSearchResult[] => { - const dynamicSelection = selectBestEvidenceRouter(development, fusion, profile) - const productionValidation = summarizeProductionRouter(validation, profile) - const validationEvidence = prepareEvidenceSamples(validation) - const holdoutProfile = unweightedProfile(profile) - const randomBaseline: SearchBaselineComparison = { - algorithm: "random-scout", - seed: RANDOM_SEARCH_SEED, - candidates: dynamicSelection.randomCandidates, - development: dynamicSelection.randomCandidate.quality, - validation: summarizeEvidenceRouter( - validationEvidence, - dynamicSelection.randomCandidate.config, - fusion, - profile, - ), - } - return dynamicSelection.selections.map((selection) => { - const staticSelection = selectBestWeights( - development, - fusion, - selection.objective, - dynamicSelection.productionQuality, - profile, - ) - return { - model, - fusion, - objective: selection.objective, - strategy, - fold, - developmentQueries: development.length, - validationQueries: validation.length, - staticWeights: staticSelection.weights, - config: benchmarkRouterConfig(fusion, selection.config), - staticDevelopment: staticSelection.quality, - staticValidation: summarize(validation, staticSelection.weights, fusion, profile), - development: selection.quality, - validation: summarizeEvidenceRouter(validationEvidence, selection.config, fusion, profile), - productionDevelopment: dynamicSelection.productionQuality, - productionValidation, - guardrailsMet: selection.guardrailsMet, - promotionStatus: selection.promotionStatus, - proxyEvaluations: dynamicSelection.proxyEvaluations, - fullEvaluations: dynamicSelection.fullEvaluations, - searchDiagnostics: dynamicSelection.searchDiagnostics, - searchBaseline: randomBaseline, - holdoutBreakdown: buildHoldoutBreakdown( - validation, - (partition) => - summarizeEvidenceRouter( - prepareEvidenceSamples(partition), + options: ParallelSearchOptions = {}, +): Promise => + withParallelEvidencePools( + development, + fusion, + profile, + options, + async (dynamicSelection, fullPool) => { + const productionValidation = summarizeProductionRouter(validation, profile) + const validationEvidence = prepareEvidenceSamples(validation) + const holdoutProfile = unweightedProfile(profile) + const randomBaseline: SearchBaselineComparison = { + algorithm: "random-scout", + seed: RANDOM_SEARCH_SEED, + candidates: dynamicSelection.randomCandidates, + development: dynamicSelection.randomCandidate.quality, + validation: summarizeEvidenceRouter( + validationEvidence, + dynamicSelection.randomCandidate.config, + fusion, + profile, + ), + } + const results: EvidenceRouterSearchResult[] = [] + for (const selection of dynamicSelection.selections) { + const staticSelection = await selectBestWeights( + development, + fusion, + fullPool, + selection.objective, + dynamicSelection.productionQuality, + profile, + ) + results.push({ + model, + fusion, + objective: selection.objective, + strategy, + fold, + developmentQueries: development.length, + validationQueries: validation.length, + staticWeights: staticSelection.weights, + config: benchmarkRouterConfig(fusion, selection.config), + staticDevelopment: staticSelection.quality, + staticValidation: summarize(validation, staticSelection.weights, fusion, profile), + development: selection.quality, + validation: summarizeEvidenceRouter( + validationEvidence, selection.config, fusion, - holdoutProfile, + profile, ), - profile, - ), - } + productionDevelopment: dynamicSelection.productionQuality, + productionValidation, + guardrailsMet: selection.guardrailsMet, + promotionStatus: selection.promotionStatus, + proxyEvaluations: dynamicSelection.proxyEvaluations, + fullEvaluations: dynamicSelection.fullEvaluations, + searchDiagnostics: dynamicSelection.searchDiagnostics, + searchBaseline: randomBaseline, + holdoutBreakdown: buildHoldoutBreakdown( + validation, + (partition) => + summarizeEvidenceRouter( + prepareEvidenceSamples(partition), + selection.config, + fusion, + holdoutProfile, + ), + profile, + ), + }) + } + return results + }, + ) + +export const optimizeEvidenceRouterParallel = optimizeEvidenceRouterWithPools + +export const optimizeEvidenceRouter = ( + model: string, + fusion: FusionMethod, + strategy: EvidenceRouterSearchResult["strategy"], + fold: string, + development: readonly WeightSearchSample[], + validation: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, +): Promise => + optimizeEvidenceRouterParallel(model, fusion, strategy, fold, development, validation, profile, { + workerCount: 0, }) -} /** Fit one deployment candidate on all samples after cross-validation has measured generalization. */ -export const fitRecommendedWeights = ( +const fitRecommendedWeightsWithPool = async ( model: string, queryKind: QueryKind, samples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): RecommendedWeights => { - const selected = selectBestWeights(samples, "rrf", "reranker-top20", undefined, profile) + pool: CandidateEvaluationPool, +): Promise => { + const selected = await selectBestWeights( + samples, + "rrf", + pool, + "reranker-top20", + undefined, + profile, + ) return { model, queryKind, @@ -1339,14 +1758,50 @@ export const fitRecommendedWeights = ( } } +const fitRecommendedWeightsWithOptions = ( + model: string, + queryKind: QueryKind, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile, + options: ParallelSearchOptions, +): Promise => + withCandidatePool(samples, "rrf", profile, options, (pool) => + fitRecommendedWeightsWithPool(model, queryKind, samples, profile, pool), + ) + +export const fitRecommendedWeights = ( + model: string, + queryKind: QueryKind, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, +): Promise => + fitRecommendedWeightsWithOptions(model, queryKind, samples, profile, { workerCount: 0 }) + +export const fitRecommendedWeightsParallel = ( + model: string, + queryKind: QueryKind, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + options: ParallelSearchOptions = {}, +): Promise => + fitRecommendedWeightsWithOptions(model, queryKind, samples, profile, options) + /** Fit one static candidate for a fusion method across all query forms. */ -export const fitRecommendedFusionWeights = ( +const fitRecommendedFusionWeightsWithPool = async ( model: string, fusion: FusionMethod, samples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): RecommendedFusionWeights => { - const selected = selectBestWeights(samples, fusion, "reranker-top20", undefined, profile) + pool: CandidateEvaluationPool, +): Promise => { + const selected = await selectBestWeights( + samples, + fusion, + pool, + "reranker-top20", + undefined, + profile, + ) const guardrailBaselines = buildGuardrailBaselines(samples, profile) const guardrailsMet = fusionGuardrailsMet( samples, @@ -1366,40 +1821,95 @@ export const fitRecommendedFusionWeights = ( } } +const fitRecommendedFusionWeightsWithOptions = ( + model: string, + fusion: FusionMethod, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile, + options: ParallelSearchOptions, +): Promise => + withCandidatePool(samples, fusion, profile, options, (pool) => + fitRecommendedFusionWeightsWithPool(model, fusion, samples, profile, pool), + ) + +export const fitRecommendedFusionWeights = ( + model: string, + fusion: FusionMethod, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, +): Promise => + fitRecommendedFusionWeightsWithOptions(model, fusion, samples, profile, { workerCount: 0 }) + +export const fitRecommendedFusionWeightsParallel = ( + model: string, + fusion: FusionMethod, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + options: ParallelSearchOptions = {}, +): Promise => + fitRecommendedFusionWeightsWithOptions(model, fusion, samples, profile, options) + /** * Fit one evidence-router candidate on all samples after cross-validation has measured * generalization. */ +const fitRecommendedEvidenceRouterWithOptions = ( + model: string, + fusion: FusionMethod, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + options: ParallelSearchOptions, +): Promise => + withParallelEvidencePools( + samples, + fusion, + profile, + options, + async (dynamicSelection, fullPool) => { + const results: RecommendedEvidenceRouter[] = [] + for (const selection of dynamicSelection.selections) { + const staticSelection = await selectBestWeights( + samples, + fusion, + fullPool, + selection.objective, + dynamicSelection.productionQuality, + profile, + ) + results.push({ + model, + fusion, + objective: selection.objective, + samples: samples.length, + staticWeights: staticSelection.weights, + config: benchmarkRouterConfig(fusion, selection.config), + staticQuality: staticSelection.quality, + fitQuality: selection.quality, + productionQuality: dynamicSelection.productionQuality, + guardrailsMet: selection.guardrailsMet, + promotionStatus: selection.promotionStatus, + proxyEvaluations: dynamicSelection.proxyEvaluations, + fullEvaluations: dynamicSelection.fullEvaluations, + searchDiagnostics: dynamicSelection.searchDiagnostics, + }) + } + return results + }, + ) + export const fitRecommendedEvidenceRouter = ( model: string, fusion: FusionMethod, samples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): readonly RecommendedEvidenceRouter[] => { - const dynamicSelection = selectBestEvidenceRouter(samples, fusion, profile) - return dynamicSelection.selections.map((selection) => { - const staticSelection = selectBestWeights( - samples, - fusion, - selection.objective, - dynamicSelection.productionQuality, - profile, - ) - return { - model, - fusion, - objective: selection.objective, - samples: samples.length, - staticWeights: staticSelection.weights, - config: benchmarkRouterConfig(fusion, selection.config), - staticQuality: staticSelection.quality, - fitQuality: selection.quality, - productionQuality: dynamicSelection.productionQuality, - guardrailsMet: selection.guardrailsMet, - promotionStatus: selection.promotionStatus, - proxyEvaluations: dynamicSelection.proxyEvaluations, - fullEvaluations: dynamicSelection.fullEvaluations, - searchDiagnostics: dynamicSelection.searchDiagnostics, - } - }) -} +): Promise => + fitRecommendedEvidenceRouterWithOptions(model, fusion, samples, profile, { workerCount: 0 }) + +export const fitRecommendedEvidenceRouterParallel = ( + model: string, + fusion: FusionMethod, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + options: ParallelSearchOptions = {}, +): Promise => + fitRecommendedEvidenceRouterWithOptions(model, fusion, samples, profile, options) diff --git a/benchmarks/retrieval/worker-pool.ts b/benchmarks/retrieval/worker-pool.ts new file mode 100644 index 0000000..c724031 --- /dev/null +++ b/benchmarks/retrieval/worker-pool.ts @@ -0,0 +1,432 @@ +import { availableParallelism } from "node:os" +import { Worker } from "node:worker_threads" + +import type { Chunk } from "../../src/domain/chunk.js" +import type { ChannelWeights } from "../../src/domain/retrieval.js" +import { evaluateCandidate } from "./fusion-core.mjs" +import { type PreparedFusionEvaluator, type PreparedFusionSnapshot } from "./fusion.js" +import type { QualitySummary } from "./types.js" + +const DEFAULT_BATCH_SIZE = 32 +const DEFAULT_WORKER_URL = new URL("./fusion-worker.mjs", import.meta.url) + +/** One prepared benchmark sample represented without source text for worker transfer. */ +export interface EvaluationSampleSnapshot { + readonly fusion: PreparedFusionSnapshot + readonly targets: readonly (readonly number[])[] + readonly contextTokens: readonly number[] + readonly sampleWeight: number +} + +/** Immutable prepared inputs shared by every candidate evaluation in one search. */ +export interface EvaluationSnapshot { + readonly samples: readonly EvaluationSampleSnapshot[] +} + +/** Main-thread inputs used to build a compact worker snapshot. */ +export interface EvaluationSampleInput { + readonly evaluator: PreparedFusionEvaluator + readonly targets: readonly ReadonlySet[] + readonly chunks: readonly Chunk[] + readonly sampleWeight: number +} + +/** One candidate's static weights or one weight vector per prepared sample. */ +export interface EvaluationCandidate { + readonly weights: ChannelWeights | readonly ChannelWeights[] +} + +/** Native-worker or serial execution mode used by a candidate pool. */ +export type EvaluationPoolMode = "parallel" | "serial" + +/** Configuration for one reusable benchmark candidate pool. */ +export interface CandidateEvaluationPoolOptions { + /** Override the default worker count; zero explicitly selects serial evaluation. */ + readonly workerCount?: number + /** Maximum number of candidate vectors sent in one worker message. */ + readonly batchSize?: number + /** Test-only or diagnostic override for the native worker entry point. */ + readonly workerUrl?: URL + /** Fall back to serial evaluation if worker startup is unavailable. */ + readonly fallbackToSerial?: boolean +} + +/** Cumulative scheduling information exposed for deterministic pool tests and reports. */ +export interface CandidateEvaluationPoolStats { + readonly mode: EvaluationPoolMode + readonly workerCount: number + readonly activeWorkerCount: number + readonly batchSize: number + readonly batches: number + readonly candidates: number +} + +/** Reusable evaluator for batched benchmark candidates. */ +export interface CandidateEvaluationPool { + readonly mode: EvaluationPoolMode + readonly workerCount: number + readonly batchSize: number + readonly evaluate: ( + candidates: readonly EvaluationCandidate[], + ) => Promise + readonly close: () => Promise + readonly stats: () => CandidateEvaluationPoolStats +} + +interface WorkerReadyMessage { + readonly type: "ready" +} + +interface WorkerResultMessage { + readonly type: "result" + readonly taskId: number + readonly results: readonly QualitySummary[] +} + +interface WorkerErrorMessage { + readonly type: "error" + readonly taskId?: number + readonly message: string +} + +type WorkerMessage = WorkerReadyMessage | WorkerResultMessage | WorkerErrorMessage + +interface WorkerTask { + readonly taskId: number + readonly start: number + readonly end: number +} + +interface WorkerSlot { + readonly worker: Worker + readonly ready: Promise + readonly resolveReady: () => void + readonly rejectReady: (error: Error) => void + busy: boolean + task?: WorkerTask +} + +interface ActiveEvaluation { + readonly candidates: readonly EvaluationCandidate[] + readonly results: Map + readonly tasks: Map + readonly resolve: (results: readonly QualitySummary[]) => void + readonly reject: (error: Error) => void + nextCandidate: number + completedCandidates: number +} + +const errorFromUnknown = (cause: unknown): Error => + cause instanceof Error ? cause : new Error(String(cause)) + +const parseInteger = (value: string | undefined): number | undefined => { + if (value === undefined || value.trim() === "") return undefined + const parsed = Number(value) + return Number.isInteger(parsed) ? parsed : undefined +} + +/** Return the default native pool size, reserving one core for the benchmark coordinator. */ +export const getDefaultWorkerCount = (): number => Math.max(1, availableParallelism() - 1) + +/** Resolve an explicit, environment, or default worker count for benchmark evaluation. */ +export const resolveWorkerCount = (requested?: number): number => { + const configured = requested ?? parseInteger(process.env.PIX_BENCH_WORKERS) + if (configured === undefined) return getDefaultWorkerCount() + return Math.max(0, configured) +} + +const resolveBatchSize = (requested?: number): number => { + const configured = requested ?? parseInteger(process.env.PIX_BENCH_WORKER_BATCH_SIZE) + return Math.max(1, configured ?? DEFAULT_BATCH_SIZE) +} + +const contextTokens = (chunk: Chunk): number => + Math.ceil( + Buffer.byteLength(`${chunk.file}:${chunk.startLine}-${chunk.endLine}\n${chunk.text}`, "utf8") / + 4, + ) + +/** Build the compact, cloneable snapshot used by serial and worker evaluators. */ +export const createEvaluationSnapshot = ( + inputs: readonly EvaluationSampleInput[], +): EvaluationSnapshot => ({ + samples: inputs.map((input) => ({ + fusion: input.evaluator.snapshot, + targets: input.targets.map((target) => [...target]), + contextTokens: input.chunks.map(contextTokens), + sampleWeight: input.sampleWeight, + })), +}) + +/** Evaluate candidates serially using the same prepared snapshot as the worker pool. */ +export const evaluateCandidatesSerial = ( + snapshot: EvaluationSnapshot, + candidates: readonly EvaluationCandidate[], +): readonly QualitySummary[] => + candidates.map((candidate) => evaluateCandidate(snapshot, candidate)) + +const isWorkerMessage = (message: unknown): message is WorkerMessage => { + if (typeof message !== "object" || message === null || !("type" in message)) return false + const type = message.type + return type === "ready" || type === "result" || type === "error" +} + +class NativeCandidateEvaluationPool implements CandidateEvaluationPool { + readonly mode = "parallel" as const + readonly workerCount: number + readonly batchSize: number + + private readonly slots: WorkerSlot[] + private readonly ready: Promise + private active: ActiveEvaluation | undefined + private closePromise: Promise | undefined + private closed = false + private nextTaskId = 0 + private batchCount = 0 + private candidateCount = 0 + + private constructor( + snapshot: EvaluationSnapshot, + workerCount: number, + batchSize: number, + workerUrl: URL, + ) { + this.workerCount = workerCount + this.batchSize = batchSize + this.slots = [] + for (let index = 0; index < workerCount; index++) { + let resolveReady: () => void = () => undefined + let rejectReady: (error: Error) => void = () => undefined + const ready = new Promise((resolve, reject) => { + resolveReady = resolve + rejectReady = reject + }) + const worker = new Worker(workerUrl) + const slot: WorkerSlot = { + worker, + ready, + resolveReady, + rejectReady, + busy: false, + } + worker.on("message", (message: unknown) => this.handleMessage(slot, message)) + worker.on("error", (cause: Error) => this.handleWorkerError(slot, cause)) + worker.on("exit", (code: number) => { + if (!this.closed && code !== 0) + this.handleWorkerError(slot, new Error(`Fusion worker exited with code ${code}`)) + }) + this.slots.push(slot) + } + this.ready = Promise.all(this.slots.map((slot) => slot.ready)).then(() => undefined) + for (const slot of this.slots) slot.worker.postMessage({ type: "init", snapshot }) + } + + /** Start a native pool after every worker has accepted its immutable snapshot. */ + static async create( + snapshot: EvaluationSnapshot, + workerCount: number, + batchSize: number, + workerUrl: URL, + ): Promise { + let pool: NativeCandidateEvaluationPool | undefined + try { + pool = new NativeCandidateEvaluationPool(snapshot, workerCount, batchSize, workerUrl) + await pool.ready + return pool + } catch (cause) { + if (pool !== undefined) await pool.close() + throw errorFromUnknown(cause) + } + } + + private handleWorkerError(slot: WorkerSlot, cause: Error): void { + if (this.closed) return + slot.rejectReady(cause) + const active = this.active + if (active !== undefined) { + this.active = undefined + active.reject(cause) + } + void this.close() + } + + private handleMessage(slot: WorkerSlot, message: unknown): void { + if (!isWorkerMessage(message)) { + this.handleWorkerError(slot, new Error("Fusion worker sent an invalid message")) + return + } + if (message.type === "ready") { + slot.resolveReady() + return + } + if (message.type === "error") { + this.handleWorkerError(slot, new Error(message.message)) + return + } + const active = this.active + const task = active?.tasks.get(message.taskId) + if (active === undefined || task === undefined) { + this.handleWorkerError( + slot, + new Error(`Fusion worker returned unknown task ${message.taskId}`), + ) + return + } + if (message.results.length !== task.end - task.start) { + this.handleWorkerError(slot, new Error(`Fusion worker returned an invalid result length`)) + return + } + for (let index = 0; index < message.results.length; index++) { + active.results.set(task.start + index, message.results[index]) + } + active.tasks.delete(message.taskId) + slot.busy = false + slot.task = undefined + active.completedCandidates += message.results.length + if (active.completedCandidates === active.candidates.length) { + const ordered: QualitySummary[] = [] + for (let index = 0; index < active.candidates.length; index++) { + const result = active.results.get(index) + if (result === undefined) { + this.handleWorkerError(slot, new Error(`Fusion worker omitted candidate ${index}`)) + return + } + ordered.push(result) + } + this.active = undefined + active.resolve(ordered) + return + } + this.dispatch() + } + + private dispatch(): void { + const active = this.active + if (active === undefined || this.closed) return + for (const slot of this.slots) { + if (slot.busy || active.nextCandidate >= active.candidates.length) continue + const start = active.nextCandidate + const end = Math.min(start + this.batchSize, active.candidates.length) + const task: WorkerTask = { taskId: this.nextTaskId++, start, end } + active.nextCandidate = end + active.tasks.set(task.taskId, task) + slot.busy = true + slot.task = task + this.batchCount++ + try { + slot.worker.postMessage({ + type: "evaluate", + taskId: task.taskId, + candidates: active.candidates.slice(start, end), + }) + } catch (cause) { + this.handleWorkerError(slot, errorFromUnknown(cause)) + return + } + } + } + + async evaluate(candidates: readonly EvaluationCandidate[]): Promise { + if (this.closed) throw new Error("Candidate evaluation pool is closed") + if (this.active !== undefined) throw new Error("Candidate evaluation pool is already busy") + await this.ready + if (candidates.length === 0) return [] + this.candidateCount += candidates.length + return new Promise((resolve, reject) => { + this.active = { + candidates, + results: new Map(), + tasks: new Map(), + resolve, + reject, + nextCandidate: 0, + completedCandidates: 0, + } + this.dispatch() + }) + } + + async close(): Promise { + if (this.closePromise !== undefined) return this.closePromise + this.closed = true + const active = this.active + this.active = undefined + active?.reject(new Error("Candidate evaluation pool closed")) + for (const slot of this.slots) { + slot.busy = false + slot.task = undefined + } + this.closePromise = Promise.all( + this.slots.map((slot) => slot.worker.terminate().catch(() => -1)), + ).then(() => undefined) + return this.closePromise + } + + stats = (): CandidateEvaluationPoolStats => ({ + mode: this.mode, + workerCount: this.workerCount, + activeWorkerCount: this.closed ? 0 : this.slots.filter((slot) => slot.busy).length, + batchSize: this.batchSize, + batches: this.batchCount, + candidates: this.candidateCount, + }) +} + +class SerialCandidateEvaluationPool implements CandidateEvaluationPool { + readonly mode = "serial" as const + readonly workerCount = 1 + readonly batchSize: number + private closed = false + private readonly snapshot: EvaluationSnapshot + private batchCount = 0 + private candidateCount = 0 + + constructor(snapshot: EvaluationSnapshot, batchSize: number) { + this.snapshot = snapshot + this.batchSize = batchSize + } + + async evaluate(candidates: readonly EvaluationCandidate[]): Promise { + if (this.closed) throw new Error("Candidate evaluation pool is closed") + this.candidateCount += candidates.length + this.batchCount += Math.ceil(candidates.length / this.batchSize) + return evaluateCandidatesSerial(this.snapshot, candidates) + } + + async close(): Promise { + this.closed = true + } + + stats = (): CandidateEvaluationPoolStats => ({ + mode: this.mode, + workerCount: this.workerCount, + activeWorkerCount: 0, + batchSize: this.batchSize, + batches: this.batchCount, + candidates: this.candidateCount, + }) +} + +/** Create a fixed benchmark evaluator pool; one worker or an explicit zero uses serial fallback. */ +export const createCandidateEvaluationPool = async ( + snapshot: EvaluationSnapshot, + options: CandidateEvaluationPoolOptions = {}, +): Promise => { + const requestedWorkerCount = resolveWorkerCount(options.workerCount) + const batchSize = resolveBatchSize(options.batchSize) + if (requestedWorkerCount <= 1 || snapshot.samples.length === 0) + return new SerialCandidateEvaluationPool(snapshot, batchSize) + + try { + return await NativeCandidateEvaluationPool.create( + snapshot, + requestedWorkerCount, + batchSize, + options.workerUrl ?? DEFAULT_WORKER_URL, + ) + } catch (cause) { + if (options.fallbackToSerial !== false) + return new SerialCandidateEvaluationPool(snapshot, batchSize) + throw errorFromUnknown(cause) + } +} diff --git a/benchmarks/tests/channels.test.ts b/benchmarks/tests/channels.test.ts index d2977f5..3d3805b 100644 --- a/benchmarks/tests/channels.test.ts +++ b/benchmarks/tests/channels.test.ts @@ -9,12 +9,19 @@ import { } from "../../src/lib/retrieval/evidence-router.js" import { fuseRankings } from "../../src/lib/retrieval/fusion.js" import { buildIdentifierIndex } from "../../src/lib/retrieval/identifier-index.js" -import { recallAt, resolveGoldTargets } from "../retrieval/metrics.js" +import { prepareFusion } from "../retrieval/fusion.js" +import { + contextRecallAtBudget, + recallAt, + reciprocalRank, + resolveGoldTargets, +} from "../retrieval/metrics.js" import type { PreparedCorpus } from "../retrieval/prepare.js" import { fuseVariant, rankLexicalChannels, RETRIEVAL_VARIANTS } from "../retrieval/ranking.js" import { ROUTER_OBJECTIVES } from "../retrieval/types.js" import { optimizeEvidenceRouter, + optimizeEvidenceRouterParallel, optimizeWeights, selectEligibleCandidate, } from "../retrieval/weight-search.js" @@ -368,6 +375,65 @@ describe("retrieval benchmark fixture", () => { expect(production[0]?.score).toBeGreaterThan(production[1]?.score ?? 0) }) + it("keeps prepared fusion quality identical on the benchmark fixture", () => { + const weights = { identity: 1, camelcase: 1, bm25: 1, dense: 1, sparse: 1 } + const quality = ( + ranked: readonly { readonly chunkIndex: number; readonly score: number }[], + sample: ReturnType[number], + ) => ({ + recallAt5: recallAt(ranked, sample.targets, 5), + recallAt10: recallAt(ranked, sample.targets, 10), + recallAt20: recallAt(ranked, sample.targets, 20), + recallAt50: recallAt(ranked, sample.targets, 50), + contextRecallAt4096: contextRecallAtBudget(ranked, sample.targets, sample.chunks, 4_096), + meanReciprocalRank: reciprocalRank(ranked, sample.targets), + }) + + for (const method of ["rrf", "relative-score", "dbsf"] as const) { + for (const sample of makeEvidenceRouterSamples()) { + const direct = fuseRankings(method, sample.rankings, weights, 200) + const prepared = prepareFusion(method, sample.rankings, 200).evaluate(weights) + + expect(prepared).toEqual(direct) + expect(quality(prepared, sample)).toEqual(quality(direct, sample)) + } + } + }) + + it("keeps prepared fusion identical for duplicate and truncated rankings", () => { + const rankings = { + identity: [ + { chunkIndex: 0, score: 1 }, + { chunkIndex: 0, score: 0.8 }, + { chunkIndex: 1, score: 0.4 }, + ], + camelcase: [ + { chunkIndex: 1, score: 0.9 }, + { chunkIndex: 2, score: 0.2 }, + ], + bm25: [ + { chunkIndex: 2, score: 9 }, + { chunkIndex: 2, score: 8 }, + { chunkIndex: 0, score: 1 }, + ], + dense: [ + { chunkIndex: 3, score: 0.7 }, + { chunkIndex: 1, score: 0.6 }, + ], + sparse: [ + { chunkIndex: 0, score: 4 }, + { chunkIndex: 3, score: 3 }, + ], + } + const weights = { identity: 0.5, camelcase: 0, bm25: 2, dense: 1.5, sparse: 0.25 } + + for (const method of ["rrf", "relative-score", "dbsf"] as const) { + expect(prepareFusion(method, rankings, 2).evaluate(weights)).toEqual( + fuseRankings(method, rankings, weights, 2), + ) + } + }) + it("fuses channel-relative scores without comparing raw score scales", () => { const ranked = fuseRankings( "relative-score", @@ -431,7 +497,7 @@ describe("retrieval benchmark fixture", () => { expect(ranked).toEqual([{ chunkIndex: 0, score: 0.5 }]) }) - it("learns weights on development and attributes holdout value with Shapley", () => { + it("learns weights on development and attributes holdout value with Shapley", async () => { const sample = { repository: "fixture", intentId: "fixture-001", @@ -448,7 +514,7 @@ describe("retrieval benchmark fixture", () => { targets: [new Set([0])], chunks, } - const result = optimizeWeights( + const result = await optimizeWeights( "fixture", "identifier", "grouped-5-fold", @@ -464,10 +530,10 @@ describe("retrieval benchmark fixture", () => { expect(result.shapleyRecallAt20.identity).toBe(1) }) - it("selects one evidence router across queries with different reliable channels", () => { + it("selects one evidence router across queries with different reliable channels", async () => { const samples = makeEvidenceRouterSamples() - const routerResults = optimizeEvidenceRouter( + const routerResults = await optimizeEvidenceRouter( "fixture", "dbsf", "grouped-5-fold", @@ -519,6 +585,27 @@ describe("retrieval benchmark fixture", () => { expect(result.validation.recallAt20).toBeGreaterThan(result.staticValidation.recallAt20) }) + it("keeps the parallel evidence-router search equal to the serial search", async () => { + const samples = makeEvidenceRouterSamples() + const serial = selectTop20Router( + await optimizeEvidenceRouter("fixture", "dbsf", "grouped-5-fold", "1", samples, samples), + ) + const parallel = selectTop20Router( + await optimizeEvidenceRouterParallel( + "fixture", + "dbsf", + "grouped-5-fold", + "1", + samples, + samples, + undefined, + { workerCount: 2, batchSize: 8 }, + ), + ) + + expect(parallel).toEqual(serial) + }) + it("does not treat a guardrail-failing fallback as promotable", () => { const selection = selectEligibleCandidate([{ name: "fallback" }], () => false) @@ -528,20 +615,34 @@ describe("retrieval benchmark fixture", () => { }) }) - it("keeps router fitting deterministic and independent of validation samples", () => { + it("keeps router fitting deterministic and independent of validation samples", async () => { const development = makeEvidenceRouterSamples() const first = selectTop20Router( - optimizeEvidenceRouter("fixture", "dbsf", "grouped-5-fold", "1", development, development), + await optimizeEvidenceRouter( + "fixture", + "dbsf", + "grouped-5-fold", + "1", + development, + development, + ), ) const repeat = selectTop20Router( - optimizeEvidenceRouter("fixture", "dbsf", "grouped-5-fold", "1", development, development), + await optimizeEvidenceRouter( + "fixture", + "dbsf", + "grouped-5-fold", + "1", + development, + development, + ), ) const alteredValidation = development.map((sample) => ({ ...sample, targets: [new Set([0])], })) const withAlteredValidation = selectTop20Router( - optimizeEvidenceRouter( + await optimizeEvidenceRouter( "fixture", "dbsf", "grouped-5-fold", diff --git a/benchmarks/tests/worker-pool.test.ts b/benchmarks/tests/worker-pool.test.ts new file mode 100644 index 0000000..f4cca7e --- /dev/null +++ b/benchmarks/tests/worker-pool.test.ts @@ -0,0 +1,202 @@ +import { availableParallelism } from "node:os" + +import { describe, expect, it } from "@effect/vitest" + +import type { Chunk } from "../../src/domain/chunk.js" +import { prepareFusion } from "../retrieval/fusion.js" +import { + fitRecommendedEvidenceRouter, + fitRecommendedFusionWeights, + fitRecommendedFusionWeightsParallel, + fitRecommendedWeights, + fitRecommendedWeightsParallel, + optimizeFusionWeights, +} from "../retrieval/weight-search.js" +import { + createCandidateEvaluationPool, + createEvaluationSnapshot, + evaluateCandidatesSerial, + getDefaultWorkerCount, + resolveWorkerCount, + type EvaluationCandidate, +} from "../retrieval/worker-pool.js" + +const chunks: readonly Chunk[] = [0, 1, 2, 3].map((index) => ({ + id: String(index), + idx: index, + file: `src/chunk-${index}.ts`, + startLine: 1, + endLine: 1, + startOffset: 0, + endOffset: 32, + text: `export function target${index}() { return ${index} }`, +})) + +const rankings = { + identity: [ + { chunkIndex: 0, score: 1 }, + { chunkIndex: 1, score: 0.5 }, + ], + camelcase: [{ chunkIndex: 1, score: 1 }], + bm25: [ + { chunkIndex: 2, score: 4 }, + { chunkIndex: 3, score: 1 }, + ], + dense: [{ chunkIndex: 3, score: 1 }], + sparse: [], +} + +const snapshot = createEvaluationSnapshot([ + { + evaluator: prepareFusion("dbsf", rankings, 10), + targets: [new Set([0])], + chunks, + sampleWeight: 1, + }, + { + evaluator: prepareFusion("dbsf", rankings, 10), + targets: [new Set([3])], + chunks, + sampleWeight: 2, + }, +]) + +const candidates: readonly EvaluationCandidate[] = [ + { weights: { identity: 1, camelcase: 0, bm25: 0, dense: 0, sparse: 0 } }, + { weights: { identity: 0, camelcase: 0, bm25: 1, dense: 0, sparse: 0 } }, + { weights: { identity: 0, camelcase: 0, bm25: 0, dense: 1, sparse: 0 } }, + { + weights: [ + { identity: 1, camelcase: 0, bm25: 0, dense: 0, sparse: 0 }, + { identity: 0, camelcase: 0, bm25: 0, dense: 1, sparse: 0 }, + ], + }, + { weights: { identity: 1, camelcase: 1, bm25: 1, dense: 1, sparse: 1 } }, +] + +const searchSample = { + repository: "fixture", + intentId: "fixture-001", + queryKind: "identifier" as const, + groupedFold: 0, + query: "target0", + rankings, + targets: [new Set([0])], + chunks, +} + +describe("benchmark candidate evaluation pool", () => { + it("derives a bounded default and honors explicit sizing", () => { + expect(getDefaultWorkerCount()).toBe(Math.max(1, availableParallelism() - 1)) + expect(resolveWorkerCount(0)).toBe(0) + expect(resolveWorkerCount(3)).toBe(3) + }) + + it("matches serial quality exactly and preserves candidate order", async () => { + const serial = evaluateCandidatesSerial(snapshot, candidates) + const pool = await createCandidateEvaluationPool(snapshot, { + workerCount: 2, + batchSize: 2, + }) + try { + await expect(pool.evaluate(candidates)).resolves.toEqual(serial) + expect(pool.stats()).toMatchObject({ + mode: "parallel", + workerCount: 2, + batchSize: 2, + batches: 3, + candidates: candidates.length, + }) + } finally { + await pool.close() + } + expect(pool.stats().activeWorkerCount).toBe(0) + }) + + it("keeps the explicit parallel search result equal to the serial search", async () => { + const serial = await fitRecommendedWeights("fixture", "identifier", [searchSample]) + const parallel = await fitRecommendedWeightsParallel( + "fixture", + "identifier", + [searchSample], + undefined, + { workerCount: 2, batchSize: 16 }, + ) + + expect(parallel).toEqual(serial) + }) + + it("keeps static fusion fitting serial and parallel paths equivalent", async () => { + const serial = await fitRecommendedFusionWeights("fixture", "dbsf", [searchSample]) + const parallel = await fitRecommendedFusionWeightsParallel( + "fixture", + "dbsf", + [searchSample], + undefined, + { workerCount: 2, batchSize: 16 }, + ) + + expect(parallel).toEqual(serial) + + const foldResult = await optimizeFusionWeights( + "fixture", + "dbsf", + "grouped-5-fold", + "1", + [searchSample], + [searchSample], + ) + expect(foldResult.developmentQueries).toBe(1) + }) + + it("retains the serial evidence-router fit path", async () => { + const result = await fitRecommendedEvidenceRouter("fixture", "dbsf", [searchSample]) + + expect(result).toHaveLength(3) + expect(result.every((candidate) => candidate.fitQuality.recallAt20 >= 0)).toBe(true) + }) + + it("uses serial fallback for a one-worker configuration", async () => { + const pool = await createCandidateEvaluationPool(snapshot, { + workerCount: 1, + batchSize: 2, + }) + try { + await expect(pool.evaluate(candidates)).resolves.toEqual( + evaluateCandidatesSerial(snapshot, candidates), + ) + expect(pool.stats()).toMatchObject({ mode: "serial", workerCount: 1, batches: 3 }) + } finally { + await pool.close() + } + }) + + it("falls back to serial evaluation when worker startup is unavailable", async () => { + const unavailableWorkerUrl = new URL("../retrieval/fusion-worker.mjs", import.meta.url) + unavailableWorkerUrl.pathname += ".missing" + const pool = await createCandidateEvaluationPool(snapshot, { + workerCount: 2, + workerUrl: unavailableWorkerUrl, + fallbackToSerial: true, + }) + try { + expect(pool.mode).toBe("serial") + await expect(pool.evaluate(candidates)).resolves.toEqual( + evaluateCandidatesSerial(snapshot, candidates), + ) + } finally { + await pool.close() + } + }) + + it("terminates the pool when a worker reports an evaluation error", async () => { + const pool = await createCandidateEvaluationPool(snapshot, { workerCount: 2, batchSize: 1 }) + try { + await expect(pool.evaluate([{ weights: [] }])).rejects.toThrow("Missing weights") + expect(pool.stats().activeWorkerCount).toBe(0) + } finally { + await pool.close() + } + await expect(pool.evaluate(candidates)).rejects.toThrow("closed") + }) +}) From 5a9077c7d75ae48ae50ac5f897b2c10ae4fc88d1 Mon Sep 17 00:00:00 2001 From: Lucas Burmeister Date: Tue, 4 Aug 2026 06:35:11 +0200 Subject: [PATCH 2/8] fix: address benchmark worker review feedback --- CONTEXT.md | 5 +- benchmarks/retrieval/fusion-core.d.mts | 6 +- benchmarks/retrieval/runner.ts | 70 +++++----- benchmarks/retrieval/weight-search.ts | 177 ++++++++++++++----------- benchmarks/retrieval/worker-pool.ts | 28 ++-- benchmarks/tests/channels.test.ts | 2 +- benchmarks/tests/worker-pool.test.ts | 40 +++++- 7 files changed, 200 insertions(+), 128 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 5441b6f..db088b6 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -245,7 +245,7 @@ current beam elites, so a later coordinate cannot regress the best development c development folds and evaluated unchanged against static weights on excluded intent folds and repositories; authored query-form labels remain informed reference strata and are not router inputs. This router remains benchmark-only until holdouts justify a production change. -Serial benchmark weight and router searches use a benchmark-only prepared evaluator: each sample and +Benchmark weight and router searches use a benchmark-only prepared evaluator: each sample and fusion method materializes per-chunk normalized or RRF contributions once, then candidate weights reuse that data. Public `fuseRankings` semantics and its existing ranking, normalization, and typed-array caches remain unchanged. The evaluator still materializes and sorts the full candidate union for every @@ -253,7 +253,8 @@ configuration; partial top-K or metric-specialized evaluation remains future wor benchmark path uses a fixed native `node:worker_threads` pool (default `max(1, availableParallelism() - 1)`), sends the prepared snapshot once per worker, batches candidate weight vectors, and keeps beam/cache/archive/ selection state on the main thread. The search APIs are async so serial and worker execution share the -same selection algorithm; `PIX_BENCH_WORKERS=0` selects serial evaluation, and the runner's +same selection algorithm; `PIX_BENCH_WORKERS=0` selects serial evaluation, `workerCount: 1` selects +the serial pool through `createCandidateEvaluationPool`, and the runner's `PIX_BENCH_SEARCH_MODE=serial` selects that serial pool mode. `PIX_BENCH_WORKER_BATCH_SIZE` bounds each worker message. Unavailable worker startup falls back to serial evaluation; task failures are surfaced and all workers are diff --git a/benchmarks/retrieval/fusion-core.d.mts b/benchmarks/retrieval/fusion-core.d.mts index 367be11..ab1a963 100644 --- a/benchmarks/retrieval/fusion-core.d.mts +++ b/benchmarks/retrieval/fusion-core.d.mts @@ -1,9 +1,11 @@ +import type { ChannelWeights } from "../../src/domain/retrieval.js" +import type { PreparedFusionSnapshot } from "./fusion.js" import type { QualitySummary } from "./types.js" import type { EvaluationCandidate, EvaluationSnapshot } from "./worker-pool.js" export function evaluatePreparedContributions( - matrix: unknown, - weights: unknown, + matrix: PreparedFusionSnapshot, + weights: ChannelWeights, ): { readonly chunkIndex: number readonly score: number diff --git a/benchmarks/retrieval/runner.ts b/benchmarks/retrieval/runner.ts index c1e04b1..3770638 100644 --- a/benchmarks/retrieval/runner.ts +++ b/benchmarks/retrieval/runner.ts @@ -115,9 +115,11 @@ const reportProgress = (message: string): void => { process.stderr.write(`[retrieval benchmark] ${message}\n`) } -const runParallelSearch = (operation: () => Promise): Effect.Effect => +const runParallelSearch = ( + operation: (signal: AbortSignal) => Promise, +): Effect.Effect => Effect.tryPromise({ - try: operation, + try: (signal) => operation(signal), catch: (cause) => (cause instanceof Error ? cause : new Error(String(cause))), }) @@ -518,17 +520,18 @@ export const runRetrievalBenchmark = ( if (config.legacyDiagnostics) { for (const group of sampleGroups.values()) { for (let fold = 0; fold < config.groupedFolds; fold++) { + const split = splitSamples(group.samples, (sample) => sample.groupedFold === fold) weightSearch.push( - yield* runParallelSearch(() => + yield* runParallelSearch((signal) => optimizeWeightsParallel( group.model, group.queryKind, groupedStrategy, String(fold + 1), - group.samples.filter((sample) => sample.groupedFold !== fold), - group.samples.filter((sample) => sample.groupedFold === fold), + split.development, + split.validation, optimizationProfile, - searchOptions, + { ...searchOptions, signal }, ), ), ) @@ -536,30 +539,31 @@ export const runRetrievalBenchmark = ( const repositories = [...new Set(group.samples.map((sample) => sample.repository))] if (config.repositoryHoldouts && repositories.length > 1) { for (const repository of repositories) { + const split = splitSamples(group.samples, (sample) => sample.repository === repository) weightSearch.push( - yield* runParallelSearch(() => + yield* runParallelSearch((signal) => optimizeWeightsParallel( group.model, group.queryKind, "leave-one-repository-out", repository, - group.samples.filter((sample) => sample.repository !== repository), - group.samples.filter((sample) => sample.repository === repository), + split.development, + split.validation, optimizationProfile, - searchOptions, + { ...searchOptions, signal }, ), ), ) } } recommendedWeights.push( - yield* runParallelSearch(() => + yield* runParallelSearch((signal) => fitRecommendedWeightsParallel( group.model, group.queryKind, group.samples, optimizationProfile, - searchOptions, + { ...searchOptions, signal }, ), ), ) @@ -606,7 +610,7 @@ export const runRetrievalBenchmark = ( for (let fold = 0; fold < config.groupedFolds; fold++) { const split = splitSamples(samples, (sample) => sample.groupedFold === fold) fusionSearch.push( - yield* runParallelSearch(() => + yield* runParallelSearch((signal) => optimizeFusionWeightsParallel( model, fusion, @@ -615,7 +619,7 @@ export const runRetrievalBenchmark = ( split.development, split.validation, optimizationProfile, - searchOptions, + { ...searchOptions, signal }, ), ), ) @@ -624,7 +628,7 @@ export const runRetrievalBenchmark = ( for (const repository of repositories) { const split = splitSamples(samples, (sample) => sample.repository === repository) fusionSearch.push( - yield* runParallelSearch(() => + yield* runParallelSearch((signal) => optimizeFusionWeightsParallel( model, fusion, @@ -633,7 +637,7 @@ export const runRetrievalBenchmark = ( split.development, split.validation, optimizationProfile, - searchOptions, + { ...searchOptions, signal }, ), ), ) @@ -645,14 +649,11 @@ export const runRetrievalBenchmark = ( for (const [model, samples] of samplesByModel) for (const fusion of config.fusionMethods) recommendedFusionWeights.push( - yield* runParallelSearch(() => - fitRecommendedFusionWeightsParallel( - model, - fusion, - samples, - optimizationProfile, - searchOptions, - ), + yield* runParallelSearch((signal) => + fitRecommendedFusionWeightsParallel(model, fusion, samples, optimizationProfile, { + ...searchOptions, + signal, + }), ), ) const fusionSearchDurationMs = performance.now() - fusionSearchStartedAt @@ -667,7 +668,7 @@ export const runRetrievalBenchmark = ( `${model}: selecting ${fusion} evidence router for grouped fold ${fold + 1}/${config.groupedFolds}`, ) evidenceRouterSearch.push( - ...(yield* runParallelSearch(() => + ...(yield* runParallelSearch((signal) => optimizeEvidenceRouterParallel( model, fusion, @@ -676,7 +677,7 @@ export const runRetrievalBenchmark = ( split.development, split.validation, optimizationProfile, - searchOptions, + { ...searchOptions, signal }, ), )), ) @@ -689,7 +690,7 @@ export const runRetrievalBenchmark = ( `${model}: selecting ${fusion} evidence router with ${repository} held out`, ) evidenceRouterSearch.push( - ...(yield* runParallelSearch(() => + ...(yield* runParallelSearch((signal) => optimizeEvidenceRouterParallel( model, fusion, @@ -698,7 +699,7 @@ export const runRetrievalBenchmark = ( split.development, split.validation, optimizationProfile, - searchOptions, + { ...searchOptions, signal }, ), )), ) @@ -711,14 +712,11 @@ export const runRetrievalBenchmark = ( for (const [model, samples] of samplesByModel) for (const fusion of config.routerFusionMethods) recommendedEvidenceRouters.push( - ...(yield* runParallelSearch(() => - fitRecommendedEvidenceRouterParallel( - model, - fusion, - samples, - optimizationProfile, - searchOptions, - ), + ...(yield* runParallelSearch((signal) => + fitRecommendedEvidenceRouterParallel(model, fusion, samples, optimizationProfile, { + ...searchOptions, + signal, + }), )), ) const evidenceRouterSearchDurationMs = performance.now() - evidenceRouterSearchStartedAt diff --git a/benchmarks/retrieval/weight-search.ts b/benchmarks/retrieval/weight-search.ts index 7a7cf4a..216a4ee 100644 --- a/benchmarks/retrieval/weight-search.ts +++ b/benchmarks/retrieval/weight-search.ts @@ -86,8 +86,10 @@ interface EvidenceSearchSample { readonly evidence: RoutingEvidence } -/** Native worker-pool controls for the explicit asynchronous benchmark search path. */ -export type ParallelSearchOptions = CandidateEvaluationPoolOptions +/** Native pool controls plus the cancellation signal owned by the benchmark Effect. */ +export interface ParallelSearchOptions extends CandidateEvaluationPoolOptions { + readonly signal?: AbortSignal +} const preparedFusionCache = new WeakMap< WeightSearchSample, @@ -199,7 +201,8 @@ const summarizeRanked = ( } } -const summarize = ( +/** Summarize one prepared benchmark candidate with the canonical quality metrics. */ +export const summarize = ( samples: readonly WeightSearchSample[], weights: ChannelWeights, fusion: FusionMethod = "rrf", @@ -617,11 +620,9 @@ const storeQualityResults = ( } const rankWeightCandidates = async ( - samples: readonly WeightSearchSample[], candidates: readonly ChannelWeights[], limit: number, qualityCache: Map, - fusion: FusionMethod, pool: CandidateEvaluationPool, objective: RouterObjective = "reranker-top20", baseline?: QualitySummary, @@ -657,8 +658,6 @@ const coordinateWeightCandidates = ( ] const selectBestWeights = async ( - samples: readonly WeightSearchSample[], - fusion: FusionMethod = "rrf", pool: CandidateEvaluationPool, objective: RouterObjective = "reranker-top20", baseline?: QualitySummary, @@ -666,11 +665,9 @@ const selectBestWeights = async ( ): Promise<{ readonly weights: ChannelWeights; readonly quality: QualitySummary }> => { const qualityCache = new Map() let beam = await rankWeightCandidates( - samples, weightCandidates(), SEARCH_BEAM_WIDTH, qualityCache, - fusion, pool, objective, baseline, @@ -679,11 +676,9 @@ const selectBestWeights = async ( for (let pass = 0; pass < SEARCH_PASSES; pass++) { for (const channel of CHANNELS) { beam = await rankWeightCandidates( - samples, coordinateWeightCandidates(beam, channel), SEARCH_BEAM_WIDTH, qualityCache, - fusion, pool, objective, baseline, @@ -731,6 +726,13 @@ const selectBestWeightsPerSubset = async ( return selectWeightSubsetCandidates(candidates, qualities, profile) } +const selectStaticWeights = ( + pool: CandidateEvaluationPool, + objective: RouterObjective, + productionQuality: QualitySummary, + profile: OptimizationProfile, +) => selectBestWeights(pool, objective, productionQuality, profile) + /** Config field containing one coefficient per retrieval channel. */ type InfluenceName = | "scoreInfluence" @@ -1193,6 +1195,7 @@ export const selectEligibleCandidate = ( } interface RouterSearchPreparation { + readonly samples: readonly WeightSearchSample[] readonly evidenceSamples: readonly EvidenceSearchSample[] readonly proxySamples: readonly EvidenceSearchSample[] readonly guardrailBaselines: GuardrailBaselines @@ -1213,6 +1216,7 @@ const prepareRouterSearch = ( profile, ) return { + samples, evidenceSamples, proxySamples, guardrailBaselines, @@ -1245,7 +1249,7 @@ const profileSeedFor = (profile: OptimizationProfile): EvidenceRouterConfig => ( }) const selectBestEvidenceRouter = async ( - samples: readonly WeightSearchSample[], + preparation: RouterSearchPreparation, fusion: FusionMethod, profile: OptimizationProfile, fullPool: CandidateEvaluationPool, @@ -1260,13 +1264,14 @@ const selectBestEvidenceRouter = async ( readonly randomCandidates: number }> => { const { + samples, evidenceSamples, proxySamples, guardrailBaselines, productionQuality, proxyBaseline, stats, - } = prepareRouterSearch(samples, profile) + } = preparation const searchContext: RouterSearchContext = { samples: evidenceSamples, proxySamples, @@ -1284,8 +1289,6 @@ const selectBestEvidenceRouter = async ( } const profileSeed = profileSeedFor(profile) const baseWeights = await selectBestWeights( - samples, - fusion, fullPool, "reranker-top20", productionQuality, @@ -1364,13 +1367,19 @@ const withCandidatePool = async ( samples: readonly WeightSearchSample[], fusion: FusionMethod, profile: OptimizationProfile, - options: CandidateEvaluationPoolOptions, + options: ParallelSearchOptions, operation: (pool: CandidateEvaluationPool) => Promise, ): Promise => { const pool = await createEvaluationPoolForSamples(samples, fusion, profile, options) + const closeOnAbort = () => { + void pool.close() + } + options.signal?.addEventListener("abort", closeOnAbort, { once: true }) + if (options.signal?.aborted) closeOnAbort() try { return await operation(pool) } finally { + options.signal?.removeEventListener("abort", closeOnAbort) await pool.close() } } @@ -1386,14 +1395,7 @@ const optimizeWeightsWithPool = async ( profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, pool: CandidateEvaluationPool, ): Promise => { - const selected = await selectBestWeights( - development, - "rrf", - pool, - "reranker-top20", - undefined, - profile, - ) + const selected = await selectBestWeights(pool, "reranker-top20", undefined, profile) return { model, queryKind, @@ -1476,14 +1478,7 @@ const optimizeFusionWeightsWithPool = async ( profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, pool: CandidateEvaluationPool, ): Promise => { - const selected = await selectBestWeights( - development, - fusion, - pool, - "reranker-top20", - undefined, - profile, - ) + const selected = await selectBestWeights(pool, "reranker-top20", undefined, profile) const guardrailBaselines = buildGuardrailBaselines(development, profile) const holdoutProfile = unweightedProfile(profile) const validationQuality = summarize(validationSamples, selected.weights, fusion, profile) @@ -1606,15 +1601,21 @@ const withParallelEvidencePools = async ( fullPool: CandidateEvaluationPool, ) => Promise, ): Promise => { - const evidenceSamples = prepareEvidenceSamples(samples) - const proxySamples = buildProxySamples(evidenceSamples) + const preparation = prepareRouterSearch(samples, profile) + const { evidenceSamples, proxySamples } = preparation const fullPool = await createEvaluationPoolForSamples( evidenceSamples.map(({ sample }) => sample), fusion, profile, options, ) + const closeFullPoolOnAbort = () => { + void fullPool.close() + } + options.signal?.addEventListener("abort", closeFullPoolOnAbort, { once: true }) + if (options.signal?.aborted) closeFullPoolOnAbort() let proxyPool: CandidateEvaluationPool | undefined + let closeProxyPoolOnAbort: (() => void) | undefined try { proxyPool = proxySamples === evidenceSamples @@ -1625,14 +1626,56 @@ const withParallelEvidencePools = async ( profile, options, ) - const selection = await selectBestEvidenceRouter(samples, fusion, profile, fullPool, proxyPool) + if (proxyPool !== undefined && proxyPool !== fullPool) { + closeProxyPoolOnAbort = () => { + void proxyPool?.close() + } + options.signal?.addEventListener("abort", closeProxyPoolOnAbort, { once: true }) + if (options.signal?.aborted) closeProxyPoolOnAbort() + } + const selection = await selectBestEvidenceRouter( + preparation, + fusion, + profile, + fullPool, + proxyPool, + ) return await operation(selection, fullPool) } finally { + options.signal?.removeEventListener("abort", closeFullPoolOnAbort) + if (closeProxyPoolOnAbort !== undefined) + options.signal?.removeEventListener("abort", closeProxyPoolOnAbort) if (proxyPool !== undefined && proxyPool !== fullPool) await proxyPool.close() await fullPool.close() } } +interface StaticRouterSelection { + readonly selection: EvidenceRouterSelection + readonly staticSelection: Awaited> +} + +const selectStaticWeightsForSelections = async ( + selections: readonly EvidenceRouterSelection[], + fullPool: CandidateEvaluationPool, + productionQuality: QualitySummary, + profile: OptimizationProfile, +): Promise => { + const selected: StaticRouterSelection[] = [] + for (const selection of selections) { + selected.push({ + selection, + staticSelection: await selectStaticWeights( + fullPool, + selection.objective, + productionQuality, + profile, + ), + }) + } + return selected +} + const optimizeEvidenceRouterWithPools = async ( model: string, fusion: FusionMethod, @@ -1664,17 +1707,14 @@ const optimizeEvidenceRouterWithPools = async ( profile, ), } - const results: EvidenceRouterSearchResult[] = [] - for (const selection of dynamicSelection.selections) { - const staticSelection = await selectBestWeights( - development, - fusion, - fullPool, - selection.objective, - dynamicSelection.productionQuality, - profile, - ) - results.push({ + const staticSelections = await selectStaticWeightsForSelections( + dynamicSelection.selections, + fullPool, + dynamicSelection.productionQuality, + profile, + ) + const results: EvidenceRouterSearchResult[] = staticSelections.map( + ({ selection, staticSelection }) => ({ model, fusion, objective: selection.objective, @@ -1712,8 +1752,8 @@ const optimizeEvidenceRouterWithPools = async ( ), profile, ), - }) - } + }), + ) return results }, ) @@ -1741,14 +1781,7 @@ const fitRecommendedWeightsWithPool = async ( profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, pool: CandidateEvaluationPool, ): Promise => { - const selected = await selectBestWeights( - samples, - "rrf", - pool, - "reranker-top20", - undefined, - profile, - ) + const selected = await selectBestWeights(pool, "reranker-top20", undefined, profile) return { model, queryKind, @@ -1794,14 +1827,7 @@ const fitRecommendedFusionWeightsWithPool = async ( profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, pool: CandidateEvaluationPool, ): Promise => { - const selected = await selectBestWeights( - samples, - fusion, - pool, - "reranker-top20", - undefined, - profile, - ) + const selected = await selectBestWeights(pool, "reranker-top20", undefined, profile) const guardrailBaselines = buildGuardrailBaselines(samples, profile) const guardrailsMet = fusionGuardrailsMet( samples, @@ -1866,17 +1892,14 @@ const fitRecommendedEvidenceRouterWithOptions = ( profile, options, async (dynamicSelection, fullPool) => { - const results: RecommendedEvidenceRouter[] = [] - for (const selection of dynamicSelection.selections) { - const staticSelection = await selectBestWeights( - samples, - fusion, - fullPool, - selection.objective, - dynamicSelection.productionQuality, - profile, - ) - results.push({ + const staticSelections = await selectStaticWeightsForSelections( + dynamicSelection.selections, + fullPool, + dynamicSelection.productionQuality, + profile, + ) + const results: RecommendedEvidenceRouter[] = staticSelections.map( + ({ selection, staticSelection }) => ({ model, fusion, objective: selection.objective, @@ -1891,8 +1914,8 @@ const fitRecommendedEvidenceRouterWithOptions = ( proxyEvaluations: dynamicSelection.proxyEvaluations, fullEvaluations: dynamicSelection.fullEvaluations, searchDiagnostics: dynamicSelection.searchDiagnostics, - }) - } + }), + ) return results }, ) diff --git a/benchmarks/retrieval/worker-pool.ts b/benchmarks/retrieval/worker-pool.ts index c724031..91798ff 100644 --- a/benchmarks/retrieval/worker-pool.ts +++ b/benchmarks/retrieval/worker-pool.ts @@ -149,14 +149,24 @@ const contextTokens = (chunk: Chunk): number => /** Build the compact, cloneable snapshot used by serial and worker evaluators. */ export const createEvaluationSnapshot = ( inputs: readonly EvaluationSampleInput[], -): EvaluationSnapshot => ({ - samples: inputs.map((input) => ({ - fusion: input.evaluator.snapshot, - targets: input.targets.map((target) => [...target]), - contextTokens: input.chunks.map(contextTokens), - sampleWeight: input.sampleWeight, - })), -}) +): EvaluationSnapshot => { + const contextTokenCache = new WeakMap() + return { + samples: inputs.map((input) => { + let cachedContextTokens = contextTokenCache.get(input.chunks) + if (cachedContextTokens === undefined) { + cachedContextTokens = input.chunks.map(contextTokens) + contextTokenCache.set(input.chunks, cachedContextTokens) + } + return { + fusion: input.evaluator.snapshot, + targets: input.targets.map((target) => [...target]), + contextTokens: cachedContextTokens, + sampleWeight: input.sampleWeight, + } + }), + } +} /** Evaluate candidates serially using the same prepared snapshot as the worker pool. */ export const evaluateCandidatesSerial = ( @@ -212,7 +222,7 @@ class NativeCandidateEvaluationPool implements CandidateEvaluationPool { worker.on("message", (message: unknown) => this.handleMessage(slot, message)) worker.on("error", (cause: Error) => this.handleWorkerError(slot, cause)) worker.on("exit", (code: number) => { - if (!this.closed && code !== 0) + if (!this.closed) this.handleWorkerError(slot, new Error(`Fusion worker exited with code ${code}`)) }) this.slots.push(slot) diff --git a/benchmarks/tests/channels.test.ts b/benchmarks/tests/channels.test.ts index 3d3805b..aca5ca5 100644 --- a/benchmarks/tests/channels.test.ts +++ b/benchmarks/tests/channels.test.ts @@ -599,7 +599,7 @@ describe("retrieval benchmark fixture", () => { samples, samples, undefined, - { workerCount: 2, batchSize: 8 }, + { workerCount: 2, batchSize: 8, fallbackToSerial: false }, ), ) diff --git a/benchmarks/tests/worker-pool.test.ts b/benchmarks/tests/worker-pool.test.ts index f4cca7e..20bb5bb 100644 --- a/benchmarks/tests/worker-pool.test.ts +++ b/benchmarks/tests/worker-pool.test.ts @@ -4,6 +4,7 @@ import { describe, expect, it } from "@effect/vitest" import type { Chunk } from "../../src/domain/chunk.js" import { prepareFusion } from "../retrieval/fusion.js" +import { SEARCH_PRIORITY_PROFILE } from "../retrieval/optimization-profiles.js" import { fitRecommendedEvidenceRouter, fitRecommendedFusionWeights, @@ -11,6 +12,7 @@ import { fitRecommendedWeights, fitRecommendedWeightsParallel, optimizeFusionWeights, + summarize, } from "../retrieval/weight-search.js" import { createCandidateEvaluationPool, @@ -87,7 +89,8 @@ const searchSample = { describe("benchmark candidate evaluation pool", () => { it("derives a bounded default and honors explicit sizing", () => { - expect(getDefaultWorkerCount()).toBe(Math.max(1, availableParallelism() - 1)) + expect(getDefaultWorkerCount()).toBeGreaterThanOrEqual(1) + expect(getDefaultWorkerCount()).toBeLessThanOrEqual(availableParallelism()) expect(resolveWorkerCount(0)).toBe(0) expect(resolveWorkerCount(3)).toBe(3) }) @@ -113,6 +116,41 @@ describe("benchmark candidate evaluation pool", () => { expect(pool.stats().activeWorkerCount).toBe(0) }) + it("keeps worker metrics aligned with canonical summarization", () => { + const parityRankings = { + ...rankings, + dense: [...rankings.dense, { chunkIndex: 99, score: 0.25 }], + } + const paritySamples = [ + { + ...searchSample, + rankings: parityRankings, + targets: [new Set([0])], + queryKind: "identifier" as const, + }, + { + ...searchSample, + intentId: "fixture-002", + rankings: parityRankings, + targets: [new Set([3])], + queryKind: "searchPhrase" as const, + }, + ] + const weights = { identity: 1, camelcase: 1, bm25: 1, dense: 1, sparse: 1 } + const snapshot = createEvaluationSnapshot( + paritySamples.map((sample) => ({ + evaluator: prepareFusion("dbsf", sample.rankings), + targets: sample.targets, + chunks: sample.chunks, + sampleWeight: SEARCH_PRIORITY_PROFILE.queryFormWeights[sample.queryKind], + })), + ) + + expect(evaluateCandidatesSerial(snapshot, [{ weights }])[0]).toEqual( + summarize(paritySamples, weights, "dbsf", SEARCH_PRIORITY_PROFILE), + ) + }) + it("keeps the explicit parallel search result equal to the serial search", async () => { const serial = await fitRecommendedWeights("fixture", "identifier", [searchSample]) const parallel = await fitRecommendedWeightsParallel( From 32e97a9053377890ddf02e9109eed7d353a53764 Mon Sep 17 00:00:00 2001 From: Lucas Burmeister Date: Tue, 4 Aug 2026 14:33:41 +0200 Subject: [PATCH 3/8] perf: parallelize complete router jobs --- .fallowrc.json | 3 +- CONTEXT.md | 9 +- benchmarks/retrieval/router-job-pool.ts | 182 ++++++++++++++++++++ benchmarks/retrieval/router-job-worker.mjs | 32 ++++ benchmarks/retrieval/runner.ts | 188 +++++++++++++++------ benchmarks/retrieval/ts-loader.mjs | 10 ++ benchmarks/retrieval/worker-message.ts | 8 + benchmarks/retrieval/worker-pool.ts | 5 +- benchmarks/tests/worker-pool.test.ts | 31 ++++ 9 files changed, 408 insertions(+), 60 deletions(-) create mode 100644 benchmarks/retrieval/router-job-pool.ts create mode 100644 benchmarks/retrieval/router-job-worker.mjs create mode 100644 benchmarks/retrieval/ts-loader.mjs create mode 100644 benchmarks/retrieval/worker-message.ts diff --git a/.fallowrc.json b/.fallowrc.json index 84576d3..e93fff5 100644 --- a/.fallowrc.json +++ b/.fallowrc.json @@ -4,7 +4,8 @@ "entry": [ "src/index.ts", "tests/scripts/bench-startup.mjs", - "tests/scripts/check-dtype-output.mjs" + "tests/scripts/check-dtype-output.mjs", + "benchmarks/retrieval/ts-loader.mjs" ], "boundaries": { "zones": [ diff --git a/CONTEXT.md b/CONTEXT.md index db088b6..8e8fc02 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -257,8 +257,13 @@ same selection algorithm; `PIX_BENCH_WORKERS=0` selects serial evaluation, `work the serial pool through `createCandidateEvaluationPool`, and the runner's `PIX_BENCH_SEARCH_MODE=serial` selects that serial pool mode. `PIX_BENCH_WORKER_BATCH_SIZE` bounds each worker message. -Unavailable worker startup falls back to serial evaluation; task failures are surfaced and all workers are -terminated before the benchmark fails. +Evidence-router grouped-fold and repository-holdout jobs run as complete native router jobs. One worker +owns the mutable candidate-generation, beam, cache, and promotion state for one job. Final deployment fits +keep their candidate-evaluation queue, and the runner splits one `PIX_BENCH_WORKERS` budget between the +holdout-job queue and that inner queue so nested pools do not oversubscribe the machine. Results retain +planner order even when jobs finish out of order. Serial mode keeps the same algorithm on the main thread +for comparison. Native router worker startup or task failures are surfaced and all workers are terminated +before the benchmark fails. Repository checkouts live under ignored `benchmarks/.cache/repos/`; generated artifacts live under ignored `benchmarks/results/`. Dense and Sparse vectors are held only in the production in-memory SQLite adapter during a benchmark run. See `benchmarks/README.md` and `benchmarks/BASELINE.md`. diff --git a/benchmarks/retrieval/router-job-pool.ts b/benchmarks/retrieval/router-job-pool.ts new file mode 100644 index 0000000..1f2b220 --- /dev/null +++ b/benchmarks/retrieval/router-job-pool.ts @@ -0,0 +1,182 @@ +import { Worker } from "node:worker_threads" + +import type { FusionMethod } from "../../src/domain/retrieval.js" +import type { OptimizationProfile } from "./optimization-profiles.js" +import type { EvidenceRouterSearchResult, ValidationStrategy } from "./types.js" +import type { WeightSearchSample } from "./weight-search.js" +import { hasWorkerMessageType } from "./worker-message.js" + +/** One independent evidence-router fold or repository holdout. */ +export interface EvidenceRouterHoldoutJob { + readonly kind: "holdout" + readonly model: string + readonly fusion: FusionMethod + readonly strategy: ValidationStrategy + readonly fold: string + readonly development: readonly WeightSearchSample[] + readonly validation: readonly WeightSearchSample[] +} + +/** Result returned by one complete native router job. */ +export type EvidenceRouterJobResult = readonly EvidenceRouterSearchResult[] + +/** Options for the native pool that runs complete router jobs. */ +export interface EvidenceRouterJobPoolOptions { + /** Maximum number of native workers, with one complete search per worker. */ + readonly workerCount: number + /** Abort signal that terminates all active router jobs. */ + readonly signal?: AbortSignal + /** Test-only override for the router job worker entry point. */ + readonly workerUrl?: URL + /** Test-only override for the TypeScript ESM loader. */ + readonly loaderUrl?: URL +} + +interface WorkerReadyMessage { + readonly type: "ready" +} + +interface WorkerResultMessage { + readonly type: "result" + readonly jobId: number + readonly results: EvidenceRouterJobResult +} + +interface WorkerErrorMessage { + readonly type: "error" + readonly jobId?: number + readonly message: string +} + +type WorkerMessage = WorkerReadyMessage | WorkerResultMessage | WorkerErrorMessage + +interface WorkerSlot { + readonly worker: Worker + busy: boolean +} + +const DEFAULT_WORKER_URL = new URL("./router-job-worker.mjs", import.meta.url) +const DEFAULT_LOADER_URL = new URL("./ts-loader.mjs", import.meta.url) + +const errorFromUnknown = (cause: unknown): Error => + cause instanceof Error ? cause : new Error(String(cause)) + +const isWorkerMessage = (message: unknown): message is WorkerMessage => { + return hasWorkerMessageType(message, ["ready", "result", "error"]) +} + +const workerExecArgv = (loaderUrl: URL): readonly string[] => [ + "--no-warnings", + "--experimental-strip-types", + "--experimental-loader", + loaderUrl.href, +] + +/** Run independent router holdout jobs in native workers and return results in input order. */ +export const runEvidenceRouterJobs = ( + jobs: readonly EvidenceRouterHoldoutJob[], + profile: OptimizationProfile, + options: EvidenceRouterJobPoolOptions, +): Promise => { + if (jobs.length === 0) return Promise.resolve([]) + if (options.signal?.aborted) + return Promise.reject(new Error("Evidence-router job pool was aborted before startup")) + + const workerCount = Math.min(jobs.length, Math.max(1, Math.floor(options.workerCount))) + const workerUrl = options.workerUrl ?? DEFAULT_WORKER_URL + const loaderUrl = options.loaderUrl ?? DEFAULT_LOADER_URL + + return new Promise((resolve, reject) => { + const slots: WorkerSlot[] = [] + const results = new Map() + let nextJob = 0 + let completedJobs = 0 + let settled = false + + const removeAbortListener = () => options.signal?.removeEventListener("abort", abort) + + const terminate = async (): Promise => { + await Promise.all(slots.map((slot) => slot.worker.terminate().catch(() => -1))) + } + + const finish = (cause?: Error): void => { + if (settled) return + settled = true + removeAbortListener() + void terminate().then(() => { + if (cause !== undefined) { + reject(cause) + return + } + const ordered: EvidenceRouterJobResult[] = [] + for (let index = 0; index < jobs.length; index++) { + const result = results.get(index) + if (result === undefined) { + reject(new Error("Evidence-router job pool returned incomplete results")) + return + } + ordered.push(result) + } + resolve(ordered) + }) + } + + const abort = () => finish(new Error("Evidence-router job pool was interrupted")) + + const dispatch = (slot: WorkerSlot): void => { + if (settled || slot.busy || nextJob >= jobs.length) return + const jobId = nextJob++ + slot.busy = true + try { + slot.worker.postMessage({ type: "run", jobId, job: jobs[jobId], profile }) + } catch (cause) { + finish(errorFromUnknown(cause)) + } + } + + const handleMessage = (slot: WorkerSlot, message: unknown): void => { + if (!isWorkerMessage(message)) { + finish(new Error("Router job worker sent an invalid message")) + return + } + if (message.type === "ready") { + dispatch(slot) + return + } + if (message.type === "error") { + finish(new Error(message.message)) + return + } + if (message.results.length === 0) { + finish(new Error(`Router job worker returned no results for job ${message.jobId}`)) + return + } + slot.busy = false + results.set(message.jobId, message.results) + completedJobs++ + if (completedJobs === jobs.length) { + finish() + return + } + dispatch(slot) + } + + options.signal?.addEventListener("abort", abort, { once: true }) + try { + for (let index = 0; index < workerCount; index++) { + const worker = new Worker(workerUrl, { + execArgv: [...workerExecArgv(loaderUrl)], + }) + const slot: WorkerSlot = { worker, busy: false } + worker.on("message", (message: unknown) => handleMessage(slot, message)) + worker.on("error", (cause: Error) => finish(cause)) + worker.on("exit", (code: number) => { + if (!settled) finish(new Error(`Router job worker exited with code ${code}`)) + }) + slots.push(slot) + } + } catch (cause) { + finish(errorFromUnknown(cause)) + } + }) +} diff --git a/benchmarks/retrieval/router-job-worker.mjs b/benchmarks/retrieval/router-job-worker.mjs new file mode 100644 index 0000000..e667f6f --- /dev/null +++ b/benchmarks/retrieval/router-job-worker.mjs @@ -0,0 +1,32 @@ +import { parentPort } from "node:worker_threads" + +import { optimizeEvidenceRouter } from "./weight-search.ts" + +if (parentPort === null) throw new Error("Router job worker requires a parent port") + +parentPort.postMessage({ type: "ready" }) + +parentPort.on("message", async (message) => { + if (message?.type !== "run") { + parentPort.postMessage({ type: "error", message: "Router job worker received an invalid task" }) + return + } + try { + const results = await optimizeEvidenceRouter( + message.job.model, + message.job.fusion, + message.job.strategy, + message.job.fold, + message.job.development, + message.job.validation, + message.profile, + ) + parentPort.postMessage({ type: "result", jobId: message.jobId, results }) + } catch (error) { + parentPort.postMessage({ + type: "error", + jobId: message.jobId, + message: error instanceof Error ? error.message : String(error), + }) + } +}) diff --git a/benchmarks/retrieval/runner.ts b/benchmarks/retrieval/runner.ts index 3770638..d5fd4e4 100644 --- a/benchmarks/retrieval/runner.ts +++ b/benchmarks/retrieval/runner.ts @@ -29,6 +29,7 @@ import { OPTIMIZATION_PROFILES, type OptimizationProfile } from "./optimization- import { prepareCorpus, type PreparedCorpus } from "./prepare.js" import { fuseVariant, rankLexicalChannels, RETRIEVAL_VARIANTS } from "./ranking.js" import { renderMarkdownReport } from "./report.js" +import { runEvidenceRouterJobs, type EvidenceRouterHoldoutJob } from "./router-job-pool.js" import { withSqliteBenchmarkStore } from "./sqlite-index.js" import { ROUTER_SEARCH_STRATEGY, @@ -37,7 +38,9 @@ import { type CorpusManifest, type QueryKind, type QueryMeasurement, + type RecommendedEvidenceRouter, type ValidationStrategy, + type EvidenceRouterSearchResult, } from "./types.js" import { fitRecommendedEvidenceRouterParallel, @@ -49,6 +52,7 @@ import { evaluateProductionRouter, type WeightSearchSample, } from "./weight-search.js" +import { getDefaultWorkerCount, resolveWorkerCount } from "./worker-pool.js" const CONTEXT_BUDGETS = [2_048, 4_096, 8_192, 16_384] as const const EMBEDDING_BATCH_SIZE = 2 @@ -134,6 +138,46 @@ const splitSamples = ( validation: samples.filter(isValidation), }) +const planEvidenceRouterJobs = ( + samplesByModel: ReadonlyMap, + config: BenchmarkProfileConfig, + groupedStrategy: ValidationStrategy, +): readonly EvidenceRouterHoldoutJob[] => { + const jobs: EvidenceRouterHoldoutJob[] = [] + for (const [model, samples] of samplesByModel) { + const repositories = [...new Set(samples.map((sample) => sample.repository))] + for (const fusion of config.routerFusionMethods) { + for (let fold = 0; fold < config.groupedFolds; fold++) { + const split = splitSamples(samples, (sample) => sample.groupedFold === fold) + jobs.push({ + kind: "holdout", + model, + fusion, + strategy: groupedStrategy, + fold: String(fold + 1), + development: split.development, + validation: split.validation, + }) + } + if (config.repositoryHoldouts && repositories.length > 1) { + for (const repository of repositories) { + const split = splitSamples(samples, (sample) => sample.repository === repository) + jobs.push({ + kind: "holdout", + model, + fusion, + strategy: "leave-one-repository-out", + fold: repository, + development: split.development, + validation: split.validation, + }) + } + } + } + } + return jobs +} + const isLongInput = (text: string): boolean => Buffer.byteLength(text, "utf8") / 4 > SINGLE_ITEM_ESTIMATED_TOKENS @@ -291,7 +335,12 @@ export const runRetrievalBenchmark = ( const benchmarkStartedAt = performance.now() const config = profileConfig(profile) const serialSearch = process.env.PIX_BENCH_SEARCH_MODE === "serial" - const searchOptions = serialSearch ? { workerCount: 0 } : {} + const searchOptions = serialSearch + ? { workerCount: 0 } + : { + workerCount: Math.min(resolveWorkerCount(), getDefaultWorkerCount()), + fallbackToSerial: false, + } const optimizationProfile = yield* selectOptimizationProfile() const groupedStrategy: ValidationStrategy = config.groupedFolds === 3 ? "grouped-3-fold" : "grouped-5-fold" @@ -659,66 +708,97 @@ export const runRetrievalBenchmark = ( const fusionSearchDurationMs = performance.now() - fusionSearchStartedAt const evidenceRouterSearchStartedAt = performance.now() - const evidenceRouterSearch: BenchmarkArtifact["evidenceRouterSearch"][number][] = [] - for (const [model, samples] of samplesByModel) { - for (const fusion of config.routerFusionMethods) { - for (let fold = 0; fold < config.groupedFolds; fold++) { - const split = splitSamples(samples, (sample) => sample.groupedFold === fold) - reportProgress( - `${model}: selecting ${fusion} evidence router for grouped fold ${fold + 1}/${config.groupedFolds}`, + const routerJobs = planEvidenceRouterJobs(samplesByModel, config, groupedStrategy) + const routerWorkerBudget = Math.min( + resolveWorkerCount(searchOptions.workerCount), + getDefaultWorkerCount(), + ) + const recommendedJobs: { + readonly model: string + readonly fusion: FusionMethod + readonly samples: readonly WeightSearchSample[] + }[] = [] + for (const [model, samples] of samplesByModel) + for (const fusion of config.routerFusionMethods) + recommendedJobs.push({ model, fusion, samples }) + const canParallelizeRouterJobs = !serialSearch && routerWorkerBudget >= 2 + // Holdouts are complete native-worker jobs; final fits retain their candidate queue. + // Split one budget between both queues so nested pools never oversubscribe the machine. + const holdoutWorkerCount = + !canParallelizeRouterJobs || routerJobs.length === 0 + ? 0 + : Math.min( + routerJobs.length, + recommendedJobs.length === 0 + ? routerWorkerBudget + : Math.max(1, Math.floor(routerWorkerBudget / 2)), ) - evidenceRouterSearch.push( - ...(yield* runParallelSearch((signal) => - optimizeEvidenceRouterParallel( - model, - fusion, - groupedStrategy, - String(fold + 1), - split.development, - split.validation, - optimizationProfile, - { ...searchOptions, signal }, - ), - )), + const recommendedWorkerCount = + !canParallelizeRouterJobs || recommendedJobs.length === 0 + ? 0 + : routerJobs.length === 0 + ? Math.min(recommendedJobs.length, routerWorkerBudget) + : Math.max(1, routerWorkerBudget - holdoutWorkerCount) + reportProgress( + `running ${routerJobs.length} complete holdout jobs with ${holdoutWorkerCount} native workers ` + + `and ${recommendedJobs.length} final fits with ${recommendedWorkerCount} pool workers`, + ) + let evidenceRouterSearch: readonly EvidenceRouterSearchResult[] = [] + let recommendedEvidenceRouters: readonly RecommendedEvidenceRouter[] = [] + const holdoutSearch = + canParallelizeRouterJobs && routerJobs.length > 0 + ? runParallelSearch((signal) => + runEvidenceRouterJobs(routerJobs, optimizationProfile, { + workerCount: holdoutWorkerCount, + signal, + }), ) - } - const repositories = [...new Set(samples.map((sample) => sample.repository))] - if (config.repositoryHoldouts && repositories.length > 1) { - for (const repository of repositories) { - const split = splitSamples(samples, (sample) => sample.repository === repository) - reportProgress( - `${model}: selecting ${fusion} evidence router with ${repository} held out`, - ) - evidenceRouterSearch.push( - ...(yield* runParallelSearch((signal) => + : Effect.forEach( + routerJobs, + (job) => + runParallelSearch((signal) => optimizeEvidenceRouterParallel( - model, - fusion, - "leave-one-repository-out", - repository, - split.development, - split.validation, + job.model, + job.fusion, + job.strategy, + job.fold, + job.development, + job.validation, optimizationProfile, - { ...searchOptions, signal }, + { ...searchOptions, workerCount: 0, signal }, ), - )), - ) - } - } - } - } - reportProgress("fitting final evidence router on all samples") - const recommendedEvidenceRouters: BenchmarkArtifact["recommendedEvidenceRouters"][number][] = [] - for (const [model, samples] of samplesByModel) - for (const fusion of config.routerFusionMethods) - recommendedEvidenceRouters.push( - ...(yield* runParallelSearch((signal) => - fitRecommendedEvidenceRouterParallel(model, fusion, samples, optimizationProfile, { + ), + { concurrency: 1 }, + ) + const recommendedSearch = Effect.forEach( + recommendedJobs, + (job) => + runParallelSearch((signal) => + fitRecommendedEvidenceRouterParallel( + job.model, + job.fusion, + job.samples, + optimizationProfile, + { ...searchOptions, + workerCount: canParallelizeRouterJobs ? recommendedWorkerCount : 0, signal, - }), - )), - ) + }, + ), + ), + { concurrency: 1 }, + ) + if (canParallelizeRouterJobs) { + const results = yield* Effect.all( + { holdouts: holdoutSearch, recommended: recommendedSearch }, + { concurrency: 2 }, + ) + evidenceRouterSearch = results.holdouts.flat() + recommendedEvidenceRouters = results.recommended.flat() + } else { + evidenceRouterSearch = (yield* holdoutSearch).flat() + recommendedEvidenceRouters = (yield* recommendedSearch).flat() + } const evidenceRouterSearchDurationMs = performance.now() - evidenceRouterSearchStartedAt const embeddingDurationMs = embeddingRuns.reduce( diff --git a/benchmarks/retrieval/ts-loader.mjs b/benchmarks/retrieval/ts-loader.mjs new file mode 100644 index 0000000..73ed509 --- /dev/null +++ b/benchmarks/retrieval/ts-loader.mjs @@ -0,0 +1,10 @@ +export async function resolve(specifier, context, defaultResolve) { + if (specifier.startsWith(".") && specifier.endsWith(".js")) { + try { + return await defaultResolve(`${specifier.slice(0, -3)}.ts`, context, defaultResolve) + } catch { + // Keep normal JavaScript resolution for actual .js modules. + } + } + return defaultResolve(specifier, context, defaultResolve) +} diff --git a/benchmarks/retrieval/worker-message.ts b/benchmarks/retrieval/worker-message.ts new file mode 100644 index 0000000..dd52042 --- /dev/null +++ b/benchmarks/retrieval/worker-message.ts @@ -0,0 +1,8 @@ +/** Check whether an unknown worker message has one of the expected protocol types. */ +export const hasWorkerMessageType = ( + message: unknown, + types: readonly string[], +): message is { readonly type: string } => { + if (typeof message !== "object" || message === null || !("type" in message)) return false + return typeof message.type === "string" && types.includes(message.type) +} diff --git a/benchmarks/retrieval/worker-pool.ts b/benchmarks/retrieval/worker-pool.ts index 91798ff..ce5b818 100644 --- a/benchmarks/retrieval/worker-pool.ts +++ b/benchmarks/retrieval/worker-pool.ts @@ -6,6 +6,7 @@ import type { ChannelWeights } from "../../src/domain/retrieval.js" import { evaluateCandidate } from "./fusion-core.mjs" import { type PreparedFusionEvaluator, type PreparedFusionSnapshot } from "./fusion.js" import type { QualitySummary } from "./types.js" +import { hasWorkerMessageType } from "./worker-message.js" const DEFAULT_BATCH_SIZE = 32 const DEFAULT_WORKER_URL = new URL("./fusion-worker.mjs", import.meta.url) @@ -176,9 +177,7 @@ export const evaluateCandidatesSerial = ( candidates.map((candidate) => evaluateCandidate(snapshot, candidate)) const isWorkerMessage = (message: unknown): message is WorkerMessage => { - if (typeof message !== "object" || message === null || !("type" in message)) return false - const type = message.type - return type === "ready" || type === "result" || type === "error" + return hasWorkerMessageType(message, ["ready", "result", "error"]) } class NativeCandidateEvaluationPool implements CandidateEvaluationPool { diff --git a/benchmarks/tests/worker-pool.test.ts b/benchmarks/tests/worker-pool.test.ts index 20bb5bb..74a9825 100644 --- a/benchmarks/tests/worker-pool.test.ts +++ b/benchmarks/tests/worker-pool.test.ts @@ -5,12 +5,17 @@ import { describe, expect, it } from "@effect/vitest" import type { Chunk } from "../../src/domain/chunk.js" import { prepareFusion } from "../retrieval/fusion.js" import { SEARCH_PRIORITY_PROFILE } from "../retrieval/optimization-profiles.js" +import { + runEvidenceRouterJobs, + type EvidenceRouterHoldoutJob, +} from "../retrieval/router-job-pool.js" import { fitRecommendedEvidenceRouter, fitRecommendedFusionWeights, fitRecommendedFusionWeightsParallel, fitRecommendedWeights, fitRecommendedWeightsParallel, + optimizeEvidenceRouter, optimizeFusionWeights, summarize, } from "../retrieval/weight-search.js" @@ -88,6 +93,32 @@ const searchSample = { } describe("benchmark candidate evaluation pool", () => { + it("runs complete router jobs in native workers and preserves results", async () => { + const job: EvidenceRouterHoldoutJob = { + kind: "holdout", + model: "fixture", + fusion: "dbsf", + strategy: "grouped-5-fold", + fold: "1", + development: [searchSample], + validation: [searchSample], + } + const [parallel] = await runEvidenceRouterJobs([job], SEARCH_PRIORITY_PROFILE, { + workerCount: 1, + }) + const serial = await optimizeEvidenceRouter( + job.model, + job.fusion, + job.strategy, + job.fold, + job.development, + job.validation, + SEARCH_PRIORITY_PROFILE, + ) + + expect(parallel).toEqual(serial) + }) + it("derives a bounded default and honors explicit sizing", () => { expect(getDefaultWorkerCount()).toBeGreaterThanOrEqual(1) expect(getDefaultWorkerCount()).toBeLessThanOrEqual(availableParallelism()) From 92d5988130c15fe8eba81ea36d92628b2d4af9de Mon Sep 17 00:00:00 2001 From: Lucas Burmeister Date: Tue, 4 Aug 2026 17:05:07 +0200 Subject: [PATCH 4/8] perf: add shared router candidate queue --- CONTEXT.md | 25 +- benchmarks/retrieval/fusion-worker.mjs | 14 +- benchmarks/retrieval/router-job-pool.ts | 164 +++++- benchmarks/retrieval/router-job-worker.mjs | 171 +++++- benchmarks/retrieval/runner.ts | 161 +++--- benchmarks/retrieval/types.ts | 19 + benchmarks/retrieval/weight-search.ts | 92 ++- benchmarks/retrieval/worker-pool.ts | 639 +++++++++++++++++++-- benchmarks/tests/worker-pool.test.ts | 122 +++- 9 files changed, 1211 insertions(+), 196 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 8e8fc02..2ad398e 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -252,18 +252,19 @@ remain unchanged. The evaluator still materializes and sorts the full candidate configuration; partial top-K or metric-specialized evaluation remains future work. The explicit parallel benchmark path uses a fixed native `node:worker_threads` pool (default `max(1, availableParallelism() - 1)`), sends the prepared snapshot once per worker, batches candidate weight vectors, and keeps beam/cache/archive/ -selection state on the main thread. The search APIs are async so serial and worker execution share the -same selection algorithm; `PIX_BENCH_WORKERS=0` selects serial evaluation, `workerCount: 1` selects -the serial pool through `createCandidateEvaluationPool`, and the runner's -`PIX_BENCH_SEARCH_MODE=serial` selects that serial pool mode. `PIX_BENCH_WORKER_BATCH_SIZE` bounds each -worker message. -Evidence-router grouped-fold and repository-holdout jobs run as complete native router jobs. One worker -owns the mutable candidate-generation, beam, cache, and promotion state for one job. Final deployment fits -keep their candidate-evaluation queue, and the runner splits one `PIX_BENCH_WORKERS` budget between the -holdout-job queue and that inner queue so nested pools do not oversubscribe the machine. Results retain -planner order even when jobs finish out of order. Serial mode keeps the same algorithm on the main thread -for comparison. Native router worker startup or task failures are surfaced and all workers are terminated -before the benchmark fails. +selection state on the controller thread. The search APIs are async so serial and worker execution share the +same selection algorithm; `PIX_BENCH_WORKERS=0` selects serial evaluation, `workerCount: 1` selects the +serial pool through `createCandidateEvaluationPool`, and the runner's `PIX_BENCH_SEARCH_MODE=serial` selects +that serial pool mode. `PIX_BENCH_WORKER_BATCH_SIZE` bounds ordinary worker messages. +Evidence-router grouped-fold, repository-holdout, and fit-all jobs use native controller workers backed by one +shared candidate-evaluation queue. Controllers keep candidate generation, beam rounds, caches, and archive +updates local and sequential; candidate batches are work-stealing tasks evaluated by the shared native queue. +The queue caches each prepared snapshot once per worker and defaults to one candidate per task because the +benchmark's candidate costs are uneven; `PIX_BENCH_WORKER_BATCH_SIZE` overrides that default. The runner +reserves roughly half of `PIX_BENCH_WORKERS` for controllers and the remainder for candidate workers. +Results retain planner order even when jobs finish out of order. Serial mode keeps the same algorithm on the +main thread for comparison. Native worker startup or task failures are surfaced and all workers are +terminated before the benchmark fails. Repository checkouts live under ignored `benchmarks/.cache/repos/`; generated artifacts live under ignored `benchmarks/results/`. Dense and Sparse vectors are held only in the production in-memory SQLite adapter during a benchmark run. See `benchmarks/README.md` and `benchmarks/BASELINE.md`. diff --git a/benchmarks/retrieval/fusion-worker.mjs b/benchmarks/retrieval/fusion-worker.mjs index 65bc473..41aa5f7 100644 --- a/benchmarks/retrieval/fusion-worker.mjs +++ b/benchmarks/retrieval/fusion-worker.mjs @@ -4,17 +4,21 @@ import { evaluateCandidate } from "./fusion-core.mjs" if (parentPort === null) throw new Error("Fusion worker requires a parent port") -let snapshot +const snapshots = new Map() + +parentPort.postMessage({ type: "ready" }) parentPort.on("message", (message) => { try { if (message.type === "init") { - snapshot = message.snapshot - parentPort.postMessage({ type: "ready" }) + snapshots.set("legacy", message.snapshot) return } - if (message.type !== "evaluate" || snapshot === undefined) - throw new Error("Fusion worker received an invalid task") + if (message.type !== "evaluate") throw new Error("Fusion worker received an invalid task") + const snapshotId = message.snapshotId ?? "legacy" + if (message.snapshot !== undefined) snapshots.set(snapshotId, message.snapshot) + const snapshot = snapshots.get(snapshotId) + if (snapshot === undefined) throw new Error(`Fusion worker has no snapshot ${snapshotId}`) parentPort.postMessage({ type: "result", taskId: message.taskId, diff --git a/benchmarks/retrieval/router-job-pool.ts b/benchmarks/retrieval/router-job-pool.ts index 1f2b220..b0ee86f 100644 --- a/benchmarks/retrieval/router-job-pool.ts +++ b/benchmarks/retrieval/router-job-pool.ts @@ -2,9 +2,17 @@ import { Worker } from "node:worker_threads" import type { FusionMethod } from "../../src/domain/retrieval.js" import type { OptimizationProfile } from "./optimization-profiles.js" -import type { EvidenceRouterSearchResult, ValidationStrategy } from "./types.js" +import type { + EvidenceRouterSearchResult, + RecommendedEvidenceRouter, + ValidationStrategy, +} from "./types.js" import type { WeightSearchSample } from "./weight-search.js" -import { hasWorkerMessageType } from "./worker-message.js" +import type { + CandidateEvaluationQueue, + EvaluationCandidate, + EvaluationSnapshot, +} from "./worker-pool.js" /** One independent evidence-router fold or repository holdout. */ export interface EvidenceRouterHoldoutJob { @@ -17,13 +25,40 @@ export interface EvidenceRouterHoldoutJob { readonly validation: readonly WeightSearchSample[] } -/** Result returned by one complete native router job. */ -export type EvidenceRouterJobResult = readonly EvidenceRouterSearchResult[] +/** One fit-all router search over every sample after holdout evaluation. */ +export interface EvidenceRouterFitAllJob { + readonly kind: "fit-all" + readonly model: string + readonly fusion: FusionMethod + readonly samples: readonly WeightSearchSample[] +} + +/** One complete router search handled by a native worker. */ +export type EvidenceRouterJob = EvidenceRouterHoldoutJob | EvidenceRouterFitAllJob + +/** Tagged result returned by one holdout router job. */ +export interface EvidenceRouterHoldoutJobResult { + readonly jobId: number + readonly kind: "holdout" + readonly results: readonly EvidenceRouterSearchResult[] +} + +/** Tagged result returned by one fit-all router job. */ +export interface EvidenceRouterFitAllJobResult { + readonly jobId: number + readonly kind: "fit-all" + readonly results: readonly RecommendedEvidenceRouter[] +} + +/** Tagged result returned by one complete native router job. */ +export type EvidenceRouterJobResult = EvidenceRouterHoldoutJobResult | EvidenceRouterFitAllJobResult /** Options for the native pool that runs complete router jobs. */ export interface EvidenceRouterJobPoolOptions { - /** Maximum number of native workers, with one complete search per worker. */ + /** Maximum number of native router-controller workers. */ readonly workerCount: number + /** Shared candidate queue used by the active router controllers. */ + readonly candidateQueue?: CandidateEvaluationQueue /** Abort signal that terminates all active router jobs. */ readonly signal?: AbortSignal /** Test-only override for the router job worker entry point. */ @@ -36,10 +71,26 @@ interface WorkerReadyMessage { readonly type: "ready" } -interface WorkerResultMessage { +interface WorkerHoldoutResultMessage { + readonly type: "result" + readonly jobId: number + readonly kind: "holdout" + readonly results: readonly EvidenceRouterSearchResult[] +} + +interface WorkerFitAllResultMessage { readonly type: "result" readonly jobId: number - readonly results: EvidenceRouterJobResult + readonly kind: "fit-all" + readonly results: readonly RecommendedEvidenceRouter[] +} + +interface WorkerCandidateEvaluateMessage { + readonly type: "candidate-evaluate" + readonly requestId: number + readonly snapshotId: string + readonly snapshot?: EvaluationSnapshot + readonly candidates: readonly EvaluationCandidate[] } interface WorkerErrorMessage { @@ -48,10 +99,17 @@ interface WorkerErrorMessage { readonly message: string } -type WorkerMessage = WorkerReadyMessage | WorkerResultMessage | WorkerErrorMessage +type WorkerMessage = + | WorkerReadyMessage + | WorkerHoldoutResultMessage + | WorkerFitAllResultMessage + | WorkerCandidateEvaluateMessage + | WorkerErrorMessage interface WorkerSlot { + readonly slotId: number readonly worker: Worker + readonly candidateSnapshots: Map busy: boolean } @@ -61,8 +119,27 @@ const DEFAULT_LOADER_URL = new URL("./ts-loader.mjs", import.meta.url) const errorFromUnknown = (cause: unknown): Error => cause instanceof Error ? cause : new Error(String(cause)) +const isRecord = (message: unknown): message is Record => + typeof message === "object" && message !== null + const isWorkerMessage = (message: unknown): message is WorkerMessage => { - return hasWorkerMessageType(message, ["ready", "result", "error"]) + if (!isRecord(message) || typeof message.type !== "string") return false + if (message.type === "ready") return true + if (message.type === "error") return typeof message.message === "string" + if (message.type === "candidate-evaluate") + return ( + Number.isInteger(message.requestId) && + typeof message.snapshotId === "string" && + (message.snapshot === undefined || + (typeof message.snapshot === "object" && message.snapshot !== null)) && + Array.isArray(message.candidates) + ) + return ( + message.type === "result" && + Number.isInteger(message.jobId) && + (message.kind === "holdout" || message.kind === "fit-all") && + Array.isArray(message.results) + ) } const workerExecArgv = (loaderUrl: URL): readonly string[] => [ @@ -72,9 +149,9 @@ const workerExecArgv = (loaderUrl: URL): readonly string[] => [ loaderUrl.href, ] -/** Run independent router holdout jobs in native workers and return results in input order. */ +/** Run independent router jobs in native workers and return tagged results in input order. */ export const runEvidenceRouterJobs = ( - jobs: readonly EvidenceRouterHoldoutJob[], + jobs: readonly EvidenceRouterJob[], profile: OptimizationProfile, options: EvidenceRouterJobPoolOptions, ): Promise => { @@ -128,7 +205,15 @@ export const runEvidenceRouterJobs = ( const jobId = nextJob++ slot.busy = true try { - slot.worker.postMessage({ type: "run", jobId, job: jobs[jobId], profile }) + slot.worker.postMessage({ + type: "run", + jobId, + job: jobs[jobId], + profile, + candidateWorkerCount: options.candidateQueue?.workerCount ?? 0, + candidateBatchSize: options.candidateQueue?.batchSize ?? 0, + useCandidateQueue: options.candidateQueue !== undefined, + }) } catch (cause) { finish(errorFromUnknown(cause)) } @@ -147,12 +232,58 @@ export const runEvidenceRouterJobs = ( finish(new Error(message.message)) return } + if (message.type === "candidate-evaluate") { + if (options.candidateQueue === undefined) { + finish(new Error("Router job worker requested a candidate queue that was not configured")) + return + } + const snapshot = message.snapshot ?? slot.candidateSnapshots.get(message.snapshotId) + if (snapshot === undefined) { + finish(new Error(`Router job worker omitted snapshot ${message.snapshotId}`)) + return + } + if (message.snapshot !== undefined) + slot.candidateSnapshots.set(message.snapshotId, message.snapshot) + void options.candidateQueue + .evaluate(snapshot, message.candidates, undefined, `${slot.slotId}:${message.snapshotId}`) + .then((results) => { + if (settled) return + try { + slot.worker.postMessage({ + type: "candidate-result", + requestId: message.requestId, + results, + }) + } catch (cause) { + finish(errorFromUnknown(cause)) + } + }) + .catch((cause) => finish(errorFromUnknown(cause))) + return + } + const job = jobs[message.jobId] + if (job === undefined || job.kind !== message.kind) { + finish(new Error(`Router job worker returned an unexpected job ${message.jobId}`)) + return + } if (message.results.length === 0) { finish(new Error(`Router job worker returned no results for job ${message.jobId}`)) return } slot.busy = false - results.set(message.jobId, message.results) + if (message.kind === "holdout") { + results.set(message.jobId, { + jobId: message.jobId, + kind: "holdout", + results: message.results, + }) + } else { + results.set(message.jobId, { + jobId: message.jobId, + kind: "fit-all", + results: message.results, + }) + } completedJobs++ if (completedJobs === jobs.length) { finish() @@ -167,7 +298,12 @@ export const runEvidenceRouterJobs = ( const worker = new Worker(workerUrl, { execArgv: [...workerExecArgv(loaderUrl)], }) - const slot: WorkerSlot = { worker, busy: false } + const slot: WorkerSlot = { + slotId: index, + worker, + candidateSnapshots: new Map(), + busy: false, + } worker.on("message", (message: unknown) => handleMessage(slot, message)) worker.on("error", (cause: Error) => finish(cause)) worker.on("exit", (code: number) => { diff --git a/benchmarks/retrieval/router-job-worker.mjs b/benchmarks/retrieval/router-job-worker.mjs index e667f6f..f85144d 100644 --- a/benchmarks/retrieval/router-job-worker.mjs +++ b/benchmarks/retrieval/router-job-worker.mjs @@ -1,32 +1,181 @@ import { parentPort } from "node:worker_threads" -import { optimizeEvidenceRouter } from "./weight-search.ts" +import { + fitRecommendedEvidenceRouter, + fitRecommendedEvidenceRouterParallel, + optimizeEvidenceRouter, + optimizeEvidenceRouterParallel, +} from "./weight-search.ts" if (parentPort === null) throw new Error("Router job worker requires a parent port") +let busy = false +let activeCandidateQueue + +const createRemoteCandidateQueue = (workerCount, batchSize) => { + const pending = new Map() + const snapshotIds = new WeakMap() + const sentSnapshots = new WeakSet() + let nextRequestId = 0 + let nextSnapshotId = 0 + let closed = false + + const queue = { + workerCount, + batchSize, + activeWorkerCount: () => 0, + evaluate(snapshot, candidates, signal, requestedSnapshotId) { + if (closed) return Promise.reject(new Error("Remote candidate queue is closed")) + if (signal?.aborted) + return Promise.reject(new Error("Remote candidate queue was interrupted")) + if (candidates.length === 0) return Promise.resolve([]) + let snapshotId = requestedSnapshotId ?? snapshotIds.get(snapshot) + if (snapshotId === undefined) { + snapshotId = String(nextSnapshotId++) + snapshotIds.set(snapshot, snapshotId) + } + const requestId = nextRequestId++ + return new Promise((resolve, reject) => { + const abort = () => { + pending.delete(requestId) + reject(new Error("Remote candidate queue was interrupted")) + } + if (signal !== undefined) { + signal.addEventListener("abort", abort, { once: true }) + if (signal.aborted) { + abort() + return + } + } + pending.set(requestId, { + resolve: (results) => { + signal?.removeEventListener("abort", abort) + resolve(results) + }, + reject: (error) => { + signal?.removeEventListener("abort", abort) + reject(error) + }, + }) + parentPort.postMessage({ + type: "candidate-evaluate", + requestId, + snapshotId, + snapshot: sentSnapshots.has(snapshot) ? undefined : snapshot, + candidates, + }) + sentSnapshots.add(snapshot) + }) + }, + handleResult(message) { + const request = pending.get(message.requestId) + if (request === undefined) return + pending.delete(message.requestId) + if (message.type === "candidate-error") request.reject(new Error(message.message)) + else request.resolve(message.results) + }, + close() { + closed = true + for (const request of pending.values()) + request.reject(new Error("Remote candidate queue closed")) + pending.clear() + }, + } + return queue +} + +const runJob = async (message) => { + const candidateQueue = message.useCandidateQueue + ? createRemoteCandidateQueue(message.candidateWorkerCount, message.candidateBatchSize) + : undefined + activeCandidateQueue = candidateQueue + const options = + candidateQueue === undefined + ? { workerCount: 0 } + : { + workerCount: candidateQueue.workerCount, + fallbackToSerial: false, + evaluationQueue: candidateQueue, + } + try { + if (message.job.kind === "holdout") { + const results = + candidateQueue === undefined + ? await optimizeEvidenceRouter( + message.job.model, + message.job.fusion, + message.job.strategy, + message.job.fold, + message.job.development, + message.job.validation, + message.profile, + ) + : await optimizeEvidenceRouterParallel( + message.job.model, + message.job.fusion, + message.job.strategy, + message.job.fold, + message.job.development, + message.job.validation, + message.profile, + options, + ) + parentPort.postMessage({ type: "result", jobId: message.jobId, kind: "holdout", results }) + } else if (message.job.kind === "fit-all") { + const results = + candidateQueue === undefined + ? await fitRecommendedEvidenceRouter( + message.job.model, + message.job.fusion, + message.job.samples, + message.profile, + ) + : await fitRecommendedEvidenceRouterParallel( + message.job.model, + message.job.fusion, + message.job.samples, + message.profile, + options, + ) + parentPort.postMessage({ type: "result", jobId: message.jobId, kind: "fit-all", results }) + } else { + throw new Error("Router job worker received an unknown job kind") + } + } finally { + candidateQueue?.close() + activeCandidateQueue = undefined + } +} + parentPort.postMessage({ type: "ready" }) parentPort.on("message", async (message) => { + if (message?.type === "candidate-result" || message?.type === "candidate-error") { + activeCandidateQueue?.handleResult(message) + return + } if (message?.type !== "run") { parentPort.postMessage({ type: "error", message: "Router job worker received an invalid task" }) return } + if (busy) { + parentPort.postMessage({ + type: "error", + jobId: message.jobId, + message: "Router job worker is busy", + }) + return + } + busy = true try { - const results = await optimizeEvidenceRouter( - message.job.model, - message.job.fusion, - message.job.strategy, - message.job.fold, - message.job.development, - message.job.validation, - message.profile, - ) - parentPort.postMessage({ type: "result", jobId: message.jobId, results }) + await runJob(message) } catch (error) { parentPort.postMessage({ type: "error", jobId: message.jobId, message: error instanceof Error ? error.message : String(error), }) + } finally { + busy = false } }) diff --git a/benchmarks/retrieval/runner.ts b/benchmarks/retrieval/runner.ts index d5fd4e4..a79c334 100644 --- a/benchmarks/retrieval/runner.ts +++ b/benchmarks/retrieval/runner.ts @@ -29,7 +29,11 @@ import { OPTIMIZATION_PROFILES, type OptimizationProfile } from "./optimization- import { prepareCorpus, type PreparedCorpus } from "./prepare.js" import { fuseVariant, rankLexicalChannels, RETRIEVAL_VARIANTS } from "./ranking.js" import { renderMarkdownReport } from "./report.js" -import { runEvidenceRouterJobs, type EvidenceRouterHoldoutJob } from "./router-job-pool.js" +import { + runEvidenceRouterJobs, + type EvidenceRouterFitAllJob, + type EvidenceRouterHoldoutJob, +} from "./router-job-pool.js" import { withSqliteBenchmarkStore } from "./sqlite-index.js" import { ROUTER_SEARCH_STRATEGY, @@ -50,9 +54,14 @@ import { optimizeFusionWeightsParallel, optimizeWeightsParallel, evaluateProductionRouter, + type ParallelSearchOptions, type WeightSearchSample, } from "./weight-search.js" -import { getDefaultWorkerCount, resolveWorkerCount } from "./worker-pool.js" +import { + createCandidateEvaluationQueue, + getDefaultWorkerCount, + resolveWorkerCount, +} from "./worker-pool.js" const CONTEXT_BUDGETS = [2_048, 4_096, 8_192, 16_384] as const const EMBEDDING_BATCH_SIZE = 2 @@ -178,6 +187,35 @@ const planEvidenceRouterJobs = ( return jobs } +type RouterSearchJobResult = + | { readonly kind: "holdout"; readonly results: readonly EvidenceRouterSearchResult[] } + | { readonly kind: "fit-all"; readonly results: readonly RecommendedEvidenceRouter[] } + +const runRouterSearchJob = ( + job: EvidenceRouterHoldoutJob | EvidenceRouterFitAllJob, + profile: OptimizationProfile, + options: ParallelSearchOptions, +): Promise => { + if (job.kind === "holdout") + return optimizeEvidenceRouterParallel( + job.model, + job.fusion, + job.strategy, + job.fold, + job.development, + job.validation, + profile, + options, + ).then((results) => ({ kind: "holdout", results })) + return fitRecommendedEvidenceRouterParallel( + job.model, + job.fusion, + job.samples, + profile, + options, + ).then((results) => ({ kind: "fit-all", results })) +} + const isLongInput = (text: string): boolean => Buffer.byteLength(text, "utf8") / 4 > SINGLE_ITEM_ESTIMATED_TOKENS @@ -713,91 +751,68 @@ export const runRetrievalBenchmark = ( resolveWorkerCount(searchOptions.workerCount), getDefaultWorkerCount(), ) - const recommendedJobs: { - readonly model: string - readonly fusion: FusionMethod - readonly samples: readonly WeightSearchSample[] - }[] = [] + const recommendedJobs: EvidenceRouterFitAllJob[] = [] for (const [model, samples] of samplesByModel) for (const fusion of config.routerFusionMethods) - recommendedJobs.push({ model, fusion, samples }) - const canParallelizeRouterJobs = !serialSearch && routerWorkerBudget >= 2 - // Holdouts are complete native-worker jobs; final fits retain their candidate queue. - // Split one budget between both queues so nested pools never oversubscribe the machine. - const holdoutWorkerCount = - !canParallelizeRouterJobs || routerJobs.length === 0 - ? 0 - : Math.min( - routerJobs.length, - recommendedJobs.length === 0 - ? routerWorkerBudget - : Math.max(1, Math.floor(routerWorkerBudget / 2)), - ) - const recommendedWorkerCount = - !canParallelizeRouterJobs || recommendedJobs.length === 0 - ? 0 - : routerJobs.length === 0 - ? Math.min(recommendedJobs.length, routerWorkerBudget) - : Math.max(1, routerWorkerBudget - holdoutWorkerCount) + recommendedJobs.push({ kind: "fit-all", model, fusion, samples }) + const allRouterJobs: readonly (EvidenceRouterHoldoutJob | EvidenceRouterFitAllJob)[] = [ + ...routerJobs, + ...recommendedJobs, + ] + const canParallelizeRouterJobs = + !serialSearch && routerWorkerBudget >= 3 && allRouterJobs.length > 0 + const routerControllerCount = canParallelizeRouterJobs + ? Math.min(allRouterJobs.length, Math.max(1, Math.floor(routerWorkerBudget / 2))) + : 0 + const candidateWorkerCount = canParallelizeRouterJobs + ? Math.max(1, routerWorkerBudget - routerControllerCount) + : 0 + const candidateQueue = canParallelizeRouterJobs + ? yield* runParallelSearch(() => + createCandidateEvaluationQueue({ workerCount: candidateWorkerCount }), + ) + : undefined reportProgress( - `running ${routerJobs.length} complete holdout jobs with ${holdoutWorkerCount} native workers ` + - `and ${recommendedJobs.length} final fits with ${recommendedWorkerCount} pool workers`, + `running ${allRouterJobs.length} evidence-router jobs with ` + + `${routerControllerCount} native controllers and ` + + `${candidateQueue?.workerCount ?? 0} shared candidate workers`, ) - let evidenceRouterSearch: readonly EvidenceRouterSearchResult[] = [] - let recommendedEvidenceRouters: readonly RecommendedEvidenceRouter[] = [] - const holdoutSearch = - canParallelizeRouterJobs && routerJobs.length > 0 + const closeCandidateQueue = + candidateQueue === undefined ? Effect.void : runParallelSearch(() => candidateQueue.close()) + const routerResults = yield* Effect.ensuring( + canParallelizeRouterJobs ? runParallelSearch((signal) => - runEvidenceRouterJobs(routerJobs, optimizationProfile, { - workerCount: holdoutWorkerCount, + runEvidenceRouterJobs(allRouterJobs, optimizationProfile, { + workerCount: routerControllerCount, + candidateQueue, signal, - }), + }).then((results) => + results.map((result): RouterSearchJobResult => { + if (result.kind === "holdout") return { kind: "holdout", results: result.results } + return { kind: "fit-all", results: result.results } + }), + ), ) : Effect.forEach( - routerJobs, + allRouterJobs, (job) => runParallelSearch((signal) => - optimizeEvidenceRouterParallel( - job.model, - job.fusion, - job.strategy, - job.fold, - job.development, - job.validation, - optimizationProfile, - { ...searchOptions, workerCount: 0, signal }, - ), + runRouterSearchJob(job, optimizationProfile, { + ...searchOptions, + workerCount: 0, + signal, + }), ), { concurrency: 1 }, - ) - const recommendedSearch = Effect.forEach( - recommendedJobs, - (job) => - runParallelSearch((signal) => - fitRecommendedEvidenceRouterParallel( - job.model, - job.fusion, - job.samples, - optimizationProfile, - { - ...searchOptions, - workerCount: canParallelizeRouterJobs ? recommendedWorkerCount : 0, - signal, - }, ), - ), - { concurrency: 1 }, + Effect.orDie(closeCandidateQueue), ) - if (canParallelizeRouterJobs) { - const results = yield* Effect.all( - { holdouts: holdoutSearch, recommended: recommendedSearch }, - { concurrency: 2 }, - ) - evidenceRouterSearch = results.holdouts.flat() - recommendedEvidenceRouters = results.recommended.flat() - } else { - evidenceRouterSearch = (yield* holdoutSearch).flat() - recommendedEvidenceRouters = (yield* recommendedSearch).flat() + let evidenceRouterSearch: readonly EvidenceRouterSearchResult[] = [] + let recommendedEvidenceRouters: readonly RecommendedEvidenceRouter[] = [] + for (const result of routerResults) { + if (result.kind === "holdout") + evidenceRouterSearch = [...evidenceRouterSearch, ...result.results] + else recommendedEvidenceRouters = [...recommendedEvidenceRouters, ...result.results] } const evidenceRouterSearchDurationMs = performance.now() - evidenceRouterSearchStartedAt diff --git a/benchmarks/retrieval/types.ts b/benchmarks/retrieval/types.ts index 69efefc..b5eff53 100644 --- a/benchmarks/retrieval/types.ts +++ b/benchmarks/retrieval/types.ts @@ -229,6 +229,24 @@ export interface RecommendedFusionWeights { } /** Search accounting needed to interpret a router candidate and its budget. */ +export interface RouterSearchTimings { + /** Time spent preparing evidence, baselines, and proxy samples. */ + readonly preparationMs: number + /** Time spent selecting the initial static/base weight seeds. */ + readonly baseWeightSearchMs: number + /** Time spent evaluating the random scout baseline. */ + readonly randomSearchMs: number + /** Time spent in initial and coordinate beam rounds. */ + readonly beamSearchMs: number + /** Time spent converting router configs into evaluator candidates. */ + readonly candidatePreparationMs: number + /** Wall time waiting for candidate evaluation results. */ + readonly candidateEvaluationMs: number + /** Time spent ranking, selecting, and archiving evaluated candidates. */ + readonly candidateSelectionMs: number +} + +/** Search diagnostics and timing breakdown for one router search. */ export interface RouterSearchDiagnostics { readonly parameterCount: number readonly parameterLevels: Readonly> @@ -241,6 +259,7 @@ export interface RouterSearchDiagnostics { readonly proxyPromotions: number readonly proxyFullAgreement: number readonly protectedEliteCount: number + readonly timings: RouterSearchTimings } /** Holdout comparison against a deterministic random-search baseline. */ diff --git a/benchmarks/retrieval/weight-search.ts b/benchmarks/retrieval/weight-search.ts index 216a4ee..eaf3d4d 100644 --- a/benchmarks/retrieval/weight-search.ts +++ b/benchmarks/retrieval/weight-search.ts @@ -42,9 +42,11 @@ import { } from "./types.js" import { createCandidateEvaluationPool, + createCandidateEvaluationPoolOnQueue, createEvaluationSnapshot, type CandidateEvaluationPool, type CandidateEvaluationPoolOptions, + type CandidateEvaluationQueue, type EvaluationCandidate, } from "./worker-pool.js" @@ -89,6 +91,8 @@ interface EvidenceSearchSample { /** Native pool controls plus the cancellation signal owned by the benchmark Effect. */ export interface ParallelSearchOptions extends CandidateEvaluationPoolOptions { readonly signal?: AbortSignal + /** Shared candidate scheduler used to interleave multiple router searches. */ + readonly evaluationQueue?: CandidateEvaluationQueue } const preparedFusionCache = new WeakMap< @@ -123,19 +127,20 @@ const createEvaluationPoolForSamples = async ( samples: readonly WeightSearchSample[], fusion: FusionMethod, profile: OptimizationProfile, - options: CandidateEvaluationPoolOptions, -): Promise => - createCandidateEvaluationPool( - createEvaluationSnapshot( - samples.map((sample) => ({ - evaluator: preparedFusionEvaluatorFor(sample, fusion), - targets: sample.targets, - chunks: sample.chunks, - sampleWeight: profile.queryFormWeights[sample.queryKind], - })), - ), - options, + options: ParallelSearchOptions, +): Promise => { + const snapshot = createEvaluationSnapshot( + samples.map((sample) => ({ + evaluator: preparedFusionEvaluatorFor(sample, fusion), + targets: sample.targets, + chunks: sample.chunks, + sampleWeight: profile.queryFormWeights[sample.queryKind], + })), ) + if (options.evaluationQueue !== undefined) + return createCandidateEvaluationPoolOnQueue(snapshot, options.evaluationQueue, options.signal) + return createCandidateEvaluationPool(snapshot, options) +} const routerEvaluationCandidate = ( samples: readonly EvidenceSearchSample[], @@ -864,6 +869,7 @@ const buildSearchDiagnostics = ( ? 1 : stats.proxyAgreementMatches / stats.proxyAgreementComparisons, protectedEliteCount: stats.protectedEliteCount, + timings: { ...stats.timings }, }) const radicalInverse = (index: number, base: number): number => { @@ -975,6 +981,17 @@ interface SearchEvaluationStats { proxyAgreementMatches: number proxyAgreementComparisons: number protectedEliteCount: number + timings: MutableRouterSearchTimings +} + +interface MutableRouterSearchTimings { + preparationMs: number + baseWeightSearchMs: number + randomSearchMs: number + beamSearchMs: number + candidatePreparationMs: number + candidateEvaluationMs: number + candidateSelectionMs: number } interface RouterSearchContext { @@ -1010,11 +1027,15 @@ const selectRandomRouter = async ( pool: CandidateEvaluationPool, baseline: QualitySummary, profile: OptimizationProfile, + stats: SearchEvaluationStats, ): Promise<{ readonly candidate: RouterCandidate; readonly candidates: number }> => { const configs = buildRandomRouterSeeds(baseSeed, parameters) - const qualities = await pool.evaluate( - configs.map((config) => routerEvaluationCandidate(samples, config)), - ) + const candidatePreparationStartedAt = performance.now() + const evaluationCandidates = configs.map((config) => routerEvaluationCandidate(samples, config)) + stats.timings.candidatePreparationMs += performance.now() - candidatePreparationStartedAt + const candidateEvaluationStartedAt = performance.now() + const qualities = await pool.evaluate(evaluationCandidates) + stats.timings.candidateEvaluationMs += performance.now() - candidateEvaluationStartedAt const candidates: RouterCandidate[] = [] for (let index = 0; index < configs.length; index++) { const config = configs[index] @@ -1061,6 +1082,8 @@ interface RouterEvaluationResult { readonly candidates: readonly RouterCandidate[] readonly cacheHits: number readonly evaluations: number + readonly candidatePreparationMs: number + readonly candidateEvaluationMs: number } const finalizeRouterCandidates = ( @@ -1103,9 +1126,14 @@ const evaluateRouterConfigs = async ( qualityCache: Map, ): Promise => { const pending = entries.filter(([key]) => qualityCache.get(key) === undefined) - const qualities = await pool.evaluate( - pending.map(([, config]) => routerEvaluationCandidate(samples, config)), + const candidatePreparationStartedAt = performance.now() + const evaluationCandidates = pending.map(([, config]) => + routerEvaluationCandidate(samples, config), ) + const candidatePreparationMs = performance.now() - candidatePreparationStartedAt + const candidateEvaluationStartedAt = performance.now() + const qualities = await pool.evaluate(evaluationCandidates) + const candidateEvaluationMs = performance.now() - candidateEvaluationStartedAt storeQualityResults( pending, qualities, @@ -1120,6 +1148,8 @@ const evaluateRouterConfigs = async ( }), cacheHits: entries.length - pending.length, evaluations: pending.length, + candidatePreparationMs, + candidateEvaluationMs, } } @@ -1146,6 +1176,9 @@ const rankRouterCandidates = async ( ) context.stats.proxyCacheHits += rankedProxy.cacheHits context.stats.proxyEvaluations += rankedProxy.evaluations + context.stats.timings.candidatePreparationMs += rankedProxy.candidatePreparationMs + context.stats.timings.candidateEvaluationMs += rankedProxy.candidateEvaluationMs + const proxySelectionStartedAt = performance.now() const selectedProxy = selectObjectiveCandidates( rankedProxy.candidates, limit * SEARCH_PROXY_PROMOTION_FACTOR, @@ -1162,6 +1195,7 @@ const rankRouterCandidates = async ( if (protectedKeys.has(key)) selected.set(key, config) } fullCandidates = [...selected] + context.stats.timings.candidateSelectionMs += performance.now() - proxySelectionStartedAt } const rankedFull = await evaluateRouterConfigs( fullCandidates, @@ -1171,7 +1205,12 @@ const rankRouterCandidates = async ( ) context.stats.fullCacheHits += rankedFull.cacheHits context.stats.fullEvaluations += rankedFull.evaluations - return finalizeRouterCandidates(context, rankedFull.candidates, proxyKeys, limit) + context.stats.timings.candidatePreparationMs += rankedFull.candidatePreparationMs + context.stats.timings.candidateEvaluationMs += rankedFull.candidateEvaluationMs + const fullSelectionStartedAt = performance.now() + const selected = finalizeRouterCandidates(context, rankedFull.candidates, proxyKeys, limit) + context.stats.timings.candidateSelectionMs += performance.now() - fullSelectionStartedAt + return selected } interface EvidenceRouterSelection { @@ -1208,6 +1247,7 @@ const prepareRouterSearch = ( samples: readonly WeightSearchSample[], profile: OptimizationProfile, ): RouterSearchPreparation => { + const preparationStartedAt = performance.now() const evidenceSamples = prepareEvidenceSamples(samples) const proxySamples = buildProxySamples(evidenceSamples) const guardrailBaselines = buildGuardrailBaselines(samples, profile) @@ -1233,6 +1273,15 @@ const prepareRouterSearch = ( proxyAgreementMatches: 0, proxyAgreementComparisons: 0, protectedEliteCount: 0, + timings: { + preparationMs: performance.now() - preparationStartedAt, + baseWeightSearchMs: 0, + randomSearchMs: 0, + beamSearchMs: 0, + candidatePreparationMs: 0, + candidateEvaluationMs: 0, + candidateSelectionMs: 0, + }, }, } } @@ -1288,6 +1337,7 @@ const selectBestEvidenceRouter = async ( proxyPool, } const profileSeed = profileSeedFor(profile) + const baseSearchStartedAt = performance.now() const baseWeights = await selectBestWeights( fullPool, "reranker-top20", @@ -1297,9 +1347,11 @@ const selectBestEvidenceRouter = async ( const subsetWeights = await selectBestWeightsPerSubset(fullPool, profile) const baseSeeds = [baseWeights.weights, ...subsetWeights].map(emptyRouterConfig) baseSeeds.unshift(profileSeed) + stats.timings.baseWeightSearchMs += performance.now() - baseSearchStartedAt const parameters = routerParameters() const randomBaseSeed = baseSeeds[0] if (randomBaseSeed === undefined) throw new Error("Evidence router search has no base seed") + const randomSearchStartedAt = performance.now() const randomSearch = await selectRandomRouter( evidenceSamples, randomBaseSeed, @@ -1307,7 +1359,10 @@ const selectBestEvidenceRouter = async ( fullPool, productionQuality, profile, + stats, ) + stats.timings.randomSearchMs += performance.now() - randomSearchStartedAt + const beamSearchStartedAt = performance.now() let beam = await rankRouterCandidates( searchContext, [...baseSeeds, ...buildGlobalRouterSeeds(baseSeeds, parameters)], @@ -1331,6 +1386,7 @@ const selectBestEvidenceRouter = async ( ) } } + stats.timings.beamSearchMs += performance.now() - beamSearchStartedAt const fallback = beam[0] const candidates = [...searchContext.archive.values()] const selections = ROUTER_OBJECTIVES.map((objective) => { diff --git a/benchmarks/retrieval/worker-pool.ts b/benchmarks/retrieval/worker-pool.ts index ce5b818..0d848ef 100644 --- a/benchmarks/retrieval/worker-pool.ts +++ b/benchmarks/retrieval/worker-pool.ts @@ -9,6 +9,7 @@ import type { QualitySummary } from "./types.js" import { hasWorkerMessageType } from "./worker-message.js" const DEFAULT_BATCH_SIZE = 32 +const DEFAULT_QUEUE_BATCH_SIZE = 1 const DEFAULT_WORKER_URL = new URL("./fusion-worker.mjs", import.meta.url) /** One prepared benchmark sample represented without source text for worker transfer. */ @@ -74,6 +75,35 @@ export interface CandidateEvaluationPool { readonly stats: () => CandidateEvaluationPoolStats } +/** Configuration for one shared candidate-evaluation queue. */ +export interface CandidateEvaluationQueueOptions { + /** Maximum number of native evaluation workers. */ + readonly workerCount?: number + /** Maximum number of candidate vectors sent in one worker message. */ + readonly batchSize?: number + /** Test-only or diagnostic override for the queue worker entry point. */ + readonly workerUrl?: URL +} + +/** Shared native executor used by multiple independent router searches. */ +export interface CandidateEvaluationQueue { + /** Number of native workers owned by the queue. */ + readonly workerCount: number + /** Maximum number of candidate vectors sent in one worker message. */ + readonly batchSize: number + /** Number of workers currently evaluating a candidate batch. */ + readonly activeWorkerCount: () => number + /** Enqueue candidates for one prepared snapshot and preserve candidate order in the result. */ + readonly evaluate: ( + snapshot: EvaluationSnapshot, + candidates: readonly EvaluationCandidate[], + signal?: AbortSignal, + snapshotId?: string, + ) => Promise + /** Stop all workers and reject queued evaluations. */ + readonly close: () => Promise +} + interface WorkerReadyMessage { readonly type: "ready" } @@ -98,11 +128,14 @@ interface WorkerTask { readonly end: number } -interface WorkerSlot { +interface ReadyWorker { readonly worker: Worker readonly ready: Promise readonly resolveReady: () => void readonly rejectReady: (error: Error) => void +} + +interface WorkerSlot extends ReadyWorker { busy: boolean task?: WorkerTask } @@ -120,6 +153,118 @@ interface ActiveEvaluation { const errorFromUnknown = (cause: unknown): Error => cause instanceof Error ? cause : new Error(String(cause)) +const createReadyWorker = (workerUrl: URL): ReadyWorker => { + let resolveReady: () => void = () => undefined + let rejectReady: (error: Error) => void = () => undefined + const ready = new Promise((resolve, reject) => { + resolveReady = resolve + rejectReady = reject + }) + return { worker: new Worker(workerUrl), ready, resolveReady, rejectReady } +} + +const attachWorkerLifecycle = ( + worker: Worker, + onMessage: (message: unknown) => void, + onError: (cause: Error) => void, + onExit: (code: number) => void, +): void => { + worker.on("message", onMessage) + worker.on("error", onError) + worker.on("exit", onExit) +} + +const attachCandidateWorkerLifecycle = ( + slot: ReadyWorker, + onMessage: (message: unknown) => void, + onError: (cause: Error) => void, + isClosed: () => boolean, + workerName: string, +): void => + attachWorkerLifecycle(slot.worker, onMessage, onError, (code) => { + if (!isClosed()) onError(new Error(`${workerName} worker exited with code ${code}`)) + }) + +const attachFusionWorkerLifecycle = ( + slot: ReadyWorker, + onMessage: (message: unknown) => void, + onError: (cause: Error) => void, + isClosed: () => boolean, +): void => attachCandidateWorkerLifecycle(slot, onMessage, onError, isClosed, "Fusion") + +const attachFusionQueueWorkerLifecycle = ( + slot: ReadyWorker, + onMessage: (message: unknown) => void, + onError: (cause: Error) => void, + isClosed: () => boolean, +): void => attachCandidateWorkerLifecycle(slot, onMessage, onError, isClosed, "Fusion queue") + +const createWorkerSlots = ( + workerCount: number, + workerUrl: URL, + createSlot: (readyWorker: ReadyWorker) => T, + attach: (slot: T) => void, +): T[] => { + const slots: T[] = [] + for (let index = 0; index < workerCount; index++) { + const slot = createSlot(createReadyWorker(workerUrl)) + attach(slot) + slots.push(slot) + } + return slots +} + +const handleWorkerControlMessage = ( + slot: ReadyWorker, + message: { readonly type: string; readonly message?: string }, + onError: (cause: Error) => void, +): boolean => { + if (message.type === "ready") { + slot.resolveReady() + return true + } + if (message.type === "error") { + onError(new Error(message.message ?? "Worker reported an error")) + return true + } + return false +} + +const failWorker = ( + slot: ReadyWorker, + cause: Error, + rejectActive: () => void, + close: () => Promise, +): void => { + slot.rejectReady(cause) + rejectActive() + void close() +} + +const candidatePoolStats = ( + mode: EvaluationPoolMode, + workerCount: number, + activeWorkerCount: number, + batchSize: number, + batches: number, + candidates: number, +): CandidateEvaluationPoolStats => ({ + mode, + workerCount, + activeWorkerCount, + batchSize, + batches, + candidates, +}) + +const candidateBatchStats = ( + candidates: readonly EvaluationCandidate[], + batchSize: number, +): { readonly batches: number; readonly candidates: number } => ({ + batches: Math.ceil(candidates.length / batchSize), + candidates: candidates.length, +}) + const parseInteger = (value: string | undefined): number | undefined => { if (value === undefined || value.trim() === "") return undefined const parsed = Number(value) @@ -141,6 +286,11 @@ const resolveBatchSize = (requested?: number): number => { return Math.max(1, configured ?? DEFAULT_BATCH_SIZE) } +const resolveQueueBatchSize = (requested?: number): number => { + const configured = requested ?? parseInteger(process.env.PIX_BENCH_WORKER_BATCH_SIZE) + return Math.max(1, configured ?? DEFAULT_QUEUE_BATCH_SIZE) +} + const contextTokens = (chunk: Chunk): number => Math.ceil( Buffer.byteLength(`${chunk.file}:${chunk.startLine}-${chunk.endLine}\n${chunk.text}`, "utf8") / @@ -202,30 +352,18 @@ class NativeCandidateEvaluationPool implements CandidateEvaluationPool { ) { this.workerCount = workerCount this.batchSize = batchSize - this.slots = [] - for (let index = 0; index < workerCount; index++) { - let resolveReady: () => void = () => undefined - let rejectReady: (error: Error) => void = () => undefined - const ready = new Promise((resolve, reject) => { - resolveReady = resolve - rejectReady = reject - }) - const worker = new Worker(workerUrl) - const slot: WorkerSlot = { - worker, - ready, - resolveReady, - rejectReady, - busy: false, - } - worker.on("message", (message: unknown) => this.handleMessage(slot, message)) - worker.on("error", (cause: Error) => this.handleWorkerError(slot, cause)) - worker.on("exit", (code: number) => { - if (!this.closed) - this.handleWorkerError(slot, new Error(`Fusion worker exited with code ${code}`)) - }) - this.slots.push(slot) - } + this.slots = createWorkerSlots( + workerCount, + workerUrl, + (readyWorker): WorkerSlot => ({ ...readyWorker, busy: false }), + (slot) => + attachFusionWorkerLifecycle( + slot, + (message) => this.handleMessage(slot, message), + (cause) => this.handleWorkerError(slot, cause), + () => this.closed, + ), + ) this.ready = Promise.all(this.slots.map((slot) => slot.ready)).then(() => undefined) for (const slot of this.slots) slot.worker.postMessage({ type: "init", snapshot }) } @@ -250,13 +388,18 @@ class NativeCandidateEvaluationPool implements CandidateEvaluationPool { private handleWorkerError(slot: WorkerSlot, cause: Error): void { if (this.closed) return - slot.rejectReady(cause) const active = this.active - if (active !== undefined) { - this.active = undefined - active.reject(cause) - } - void this.close() + failWorker( + slot, + cause, + () => { + if (active !== undefined) { + this.active = undefined + active.reject(cause) + } + }, + () => this.close(), + ) } private handleMessage(slot: WorkerSlot, message: unknown): void { @@ -264,14 +407,9 @@ class NativeCandidateEvaluationPool implements CandidateEvaluationPool { this.handleWorkerError(slot, new Error("Fusion worker sent an invalid message")) return } - if (message.type === "ready") { - slot.resolveReady() + if (handleWorkerControlMessage(slot, message, (cause) => this.handleWorkerError(slot, cause))) return - } - if (message.type === "error") { - this.handleWorkerError(slot, new Error(message.message)) - return - } + if (message.type !== "result") return const active = this.active const task = active?.tasks.get(message.taskId) if (active === undefined || task === undefined) { @@ -381,6 +519,409 @@ class NativeCandidateEvaluationPool implements CandidateEvaluationPool { }) } +interface QueueWorkerReadyMessage { + readonly type: "ready" +} + +interface QueueWorkerResultMessage { + readonly type: "result" + readonly taskId: number + readonly results: readonly QualitySummary[] +} + +interface QueueWorkerErrorMessage { + readonly type: "error" + readonly taskId?: number + readonly message: string +} + +type QueueWorkerMessage = + | QueueWorkerReadyMessage + | QueueWorkerResultMessage + | QueueWorkerErrorMessage + +interface QueueWorkerSlot extends ReadyWorker { + readonly knownSnapshots: Set + busy: boolean +} + +interface CandidateQueueRequest { + readonly requestId: number + readonly snapshotId: string + readonly snapshot: EvaluationSnapshot + readonly candidates: readonly EvaluationCandidate[] + readonly results: Map + readonly resolve: (results: readonly QualitySummary[]) => void + readonly reject: (error: Error) => void + readonly removeAbortListener: () => void + nextCandidate: number + completedCandidates: number + enqueued: boolean + settled: boolean +} + +interface CandidateQueueTask { + readonly taskId: number + readonly request: CandidateQueueRequest + readonly slot: QueueWorkerSlot + readonly start: number + readonly end: number +} + +const isQueueRecord = (message: unknown): message is Record => + typeof message === "object" && message !== null + +const isQueueWorkerMessage = (message: unknown): message is QueueWorkerMessage => { + if (!isQueueRecord(message) || typeof message.type !== "string") return false + if (typeof message.type !== "string") return false + if (message.type === "ready") return true + if (message.type === "error") return typeof message.message === "string" + return ( + Number.isInteger(message.taskId) && Array.isArray(message.results) && message.results.length > 0 + ) +} + +class NativeCandidateEvaluationQueue implements CandidateEvaluationQueue { + readonly workerCount: number + readonly batchSize: number + + private readonly slots: QueueWorkerSlot[] + private readonly ready: Promise + private readonly requests = new Set() + private readonly readyRequests: CandidateQueueRequest[] = [] + private readonly activeTasks = new Map() + private readonly snapshotIds = new WeakMap() + private closePromise: Promise | undefined + private closed = false + private nextRequestId = 0 + private nextSnapshotId = 0 + private nextTaskId = 0 + + private constructor(workerCount: number, batchSize: number, workerUrl: URL) { + this.workerCount = workerCount + this.batchSize = batchSize + this.slots = createWorkerSlots( + workerCount, + workerUrl, + (readyWorker): QueueWorkerSlot => ({ + ...readyWorker, + knownSnapshots: new Set(), + busy: false, + }), + (slot) => + attachFusionQueueWorkerLifecycle( + slot, + (message) => this.handleMessage(slot, message), + (cause) => this.handleWorkerError(slot, cause), + () => this.closed, + ), + ) + this.ready = Promise.all(this.slots.map((slot) => slot.ready)).then(() => undefined) + } + + /** Start a queue after every native worker has completed protocol startup. */ + static async create( + workerCount: number, + batchSize: number, + workerUrl: URL, + ): Promise { + let queue: NativeCandidateEvaluationQueue | undefined + try { + queue = new NativeCandidateEvaluationQueue(workerCount, batchSize, workerUrl) + await queue.ready + return queue + } catch (cause) { + if (queue !== undefined) await queue.close() + throw errorFromUnknown(cause) + } + } + + private handleWorkerError(slot: QueueWorkerSlot, cause: Error): void { + if (this.closed) return + failWorker( + slot, + cause, + () => { + for (const request of this.requests) this.settleRequest(request, cause) + }, + () => this.close(), + ) + } + + private snapshotIdFor(snapshot: EvaluationSnapshot, requestedId?: string): string { + if (requestedId !== undefined) return requestedId + const existing = this.snapshotIds.get(snapshot) + if (existing !== undefined) return existing + const snapshotId = String(this.nextSnapshotId++) + this.snapshotIds.set(snapshot, snapshotId) + return snapshotId + } + + private enqueueRequest(request: CandidateQueueRequest): void { + if (request.settled || request.enqueued || request.nextCandidate >= request.candidates.length) + return + request.enqueued = true + this.readyRequests.push(request) + } + + private nextRequest(): CandidateQueueRequest | undefined { + while (this.readyRequests.length > 0) { + const request = this.readyRequests.shift() + if (request === undefined) return undefined + request.enqueued = false + if (request.settled || request.nextCandidate >= request.candidates.length) continue + return request + } + return undefined + } + + private settleRequest(request: CandidateQueueRequest, cause?: Error): void { + if (request.settled) return + request.settled = true + this.requests.delete(request) + request.removeAbortListener() + if (cause !== undefined) { + request.reject(cause) + return + } + const ordered: QualitySummary[] = [] + for (let index = 0; index < request.candidates.length; index++) { + const result = request.results.get(index) + if (result === undefined) { + request.reject(new Error(`Fusion queue omitted candidate ${index}`)) + return + } + ordered.push(result) + } + request.resolve(ordered) + } + + private dispatch(): void { + if (this.closed) return + for (const slot of this.slots) { + if (slot.busy) continue + const request = this.nextRequest() + if (request === undefined) return + const start = request.nextCandidate + const end = Math.min(start + this.batchSize, request.candidates.length) + const task: CandidateQueueTask = { + taskId: this.nextTaskId++, + request, + slot, + start, + end, + } + request.nextCandidate = end + this.activeTasks.set(task.taskId, task) + slot.busy = true + const includeSnapshot = !slot.knownSnapshots.has(request.snapshotId) + try { + slot.worker.postMessage({ + type: "evaluate", + taskId: task.taskId, + snapshotId: request.snapshotId, + snapshot: includeSnapshot ? request.snapshot : undefined, + candidates: request.candidates.slice(start, end), + }) + if (includeSnapshot) slot.knownSnapshots.add(request.snapshotId) + } catch (cause) { + this.activeTasks.delete(task.taskId) + slot.busy = false + this.handleWorkerError(slot, errorFromUnknown(cause)) + return + } + this.enqueueRequest(request) + } + } + + private handleMessage(slot: QueueWorkerSlot, message: unknown): void { + if (this.closed) return + if (!isQueueWorkerMessage(message)) { + this.handleWorkerError(slot, new Error("Fusion queue worker sent an invalid message")) + return + } + if (handleWorkerControlMessage(slot, message, (cause) => this.handleWorkerError(slot, cause))) + return + if (message.type !== "result") return + const task = this.activeTasks.get(message.taskId) + if (task === undefined || task.slot !== slot) { + this.handleWorkerError( + slot, + new Error(`Fusion queue returned unknown task ${message.taskId}`), + ) + return + } + this.activeTasks.delete(message.taskId) + slot.busy = false + if (message.results.length !== task.end - task.start) { + this.handleWorkerError(slot, new Error("Fusion queue returned an invalid result length")) + return + } + if (!task.request.settled) { + for (let index = 0; index < message.results.length; index++) + task.request.results.set(task.start + index, message.results[index]) + task.request.completedCandidates += message.results.length + if (task.request.completedCandidates === task.request.candidates.length) + this.settleRequest(task.request) + else this.enqueueRequest(task.request) + } + this.dispatch() + } + + activeWorkerCount = (): number => + this.closed ? 0 : this.slots.filter((slot) => slot.busy).length + + async evaluate( + snapshot: EvaluationSnapshot, + candidates: readonly EvaluationCandidate[], + signal?: AbortSignal, + snapshotId?: string, + ): Promise { + if (this.closed) throw new Error("Candidate evaluation queue is closed") + if (signal?.aborted) throw new Error("Candidate evaluation queue was interrupted") + if (candidates.length === 0) return [] + await this.ready + if (this.closed) throw new Error("Candidate evaluation queue is closed") + const resolvedSnapshotId = this.snapshotIdFor(snapshot, snapshotId) + return new Promise((resolve, reject) => { + let request: CandidateQueueRequest | undefined + const abort = () => { + if (request === undefined) return + this.settleRequest(request, new Error("Candidate evaluation queue was interrupted")) + this.dispatch() + } + const removeAbortListener = () => signal?.removeEventListener("abort", abort) + request = { + requestId: this.nextRequestId++, + snapshotId: resolvedSnapshotId, + snapshot, + candidates, + results: new Map(), + resolve, + reject, + removeAbortListener, + nextCandidate: 0, + completedCandidates: 0, + enqueued: false, + settled: false, + } + if (signal !== undefined) { + signal.addEventListener("abort", abort, { once: true }) + if (signal.aborted) { + abort() + return + } + } + this.requests.add(request) + this.enqueueRequest(request) + this.dispatch() + }) + } + + async close(): Promise { + if (this.closePromise !== undefined) return this.closePromise + this.closed = true + for (const request of this.requests) + this.settleRequest(request, new Error("Candidate evaluation queue closed")) + this.readyRequests.length = 0 + for (const slot of this.slots) slot.busy = false + this.closePromise = Promise.all( + this.slots.map((slot) => slot.worker.terminate().catch(() => -1)), + ).then(() => undefined) + return this.closePromise + } +} + +class SerialCandidateEvaluationQueue implements CandidateEvaluationQueue { + readonly workerCount = 1 + readonly batchSize: number + private closed = false + + constructor(batchSize: number) { + this.batchSize = batchSize + } + + activeWorkerCount = (): number => 0 + + async evaluate( + snapshot: EvaluationSnapshot, + candidates: readonly EvaluationCandidate[], + signal?: AbortSignal, + _snapshotId?: string, + ): Promise { + if (this.closed) throw new Error("Candidate evaluation queue is closed") + if (signal?.aborted) throw new Error("Candidate evaluation queue was interrupted") + return evaluateCandidatesSerial(snapshot, candidates) + } + + async close(): Promise { + this.closed = true + } +} + +class QueuedCandidateEvaluationPool implements CandidateEvaluationPool { + readonly mode = "parallel" as const + readonly workerCount: number + readonly batchSize: number + private readonly snapshot: EvaluationSnapshot + private readonly queue: CandidateEvaluationQueue + private readonly signal: AbortSignal | undefined + private closed = false + private batchCount = 0 + private candidateCount = 0 + + constructor(snapshot: EvaluationSnapshot, queue: CandidateEvaluationQueue, signal?: AbortSignal) { + this.snapshot = snapshot + this.queue = queue + this.signal = signal + this.workerCount = queue.workerCount + this.batchSize = queue.batchSize + } + + async evaluate(candidates: readonly EvaluationCandidate[]): Promise { + if (this.closed) throw new Error("Candidate evaluation pool is closed") + const counts = candidateBatchStats(candidates, this.batchSize) + this.candidateCount += counts.candidates + this.batchCount += counts.batches + return this.queue.evaluate(this.snapshot, candidates, this.signal) + } + + async close(): Promise { + this.closed = true + } + + stats = (): CandidateEvaluationPoolStats => + candidatePoolStats( + this.mode, + this.workerCount, + this.closed ? 0 : this.queue.activeWorkerCount(), + this.batchSize, + this.batchCount, + this.candidateCount, + ) +} + +/** Create one shared candidate queue for multiple independent router searches. */ +export const createCandidateEvaluationQueue = async ( + options: CandidateEvaluationQueueOptions = {}, +): Promise => { + const workerCount = resolveWorkerCount(options.workerCount) + const batchSize = resolveQueueBatchSize(options.batchSize) + if (workerCount <= 1) return new SerialCandidateEvaluationQueue(batchSize) + return NativeCandidateEvaluationQueue.create( + workerCount, + batchSize, + options.workerUrl ?? DEFAULT_WORKER_URL, + ) +} + +/** Create a per-search pool facade that submits work to a shared candidate queue. */ +export const createCandidateEvaluationPoolOnQueue = ( + snapshot: EvaluationSnapshot, + queue: CandidateEvaluationQueue, + signal?: AbortSignal, +): CandidateEvaluationPool => new QueuedCandidateEvaluationPool(snapshot, queue, signal) + class SerialCandidateEvaluationPool implements CandidateEvaluationPool { readonly mode = "serial" as const readonly workerCount = 1 @@ -397,8 +938,9 @@ class SerialCandidateEvaluationPool implements CandidateEvaluationPool { async evaluate(candidates: readonly EvaluationCandidate[]): Promise { if (this.closed) throw new Error("Candidate evaluation pool is closed") - this.candidateCount += candidates.length - this.batchCount += Math.ceil(candidates.length / this.batchSize) + const counts = candidateBatchStats(candidates, this.batchSize) + this.candidateCount += counts.candidates + this.batchCount += counts.batches return evaluateCandidatesSerial(this.snapshot, candidates) } @@ -406,14 +948,15 @@ class SerialCandidateEvaluationPool implements CandidateEvaluationPool { this.closed = true } - stats = (): CandidateEvaluationPoolStats => ({ - mode: this.mode, - workerCount: this.workerCount, - activeWorkerCount: 0, - batchSize: this.batchSize, - batches: this.batchCount, - candidates: this.candidateCount, - }) + stats = (): CandidateEvaluationPoolStats => + candidatePoolStats( + this.mode, + this.workerCount, + 0, + this.batchSize, + this.batchCount, + this.candidateCount, + ) } /** Create a fixed benchmark evaluator pool; one worker or an explicit zero uses serial fallback. */ diff --git a/benchmarks/tests/worker-pool.test.ts b/benchmarks/tests/worker-pool.test.ts index 74a9825..5592d02 100644 --- a/benchmarks/tests/worker-pool.test.ts +++ b/benchmarks/tests/worker-pool.test.ts @@ -7,6 +7,7 @@ import { prepareFusion } from "../retrieval/fusion.js" import { SEARCH_PRIORITY_PROFILE } from "../retrieval/optimization-profiles.js" import { runEvidenceRouterJobs, + type EvidenceRouterFitAllJob, type EvidenceRouterHoldoutJob, } from "../retrieval/router-job-pool.js" import { @@ -21,6 +22,8 @@ import { } from "../retrieval/weight-search.js" import { createCandidateEvaluationPool, + createCandidateEvaluationPoolOnQueue, + createCandidateEvaluationQueue, createEvaluationSnapshot, evaluateCandidatesSerial, getDefaultWorkerCount, @@ -92,9 +95,18 @@ const searchSample = { chunks, } +const withoutSearchTimings = < + T extends { readonly searchDiagnostics: { readonly timings: object } }, +>( + result: T, +) => { + const { timings: _timings, ...searchDiagnostics } = result.searchDiagnostics + return { ...result, searchDiagnostics } +} + describe("benchmark candidate evaluation pool", () => { - it("runs complete router jobs in native workers and preserves results", async () => { - const job: EvidenceRouterHoldoutJob = { + it("runs holdout and fit-all jobs through one native queue", async () => { + const holdoutJob: EvidenceRouterHoldoutJob = { kind: "holdout", model: "fixture", fusion: "dbsf", @@ -103,20 +115,54 @@ describe("benchmark candidate evaluation pool", () => { development: [searchSample], validation: [searchSample], } - const [parallel] = await runEvidenceRouterJobs([job], SEARCH_PRIORITY_PROFILE, { - workerCount: 1, - }) - const serial = await optimizeEvidenceRouter( - job.model, - job.fusion, - job.strategy, - job.fold, - job.development, - job.validation, - SEARCH_PRIORITY_PROFILE, - ) + const fitAllJob: EvidenceRouterFitAllJob = { + kind: "fit-all", + model: "fixture", + fusion: "dbsf", + samples: [searchSample], + } + const candidateQueue = await createCandidateEvaluationQueue({ workerCount: 2 }) + try { + const [parallelHoldout, parallelFitAll] = await runEvidenceRouterJobs( + [holdoutJob, fitAllJob], + SEARCH_PRIORITY_PROFILE, + { + workerCount: 2, + candidateQueue, + }, + ) + const serialHoldout = await optimizeEvidenceRouter( + holdoutJob.model, + holdoutJob.fusion, + holdoutJob.strategy, + holdoutJob.fold, + holdoutJob.development, + holdoutJob.validation, + SEARCH_PRIORITY_PROFILE, + ) + const serialFitAll = await fitRecommendedEvidenceRouter( + fitAllJob.model, + fitAllJob.fusion, + fitAllJob.samples, + SEARCH_PRIORITY_PROFILE, + ) - expect(parallel).toEqual(serial) + expect(parallelHoldout.jobId).toBe(0) + expect(parallelHoldout.kind).toBe("holdout") + expect(parallelHoldout.results.map(withoutSearchTimings)).toEqual( + serialHoldout.map(withoutSearchTimings), + ) + expect(parallelFitAll.jobId).toBe(1) + expect(parallelFitAll.kind).toBe("fit-all") + expect(parallelFitAll.results.map(withoutSearchTimings)).toEqual( + serialFitAll.map(withoutSearchTimings), + ) + expect( + parallelHoldout.results[0]?.searchDiagnostics.timings.candidateEvaluationMs, + ).toBeGreaterThan(0) + } finally { + await candidateQueue.close() + } }) it("derives a bounded default and honors explicit sizing", () => { @@ -147,6 +193,52 @@ describe("benchmark candidate evaluation pool", () => { expect(pool.stats().activeWorkerCount).toBe(0) }) + it("shares one queue across independent snapshots without changing results", async () => { + const queue = await createCandidateEvaluationQueue({ workerCount: 2, batchSize: 1 }) + const secondSnapshot = createEvaluationSnapshot([ + { + evaluator: prepareFusion("dbsf", rankings, 10), + targets: [new Set([0])], + chunks, + sampleWeight: 1, + }, + { + evaluator: prepareFusion("dbsf", rankings, 10), + targets: [new Set([3])], + chunks, + sampleWeight: 2, + }, + ]) + const firstPool = createCandidateEvaluationPoolOnQueue(snapshot, queue) + const secondPool = createCandidateEvaluationPoolOnQueue(secondSnapshot, queue) + try { + const [firstResults, secondResults] = await Promise.all([ + firstPool.evaluate(candidates), + secondPool.evaluate(candidates), + ]) + expect(firstResults).toEqual(evaluateCandidatesSerial(snapshot, candidates)) + expect(secondResults).toEqual(evaluateCandidatesSerial(secondSnapshot, candidates)) + expect(queue.activeWorkerCount()).toBe(0) + } finally { + await firstPool.close() + await secondPool.close() + await queue.close() + } + }) + + it("rejects an aborted queue request without closing the shared queue", async () => { + const queue = await createCandidateEvaluationQueue({ workerCount: 2, batchSize: 1 }) + const controller = new AbortController() + try { + const pending = queue.evaluate(snapshot, candidates, controller.signal) + controller.abort() + await expect(pending).rejects.toThrow("interrupted") + await expect(queue.evaluate(snapshot, candidates.slice(0, 1))).resolves.toHaveLength(1) + } finally { + await queue.close() + } + }) + it("keeps worker metrics aligned with canonical summarization", () => { const parityRankings = { ...rankings, From 67529d2b5fcf5aee287fecf8cbd672f23eb7fb4e Mon Sep 17 00:00:00 2001 From: Lucas Burmeister Date: Tue, 4 Aug 2026 17:51:03 +0200 Subject: [PATCH 5/8] perf: record benchmark pool timings --- CONTEXT.md | 4 +++- benchmarks/README.md | 3 ++- benchmarks/retrieval/report.ts | 6 +++--- benchmarks/retrieval/runner.ts | 28 ++++++++++++++++++++------- benchmarks/retrieval/types.ts | 8 +++++++- benchmarks/retrieval/weight-search.ts | 18 ++++++++++++++--- benchmarks/tests/retrieval.test.ts | 7 ++++++- 7 files changed, 57 insertions(+), 17 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 2ad398e..1717730 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -255,7 +255,9 @@ sends the prepared snapshot once per worker, batches candidate weight vectors, a selection state on the controller thread. The search APIs are async so serial and worker execution share the same selection algorithm; `PIX_BENCH_WORKERS=0` selects serial evaluation, `workerCount: 1` selects the serial pool through `createCandidateEvaluationPool`, and the runner's `PIX_BENCH_SEARCH_MODE=serial` selects -that serial pool mode. `PIX_BENCH_WORKER_BATCH_SIZE` bounds ordinary worker messages. +that serial pool mode. `PIX_BENCH_WORKER_BATCH_SIZE` bounds ordinary worker messages. Router search +diagnostics retain candidate snapshot and pool initialization time separately from candidate evaluation time; +benchmark timings also record shared candidate-queue startup and shutdown. Evidence-router grouped-fold, repository-holdout, and fit-all jobs use native controller workers backed by one shared candidate-evaluation queue. Controllers keep candidate generation, beam rounds, caches, and archive updates local and sequential; candidate batches are work-stealing tasks evaluated by the shared native queue. diff --git a/benchmarks/README.md b/benchmarks/README.md index 7c1f1d8..63d7317 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -339,7 +339,8 @@ output size without introducing an LLM or provider-specific tokenizer. Each run writes ignored JSON and Markdown artifacts under `benchmarks/results`. JSON rows retain the repository, revision, language, size, category, difficulty, query form, grouped fold, model, variant, -individual gold ranks, timing, and every metric. Schema 22 stores each authored query and its exact +individual gold ranks, timing, and every metric. Schema 23 adds shared candidate-queue lifecycle and +per-router candidate-pool initialization timings. Each artifact stores each authored query and its exact file-qualified ground truth once, records productive Sparse timings, and adds the fixed equal-weight RRF baseline. The Markdown report includes quality by query form, marginal leave-one-channel-out contribution, cross-validation folds, Shapley values, and final fitted diff --git a/benchmarks/retrieval/report.ts b/benchmarks/retrieval/report.ts index 02c994a..1843a79 100644 --- a/benchmarks/retrieval/report.ts +++ b/benchmarks/retrieval/report.ts @@ -86,9 +86,9 @@ export const renderMarkdownReport = (artifact: BenchmarkArtifact): string => { "", "Compute timings exclude JSON and Markdown artifact serialization.", "", - "| Total | Corpus preparation | Embedding | Retrieval | Weight search | Fusion search | Router search |", - "| ---: | ---: | ---: | ---: | ---: | ---: | ---: |", - `| ${duration(artifact.timings.totalDurationMs)} | ${duration(artifact.timings.corpusPreparationDurationMs)} | ${duration(artifact.timings.embeddingDurationMs)} | ${duration(artifact.timings.retrievalDurationMs)} | ${duration(artifact.timings.weightSearchDurationMs)} | ${duration(artifact.timings.fusionSearchDurationMs)} | ${duration(artifact.timings.evidenceRouterSearchDurationMs)} |`, + "| Total | Corpus preparation | Embedding | Retrieval | Weight search | Fusion search | Router search | Candidate queue startup | Candidate queue shutdown |", + "| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", + `| ${duration(artifact.timings.totalDurationMs)} | ${duration(artifact.timings.corpusPreparationDurationMs)} | ${duration(artifact.timings.embeddingDurationMs)} | ${duration(artifact.timings.retrievalDurationMs)} | ${duration(artifact.timings.weightSearchDurationMs)} | ${duration(artifact.timings.fusionSearchDurationMs)} | ${duration(artifact.timings.evidenceRouterSearchDurationMs)} | ${duration(artifact.timings.candidateQueueStartupDurationMs)} | ${duration(artifact.timings.candidateQueueShutdownDurationMs)} |`, "", "## Sparse Encoder", "", diff --git a/benchmarks/retrieval/runner.ts b/benchmarks/retrieval/runner.ts index a79c334..e95a96b 100644 --- a/benchmarks/retrieval/runner.ts +++ b/benchmarks/retrieval/runner.ts @@ -61,6 +61,7 @@ import { createCandidateEvaluationQueue, getDefaultWorkerCount, resolveWorkerCount, + type CandidateEvaluationQueue, } from "./worker-pool.js" const CONTEXT_BUDGETS = [2_048, 4_096, 8_192, 16_384] as const @@ -399,6 +400,8 @@ export const runRetrievalBenchmark = ( >() const samplesByModel = new Map() let retrievalDurationMs = 0 + let candidateQueueStartupDurationMs = 0 + let candidateQueueShutdownDurationMs = 0 for (const manifest of manifests) { const repositoryPath = yield* prepareRepository(manifest) @@ -767,18 +770,27 @@ export const runRetrievalBenchmark = ( const candidateWorkerCount = canParallelizeRouterJobs ? Math.max(1, routerWorkerBudget - routerControllerCount) : 0 - const candidateQueue = canParallelizeRouterJobs - ? yield* runParallelSearch(() => - createCandidateEvaluationQueue({ workerCount: candidateWorkerCount }), - ) - : undefined + let candidateQueue: CandidateEvaluationQueue | undefined + if (canParallelizeRouterJobs) { + const candidateQueueStartedAt = performance.now() + candidateQueue = yield* runParallelSearch(() => + createCandidateEvaluationQueue({ workerCount: candidateWorkerCount }), + ) + candidateQueueStartupDurationMs = performance.now() - candidateQueueStartedAt + } reportProgress( `running ${allRouterJobs.length} evidence-router jobs with ` + `${routerControllerCount} native controllers and ` + `${candidateQueue?.workerCount ?? 0} shared candidate workers`, ) const closeCandidateQueue = - candidateQueue === undefined ? Effect.void : runParallelSearch(() => candidateQueue.close()) + candidateQueue === undefined + ? Effect.void + : Effect.gen(function* () { + const candidateQueueStartedAt = performance.now() + yield* runParallelSearch(() => candidateQueue!.close()) + candidateQueueShutdownDurationMs = performance.now() - candidateQueueStartedAt + }) const routerResults = yield* Effect.ensuring( canParallelizeRouterJobs ? runParallelSearch((signal) => @@ -826,7 +838,7 @@ export const runRetrievalBenchmark = ( ) const artifact: BenchmarkArtifact = { - schemaVersion: 22, + schemaVersion: 23, benchmarkProfile: profile, optimizationProfile, validationProtocol: { @@ -849,6 +861,8 @@ export const runRetrievalBenchmark = ( weightSearchDurationMs, fusionSearchDurationMs, evidenceRouterSearchDurationMs, + candidateQueueStartupDurationMs, + candidateQueueShutdownDurationMs, }, chunkConfig: { chunkLines: DEFAULT_CONFIG.chunkLines, diff --git a/benchmarks/retrieval/types.ts b/benchmarks/retrieval/types.ts index b5eff53..72f12d6 100644 --- a/benchmarks/retrieval/types.ts +++ b/benchmarks/retrieval/types.ts @@ -232,6 +232,8 @@ export interface RecommendedFusionWeights { export interface RouterSearchTimings { /** Time spent preparing evidence, baselines, and proxy samples. */ readonly preparationMs: number + /** Time spent preparing candidate snapshots and starting their evaluation pool. */ + readonly candidatePoolInitializationMs: number /** Time spent selecting the initial static/base weight seeds. */ readonly baseWeightSearchMs: number /** Time spent evaluating the random scout baseline. */ @@ -333,11 +335,15 @@ export interface BenchmarkTimings { readonly weightSearchDurationMs: number readonly fusionSearchDurationMs: number readonly evidenceRouterSearchDurationMs: number + /** Time spent starting the shared native candidate queue. */ + readonly candidateQueueStartupDurationMs: number + /** Time spent shutting down the shared native candidate queue. */ + readonly candidateQueueShutdownDurationMs: number } /** Reproducible machine-readable output of one complete benchmark run. */ export interface BenchmarkArtifact { - readonly schemaVersion: 22 + readonly schemaVersion: 23 /** Profile controlling benchmark coverage without changing retrieval behavior. */ readonly benchmarkProfile: BenchmarkProfile /** Versioned objective profile used for candidate selection and aggregate metrics. */ diff --git a/benchmarks/retrieval/weight-search.ts b/benchmarks/retrieval/weight-search.ts index eaf3d4d..3411c99 100644 --- a/benchmarks/retrieval/weight-search.ts +++ b/benchmarks/retrieval/weight-search.ts @@ -99,6 +99,7 @@ const preparedFusionCache = new WeakMap< WeightSearchSample, Map >() +const candidatePoolInitializationMs = new WeakMap() const preparedFusionEvaluatorFor = ( sample: WeightSearchSample, @@ -129,6 +130,7 @@ const createEvaluationPoolForSamples = async ( profile: OptimizationProfile, options: ParallelSearchOptions, ): Promise => { + const snapshotPreparationStartedAt = performance.now() const snapshot = createEvaluationSnapshot( samples.map((sample) => ({ evaluator: preparedFusionEvaluatorFor(sample, fusion), @@ -137,9 +139,12 @@ const createEvaluationPoolForSamples = async ( sampleWeight: profile.queryFormWeights[sample.queryKind], })), ) - if (options.evaluationQueue !== undefined) - return createCandidateEvaluationPoolOnQueue(snapshot, options.evaluationQueue, options.signal) - return createCandidateEvaluationPool(snapshot, options) + const pool = + options.evaluationQueue !== undefined + ? createCandidateEvaluationPoolOnQueue(snapshot, options.evaluationQueue, options.signal) + : await createCandidateEvaluationPool(snapshot, options) + candidatePoolInitializationMs.set(pool, performance.now() - snapshotPreparationStartedAt) + return pool } const routerEvaluationCandidate = ( @@ -986,6 +991,7 @@ interface SearchEvaluationStats { interface MutableRouterSearchTimings { preparationMs: number + candidatePoolInitializationMs: number baseWeightSearchMs: number randomSearchMs: number beamSearchMs: number @@ -1275,6 +1281,7 @@ const prepareRouterSearch = ( protectedEliteCount: 0, timings: { preparationMs: performance.now() - preparationStartedAt, + candidatePoolInitializationMs: 0, baseWeightSearchMs: 0, randomSearchMs: 0, beamSearchMs: 0, @@ -1665,6 +1672,8 @@ const withParallelEvidencePools = async ( profile, options, ) + preparation.stats.timings.candidatePoolInitializationMs += + candidatePoolInitializationMs.get(fullPool) ?? 0 const closeFullPoolOnAbort = () => { void fullPool.close() } @@ -1682,6 +1691,9 @@ const withParallelEvidencePools = async ( profile, options, ) + if (proxyPool !== undefined && proxyPool !== fullPool) + preparation.stats.timings.candidatePoolInitializationMs += + candidatePoolInitializationMs.get(proxyPool) ?? 0 if (proxyPool !== undefined && proxyPool !== fullPool) { closeProxyPoolOnAbort = () => { void proxyPool?.close() diff --git a/benchmarks/tests/retrieval.test.ts b/benchmarks/tests/retrieval.test.ts index 269ffb9..3d723ad 100644 --- a/benchmarks/tests/retrieval.test.ts +++ b/benchmarks/tests/retrieval.test.ts @@ -99,10 +99,15 @@ const runProfile = (profile: BenchmarkProfile, groupedFolds: number, fusionMetho expect(artifact.evaluationCases.every(({ groundTruth }) => groundTruth.length > 0)).toBe(true) expect(artifact.models.length).toBeGreaterThan(0) expect(artifact.measurements.length).toBeGreaterThan(0) - expect(artifact.schemaVersion).toBe(22) + expect(artifact.schemaVersion).toBe(23) expect(artifact.searchStrategy).toEqual(ROUTER_SEARCH_STRATEGY) expect(artifact.timings.totalDurationMs).toBeGreaterThan(0) expect(Object.values(artifact.timings).every((duration) => duration >= 0)).toBe(true) + expect( + artifact.evidenceRouterSearch.every( + ({ searchDiagnostics }) => searchDiagnostics.timings.candidatePoolInitializationMs >= 0, + ), + ).toBe(true) expect(artifact.embeddingRuns.every((run) => run.queryEmbeddingDurationMs >= 0)).toBe(true) expect(artifact.sparseEmbeddingRuns.length).toBe(artifact.repositories.length) expect(artifact.sparseEmbeddingRuns.every((run) => run.queryTokenizationDurationMs >= 0)).toBe( From 93131c11b228efc0eedcbc45fe431fed26f2838b Mon Sep 17 00:00:00 2001 From: Lucas Burmeister Date: Tue, 4 Aug 2026 19:46:22 +0200 Subject: [PATCH 6/8] refactor: split retrieval benchmark modules --- .fallowrc.json | 3 +- CONTEXT.md | 40 +- benchmarks/BASELINE.md | 4 +- benchmarks/README.md | 12 +- benchmarks/retrieval/{ => corpus}/prepare.ts | 26 +- .../{corpus.ts => corpus/repository.ts} | 2 +- .../retrieval/{ => evaluation}/baseline.ts | 2 +- benchmarks/retrieval/evaluation/collect.ts | 525 +++++++++++++ .../retrieval/{ => evaluation}/folds.ts | 0 .../retrieval/{ => evaluation}/metrics.ts | 4 +- .../{ => evaluation}/optimization-profiles.ts | 2 +- .../prepared-fusion-core.d.mts} | 9 +- .../prepared-fusion-core.mjs} | 0 .../prepared-fusion.ts} | 8 +- .../retrieval/{ => evaluation}/ranking.ts | 19 +- .../retrieval/{ => evaluation}/report.ts | 19 +- benchmarks/retrieval/evaluation/search.ts | 555 +++++++++++++ .../retrieval/{ => evaluation}/types.ts | 2 +- .../{ => evaluation}/weight-search.ts | 330 +++----- .../candidate-evaluation-pool.ts} | 369 ++------- .../candidate-evaluation-worker.mjs} | 14 +- .../retrieval/{ => execution}/sqlite-index.ts | 12 +- benchmarks/retrieval/router-job-pool.ts | 318 -------- benchmarks/retrieval/router-job-worker.mjs | 181 ----- benchmarks/retrieval/runner.ts | 739 +----------------- benchmarks/retrieval/ts-loader.mjs | 10 - benchmarks/retrieval/worker-message.ts | 8 - benchmarks/tests/channels.test.ts | 30 +- benchmarks/tests/corpus.test.ts | 6 +- .../tests/optimization-profiles.test.ts | 2 +- benchmarks/tests/retrieval.test.ts | 6 +- benchmarks/tests/worker-pool.test.ts | 93 +-- 32 files changed, 1459 insertions(+), 1891 deletions(-) rename benchmarks/retrieval/{ => corpus}/prepare.ts (74%) rename benchmarks/retrieval/{corpus.ts => corpus/repository.ts} (97%) rename benchmarks/retrieval/{ => evaluation}/baseline.ts (95%) create mode 100644 benchmarks/retrieval/evaluation/collect.ts rename benchmarks/retrieval/{ => evaluation}/folds.ts (100%) rename benchmarks/retrieval/{ => evaluation}/metrics.ts (96%) rename benchmarks/retrieval/{ => evaluation}/optimization-profiles.ts (99%) rename benchmarks/retrieval/{fusion-core.d.mts => evaluation/prepared-fusion-core.d.mts} (57%) rename benchmarks/retrieval/{fusion-core.mjs => evaluation/prepared-fusion-core.mjs} (100%) rename benchmarks/retrieval/{fusion.ts => evaluation/prepared-fusion.ts} (91%) rename benchmarks/retrieval/{ => evaluation}/ranking.ts (87%) rename benchmarks/retrieval/{ => evaluation}/report.ts (97%) create mode 100644 benchmarks/retrieval/evaluation/search.ts rename benchmarks/retrieval/{ => evaluation}/types.ts (99%) rename benchmarks/retrieval/{ => evaluation}/weight-search.ts (89%) rename benchmarks/retrieval/{worker-pool.ts => execution/candidate-evaluation-pool.ts} (67%) rename benchmarks/retrieval/{fusion-worker.mjs => execution/candidate-evaluation-worker.mjs} (65%) rename benchmarks/retrieval/{ => execution}/sqlite-index.ts (73%) delete mode 100644 benchmarks/retrieval/router-job-pool.ts delete mode 100644 benchmarks/retrieval/router-job-worker.mjs delete mode 100644 benchmarks/retrieval/ts-loader.mjs delete mode 100644 benchmarks/retrieval/worker-message.ts diff --git a/.fallowrc.json b/.fallowrc.json index e93fff5..84576d3 100644 --- a/.fallowrc.json +++ b/.fallowrc.json @@ -4,8 +4,7 @@ "entry": [ "src/index.ts", "tests/scripts/bench-startup.mjs", - "tests/scripts/check-dtype-output.mjs", - "benchmarks/retrieval/ts-loader.mjs" + "tests/scripts/check-dtype-output.mjs" ], "boundaries": { "zones": [ diff --git a/CONTEXT.md b/CONTEXT.md index 1717730..384f3b5 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -168,10 +168,12 @@ Production queries accept the named profile selection `compatibility`, `balanced `natural-language`. Only `compatibility` currently has matrix-derived values; the remaining names are runtime placeholders until their configurations are selected from the benchmark matrix. -Benchmark-owned profile seeds and optimizer search live under `benchmarks/retrieval/`; current profile -values are marked `authored-seed` and are not presented as benchmark-derived weights. Production keeps -only the active router configuration and reusable fusion/evidence seams. A benchmark result is promoted -explicitly rather than making the production package responsible for discovering its own profile. +Benchmark-owned corpus checkout and preparation live under `benchmarks/retrieval/corpus/`; quality +evaluation, profile seeds, and optimizer search live under `benchmarks/retrieval/evaluation/`; native +SQLite and worker execution live under `benchmarks/retrieval/execution/`. Current profile values are +marked `authored-seed` and are not presented as benchmark-derived weights. Production keeps only the +active router configuration and reusable fusion/evidence seams. A benchmark result is promoted explicitly +rather than making the production package responsible for discovering its own profile. ### Retrieval Quality Benchmark @@ -251,22 +253,20 @@ that data. Public `fuseRankings` semantics and its existing ranking, normalizati remain unchanged. The evaluator still materializes and sorts the full candidate union for every configuration; partial top-K or metric-specialized evaluation remains future work. The explicit parallel benchmark path uses a fixed native `node:worker_threads` pool (default `max(1, availableParallelism() - 1)`), -sends the prepared snapshot once per worker, batches candidate weight vectors, and keeps beam/cache/archive/ -selection state on the controller thread. The search APIs are async so serial and worker execution share the -same selection algorithm; `PIX_BENCH_WORKERS=0` selects serial evaluation, `workerCount: 1` selects the -serial pool through `createCandidateEvaluationPool`, and the runner's `PIX_BENCH_SEARCH_MODE=serial` selects -that serial pool mode. `PIX_BENCH_WORKER_BATCH_SIZE` bounds ordinary worker messages. Router search -diagnostics retain candidate snapshot and pool initialization time separately from candidate evaluation time; -benchmark timings also record shared candidate-queue startup and shutdown. -Evidence-router grouped-fold, repository-holdout, and fit-all jobs use native controller workers backed by one -shared candidate-evaluation queue. Controllers keep candidate generation, beam rounds, caches, and archive -updates local and sequential; candidate batches are work-stealing tasks evaluated by the shared native queue. -The queue caches each prepared snapshot once per worker and defaults to one candidate per task because the -benchmark's candidate costs are uneven; `PIX_BENCH_WORKER_BATCH_SIZE` overrides that default. The runner -reserves roughly half of `PIX_BENCH_WORKERS` for controllers and the remainder for candidate workers. -Results retain planner order even when jobs finish out of order. Serial mode keeps the same algorithm on the -main thread for comparison. Native worker startup or task failures are surfaced and all workers are -terminated before the benchmark fails. +sends the prepared snapshot once per worker, and batches candidate weight vectors. Beam, cache, archive, and +selection state stay on the main thread; async candidate evaluations let independent router jobs share one +queue without a second controller-worker layer. The search APIs are async so serial and worker execution +share the same selection algorithm; `PIX_BENCH_WORKERS=0` selects serial evaluation, `workerCount: 1` +selects the serial pool through `createCandidateEvaluationPool`, and the runner's +`PIX_BENCH_SEARCH_MODE=serial` selects that serial pool mode. `PIX_BENCH_WORKER_BATCH_SIZE` bounds worker +messages. Router search diagnostics retain candidate snapshot and pool initialization time separately from +candidate evaluation time; benchmark timings also record shared candidate-queue startup and shutdown. +Evidence-router grouped-fold, repository-holdout, and fit-all jobs run concurrently on the main thread while +candidate batches use work-stealing tasks in the shared native queue. The queue caches each prepared snapshot +once per worker and defaults to one candidate per task because the benchmark's candidate costs are uneven; +`PIX_BENCH_WORKER_BATCH_SIZE` overrides that default. Results retain planner order through `Promise.all`. +Serial mode keeps the same algorithm on the main thread for comparison. Native worker startup or task +failures are surfaced and all workers are terminated before the benchmark fails. Repository checkouts live under ignored `benchmarks/.cache/repos/`; generated artifacts live under ignored `benchmarks/results/`. Dense and Sparse vectors are held only in the production in-memory SQLite adapter during a benchmark run. See `benchmarks/README.md` and `benchmarks/BASELINE.md`. diff --git a/benchmarks/BASELINE.md b/benchmarks/BASELINE.md index eedf1f7..ea8b89b 100644 --- a/benchmarks/BASELINE.md +++ b/benchmarks/BASELINE.md @@ -1,14 +1,14 @@ # Preliminary Retrieval Baseline The schema-17 entries below are historical artifacts from the benchmark-owned Sparse implementation. -Current schema-22 runs use the production SparseEmbedder and IndexStore without benchmark vector caches. +Current schema-23 runs use the production SparseEmbedder and IndexStore without benchmark vector caches. ## Schema 20: Search-Priority DBSF Selection The historical schema-20 `search-priority` full-profile `fd` run is `benchmarks/results/retrieval-2026-08-03T23-34-02.311Z.json`. Its authored query-form objective is `identifier/agentTask/naturalQuestion/searchPhrase = 1/2/3/4`; the profile's channel weights are -authored seeds, not benchmark-derived deployment weights. This evidence predates the current schema-22 +authored seeds, not benchmark-derived deployment weights. This evidence predates the current schema-23 artifact format. Result artifacts are ignored and the referenced JSON is available only in the local worktree that produced that promotion evidence. diff --git a/benchmarks/README.md b/benchmarks/README.md index 63d7317..5c0eb81 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -95,7 +95,7 @@ vp run bench:retrieval:full `bench:retrieval` aliases `bench:retrieval:validate`. Every profile measures the same physical rankings and retrieval variants; profiles only control matrix size, holdout coverage, and expensive -diagnostics. The selected profile is recorded in schema-22 artifacts without changing retrieval +diagnostics. The selected profile is recorded in schema-23 artifacts without changing retrieval semantics. The full profile includes all three fusion methods; short profiles intentionally omit RRF to keep development runs fast. @@ -310,9 +310,13 @@ Score and DBSF consume the same `ChannelRankings` interface and are evaluated wi encoders, persistence, or scoring. `src/lib/retrieval/evidence-router.ts` is likewise shared by production configuration and benchmark evidence evaluation. -`benchmarks/retrieval/optimization-profiles.ts` owns authored (`authored-seed`) profile seeds, and -`benchmarks/retrieval/weight-search.ts` owns candidate search. A validated benchmark result is promoted -to an explicit production configuration; production does not discover or optimize its own profile. +`benchmarks/retrieval/evaluation/optimization-profiles.ts` owns authored (`authored-seed`) profile seeds, and +`benchmarks/retrieval/evaluation/weight-search.ts` owns candidate search. Corpus checkout and preparation +live under `benchmarks/retrieval/corpus/`; native SQLite and worker execution live under +`benchmarks/retrieval/execution/`. A validated benchmark result is promoted to an explicit production +configuration; production does not discover or optimize its own profile. Router searches keep their beam +and archive state on the main thread; only candidate scoring crosses the worker seam, so independent jobs +share one queue without a second controller-worker protocol. The remaining architectural follow-up is a diagnostic retrieval snapshot from `IndexStore` if future benchmark work needs to inspect persisted channel evidence through the application boundary. Current diff --git a/benchmarks/retrieval/prepare.ts b/benchmarks/retrieval/corpus/prepare.ts similarity index 74% rename from benchmarks/retrieval/prepare.ts rename to benchmarks/retrieval/corpus/prepare.ts index 76612a2..22725a6 100644 --- a/benchmarks/retrieval/prepare.ts +++ b/benchmarks/retrieval/corpus/prepare.ts @@ -3,19 +3,19 @@ import path from "node:path" import { Effect } from "effect" -import type { Chunk } from "../../src/domain/chunk.js" -import { DEFAULT_CONFIG } from "../../src/domain/config.js" -import type { Identifier } from "../../src/domain/identifier.js" -import type { Bm25Index } from "../../src/domain/ports.js" -import { getExtension } from "../../src/lib/config/extension.js" -import { extractIdentifiers } from "../../src/lib/parsing/identifier-extractor.js" -import { buildExtensionRegistry } from "../../src/lib/registry.js" -import { buildBm25Index } from "../../src/lib/retrieval/bm25.js" -import { buildIdentifierIndex } from "../../src/lib/retrieval/identifier-index.js" -import { chunkTextWithConfig } from "../../src/services/chunker.js" -import { listCorpusFiles } from "./corpus.js" -import type { ChunkIdentifiers } from "./metrics.js" -import type { CorpusManifest } from "./types.js" +import type { Chunk } from "../../../src/domain/chunk.js" +import { DEFAULT_CONFIG } from "../../../src/domain/config.js" +import type { Identifier } from "../../../src/domain/identifier.js" +import type { Bm25Index } from "../../../src/domain/ports.js" +import { getExtension } from "../../../src/lib/config/extension.js" +import { extractIdentifiers } from "../../../src/lib/parsing/identifier-extractor.js" +import { buildExtensionRegistry } from "../../../src/lib/registry.js" +import { buildBm25Index } from "../../../src/lib/retrieval/bm25.js" +import { buildIdentifierIndex } from "../../../src/lib/retrieval/identifier-index.js" +import { chunkTextWithConfig } from "../../../src/services/chunker.js" +import type { ChunkIdentifiers } from "../evaluation/metrics.js" +import type { CorpusManifest } from "../evaluation/types.js" +import { listCorpusFiles } from "./repository.js" /** Query-independent indexes and source chunks prepared once per repository. */ export interface PreparedCorpus { diff --git a/benchmarks/retrieval/corpus.ts b/benchmarks/retrieval/corpus/repository.ts similarity index 97% rename from benchmarks/retrieval/corpus.ts rename to benchmarks/retrieval/corpus/repository.ts index d2a108d..35886e2 100644 --- a/benchmarks/retrieval/corpus.ts +++ b/benchmarks/retrieval/corpus/repository.ts @@ -5,7 +5,7 @@ import { promisify } from "node:util" import { Effect, Schema } from "effect" -import { CorpusManifestSchema, type CorpusManifest } from "./types.js" +import { CorpusManifestSchema, type CorpusManifest } from "../evaluation/types.js" const execFilePromise = promisify(execFile) const CACHE_ROOT = path.resolve("benchmarks/.cache/repos") diff --git a/benchmarks/retrieval/baseline.ts b/benchmarks/retrieval/evaluation/baseline.ts similarity index 95% rename from benchmarks/retrieval/baseline.ts rename to benchmarks/retrieval/evaluation/baseline.ts index dc5b649..1b2207d 100644 --- a/benchmarks/retrieval/baseline.ts +++ b/benchmarks/retrieval/evaluation/baseline.ts @@ -2,7 +2,7 @@ import { decodeEvidenceRouterConfig, ZERO_CHANNEL_COEFFICIENTS, type EvidenceRouterConfig, -} from "../../src/domain/retrieval.js" +} from "../../../src/domain/retrieval.js" /** Historical RRF configuration used only for explicit benchmark and rollback comparisons. */ export const HISTORICAL_RRF_BASELINE_CONFIG: EvidenceRouterConfig = decodeEvidenceRouterConfig({ diff --git a/benchmarks/retrieval/evaluation/collect.ts b/benchmarks/retrieval/evaluation/collect.ts new file mode 100644 index 0000000..458254f --- /dev/null +++ b/benchmarks/retrieval/evaluation/collect.ts @@ -0,0 +1,525 @@ +import { Effect, Stream } from "effect" + +import type { Embedding } from "../../../src/domain/chunk.js" +import { DEFAULT_CONFIG } from "../../../src/domain/config.js" +import type { EmbeddingDtype } from "../../../src/domain/dtype.js" +import type { StoredChunk } from "../../../src/domain/index-data.js" +import { MODEL_REGISTRY } from "../../../src/domain/models.js" +import type { BoundEmbedder, SearchData } from "../../../src/domain/ports.js" +import { IndexStore, SparseEmbedder } from "../../../src/domain/ports.js" +import type { + SparseContract, + SparseQuery, + SparseTerm, + SparseVector, +} from "../../../src/domain/sparse.js" +import { contentHash } from "../../../src/lib/content-hash.js" +import { buildQueryTermCoverage } from "../../../src/lib/retrieval/evidence-router.js" +import { createAutoBoundEmbedder } from "../../../src/services/embedder.js" +import { prepareCorpus, type PreparedCorpus } from "../corpus/prepare.js" +import { prepareRepository } from "../corpus/repository.js" +import { withSqliteBenchmarkStore } from "../execution/sqlite-index.js" +import { foldKey } from "./folds.js" +import { + contextRecallAtBudget, + goldTargetRanks, + recallAt, + reciprocalRank, + resolveGoldTargets, + successAt, +} from "./metrics.js" +import { fuseVariant, rankLexicalChannels, RETRIEVAL_VARIANTS } from "./ranking.js" +import type { BenchmarkArtifact, CorpusManifest, QueryKind, QueryMeasurement } from "./types.js" +import type { WeightSearchSample } from "./weight-search.js" + +const CONTEXT_BUDGETS = [2_048, 4_096, 8_192, 16_384] as const +const EMBEDDING_BATCH_SIZE = 2 +const SINGLE_ITEM_ESTIMATED_TOKENS = 2_048 +const QUERY_KINDS: readonly QueryKind[] = [ + "identifier", + "searchPhrase", + "naturalQuestion", + "agentTask", +] + +interface BenchmarkQuery { + readonly questionIndex: number + readonly queryKind: QueryKind + readonly query: string +} + +interface ModelMeasurements { + readonly sparseEmbeddingRun: BenchmarkArtifact["sparseEmbeddingRuns"][number] + readonly measurements: readonly QueryMeasurement[] + readonly samples: readonly WeightSearchSample[] + readonly samplesByQueryKind: ReadonlyMap + readonly retrievalDurationMs: number +} + +interface RepositoryMeasurements { + readonly repository: BenchmarkArtifact["repositories"][number] + readonly embeddingRuns: readonly BenchmarkArtifact["embeddingRuns"][number][] + readonly sparseEmbeddingRuns: readonly BenchmarkArtifact["sparseEmbeddingRuns"][number][] + readonly measurements: readonly QueryMeasurement[] + readonly samplesByModel: ReadonlyMap + readonly sampleGroups: ReadonlyMap< + string, + { + readonly model: string + readonly queryKind: QueryKind + readonly samples: readonly WeightSearchSample[] + } + > + readonly retrievalDurationMs: number +} + +/** Collected corpus, channel, and quality samples reused by all search stages. */ +export interface CollectedBenchmarkData { + readonly repositories: readonly BenchmarkArtifact["repositories"][number][] + readonly embeddingRuns: readonly BenchmarkArtifact["embeddingRuns"][number][] + readonly sparseEmbeddingRuns: readonly BenchmarkArtifact["sparseEmbeddingRuns"][number][] + readonly measurements: readonly QueryMeasurement[] + readonly sampleGroups: ReadonlyMap< + string, + { + readonly model: string + readonly queryKind: QueryKind + readonly samples: readonly WeightSearchSample[] + } + > + readonly samplesByModel: ReadonlyMap + readonly retrievalDurationMs: number +} + +const isLongInput = (text: string): boolean => + Buffer.byteLength(text, "utf8") / 4 > SINGLE_ITEM_ESTIMATED_TOKENS + +const embedTexts = ( + texts: readonly string[], + model: string, + embedder: BoundEmbedder, +): Effect.Effect => + Effect.gen(function* () { + const vectors: Float32Array[] = [] + let start = 0 + while (start < texts.length) { + const currentIsLong = isLongInput(texts[start]) + const next = texts[start + 1] + const nextIsLong = next !== undefined && isLongInput(next) + const batchSize = currentIsLong || nextIsLong ? 1 : EMBEDDING_BATCH_SIZE + const batch = texts.slice(start, start + batchSize) + const embeddings = yield* embedder.batch(batch) + vectors.push(...embeddings.map((embedding) => embedding.vector)) + start += batchSize + } + return vectors + }).pipe(Effect.mapError((cause) => new Error(`Embedding failed for ${model}`, { cause }))) + +const embedSparseTexts = ( + texts: readonly string[], + embedder: typeof SparseEmbedder.Service, +): Effect.Effect => + Effect.gen(function* () { + const vectors: SparseVector[] = [] + let start = 0 + while (start < texts.length) { + const batch = texts.slice(start, start + DEFAULT_CONFIG.sparseEmbedder.batchSize) + vectors.push(...(yield* embedder.batch(batch))) + start += DEFAULT_CONFIG.sparseEmbedder.batchSize + } + return vectors + }).pipe(Effect.mapError((cause) => new Error("Sparse document embedding failed", { cause }))) + +const toStoredChunk = (chunk: PreparedCorpus["chunks"][number]): StoredChunk => { + const { text, ...location } = chunk + return { ...location, contentHash: contentHash(text) } +} + +const persistBenchmarkCorpus = ( + store: typeof IndexStore.Service, + corpus: PreparedCorpus, + vectors: readonly Float32Array[], + sparseVectors: readonly SparseVector[], + dims: number, + dtype: EmbeddingDtype, + sparseContract: SparseContract, + sparseIdf: readonly SparseTerm[], +): Effect.Effect => + Effect.gen(function* () { + const pairs = corpus.chunks.map((chunk, index): readonly [StoredChunk, Embedding] => [ + toStoredChunk(chunk), + { vector: vectors[index]!, dims, dtype }, + ]) + yield* store.persistIndex({ + chunks: Stream.succeed( + pairs.map( + ([chunk, embedding], index) => [chunk, embedding, sparseVectors[index]!] as const, + ), + ), + identifierIndex: corpus.identifierIndex, + bm25Index: corpus.bm25Index, + files: [], + dims, + dtype, + embeddingCache: [], + sparseEmbeddingCache: [], + sparseContract, + sparseIdf, + }) + }) + +const collectModelSamples = ( + manifest: CorpusManifest, + corpus: PreparedCorpus, + queries: readonly BenchmarkQuery[], + targetsByQuestion: readonly (readonly ReadonlySet[])[], + groupedFoldAssignments: ReadonlyMap, + model: string, + queryVectors: readonly Float32Array[], + sparseQueries: readonly SparseQuery[], + searchData: SearchData, + store: typeof IndexStore.Service, + dims: number, + dtype: EmbeddingDtype, +): Effect.Effect< + { + readonly measurements: readonly QueryMeasurement[] + readonly samples: readonly WeightSearchSample[] + readonly samplesByQueryKind: ReadonlyMap + }, + Error +> => + Effect.gen(function* () { + const modelMeasurements: QueryMeasurement[] = [] + const modelSamples: WeightSearchSample[] = [] + const samplesByQueryKind = new Map() + + for (let queryIndex = 0; queryIndex < queries.length; queryIndex++) { + const entry = queries[queryIndex] + const question = manifest.questions[entry.questionIndex] + const targets = targetsByQuestion[entry.questionIndex] + const groupedFold = groupedFoldAssignments.get(foldKey(manifest.id, question.id)) + if (groupedFold === undefined) + return yield* Effect.fail( + new Error(`No grouped fold assignment for ${manifest.id}/${question.id}`), + ) + const channelStartedAt = performance.now() + const lexicalRankings = rankLexicalChannels(entry.query, searchData) + const dense = yield* store.searchDense({ + vector: queryVectors[queryIndex]!, + dims, + dtype, + }) + const sparse = yield* store.searchSparse(sparseQueries[queryIndex]!) + const rankings = { ...lexicalRankings, dense, sparse } + const channelDurationMs = performance.now() - channelStartedAt + const sample: WeightSearchSample = { + repository: manifest.id, + intentId: question.id, + queryKind: entry.queryKind, + groupedFold, + query: entry.query, + rankings, + targets, + chunks: corpus.chunks, + termCoverage: buildQueryTermCoverage( + entry.query, + searchData.bm25Index, + searchData.identifierIndex, + ), + } + modelSamples.push(sample) + samplesByQueryKind.set(entry.queryKind, [ + ...(samplesByQueryKind.get(entry.queryKind) ?? []), + sample, + ]) + for (const variant of RETRIEVAL_VARIANTS) { + const variantStartedAt = performance.now() + const ranked = fuseVariant(variant, entry.query, rankings) + const queryDurationMs = channelDurationMs + performance.now() - variantStartedAt + modelMeasurements.push({ + repository: manifest.id, + language: manifest.language, + size: manifest.size, + revision: manifest.revision, + model, + variant, + questionId: question.id, + queryKind: entry.queryKind, + query: entry.query, + category: question.category, + difficulty: question.difficulty, + groupedFold, + recallAt5: recallAt(ranked, targets, 5), + recallAt10: recallAt(ranked, targets, 10), + recallAt20: recallAt(ranked, targets, 20), + recallAt50: recallAt(ranked, targets, 50), + successAt10: successAt(ranked, targets, 10), + successAt20: successAt(ranked, targets, 20), + reciprocalRank: reciprocalRank(ranked, targets), + goldRanks: goldTargetRanks(ranked, targets), + contextRecall: Object.fromEntries( + CONTEXT_BUDGETS.map((budget) => [ + String(budget), + contextRecallAtBudget(ranked, targets, corpus.chunks, budget), + ]), + ), + queryDurationMs, + }) + } + } + + return { measurements: modelMeasurements, samples: modelSamples, samplesByQueryKind } + }) + +const collectModelMeasurements = ( + manifest: CorpusManifest, + corpus: PreparedCorpus, + queries: readonly BenchmarkQuery[], + targetsByQuestion: readonly (readonly ReadonlySet[])[], + groupedFoldAssignments: ReadonlyMap, + model: string, + info: (typeof MODEL_REGISTRY)[string] & object, + chunkVectors: readonly Float32Array[], + queryVectors: readonly Float32Array[], +): Effect.Effect => + Effect.gen(function* () { + const retrievalStartedAt = performance.now() + const modelRun = yield* withSqliteBenchmarkStore( + model, + info.defaultDtype, + Effect.gen(function* () { + const store = yield* IndexStore + const sparseEmbedder = yield* SparseEmbedder + const sparseStartedAt = performance.now() + const sparseVectors = yield* embedSparseTexts( + corpus.chunks.map((chunk) => chunk.text), + sparseEmbedder, + ) + const sparseChunkEmbeddingDurationMs = performance.now() - sparseStartedAt + const sparseIdf = yield* sparseEmbedder.loadIdf() + const sparseQueryStartedAt = performance.now() + const sparseQueries = yield* Effect.forEach(queries, ({ query }) => + sparseEmbedder.tokenizeQuery(query), + ) + const sparseQueryTokenizationDurationMs = performance.now() - sparseQueryStartedAt + yield* persistBenchmarkCorpus( + store, + corpus, + chunkVectors, + sparseVectors, + info.dims, + info.defaultDtype, + sparseEmbedder.contract, + sparseIdf, + ) + const searchData = yield* store.loadSearchData() + const collected = yield* collectModelSamples( + manifest, + corpus, + queries, + targetsByQuestion, + groupedFoldAssignments, + model, + queryVectors, + sparseQueries, + searchData, + store, + info.dims, + info.defaultDtype, + ) + return { + sparseEmbeddingRun: { + repository: manifest.id, + model: sparseEmbedder.contract.model, + tokenizerModel: sparseEmbedder.contract.tokenizer, + batchSize: DEFAULT_CONFIG.sparseEmbedder.batchSize, + chunkEmbeddingDurationMs: sparseChunkEmbeddingDurationMs, + queryTokenizationDurationMs: sparseQueryTokenizationDurationMs, + }, + ...collected, + } + }), + ) + return { + sparseEmbeddingRun: modelRun.sparseEmbeddingRun, + measurements: modelRun.measurements, + samples: modelRun.samples, + samplesByQueryKind: modelRun.samplesByQueryKind, + retrievalDurationMs: performance.now() - retrievalStartedAt, + } + }) + +const collectRepositoryMeasurements = ( + manifest: CorpusManifest, + models: readonly string[], + groupedFoldAssignments: ReadonlyMap, +): Effect.Effect => + Effect.gen(function* () { + const repositoryPath = yield* prepareRepository(manifest) + const corpus = yield* prepareCorpus(repositoryPath, manifest) + const targetsByQuestion: (readonly ReadonlySet[])[] = [] + for (const question of manifest.questions) { + const targets = resolveGoldTargets( + question.groundTruth, + corpus.chunks, + corpus.identifiersByChunk, + ) + const unresolved = question.groundTruth.filter((_, index) => targets[index].size === 0) + if (unresolved.length > 0) + return yield* Effect.fail( + new Error( + `${question.id} has unresolved gold targets: ${unresolved.map((target) => `${target.file}::${target.symbol}`).join(", ")}`, + ), + ) + targetsByQuestion.push(targets) + } + + const queries = manifest.questions.flatMap((question, questionIndex) => + QUERY_KINDS.map((queryKind) => ({ + questionIndex, + queryKind, + query: question.queries[queryKind], + })), + ) + const embeddingRuns: BenchmarkArtifact["embeddingRuns"][number][] = [] + const sparseEmbeddingRuns: BenchmarkArtifact["sparseEmbeddingRuns"][number][] = [] + const measurements: QueryMeasurement[] = [] + const samplesByModel = new Map() + const sampleGroups = new Map< + string, + { + readonly model: string + readonly queryKind: QueryKind + readonly samples: readonly WeightSearchSample[] + } + >() + let retrievalDurationMs = 0 + + for (const model of models) { + const info = MODEL_REGISTRY[model] + if (info === undefined) + return yield* Effect.fail(new Error(`Unknown embedding model ${model}`)) + const bound = yield* createAutoBoundEmbedder({ + model, + dtype: info.defaultDtype, + dims: info.dims, + }).pipe( + Effect.mapError( + (cause) => new Error(`Could not auto-select a device for ${model}`, { cause }), + ), + ) + const embeddingStartedAt = performance.now() + const chunkVectors = yield* embedTexts( + corpus.chunks.map((chunk) => chunk.text), + model, + bound.embedder, + ) + const chunkEmbeddingDurationMs = performance.now() - embeddingStartedAt + const queryEmbeddingStartedAt = performance.now() + const queryVectors = yield* embedTexts( + queries.map((entry) => entry.query), + model, + bound.embedder, + ) + const modelData = yield* collectModelMeasurements( + manifest, + corpus, + queries, + targetsByQuestion, + groupedFoldAssignments, + model, + info, + chunkVectors, + queryVectors, + ) + embeddingRuns.push({ + repository: manifest.id, + model, + device: bound.device, + batchSize: EMBEDDING_BATCH_SIZE, + chunkEmbeddingDurationMs, + queryEmbeddingDurationMs: performance.now() - queryEmbeddingStartedAt, + }) + sparseEmbeddingRuns.push(modelData.sparseEmbeddingRun) + measurements.push(...modelData.measurements) + retrievalDurationMs += modelData.retrievalDurationMs + samplesByModel.set(model, modelData.samples) + for (const [queryKind, samples] of modelData.samplesByQueryKind) { + sampleGroups.set(`${model}\0${queryKind}`, { model, queryKind, samples }) + } + } + + return { + repository: { + id: manifest.id, + repository: manifest.repository, + revision: manifest.revision, + chunks: corpus.chunks.length, + preparationDurationMs: corpus.preparationDurationMs, + }, + embeddingRuns, + sparseEmbeddingRuns, + measurements, + samplesByModel, + sampleGroups, + retrievalDurationMs, + } + }) + +/** Prepare every selected corpus and collect the physical rankings once. */ +export const collectBenchmarkData = ( + manifests: readonly CorpusManifest[], + models: readonly string[], + groupedFoldAssignments: ReadonlyMap, +): Effect.Effect => + Effect.gen(function* () { + const repositories: BenchmarkArtifact["repositories"][number][] = [] + const embeddingRuns: BenchmarkArtifact["embeddingRuns"][number][] = [] + const sparseEmbeddingRuns: BenchmarkArtifact["sparseEmbeddingRuns"][number][] = [] + const measurements: QueryMeasurement[] = [] + const sampleGroups = new Map< + string, + { + readonly model: string + readonly queryKind: QueryKind + readonly samples: readonly WeightSearchSample[] + } + >() + const samplesByModel = new Map() + let retrievalDurationMs = 0 + + for (const manifest of manifests) { + const repositoryData = yield* collectRepositoryMeasurements( + manifest, + models, + groupedFoldAssignments, + ) + repositories.push(repositoryData.repository) + embeddingRuns.push(...repositoryData.embeddingRuns) + sparseEmbeddingRuns.push(...repositoryData.sparseEmbeddingRuns) + measurements.push(...repositoryData.measurements) + retrievalDurationMs += repositoryData.retrievalDurationMs + for (const [model, samples] of repositoryData.samplesByModel) { + samplesByModel.set(model, [...(samplesByModel.get(model) ?? []), ...samples]) + } + for (const [key, group] of repositoryData.sampleGroups) { + const current = sampleGroups.get(key) + sampleGroups.set(key, { + model: group.model, + queryKind: group.queryKind, + samples: [...(current?.samples ?? []), ...group.samples], + }) + } + } + + return { + repositories, + embeddingRuns, + sparseEmbeddingRuns, + measurements, + sampleGroups, + samplesByModel, + retrievalDurationMs, + } + }) diff --git a/benchmarks/retrieval/folds.ts b/benchmarks/retrieval/evaluation/folds.ts similarity index 100% rename from benchmarks/retrieval/folds.ts rename to benchmarks/retrieval/evaluation/folds.ts diff --git a/benchmarks/retrieval/metrics.ts b/benchmarks/retrieval/evaluation/metrics.ts similarity index 96% rename from benchmarks/retrieval/metrics.ts rename to benchmarks/retrieval/evaluation/metrics.ts index a6bca66..24f19aa 100644 --- a/benchmarks/retrieval/metrics.ts +++ b/benchmarks/retrieval/evaluation/metrics.ts @@ -1,5 +1,5 @@ -import type { Chunk } from "../../src/domain/chunk.js" -import type { RankedChunk } from "../../src/domain/ports.js" +import type { Chunk } from "../../../src/domain/chunk.js" +import type { RankedChunk } from "../../../src/domain/ports.js" import type { GoldLocation } from "./types.js" /** Indexed identifiers retained per chunk so gold symbols can be matched exactly. */ diff --git a/benchmarks/retrieval/optimization-profiles.ts b/benchmarks/retrieval/evaluation/optimization-profiles.ts similarity index 99% rename from benchmarks/retrieval/optimization-profiles.ts rename to benchmarks/retrieval/evaluation/optimization-profiles.ts index cc91ec6..3c46402 100644 --- a/benchmarks/retrieval/optimization-profiles.ts +++ b/benchmarks/retrieval/evaluation/optimization-profiles.ts @@ -5,7 +5,7 @@ import { ZERO_CHANNEL_COEFFICIENTS, type ChannelWeights, type EvidenceRouterParameters, -} from "../../src/domain/retrieval.js" +} from "../../../src/domain/retrieval.js" const NonNegativeNumber = Schema.Number.check(Schema.isGreaterThanOrEqualTo(0)) const QualityMetrics = [ diff --git a/benchmarks/retrieval/fusion-core.d.mts b/benchmarks/retrieval/evaluation/prepared-fusion-core.d.mts similarity index 57% rename from benchmarks/retrieval/fusion-core.d.mts rename to benchmarks/retrieval/evaluation/prepared-fusion-core.d.mts index ab1a963..591ec9e 100644 --- a/benchmarks/retrieval/fusion-core.d.mts +++ b/benchmarks/retrieval/evaluation/prepared-fusion-core.d.mts @@ -1,7 +1,10 @@ -import type { ChannelWeights } from "../../src/domain/retrieval.js" -import type { PreparedFusionSnapshot } from "./fusion.js" +import type { ChannelWeights } from "../../../src/domain/retrieval.js" +import type { + EvaluationCandidate, + EvaluationSnapshot, +} from "../execution/candidate-evaluation-pool.js" +import type { PreparedFusionSnapshot } from "./prepared-fusion.js" import type { QualitySummary } from "./types.js" -import type { EvaluationCandidate, EvaluationSnapshot } from "./worker-pool.js" export function evaluatePreparedContributions( matrix: PreparedFusionSnapshot, diff --git a/benchmarks/retrieval/fusion-core.mjs b/benchmarks/retrieval/evaluation/prepared-fusion-core.mjs similarity index 100% rename from benchmarks/retrieval/fusion-core.mjs rename to benchmarks/retrieval/evaluation/prepared-fusion-core.mjs diff --git a/benchmarks/retrieval/fusion.ts b/benchmarks/retrieval/evaluation/prepared-fusion.ts similarity index 91% rename from benchmarks/retrieval/fusion.ts rename to benchmarks/retrieval/evaluation/prepared-fusion.ts index 1f6e12b..ce88150 100644 --- a/benchmarks/retrieval/fusion.ts +++ b/benchmarks/retrieval/evaluation/prepared-fusion.ts @@ -1,13 +1,13 @@ -import type { RankedChunk } from "../../src/domain/ports.js" +import type { RankedChunk } from "../../../src/domain/ports.js" import { CHANNEL_NAMES, type ChannelName, type ChannelRankings, type ChannelWeights, type FusionMethod, -} from "../../src/domain/retrieval.js" -import { fuseRankings as productionFuseRankings } from "../../src/lib/retrieval/fusion.js" -import { evaluatePreparedContributions } from "./fusion-core.mjs" +} from "../../../src/domain/retrieval.js" +import { fuseRankings as productionFuseRankings } from "../../../src/lib/retrieval/fusion.js" +import { evaluatePreparedContributions } from "./prepared-fusion-core.mjs" const DEFAULT_CANDIDATE_DEPTH = Number.POSITIVE_INFINITY diff --git a/benchmarks/retrieval/ranking.ts b/benchmarks/retrieval/evaluation/ranking.ts similarity index 87% rename from benchmarks/retrieval/ranking.ts rename to benchmarks/retrieval/evaluation/ranking.ts index c5405cc..4c4de52 100644 --- a/benchmarks/retrieval/ranking.ts +++ b/benchmarks/retrieval/evaluation/ranking.ts @@ -1,17 +1,20 @@ -import type { RankedChunk } from "../../src/domain/ports.js" +import type { RankedChunk } from "../../../src/domain/ports.js" import { CHANNEL_NAMES, type ChannelName, type ChannelRankings, type ChannelWeights, -} from "../../src/domain/retrieval.js" -import { rankBm25 } from "../../src/lib/retrieval/bm25.js" -import { rankCamelCase } from "../../src/lib/retrieval/camelcase.js" -import { buildRoutingEvidence, routeWithEvidence } from "../../src/lib/retrieval/evidence-router.js" -import { fuseRankings } from "../../src/lib/retrieval/fusion.js" -import { rankIdentity } from "../../src/lib/retrieval/identity.js" +} from "../../../src/domain/retrieval.js" +import { rankBm25 } from "../../../src/lib/retrieval/bm25.js" +import { rankCamelCase } from "../../../src/lib/retrieval/camelcase.js" +import { + buildRoutingEvidence, + routeWithEvidence, +} from "../../../src/lib/retrieval/evidence-router.js" +import { fuseRankings } from "../../../src/lib/retrieval/fusion.js" +import { rankIdentity } from "../../../src/lib/retrieval/identity.js" +import type { PreparedCorpus } from "../corpus/prepare.js" import { HISTORICAL_RRF_BASELINE_CONFIG } from "./baseline.js" -import type { PreparedCorpus } from "./prepare.js" import type { RetrievalVariant } from "./types.js" /** Rankings produced by the lexical channels before dense search is delegated to SQLite. */ diff --git a/benchmarks/retrieval/report.ts b/benchmarks/retrieval/evaluation/report.ts similarity index 97% rename from benchmarks/retrieval/report.ts rename to benchmarks/retrieval/evaluation/report.ts index 1843a79..0177ae5 100644 --- a/benchmarks/retrieval/report.ts +++ b/benchmarks/retrieval/evaluation/report.ts @@ -3,7 +3,7 @@ import { type ChannelCoefficients, type ChannelWeights, type EvidenceRouterConfig, -} from "../../src/domain/retrieval.js" +} from "../../../src/domain/retrieval.js" import type { BenchmarkArtifact, EvidenceRouterSearchResult, @@ -53,6 +53,15 @@ const formatInfluences = (config: EvidenceRouterConfig): string => `L:${formatCoefficients(config.queryLengthInfluence)}`, ].join("; ") +const formatRouterWeightColumns = (result: { + readonly config: EvidenceRouterConfig + readonly staticWeights: ChannelWeights +}) => ({ + baseWeights: formatWeights(result.config.baseWeights), + staticWeights: formatWeights(result.staticWeights), + influences: formatInfluences(result.config), +}) + /** Render quality and marginal channel contribution grouped by query representation. */ export const renderMarkdownReport = (artifact: BenchmarkArtifact): string => { const groups = new Map() @@ -242,9 +251,7 @@ export const renderMarkdownReport = (artifact: BenchmarkArtifact): string => { "| --- | --- | --- | --- | --- | --- | ---: | ---: | ---: | ---: | --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", ) for (const result of artifact.evidenceRouterSearch) { - const baseWeights = formatWeights(result.config.baseWeights) - const staticWeights = formatWeights(result.staticWeights) - const influences = formatInfluences(result.config) + const { baseWeights, staticWeights, influences } = formatRouterWeightColumns(result) lines.push( `| ${result.model} | ${result.fusion} | ${result.objective} | ${result.strategy} | ${result.fold} | ${promotionLabel(result.promotionStatus)} | ${result.searchDiagnostics.parameterCount} | ${result.proxyEvaluations} | ${result.fullEvaluations} | ${percent(result.searchDiagnostics.proxyFullAgreement)} | ${staticWeights} | ${baseWeights} | ${influences} | ${percent(result.staticValidation.recallAt5)} | ${percent(result.validation.recallAt5)} | ${percent(result.staticValidation.recallAt10)} | ${percent(result.validation.recallAt10)} | ${percent(result.staticValidation.recallAt20)} | ${percent(result.validation.recallAt20)} | ${percent(result.validation.recallAt50)} | ${percent(result.staticValidation.contextRecallAt4096)} | ${percent(result.validation.contextRecallAt4096)} |`, ) @@ -319,9 +326,7 @@ export const renderMarkdownReport = (artifact: BenchmarkArtifact): string => { "| --- | --- | --- | --- | ---: | ---: | ---: | --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", ) for (const result of artifact.recommendedEvidenceRouters) { - const baseWeights = formatWeights(result.config.baseWeights) - const staticWeights = formatWeights(result.staticWeights) - const influences = formatInfluences(result.config) + const { baseWeights, staticWeights, influences } = formatRouterWeightColumns(result) lines.push( `| ${result.model} | ${result.fusion} | ${result.objective} | ${promotionLabel(result.promotionStatus)} | ${result.samples} | ${result.proxyEvaluations} | ${result.fullEvaluations} | ${staticWeights} | ${baseWeights} | ${influences} | ${percent(result.staticQuality.recallAt5)} | ${percent(result.fitQuality.recallAt5)} | ${percent(result.staticQuality.recallAt10)} | ${percent(result.fitQuality.recallAt10)} | ${percent(result.staticQuality.recallAt20)} | ${percent(result.fitQuality.recallAt20)} | ${percent(result.fitQuality.recallAt50)} | ${percent(result.staticQuality.contextRecallAt4096)} | ${percent(result.fitQuality.contextRecallAt4096)} |`, ) diff --git a/benchmarks/retrieval/evaluation/search.ts b/benchmarks/retrieval/evaluation/search.ts new file mode 100644 index 0000000..3610d8c --- /dev/null +++ b/benchmarks/retrieval/evaluation/search.ts @@ -0,0 +1,555 @@ +import { Effect } from "effect" + +import type { FusionMethod } from "../../../src/domain/retrieval.js" +import { + createCandidateEvaluationQueue, + getDefaultWorkerCount, + resolveWorkerCount, + type CandidateEvaluationQueue, +} from "../execution/candidate-evaluation-pool.js" +import type { OptimizationProfile } from "./optimization-profiles.js" +import type { + BenchmarkArtifact, + EvidenceRouterSearchResult, + RecommendedEvidenceRouter, + ValidationStrategy, +} from "./types.js" +import { + fitRecommendedEvidenceRouter, + fitRecommendedFusionWeights, + fitRecommendedWeights, + optimizeEvidenceRouter, + optimizeFusionWeights, + optimizeWeights, + evaluateProductionRouter, + type BenchmarkSearchOptions, + type WeightSearchSample, +} from "./weight-search.js" + +/** Search and validation stages enabled by one benchmark profile. */ +export interface BenchmarkSearchConfig { + /** Number of intent-grouped cross-validation folds. */ + readonly groupedFolds: number + /** Whether each selected repository is evaluated as an excluded holdout. */ + readonly repositoryHoldouts: boolean + /** Whether historical query-kind RRF grids and Shapley diagnostics run. */ + readonly legacyDiagnostics: boolean + /** Static fusion formulas evaluated by this profile. */ + readonly fusionMethods: readonly FusionMethod[] + /** Fusion formulas used when evaluating the evidence router. */ + readonly routerFusionMethods: readonly FusionMethod[] +} + +/** Quality search outputs and timing fields assembled for one benchmark artifact. */ +export interface BenchmarkSearchResults { + readonly weightSearch: readonly BenchmarkArtifact["weightSearch"][number][] + readonly recommendedWeights: readonly BenchmarkArtifact["recommendedWeights"][number][] + readonly productionRouterSearch: readonly BenchmarkArtifact["productionRouterSearch"][number][] + readonly fusionSearch: readonly BenchmarkArtifact["fusionSearch"][number][] + readonly recommendedFusionWeights: readonly BenchmarkArtifact["recommendedFusionWeights"][number][] + readonly evidenceRouterSearch: readonly BenchmarkArtifact["evidenceRouterSearch"][number][] + readonly recommendedEvidenceRouters: readonly BenchmarkArtifact["recommendedEvidenceRouters"][number][] + readonly weightSearchDurationMs: number + readonly fusionSearchDurationMs: number + readonly evidenceRouterSearchDurationMs: number + readonly candidateQueueStartupDurationMs: number + readonly candidateQueueShutdownDurationMs: number +} + +type SampleGroup = { + readonly model: string + readonly queryKind: WeightSearchSample["queryKind"] + readonly samples: readonly WeightSearchSample[] +} + +const reportProgress = (message: string): void => { + process.stderr.write(`[retrieval benchmark] ${message}\n`) +} + +const runParallelSearch = ( + operation: (signal: AbortSignal) => Promise, +): Effect.Effect => + Effect.tryPromise({ + try: (signal) => operation(signal), + catch: (cause) => (cause instanceof Error ? cause : new Error(String(cause))), + }) + +const splitSamples = ( + samples: readonly WeightSearchSample[], + isValidation: (sample: WeightSearchSample) => boolean, +): { + readonly development: readonly WeightSearchSample[] + readonly validation: readonly WeightSearchSample[] +} => ({ + development: samples.filter((sample) => !isValidation(sample)), + validation: samples.filter(isValidation), +}) + +interface EvidenceRouterHoldoutJob { + readonly kind: "holdout" + readonly model: string + readonly fusion: FusionMethod + readonly strategy: ValidationStrategy + readonly fold: string + readonly development: readonly WeightSearchSample[] + readonly validation: readonly WeightSearchSample[] +} + +interface EvidenceRouterFitAllJob { + readonly kind: "fit-all" + readonly model: string + readonly fusion: FusionMethod + readonly samples: readonly WeightSearchSample[] +} + +const planEvidenceRouterJobs = ( + samplesByModel: ReadonlyMap, + config: BenchmarkSearchConfig, + groupedStrategy: ValidationStrategy, +): readonly EvidenceRouterHoldoutJob[] => { + const jobs: EvidenceRouterHoldoutJob[] = [] + for (const [model, samples] of samplesByModel) { + const repositories = [...new Set(samples.map((sample) => sample.repository))] + for (const fusion of config.routerFusionMethods) { + for (let fold = 0; fold < config.groupedFolds; fold++) { + const split = splitSamples(samples, (sample) => sample.groupedFold === fold) + jobs.push({ + kind: "holdout", + model, + fusion, + strategy: groupedStrategy, + fold: String(fold + 1), + development: split.development, + validation: split.validation, + }) + } + if (config.repositoryHoldouts && repositories.length > 1) { + for (const repository of repositories) { + const split = splitSamples(samples, (sample) => sample.repository === repository) + jobs.push({ + kind: "holdout", + model, + fusion, + strategy: "leave-one-repository-out", + fold: repository, + development: split.development, + validation: split.validation, + }) + } + } + } + } + return jobs +} + +type RouterSearchJob = EvidenceRouterHoldoutJob | EvidenceRouterFitAllJob +type RouterSearchJobResult = + | { readonly kind: "holdout"; readonly results: readonly EvidenceRouterSearchResult[] } + | { readonly kind: "fit-all"; readonly results: readonly RecommendedEvidenceRouter[] } + +const runRouterSearchJob = ( + job: RouterSearchJob, + profile: OptimizationProfile, + options: BenchmarkSearchOptions, +): Promise => { + if (job.kind === "holdout") + return optimizeEvidenceRouter( + job.model, + job.fusion, + job.strategy, + job.fold, + job.development, + job.validation, + profile, + options, + ).then((results) => ({ kind: "holdout", results })) + return fitRecommendedEvidenceRouter(job.model, job.fusion, job.samples, profile, options).then( + (results) => ({ kind: "fit-all", results }), + ) +} + +interface SearchSplit { + readonly strategy: ValidationStrategy + readonly fold: string + readonly development: readonly WeightSearchSample[] + readonly validation: readonly WeightSearchSample[] +} + +const planSearchSplits = ( + samples: readonly WeightSearchSample[], + config: BenchmarkSearchConfig, + groupedStrategy: ValidationStrategy, +): readonly SearchSplit[] => { + const splits: SearchSplit[] = [] + for (let fold = 0; fold < config.groupedFolds; fold++) { + const split = splitSamples(samples, (sample) => sample.groupedFold === fold) + splits.push({ + strategy: groupedStrategy, + fold: String(fold + 1), + development: split.development, + validation: split.validation, + }) + } + const repositories = [...new Set(samples.map((sample) => sample.repository))] + if (config.repositoryHoldouts && repositories.length > 1) + for (const repository of repositories) { + const split = splitSamples(samples, (sample) => sample.repository === repository) + splits.push({ + strategy: "leave-one-repository-out", + fold: repository, + development: split.development, + validation: split.validation, + }) + } + return splits +} + +interface WeightSearchGroupResult { + readonly weightSearch: readonly BenchmarkArtifact["weightSearch"][number][] + readonly recommendedWeights: readonly BenchmarkArtifact["recommendedWeights"][number][] +} + +const runWeightSearchForGroup = ( + group: SampleGroup, + config: BenchmarkSearchConfig, + groupedStrategy: ValidationStrategy, + optimizationProfile: OptimizationProfile, + searchOptions: BenchmarkSearchOptions, +): Effect.Effect => + Effect.gen(function* () { + const weightSearch: BenchmarkArtifact["weightSearch"][number][] = [] + for (const split of planSearchSplits(group.samples, config, groupedStrategy)) + weightSearch.push( + yield* runParallelSearch((signal) => + optimizeWeights( + group.model, + group.queryKind, + split.strategy, + split.fold, + split.development, + split.validation, + optimizationProfile, + { ...searchOptions, signal }, + ), + ), + ) + const recommendedWeights = [ + yield* runParallelSearch((signal) => + fitRecommendedWeights(group.model, group.queryKind, group.samples, optimizationProfile, { + ...searchOptions, + signal, + }), + ), + ] + return { weightSearch, recommendedWeights } + }) + +const runWeightSearchStage = ( + config: BenchmarkSearchConfig, + sampleGroups: ReadonlyMap, + groupedStrategy: ValidationStrategy, + optimizationProfile: OptimizationProfile, + searchOptions: BenchmarkSearchOptions, +): Effect.Effect< + Pick, + Error +> => + Effect.gen(function* () { + const startedAt = performance.now() + const weightSearch: BenchmarkArtifact["weightSearch"][number][] = [] + const recommendedWeights: BenchmarkArtifact["recommendedWeights"][number][] = [] + if (config.legacyDiagnostics) + for (const group of sampleGroups.values()) { + const result = yield* runWeightSearchForGroup( + group, + config, + groupedStrategy, + optimizationProfile, + searchOptions, + ) + weightSearch.push(...result.weightSearch) + recommendedWeights.push(...result.recommendedWeights) + } + return { + weightSearch, + recommendedWeights, + weightSearchDurationMs: performance.now() - startedAt, + } + }) + +const runProductionRouterSearch = ( + config: BenchmarkSearchConfig, + samplesByModel: ReadonlyMap, + groupedStrategy: ValidationStrategy, + optimizationProfile: OptimizationProfile, +): readonly BenchmarkArtifact["productionRouterSearch"][number][] => { + const results: BenchmarkArtifact["productionRouterSearch"][number][] = [] + for (const [model, samples] of samplesByModel) + for (const split of planSearchSplits(samples, config, groupedStrategy)) + results.push( + evaluateProductionRouter( + model, + split.strategy, + split.fold, + split.development, + split.validation, + optimizationProfile, + ), + ) + return results +} + +const runStaticFusionSearchForModel = ( + model: string, + samples: readonly WeightSearchSample[], + config: BenchmarkSearchConfig, + groupedStrategy: ValidationStrategy, + optimizationProfile: OptimizationProfile, + searchOptions: BenchmarkSearchOptions, +): Effect.Effect => + Effect.gen(function* () { + const results: BenchmarkArtifact["fusionSearch"][number][] = [] + for (const fusion of config.fusionMethods) { + reportProgress(`${model}: selecting static ${fusion} fusion weights`) + for (const split of planSearchSplits(samples, config, groupedStrategy)) + results.push( + yield* runParallelSearch((signal) => + optimizeFusionWeights( + model, + fusion, + split.strategy, + split.fold, + split.development, + split.validation, + optimizationProfile, + { ...searchOptions, signal }, + ), + ), + ) + } + return results + }) + +const runRecommendedFusionSearchForModel = ( + model: string, + samples: readonly WeightSearchSample[], + config: BenchmarkSearchConfig, + optimizationProfile: OptimizationProfile, + searchOptions: BenchmarkSearchOptions, +): Effect.Effect => + Effect.gen(function* () { + const results: BenchmarkArtifact["recommendedFusionWeights"][number][] = [] + for (const fusion of config.fusionMethods) + results.push( + yield* runParallelSearch((signal) => + fitRecommendedFusionWeights(model, fusion, samples, optimizationProfile, { + ...searchOptions, + signal, + }), + ), + ) + return results + }) + +const runFusionSearchStage = ( + config: BenchmarkSearchConfig, + samplesByModel: ReadonlyMap, + groupedStrategy: ValidationStrategy, + optimizationProfile: OptimizationProfile, + searchOptions: BenchmarkSearchOptions, +): Effect.Effect< + Pick< + BenchmarkSearchResults, + | "productionRouterSearch" + | "fusionSearch" + | "recommendedFusionWeights" + | "fusionSearchDurationMs" + >, + Error +> => + Effect.gen(function* () { + const startedAt = performance.now() + const productionRouterSearch = runProductionRouterSearch( + config, + samplesByModel, + groupedStrategy, + optimizationProfile, + ) + const fusionSearch: BenchmarkArtifact["fusionSearch"][number][] = [] + for (const [model, samples] of samplesByModel) + fusionSearch.push( + ...(yield* runStaticFusionSearchForModel( + model, + samples, + config, + groupedStrategy, + optimizationProfile, + searchOptions, + )), + ) + const recommendedFusionWeights: BenchmarkArtifact["recommendedFusionWeights"][number][] = [] + for (const [model, samples] of samplesByModel) + recommendedFusionWeights.push( + ...(yield* runRecommendedFusionSearchForModel( + model, + samples, + config, + optimizationProfile, + searchOptions, + )), + ) + return { + productionRouterSearch, + fusionSearch, + recommendedFusionWeights, + fusionSearchDurationMs: performance.now() - startedAt, + } + }) + +const runRouterSearchJobs = ( + allRouterJobs: readonly RouterSearchJob[], + optimizationProfile: OptimizationProfile, + searchOptions: BenchmarkSearchOptions, + candidateQueue: CandidateEvaluationQueue | undefined, + canParallelize: boolean, +): Effect.Effect => + canParallelize + ? runParallelSearch((signal) => + Promise.all( + allRouterJobs.map((job) => + runRouterSearchJob(job, optimizationProfile, { + ...searchOptions, + workerCount: 0, + evaluationQueue: candidateQueue, + signal, + }), + ), + ), + ) + : Effect.forEach( + allRouterJobs, + (job) => + runParallelSearch((signal) => + runRouterSearchJob(job, optimizationProfile, { + ...searchOptions, + workerCount: 0, + signal, + }), + ), + { concurrency: 1 }, + ) + +const runEvidenceRouterSearchStage = ( + config: BenchmarkSearchConfig, + samplesByModel: ReadonlyMap, + groupedStrategy: ValidationStrategy, + optimizationProfile: OptimizationProfile, + serialSearch: boolean, + searchOptions: BenchmarkSearchOptions, +): Effect.Effect< + Pick< + BenchmarkSearchResults, + | "evidenceRouterSearch" + | "recommendedEvidenceRouters" + | "evidenceRouterSearchDurationMs" + | "candidateQueueStartupDurationMs" + | "candidateQueueShutdownDurationMs" + >, + Error +> => + Effect.gen(function* () { + const startedAt = performance.now() + const routerJobs = planEvidenceRouterJobs(samplesByModel, config, groupedStrategy) + const routerWorkerBudget = Math.min( + resolveWorkerCount(searchOptions.workerCount), + getDefaultWorkerCount(), + ) + const recommendedJobs: EvidenceRouterFitAllJob[] = [] + for (const [model, samples] of samplesByModel) + for (const fusion of config.routerFusionMethods) + recommendedJobs.push({ kind: "fit-all", model, fusion, samples }) + const allRouterJobs: readonly RouterSearchJob[] = [...routerJobs, ...recommendedJobs] + const canParallelize = !serialSearch && routerWorkerBudget >= 2 && allRouterJobs.length > 0 + const candidateWorkerCount = canParallelize ? routerWorkerBudget : 0 + let candidateQueueStartupDurationMs = 0 + let candidateQueueShutdownDurationMs = 0 + let candidateQueue: CandidateEvaluationQueue | undefined + if (canParallelize) { + const queueStartedAt = performance.now() + candidateQueue = yield* runParallelSearch(() => + createCandidateEvaluationQueue({ workerCount: candidateWorkerCount }), + ) + candidateQueueStartupDurationMs = performance.now() - queueStartedAt + } + reportProgress( + `running ${allRouterJobs.length} evidence-router jobs with ` + + `${candidateQueue?.workerCount ?? 0} shared candidate workers`, + ) + const routerResults = yield* Effect.ensuring( + runRouterSearchJobs( + allRouterJobs, + optimizationProfile, + searchOptions, + candidateQueue, + canParallelize, + ), + candidateQueue === undefined + ? Effect.void + : Effect.orDie( + Effect.gen(function* () { + const queueStartedAt = performance.now() + yield* runParallelSearch(() => candidateQueue!.close()) + candidateQueueShutdownDurationMs = performance.now() - queueStartedAt + }), + ), + ) + const evidenceRouterSearch: EvidenceRouterSearchResult[] = [] + const recommendedEvidenceRouters: RecommendedEvidenceRouter[] = [] + for (const result of routerResults) { + if (result.kind === "holdout") evidenceRouterSearch.push(...result.results) + else recommendedEvidenceRouters.push(...result.results) + } + return { + evidenceRouterSearch, + recommendedEvidenceRouters, + evidenceRouterSearchDurationMs: performance.now() - startedAt, + candidateQueueStartupDurationMs, + candidateQueueShutdownDurationMs, + } + }) + +/** Run all static and evidence-router quality searches over prepared samples. */ +export const runBenchmarkSearch = ( + config: BenchmarkSearchConfig, + sampleGroups: ReadonlyMap, + samplesByModel: ReadonlyMap, + groupedStrategy: ValidationStrategy, + optimizationProfile: OptimizationProfile, + serialSearch: boolean, + searchOptions: BenchmarkSearchOptions, +): Effect.Effect => + Effect.gen(function* () { + const weight = yield* runWeightSearchStage( + config, + sampleGroups, + groupedStrategy, + optimizationProfile, + searchOptions, + ) + const fusion = yield* runFusionSearchStage( + config, + samplesByModel, + groupedStrategy, + optimizationProfile, + searchOptions, + ) + const evidenceRouter = yield* runEvidenceRouterSearchStage( + config, + samplesByModel, + groupedStrategy, + optimizationProfile, + serialSearch, + searchOptions, + ) + return { ...weight, ...fusion, ...evidenceRouter } + }) diff --git a/benchmarks/retrieval/types.ts b/benchmarks/retrieval/evaluation/types.ts similarity index 99% rename from benchmarks/retrieval/types.ts rename to benchmarks/retrieval/evaluation/types.ts index 72f12d6..5f56cff 100644 --- a/benchmarks/retrieval/types.ts +++ b/benchmarks/retrieval/evaluation/types.ts @@ -3,7 +3,7 @@ import { Schema } from "effect" import type { EvidenceRouterConfig, FusionMethod as ProductionFusionMethod, -} from "../../src/domain/retrieval.js" +} from "../../../src/domain/retrieval.js" import type { OptimizationProfile } from "./optimization-profiles.js" /** Benchmark repository size band used for report segmentation. */ diff --git a/benchmarks/retrieval/weight-search.ts b/benchmarks/retrieval/evaluation/weight-search.ts similarity index 89% rename from benchmarks/retrieval/weight-search.ts rename to benchmarks/retrieval/evaluation/weight-search.ts index 3411c99..19ca1c3 100644 --- a/benchmarks/retrieval/weight-search.ts +++ b/benchmarks/retrieval/evaluation/weight-search.ts @@ -1,5 +1,5 @@ -import type { Chunk } from "../../src/domain/chunk.js" -import type { RankedChunk } from "../../src/domain/ports.js" +import type { Chunk } from "../../../src/domain/chunk.js" +import type { RankedChunk } from "../../../src/domain/ports.js" import { decodeEvidenceRouterConfig, PRODUCTION_COMPATIBILITY_CONFIG, @@ -12,16 +12,25 @@ import { CHANNEL_NAMES, type ChannelName, type ChannelRankings, -} from "../../src/domain/retrieval.js" +} from "../../../src/domain/retrieval.js" import { buildRoutingEvidence, routeWithEvidence, type QueryTermCoverage, type RoutingEvidence, -} from "../../src/lib/retrieval/evidence-router.js" -import { prepareFusion, type PreparedFusionEvaluator } from "./fusion.js" +} from "../../../src/lib/retrieval/evidence-router.js" +import { + createCandidateEvaluationPool, + createCandidateEvaluationPoolOnQueue, + createEvaluationSnapshot, + type CandidateEvaluationPool, + type CandidateEvaluationPoolOptions, + type CandidateEvaluationQueue, + type EvaluationCandidate, +} from "../execution/candidate-evaluation-pool.js" import { contextRecallAtBudget, recallAt, reciprocalRank } from "./metrics.js" import { SEARCH_PRIORITY_PROFILE, type OptimizationProfile } from "./optimization-profiles.js" +import { prepareFusion, type PreparedFusionEvaluator } from "./prepared-fusion.js" import { ROUTER_OBJECTIVES, ROUTER_SEARCH_STRATEGY, @@ -40,15 +49,6 @@ import { type SearchBaselineComparison, type WeightSearchResult, } from "./types.js" -import { - createCandidateEvaluationPool, - createCandidateEvaluationPoolOnQueue, - createEvaluationSnapshot, - type CandidateEvaluationPool, - type CandidateEvaluationPoolOptions, - type CandidateEvaluationQueue, - type EvaluationCandidate, -} from "./worker-pool.js" const CHANNELS: readonly ChannelName[] = CHANNEL_NAMES const WEIGHT_LEVELS = [0, 0.5, 1, 2] as const @@ -89,12 +89,15 @@ interface EvidenceSearchSample { } /** Native pool controls plus the cancellation signal owned by the benchmark Effect. */ -export interface ParallelSearchOptions extends CandidateEvaluationPoolOptions { +interface SearchOptions extends CandidateEvaluationPoolOptions { readonly signal?: AbortSignal /** Shared candidate scheduler used to interleave multiple router searches. */ readonly evaluationQueue?: CandidateEvaluationQueue } +/** Options for benchmark search APIs without colliding with the product query options type. */ +export type BenchmarkSearchOptions = SearchOptions + const preparedFusionCache = new WeakMap< WeightSearchSample, Map @@ -128,7 +131,7 @@ const createEvaluationPoolForSamples = async ( samples: readonly WeightSearchSample[], fusion: FusionMethod, profile: OptimizationProfile, - options: ParallelSearchOptions, + options: SearchOptions, ): Promise => { const snapshotPreparationStartedAt = performance.now() const snapshot = createEvaluationSnapshot( @@ -645,7 +648,7 @@ const rankWeightCandidates = async ( pending, qualities, qualityCache, - "Parallel weight evaluation returned an incomplete result", + "Candidate evaluation returned an incomplete weight result", ) return sortWeightCandidates(unique, qualityCache, limit, objective, baseline, profile) } @@ -1047,14 +1050,14 @@ const selectRandomRouter = async ( const config = configs[index] const quality = qualities[index] if (config === undefined || quality === undefined) - throw new Error("Parallel random router search returned an incomplete result") + throw new Error("Candidate evaluation returned an incomplete random router result") candidates.push({ config, quality }) } const candidate = [...candidates].sort((left, right) => compareRouterCandidates(left, right, "reranker-top20", baseline, profile), )[0] if (candidate === undefined) - throw new Error("Parallel random router search produced no candidate") + throw new Error("Candidate evaluation produced no random router candidate") return { candidate, candidates: candidates.length } } @@ -1144,7 +1147,7 @@ const evaluateRouterConfigs = async ( pending, qualities, qualityCache, - "Parallel router evaluation returned an incomplete result", + "Candidate evaluation returned an incomplete router result", ) return { candidates: entries.map(([key, config]) => { @@ -1430,7 +1433,7 @@ const withCandidatePool = async ( samples: readonly WeightSearchSample[], fusion: FusionMethod, profile: OptimizationProfile, - options: ParallelSearchOptions, + options: SearchOptions, operation: (pool: CandidateEvaluationPool) => Promise, ): Promise => { const pool = await createEvaluationPoolForSamples(samples, fusion, profile, options) @@ -1481,7 +1484,7 @@ const optimizeWeightsWithOptions = ( development: readonly WeightSearchSample[], validation: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, - options: ParallelSearchOptions, + options: SearchOptions, ): Promise => withCandidatePool(development, "rrf", profile, options, (pool) => optimizeWeightsWithPool( @@ -1504,20 +1507,7 @@ export const optimizeWeights = ( development: readonly WeightSearchSample[], validation: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): Promise => - optimizeWeightsWithOptions(model, queryKind, strategy, fold, development, validation, profile, { - workerCount: 0, - }) - -export const optimizeWeightsParallel = ( - model: string, - queryKind: QueryKind, - strategy: WeightSearchResult["strategy"], - fold: string, - development: readonly WeightSearchSample[], - validation: readonly WeightSearchSample[], - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, - options: ParallelSearchOptions = {}, + options: SearchOptions = { workerCount: 0 }, ): Promise => optimizeWeightsWithOptions( model, @@ -1580,7 +1570,7 @@ const optimizeFusionWeightsWithOptions = ( development: readonly WeightSearchSample[], validationSamples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, - options: ParallelSearchOptions, + options: SearchOptions, ): Promise => withCandidatePool(development, fusion, profile, options, (pool) => optimizeFusionWeightsWithPool( @@ -1603,27 +1593,7 @@ export const optimizeFusionWeights = ( development: readonly WeightSearchSample[], validationSamples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): Promise => - optimizeFusionWeightsWithOptions( - model, - fusion, - strategy, - fold, - development, - validationSamples, - profile, - { workerCount: 0 }, - ) - -export const optimizeFusionWeightsParallel = ( - model: string, - fusion: FusionMethod, - strategy: FusionSearchResult["strategy"], - fold: string, - development: readonly WeightSearchSample[], - validationSamples: readonly WeightSearchSample[], - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, - options: ParallelSearchOptions = {}, + options: SearchOptions = { workerCount: 0 }, ): Promise => optimizeFusionWeightsWithOptions( model, @@ -1654,11 +1624,11 @@ export const evaluateProductionRouter = ( validation: summarizeProductionRouter(validation, profile), }) -const withParallelEvidencePools = async ( +const withEvidencePools = async ( samples: readonly WeightSearchSample[], fusion: FusionMethod, profile: OptimizationProfile, - options: ParallelSearchOptions, + options: SearchOptions, operation: ( selection: Awaited>, fullPool: CandidateEvaluationPool, @@ -1744,7 +1714,7 @@ const selectStaticWeightsForSelections = async ( return selected } -const optimizeEvidenceRouterWithPools = async ( +export const optimizeEvidenceRouter = async ( model: string, fusion: FusionMethod, strategy: EvidenceRouterSearchResult["strategy"], @@ -1752,93 +1722,67 @@ const optimizeEvidenceRouterWithPools = async ( development: readonly WeightSearchSample[], validation: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, - options: ParallelSearchOptions = {}, + options: SearchOptions = { workerCount: 0 }, ): Promise => - withParallelEvidencePools( - development, - fusion, - profile, - options, - async (dynamicSelection, fullPool) => { - const productionValidation = summarizeProductionRouter(validation, profile) - const validationEvidence = prepareEvidenceSamples(validation) - const holdoutProfile = unweightedProfile(profile) - const randomBaseline: SearchBaselineComparison = { - algorithm: "random-scout", - seed: RANDOM_SEARCH_SEED, - candidates: dynamicSelection.randomCandidates, - development: dynamicSelection.randomCandidate.quality, - validation: summarizeEvidenceRouter( - validationEvidence, - dynamicSelection.randomCandidate.config, - fusion, + withEvidencePools(development, fusion, profile, options, async (dynamicSelection, fullPool) => { + const productionValidation = summarizeProductionRouter(validation, profile) + const validationEvidence = prepareEvidenceSamples(validation) + const holdoutProfile = unweightedProfile(profile) + const randomBaseline: SearchBaselineComparison = { + algorithm: "random-scout", + seed: RANDOM_SEARCH_SEED, + candidates: dynamicSelection.randomCandidates, + development: dynamicSelection.randomCandidate.quality, + validation: summarizeEvidenceRouter( + validationEvidence, + dynamicSelection.randomCandidate.config, + fusion, + profile, + ), + } + const staticSelections = await selectStaticWeightsForSelections( + dynamicSelection.selections, + fullPool, + dynamicSelection.productionQuality, + profile, + ) + const results: EvidenceRouterSearchResult[] = staticSelections.map( + ({ selection, staticSelection }) => ({ + model, + fusion, + objective: selection.objective, + strategy, + fold, + developmentQueries: development.length, + validationQueries: validation.length, + staticWeights: staticSelection.weights, + config: benchmarkRouterConfig(fusion, selection.config), + staticDevelopment: staticSelection.quality, + staticValidation: summarize(validation, staticSelection.weights, fusion, profile), + development: selection.quality, + validation: summarizeEvidenceRouter(validationEvidence, selection.config, fusion, profile), + productionDevelopment: dynamicSelection.productionQuality, + productionValidation, + guardrailsMet: selection.guardrailsMet, + promotionStatus: selection.promotionStatus, + proxyEvaluations: dynamicSelection.proxyEvaluations, + fullEvaluations: dynamicSelection.fullEvaluations, + searchDiagnostics: dynamicSelection.searchDiagnostics, + searchBaseline: randomBaseline, + holdoutBreakdown: buildHoldoutBreakdown( + validation, + (partition) => + summarizeEvidenceRouter( + prepareEvidenceSamples(partition), + selection.config, + fusion, + holdoutProfile, + ), profile, ), - } - const staticSelections = await selectStaticWeightsForSelections( - dynamicSelection.selections, - fullPool, - dynamicSelection.productionQuality, - profile, - ) - const results: EvidenceRouterSearchResult[] = staticSelections.map( - ({ selection, staticSelection }) => ({ - model, - fusion, - objective: selection.objective, - strategy, - fold, - developmentQueries: development.length, - validationQueries: validation.length, - staticWeights: staticSelection.weights, - config: benchmarkRouterConfig(fusion, selection.config), - staticDevelopment: staticSelection.quality, - staticValidation: summarize(validation, staticSelection.weights, fusion, profile), - development: selection.quality, - validation: summarizeEvidenceRouter( - validationEvidence, - selection.config, - fusion, - profile, - ), - productionDevelopment: dynamicSelection.productionQuality, - productionValidation, - guardrailsMet: selection.guardrailsMet, - promotionStatus: selection.promotionStatus, - proxyEvaluations: dynamicSelection.proxyEvaluations, - fullEvaluations: dynamicSelection.fullEvaluations, - searchDiagnostics: dynamicSelection.searchDiagnostics, - searchBaseline: randomBaseline, - holdoutBreakdown: buildHoldoutBreakdown( - validation, - (partition) => - summarizeEvidenceRouter( - prepareEvidenceSamples(partition), - selection.config, - fusion, - holdoutProfile, - ), - profile, - ), - }), - ) - return results - }, - ) - -export const optimizeEvidenceRouterParallel = optimizeEvidenceRouterWithPools - -export const optimizeEvidenceRouter = ( - model: string, - fusion: FusionMethod, - strategy: EvidenceRouterSearchResult["strategy"], - fold: string, - development: readonly WeightSearchSample[], - validation: readonly WeightSearchSample[], - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): Promise => - optimizeEvidenceRouterParallel(model, fusion, strategy, fold, development, validation, profile, { - workerCount: 0, + }), + ) + return results }) /** Fit one deployment candidate on all samples after cross-validation has measured generalization. */ @@ -1864,7 +1808,7 @@ const fitRecommendedWeightsWithOptions = ( queryKind: QueryKind, samples: readonly WeightSearchSample[], profile: OptimizationProfile, - options: ParallelSearchOptions, + options: SearchOptions, ): Promise => withCandidatePool(samples, "rrf", profile, options, (pool) => fitRecommendedWeightsWithPool(model, queryKind, samples, profile, pool), @@ -1875,15 +1819,7 @@ export const fitRecommendedWeights = ( queryKind: QueryKind, samples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): Promise => - fitRecommendedWeightsWithOptions(model, queryKind, samples, profile, { workerCount: 0 }) - -export const fitRecommendedWeightsParallel = ( - model: string, - queryKind: QueryKind, - samples: readonly WeightSearchSample[], - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, - options: ParallelSearchOptions = {}, + options: SearchOptions = { workerCount: 0 }, ): Promise => fitRecommendedWeightsWithOptions(model, queryKind, samples, profile, options) @@ -1920,7 +1856,7 @@ const fitRecommendedFusionWeightsWithOptions = ( fusion: FusionMethod, samples: readonly WeightSearchSample[], profile: OptimizationProfile, - options: ParallelSearchOptions, + options: SearchOptions, ): Promise => withCandidatePool(samples, fusion, profile, options, (pool) => fitRecommendedFusionWeightsWithPool(model, fusion, samples, profile, pool), @@ -1931,15 +1867,7 @@ export const fitRecommendedFusionWeights = ( fusion: FusionMethod, samples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): Promise => - fitRecommendedFusionWeightsWithOptions(model, fusion, samples, profile, { workerCount: 0 }) - -export const fitRecommendedFusionWeightsParallel = ( - model: string, - fusion: FusionMethod, - samples: readonly WeightSearchSample[], - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, - options: ParallelSearchOptions = {}, + options: SearchOptions = { workerCount: 0 }, ): Promise => fitRecommendedFusionWeightsWithOptions(model, fusion, samples, profile, options) @@ -1952,55 +1880,41 @@ const fitRecommendedEvidenceRouterWithOptions = ( fusion: FusionMethod, samples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, - options: ParallelSearchOptions, + options: SearchOptions, ): Promise => - withParallelEvidencePools( - samples, - fusion, - profile, - options, - async (dynamicSelection, fullPool) => { - const staticSelections = await selectStaticWeightsForSelections( - dynamicSelection.selections, - fullPool, - dynamicSelection.productionQuality, - profile, - ) - const results: RecommendedEvidenceRouter[] = staticSelections.map( - ({ selection, staticSelection }) => ({ - model, - fusion, - objective: selection.objective, - samples: samples.length, - staticWeights: staticSelection.weights, - config: benchmarkRouterConfig(fusion, selection.config), - staticQuality: staticSelection.quality, - fitQuality: selection.quality, - productionQuality: dynamicSelection.productionQuality, - guardrailsMet: selection.guardrailsMet, - promotionStatus: selection.promotionStatus, - proxyEvaluations: dynamicSelection.proxyEvaluations, - fullEvaluations: dynamicSelection.fullEvaluations, - searchDiagnostics: dynamicSelection.searchDiagnostics, - }), - ) - return results - }, - ) + withEvidencePools(samples, fusion, profile, options, async (dynamicSelection, fullPool) => { + const staticSelections = await selectStaticWeightsForSelections( + dynamicSelection.selections, + fullPool, + dynamicSelection.productionQuality, + profile, + ) + const results: RecommendedEvidenceRouter[] = staticSelections.map( + ({ selection, staticSelection }) => ({ + model, + fusion, + objective: selection.objective, + samples: samples.length, + staticWeights: staticSelection.weights, + config: benchmarkRouterConfig(fusion, selection.config), + staticQuality: staticSelection.quality, + fitQuality: selection.quality, + productionQuality: dynamicSelection.productionQuality, + guardrailsMet: selection.guardrailsMet, + promotionStatus: selection.promotionStatus, + proxyEvaluations: dynamicSelection.proxyEvaluations, + fullEvaluations: dynamicSelection.fullEvaluations, + searchDiagnostics: dynamicSelection.searchDiagnostics, + }), + ) + return results + }) export const fitRecommendedEvidenceRouter = ( model: string, fusion: FusionMethod, samples: readonly WeightSearchSample[], profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): Promise => - fitRecommendedEvidenceRouterWithOptions(model, fusion, samples, profile, { workerCount: 0 }) - -export const fitRecommendedEvidenceRouterParallel = ( - model: string, - fusion: FusionMethod, - samples: readonly WeightSearchSample[], - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, - options: ParallelSearchOptions = {}, + options: SearchOptions = { workerCount: 0 }, ): Promise => fitRecommendedEvidenceRouterWithOptions(model, fusion, samples, profile, options) diff --git a/benchmarks/retrieval/worker-pool.ts b/benchmarks/retrieval/execution/candidate-evaluation-pool.ts similarity index 67% rename from benchmarks/retrieval/worker-pool.ts rename to benchmarks/retrieval/execution/candidate-evaluation-pool.ts index 0d848ef..f78039e 100644 --- a/benchmarks/retrieval/worker-pool.ts +++ b/benchmarks/retrieval/execution/candidate-evaluation-pool.ts @@ -1,16 +1,18 @@ import { availableParallelism } from "node:os" import { Worker } from "node:worker_threads" -import type { Chunk } from "../../src/domain/chunk.js" -import type { ChannelWeights } from "../../src/domain/retrieval.js" -import { evaluateCandidate } from "./fusion-core.mjs" -import { type PreparedFusionEvaluator, type PreparedFusionSnapshot } from "./fusion.js" -import type { QualitySummary } from "./types.js" -import { hasWorkerMessageType } from "./worker-message.js" +import type { Chunk } from "../../../src/domain/chunk.js" +import type { ChannelWeights } from "../../../src/domain/retrieval.js" +import { evaluateCandidate } from "../evaluation/prepared-fusion-core.mjs" +import { + type PreparedFusionEvaluator, + type PreparedFusionSnapshot, +} from "../evaluation/prepared-fusion.js" +import type { QualitySummary } from "../evaluation/types.js" const DEFAULT_BATCH_SIZE = 32 const DEFAULT_QUEUE_BATCH_SIZE = 1 -const DEFAULT_WORKER_URL = new URL("./fusion-worker.mjs", import.meta.url) +const DEFAULT_WORKER_URL = new URL("./candidate-evaluation-worker.mjs", import.meta.url) /** One prepared benchmark sample represented without source text for worker transfer. */ export interface EvaluationSampleSnapshot { @@ -122,12 +124,6 @@ interface WorkerErrorMessage { type WorkerMessage = WorkerReadyMessage | WorkerResultMessage | WorkerErrorMessage -interface WorkerTask { - readonly taskId: number - readonly start: number - readonly end: number -} - interface ReadyWorker { readonly worker: Worker readonly ready: Promise @@ -135,21 +131,6 @@ interface ReadyWorker { readonly rejectReady: (error: Error) => void } -interface WorkerSlot extends ReadyWorker { - busy: boolean - task?: WorkerTask -} - -interface ActiveEvaluation { - readonly candidates: readonly EvaluationCandidate[] - readonly results: Map - readonly tasks: Map - readonly resolve: (results: readonly QualitySummary[]) => void - readonly reject: (error: Error) => void - nextCandidate: number - completedCandidates: number -} - const errorFromUnknown = (cause: unknown): Error => cause instanceof Error ? cause : new Error(String(cause)) @@ -174,31 +155,16 @@ const attachWorkerLifecycle = ( worker.on("exit", onExit) } -const attachCandidateWorkerLifecycle = ( +const attachFusionWorkerLifecycle = ( slot: ReadyWorker, onMessage: (message: unknown) => void, onError: (cause: Error) => void, isClosed: () => boolean, - workerName: string, ): void => attachWorkerLifecycle(slot.worker, onMessage, onError, (code) => { - if (!isClosed()) onError(new Error(`${workerName} worker exited with code ${code}`)) + if (!isClosed()) onError(new Error(`Candidate evaluation worker exited with code ${code}`)) }) -const attachFusionWorkerLifecycle = ( - slot: ReadyWorker, - onMessage: (message: unknown) => void, - onError: (cause: Error) => void, - isClosed: () => boolean, -): void => attachCandidateWorkerLifecycle(slot, onMessage, onError, isClosed, "Fusion") - -const attachFusionQueueWorkerLifecycle = ( - slot: ReadyWorker, - onMessage: (message: unknown) => void, - onError: (cause: Error) => void, - isClosed: () => boolean, -): void => attachCandidateWorkerLifecycle(slot, onMessage, onError, isClosed, "Fusion queue") - const createWorkerSlots = ( workerCount: number, workerUrl: URL, @@ -326,227 +292,12 @@ export const evaluateCandidatesSerial = ( ): readonly QualitySummary[] => candidates.map((candidate) => evaluateCandidate(snapshot, candidate)) -const isWorkerMessage = (message: unknown): message is WorkerMessage => { - return hasWorkerMessageType(message, ["ready", "result", "error"]) -} - -class NativeCandidateEvaluationPool implements CandidateEvaluationPool { - readonly mode = "parallel" as const - readonly workerCount: number - readonly batchSize: number - - private readonly slots: WorkerSlot[] - private readonly ready: Promise - private active: ActiveEvaluation | undefined - private closePromise: Promise | undefined - private closed = false - private nextTaskId = 0 - private batchCount = 0 - private candidateCount = 0 - - private constructor( - snapshot: EvaluationSnapshot, - workerCount: number, - batchSize: number, - workerUrl: URL, - ) { - this.workerCount = workerCount - this.batchSize = batchSize - this.slots = createWorkerSlots( - workerCount, - workerUrl, - (readyWorker): WorkerSlot => ({ ...readyWorker, busy: false }), - (slot) => - attachFusionWorkerLifecycle( - slot, - (message) => this.handleMessage(slot, message), - (cause) => this.handleWorkerError(slot, cause), - () => this.closed, - ), - ) - this.ready = Promise.all(this.slots.map((slot) => slot.ready)).then(() => undefined) - for (const slot of this.slots) slot.worker.postMessage({ type: "init", snapshot }) - } - - /** Start a native pool after every worker has accepted its immutable snapshot. */ - static async create( - snapshot: EvaluationSnapshot, - workerCount: number, - batchSize: number, - workerUrl: URL, - ): Promise { - let pool: NativeCandidateEvaluationPool | undefined - try { - pool = new NativeCandidateEvaluationPool(snapshot, workerCount, batchSize, workerUrl) - await pool.ready - return pool - } catch (cause) { - if (pool !== undefined) await pool.close() - throw errorFromUnknown(cause) - } - } - - private handleWorkerError(slot: WorkerSlot, cause: Error): void { - if (this.closed) return - const active = this.active - failWorker( - slot, - cause, - () => { - if (active !== undefined) { - this.active = undefined - active.reject(cause) - } - }, - () => this.close(), - ) - } - - private handleMessage(slot: WorkerSlot, message: unknown): void { - if (!isWorkerMessage(message)) { - this.handleWorkerError(slot, new Error("Fusion worker sent an invalid message")) - return - } - if (handleWorkerControlMessage(slot, message, (cause) => this.handleWorkerError(slot, cause))) - return - if (message.type !== "result") return - const active = this.active - const task = active?.tasks.get(message.taskId) - if (active === undefined || task === undefined) { - this.handleWorkerError( - slot, - new Error(`Fusion worker returned unknown task ${message.taskId}`), - ) - return - } - if (message.results.length !== task.end - task.start) { - this.handleWorkerError(slot, new Error(`Fusion worker returned an invalid result length`)) - return - } - for (let index = 0; index < message.results.length; index++) { - active.results.set(task.start + index, message.results[index]) - } - active.tasks.delete(message.taskId) - slot.busy = false - slot.task = undefined - active.completedCandidates += message.results.length - if (active.completedCandidates === active.candidates.length) { - const ordered: QualitySummary[] = [] - for (let index = 0; index < active.candidates.length; index++) { - const result = active.results.get(index) - if (result === undefined) { - this.handleWorkerError(slot, new Error(`Fusion worker omitted candidate ${index}`)) - return - } - ordered.push(result) - } - this.active = undefined - active.resolve(ordered) - return - } - this.dispatch() - } - - private dispatch(): void { - const active = this.active - if (active === undefined || this.closed) return - for (const slot of this.slots) { - if (slot.busy || active.nextCandidate >= active.candidates.length) continue - const start = active.nextCandidate - const end = Math.min(start + this.batchSize, active.candidates.length) - const task: WorkerTask = { taskId: this.nextTaskId++, start, end } - active.nextCandidate = end - active.tasks.set(task.taskId, task) - slot.busy = true - slot.task = task - this.batchCount++ - try { - slot.worker.postMessage({ - type: "evaluate", - taskId: task.taskId, - candidates: active.candidates.slice(start, end), - }) - } catch (cause) { - this.handleWorkerError(slot, errorFromUnknown(cause)) - return - } - } - } - - async evaluate(candidates: readonly EvaluationCandidate[]): Promise { - if (this.closed) throw new Error("Candidate evaluation pool is closed") - if (this.active !== undefined) throw new Error("Candidate evaluation pool is already busy") - await this.ready - if (candidates.length === 0) return [] - this.candidateCount += candidates.length - return new Promise((resolve, reject) => { - this.active = { - candidates, - results: new Map(), - tasks: new Map(), - resolve, - reject, - nextCandidate: 0, - completedCandidates: 0, - } - this.dispatch() - }) - } - - async close(): Promise { - if (this.closePromise !== undefined) return this.closePromise - this.closed = true - const active = this.active - this.active = undefined - active?.reject(new Error("Candidate evaluation pool closed")) - for (const slot of this.slots) { - slot.busy = false - slot.task = undefined - } - this.closePromise = Promise.all( - this.slots.map((slot) => slot.worker.terminate().catch(() => -1)), - ).then(() => undefined) - return this.closePromise - } - - stats = (): CandidateEvaluationPoolStats => ({ - mode: this.mode, - workerCount: this.workerCount, - activeWorkerCount: this.closed ? 0 : this.slots.filter((slot) => slot.busy).length, - batchSize: this.batchSize, - batches: this.batchCount, - candidates: this.candidateCount, - }) -} - -interface QueueWorkerReadyMessage { - readonly type: "ready" -} - -interface QueueWorkerResultMessage { - readonly type: "result" - readonly taskId: number - readonly results: readonly QualitySummary[] -} - -interface QueueWorkerErrorMessage { - readonly type: "error" - readonly taskId?: number - readonly message: string -} - -type QueueWorkerMessage = - | QueueWorkerReadyMessage - | QueueWorkerResultMessage - | QueueWorkerErrorMessage - interface QueueWorkerSlot extends ReadyWorker { readonly knownSnapshots: Set busy: boolean } interface CandidateQueueRequest { - readonly requestId: number readonly snapshotId: string readonly snapshot: EvaluationSnapshot readonly candidates: readonly EvaluationCandidate[] @@ -571,13 +322,15 @@ interface CandidateQueueTask { const isQueueRecord = (message: unknown): message is Record => typeof message === "object" && message !== null -const isQueueWorkerMessage = (message: unknown): message is QueueWorkerMessage => { +const isQueueWorkerMessage = (message: unknown): message is WorkerMessage => { if (!isQueueRecord(message) || typeof message.type !== "string") return false - if (typeof message.type !== "string") return false if (message.type === "ready") return true if (message.type === "error") return typeof message.message === "string" return ( - Number.isInteger(message.taskId) && Array.isArray(message.results) && message.results.length > 0 + message.type === "result" && + Number.isInteger(message.taskId) && + Array.isArray(message.results) && + message.results.length > 0 ) } @@ -593,7 +346,6 @@ class NativeCandidateEvaluationQueue implements CandidateEvaluationQueue { private readonly snapshotIds = new WeakMap() private closePromise: Promise | undefined private closed = false - private nextRequestId = 0 private nextSnapshotId = 0 private nextTaskId = 0 @@ -609,7 +361,7 @@ class NativeCandidateEvaluationQueue implements CandidateEvaluationQueue { busy: false, }), (slot) => - attachFusionQueueWorkerLifecycle( + attachFusionWorkerLifecycle( slot, (message) => this.handleMessage(slot, message), (cause) => this.handleWorkerError(slot, cause), @@ -688,7 +440,7 @@ class NativeCandidateEvaluationQueue implements CandidateEvaluationQueue { for (let index = 0; index < request.candidates.length; index++) { const result = request.results.get(index) if (result === undefined) { - request.reject(new Error(`Fusion queue omitted candidate ${index}`)) + request.reject(new Error(`Candidate evaluation queue omitted candidate ${index}`)) return } ordered.push(result) @@ -737,7 +489,7 @@ class NativeCandidateEvaluationQueue implements CandidateEvaluationQueue { private handleMessage(slot: QueueWorkerSlot, message: unknown): void { if (this.closed) return if (!isQueueWorkerMessage(message)) { - this.handleWorkerError(slot, new Error("Fusion queue worker sent an invalid message")) + this.handleWorkerError(slot, new Error("Candidate evaluation worker sent an invalid message")) return } if (handleWorkerControlMessage(slot, message, (cause) => this.handleWorkerError(slot, cause))) @@ -747,14 +499,17 @@ class NativeCandidateEvaluationQueue implements CandidateEvaluationQueue { if (task === undefined || task.slot !== slot) { this.handleWorkerError( slot, - new Error(`Fusion queue returned unknown task ${message.taskId}`), + new Error(`Candidate evaluation queue returned unknown task ${message.taskId}`), ) return } this.activeTasks.delete(message.taskId) slot.busy = false if (message.results.length !== task.end - task.start) { - this.handleWorkerError(slot, new Error("Fusion queue returned an invalid result length")) + this.handleWorkerError( + slot, + new Error("Candidate evaluation queue returned an invalid result length"), + ) return } if (!task.request.settled) { @@ -792,7 +547,6 @@ class NativeCandidateEvaluationQueue implements CandidateEvaluationQueue { } const removeAbortListener = () => signal?.removeEventListener("abort", abort) request = { - requestId: this.nextRequestId++, snapshotId: resolvedSnapshotId, snapshot, candidates, @@ -860,20 +614,28 @@ class SerialCandidateEvaluationQueue implements CandidateEvaluationQueue { } class QueuedCandidateEvaluationPool implements CandidateEvaluationPool { - readonly mode = "parallel" as const + readonly mode: EvaluationPoolMode readonly workerCount: number readonly batchSize: number private readonly snapshot: EvaluationSnapshot private readonly queue: CandidateEvaluationQueue private readonly signal: AbortSignal | undefined + private readonly ownsQueue: boolean private closed = false private batchCount = 0 private candidateCount = 0 - constructor(snapshot: EvaluationSnapshot, queue: CandidateEvaluationQueue, signal?: AbortSignal) { + constructor( + snapshot: EvaluationSnapshot, + queue: CandidateEvaluationQueue, + signal?: AbortSignal, + ownsQueue = false, + ) { this.snapshot = snapshot this.queue = queue this.signal = signal + this.ownsQueue = ownsQueue + this.mode = queue.workerCount > 1 ? "parallel" : "serial" this.workerCount = queue.workerCount this.batchSize = queue.batchSize } @@ -887,7 +649,9 @@ class QueuedCandidateEvaluationPool implements CandidateEvaluationPool { } async close(): Promise { + if (this.closed) return this.closed = true + if (this.ownsQueue) await this.queue.close() } stats = (): CandidateEvaluationPoolStats => @@ -922,42 +686,16 @@ export const createCandidateEvaluationPoolOnQueue = ( signal?: AbortSignal, ): CandidateEvaluationPool => new QueuedCandidateEvaluationPool(snapshot, queue, signal) -class SerialCandidateEvaluationPool implements CandidateEvaluationPool { - readonly mode = "serial" as const - readonly workerCount = 1 - readonly batchSize: number - private closed = false - private readonly snapshot: EvaluationSnapshot - private batchCount = 0 - private candidateCount = 0 - - constructor(snapshot: EvaluationSnapshot, batchSize: number) { - this.snapshot = snapshot - this.batchSize = batchSize - } - - async evaluate(candidates: readonly EvaluationCandidate[]): Promise { - if (this.closed) throw new Error("Candidate evaluation pool is closed") - const counts = candidateBatchStats(candidates, this.batchSize) - this.candidateCount += counts.candidates - this.batchCount += counts.batches - return evaluateCandidatesSerial(this.snapshot, candidates) - } - - async close(): Promise { - this.closed = true - } - - stats = (): CandidateEvaluationPoolStats => - candidatePoolStats( - this.mode, - this.workerCount, - 0, - this.batchSize, - this.batchCount, - this.candidateCount, - ) -} +const createOwnedSerialPool = ( + snapshot: EvaluationSnapshot, + batchSize: number, +): CandidateEvaluationPool => + new QueuedCandidateEvaluationPool( + snapshot, + new SerialCandidateEvaluationQueue(batchSize), + undefined, + true, + ) /** Create a fixed benchmark evaluator pool; one worker or an explicit zero uses serial fallback. */ export const createCandidateEvaluationPool = async ( @@ -967,18 +705,17 @@ export const createCandidateEvaluationPool = async ( const requestedWorkerCount = resolveWorkerCount(options.workerCount) const batchSize = resolveBatchSize(options.batchSize) if (requestedWorkerCount <= 1 || snapshot.samples.length === 0) - return new SerialCandidateEvaluationPool(snapshot, batchSize) + return createOwnedSerialPool(snapshot, batchSize) try { - return await NativeCandidateEvaluationPool.create( - snapshot, - requestedWorkerCount, + const queue = await createCandidateEvaluationQueue({ + workerCount: requestedWorkerCount, batchSize, - options.workerUrl ?? DEFAULT_WORKER_URL, - ) + workerUrl: options.workerUrl ?? DEFAULT_WORKER_URL, + }) + return new QueuedCandidateEvaluationPool(snapshot, queue, undefined, true) } catch (cause) { - if (options.fallbackToSerial !== false) - return new SerialCandidateEvaluationPool(snapshot, batchSize) + if (options.fallbackToSerial !== false) return createOwnedSerialPool(snapshot, batchSize) throw errorFromUnknown(cause) } } diff --git a/benchmarks/retrieval/fusion-worker.mjs b/benchmarks/retrieval/execution/candidate-evaluation-worker.mjs similarity index 65% rename from benchmarks/retrieval/fusion-worker.mjs rename to benchmarks/retrieval/execution/candidate-evaluation-worker.mjs index 41aa5f7..5c401e5 100644 --- a/benchmarks/retrieval/fusion-worker.mjs +++ b/benchmarks/retrieval/execution/candidate-evaluation-worker.mjs @@ -1,6 +1,6 @@ import { parentPort } from "node:worker_threads" -import { evaluateCandidate } from "./fusion-core.mjs" +import { evaluateCandidate } from "../evaluation/prepared-fusion-core.mjs" if (parentPort === null) throw new Error("Fusion worker requires a parent port") @@ -10,15 +10,13 @@ parentPort.postMessage({ type: "ready" }) parentPort.on("message", (message) => { try { - if (message.type === "init") { - snapshots.set("legacy", message.snapshot) - return - } - if (message.type !== "evaluate") throw new Error("Fusion worker received an invalid task") - const snapshotId = message.snapshotId ?? "legacy" + if (message.type !== "evaluate") + throw new Error("Candidate evaluation worker received an invalid task") + const snapshotId = message.snapshotId if (message.snapshot !== undefined) snapshots.set(snapshotId, message.snapshot) const snapshot = snapshots.get(snapshotId) - if (snapshot === undefined) throw new Error(`Fusion worker has no snapshot ${snapshotId}`) + if (snapshot === undefined) + throw new Error(`Candidate evaluation worker has no snapshot ${snapshotId}`) parentPort.postMessage({ type: "result", taskId: message.taskId, diff --git a/benchmarks/retrieval/sqlite-index.ts b/benchmarks/retrieval/execution/sqlite-index.ts similarity index 73% rename from benchmarks/retrieval/sqlite-index.ts rename to benchmarks/retrieval/execution/sqlite-index.ts index 599c67c..15408ec 100644 --- a/benchmarks/retrieval/sqlite-index.ts +++ b/benchmarks/retrieval/execution/sqlite-index.ts @@ -1,12 +1,12 @@ import { Effect, Layer } from "effect" import { layerNoop } from "effect/FileSystem" -import { DEFAULT_CONFIG, type Config } from "../../src/domain/config.js" -import type { EmbeddingDtype } from "../../src/domain/dtype.js" -import { ConfigStore, IndexStore, SparseEmbedder } from "../../src/domain/ports.js" -import { SparseEmbedderBase } from "../../src/services/sparse-embedder.js" -import { SqliteIndexStoreBase } from "../../src/services/sqlite-index-store.js" -import { sqliteIndexDatabaseLayer } from "../../src/services/sqlite-index-store/client.js" +import { DEFAULT_CONFIG, type Config } from "../../../src/domain/config.js" +import type { EmbeddingDtype } from "../../../src/domain/dtype.js" +import { ConfigStore, IndexStore, SparseEmbedder } from "../../../src/domain/ports.js" +import { SparseEmbedderBase } from "../../../src/services/sparse-embedder.js" +import { SqliteIndexStoreBase } from "../../../src/services/sqlite-index-store.js" +import { sqliteIndexDatabaseLayer } from "../../../src/services/sqlite-index-store/client.js" const benchmarkConfig = (model: string, dtype: EmbeddingDtype): Config => ({ ...DEFAULT_CONFIG, diff --git a/benchmarks/retrieval/router-job-pool.ts b/benchmarks/retrieval/router-job-pool.ts deleted file mode 100644 index b0ee86f..0000000 --- a/benchmarks/retrieval/router-job-pool.ts +++ /dev/null @@ -1,318 +0,0 @@ -import { Worker } from "node:worker_threads" - -import type { FusionMethod } from "../../src/domain/retrieval.js" -import type { OptimizationProfile } from "./optimization-profiles.js" -import type { - EvidenceRouterSearchResult, - RecommendedEvidenceRouter, - ValidationStrategy, -} from "./types.js" -import type { WeightSearchSample } from "./weight-search.js" -import type { - CandidateEvaluationQueue, - EvaluationCandidate, - EvaluationSnapshot, -} from "./worker-pool.js" - -/** One independent evidence-router fold or repository holdout. */ -export interface EvidenceRouterHoldoutJob { - readonly kind: "holdout" - readonly model: string - readonly fusion: FusionMethod - readonly strategy: ValidationStrategy - readonly fold: string - readonly development: readonly WeightSearchSample[] - readonly validation: readonly WeightSearchSample[] -} - -/** One fit-all router search over every sample after holdout evaluation. */ -export interface EvidenceRouterFitAllJob { - readonly kind: "fit-all" - readonly model: string - readonly fusion: FusionMethod - readonly samples: readonly WeightSearchSample[] -} - -/** One complete router search handled by a native worker. */ -export type EvidenceRouterJob = EvidenceRouterHoldoutJob | EvidenceRouterFitAllJob - -/** Tagged result returned by one holdout router job. */ -export interface EvidenceRouterHoldoutJobResult { - readonly jobId: number - readonly kind: "holdout" - readonly results: readonly EvidenceRouterSearchResult[] -} - -/** Tagged result returned by one fit-all router job. */ -export interface EvidenceRouterFitAllJobResult { - readonly jobId: number - readonly kind: "fit-all" - readonly results: readonly RecommendedEvidenceRouter[] -} - -/** Tagged result returned by one complete native router job. */ -export type EvidenceRouterJobResult = EvidenceRouterHoldoutJobResult | EvidenceRouterFitAllJobResult - -/** Options for the native pool that runs complete router jobs. */ -export interface EvidenceRouterJobPoolOptions { - /** Maximum number of native router-controller workers. */ - readonly workerCount: number - /** Shared candidate queue used by the active router controllers. */ - readonly candidateQueue?: CandidateEvaluationQueue - /** Abort signal that terminates all active router jobs. */ - readonly signal?: AbortSignal - /** Test-only override for the router job worker entry point. */ - readonly workerUrl?: URL - /** Test-only override for the TypeScript ESM loader. */ - readonly loaderUrl?: URL -} - -interface WorkerReadyMessage { - readonly type: "ready" -} - -interface WorkerHoldoutResultMessage { - readonly type: "result" - readonly jobId: number - readonly kind: "holdout" - readonly results: readonly EvidenceRouterSearchResult[] -} - -interface WorkerFitAllResultMessage { - readonly type: "result" - readonly jobId: number - readonly kind: "fit-all" - readonly results: readonly RecommendedEvidenceRouter[] -} - -interface WorkerCandidateEvaluateMessage { - readonly type: "candidate-evaluate" - readonly requestId: number - readonly snapshotId: string - readonly snapshot?: EvaluationSnapshot - readonly candidates: readonly EvaluationCandidate[] -} - -interface WorkerErrorMessage { - readonly type: "error" - readonly jobId?: number - readonly message: string -} - -type WorkerMessage = - | WorkerReadyMessage - | WorkerHoldoutResultMessage - | WorkerFitAllResultMessage - | WorkerCandidateEvaluateMessage - | WorkerErrorMessage - -interface WorkerSlot { - readonly slotId: number - readonly worker: Worker - readonly candidateSnapshots: Map - busy: boolean -} - -const DEFAULT_WORKER_URL = new URL("./router-job-worker.mjs", import.meta.url) -const DEFAULT_LOADER_URL = new URL("./ts-loader.mjs", import.meta.url) - -const errorFromUnknown = (cause: unknown): Error => - cause instanceof Error ? cause : new Error(String(cause)) - -const isRecord = (message: unknown): message is Record => - typeof message === "object" && message !== null - -const isWorkerMessage = (message: unknown): message is WorkerMessage => { - if (!isRecord(message) || typeof message.type !== "string") return false - if (message.type === "ready") return true - if (message.type === "error") return typeof message.message === "string" - if (message.type === "candidate-evaluate") - return ( - Number.isInteger(message.requestId) && - typeof message.snapshotId === "string" && - (message.snapshot === undefined || - (typeof message.snapshot === "object" && message.snapshot !== null)) && - Array.isArray(message.candidates) - ) - return ( - message.type === "result" && - Number.isInteger(message.jobId) && - (message.kind === "holdout" || message.kind === "fit-all") && - Array.isArray(message.results) - ) -} - -const workerExecArgv = (loaderUrl: URL): readonly string[] => [ - "--no-warnings", - "--experimental-strip-types", - "--experimental-loader", - loaderUrl.href, -] - -/** Run independent router jobs in native workers and return tagged results in input order. */ -export const runEvidenceRouterJobs = ( - jobs: readonly EvidenceRouterJob[], - profile: OptimizationProfile, - options: EvidenceRouterJobPoolOptions, -): Promise => { - if (jobs.length === 0) return Promise.resolve([]) - if (options.signal?.aborted) - return Promise.reject(new Error("Evidence-router job pool was aborted before startup")) - - const workerCount = Math.min(jobs.length, Math.max(1, Math.floor(options.workerCount))) - const workerUrl = options.workerUrl ?? DEFAULT_WORKER_URL - const loaderUrl = options.loaderUrl ?? DEFAULT_LOADER_URL - - return new Promise((resolve, reject) => { - const slots: WorkerSlot[] = [] - const results = new Map() - let nextJob = 0 - let completedJobs = 0 - let settled = false - - const removeAbortListener = () => options.signal?.removeEventListener("abort", abort) - - const terminate = async (): Promise => { - await Promise.all(slots.map((slot) => slot.worker.terminate().catch(() => -1))) - } - - const finish = (cause?: Error): void => { - if (settled) return - settled = true - removeAbortListener() - void terminate().then(() => { - if (cause !== undefined) { - reject(cause) - return - } - const ordered: EvidenceRouterJobResult[] = [] - for (let index = 0; index < jobs.length; index++) { - const result = results.get(index) - if (result === undefined) { - reject(new Error("Evidence-router job pool returned incomplete results")) - return - } - ordered.push(result) - } - resolve(ordered) - }) - } - - const abort = () => finish(new Error("Evidence-router job pool was interrupted")) - - const dispatch = (slot: WorkerSlot): void => { - if (settled || slot.busy || nextJob >= jobs.length) return - const jobId = nextJob++ - slot.busy = true - try { - slot.worker.postMessage({ - type: "run", - jobId, - job: jobs[jobId], - profile, - candidateWorkerCount: options.candidateQueue?.workerCount ?? 0, - candidateBatchSize: options.candidateQueue?.batchSize ?? 0, - useCandidateQueue: options.candidateQueue !== undefined, - }) - } catch (cause) { - finish(errorFromUnknown(cause)) - } - } - - const handleMessage = (slot: WorkerSlot, message: unknown): void => { - if (!isWorkerMessage(message)) { - finish(new Error("Router job worker sent an invalid message")) - return - } - if (message.type === "ready") { - dispatch(slot) - return - } - if (message.type === "error") { - finish(new Error(message.message)) - return - } - if (message.type === "candidate-evaluate") { - if (options.candidateQueue === undefined) { - finish(new Error("Router job worker requested a candidate queue that was not configured")) - return - } - const snapshot = message.snapshot ?? slot.candidateSnapshots.get(message.snapshotId) - if (snapshot === undefined) { - finish(new Error(`Router job worker omitted snapshot ${message.snapshotId}`)) - return - } - if (message.snapshot !== undefined) - slot.candidateSnapshots.set(message.snapshotId, message.snapshot) - void options.candidateQueue - .evaluate(snapshot, message.candidates, undefined, `${slot.slotId}:${message.snapshotId}`) - .then((results) => { - if (settled) return - try { - slot.worker.postMessage({ - type: "candidate-result", - requestId: message.requestId, - results, - }) - } catch (cause) { - finish(errorFromUnknown(cause)) - } - }) - .catch((cause) => finish(errorFromUnknown(cause))) - return - } - const job = jobs[message.jobId] - if (job === undefined || job.kind !== message.kind) { - finish(new Error(`Router job worker returned an unexpected job ${message.jobId}`)) - return - } - if (message.results.length === 0) { - finish(new Error(`Router job worker returned no results for job ${message.jobId}`)) - return - } - slot.busy = false - if (message.kind === "holdout") { - results.set(message.jobId, { - jobId: message.jobId, - kind: "holdout", - results: message.results, - }) - } else { - results.set(message.jobId, { - jobId: message.jobId, - kind: "fit-all", - results: message.results, - }) - } - completedJobs++ - if (completedJobs === jobs.length) { - finish() - return - } - dispatch(slot) - } - - options.signal?.addEventListener("abort", abort, { once: true }) - try { - for (let index = 0; index < workerCount; index++) { - const worker = new Worker(workerUrl, { - execArgv: [...workerExecArgv(loaderUrl)], - }) - const slot: WorkerSlot = { - slotId: index, - worker, - candidateSnapshots: new Map(), - busy: false, - } - worker.on("message", (message: unknown) => handleMessage(slot, message)) - worker.on("error", (cause: Error) => finish(cause)) - worker.on("exit", (code: number) => { - if (!settled) finish(new Error(`Router job worker exited with code ${code}`)) - }) - slots.push(slot) - } - } catch (cause) { - finish(errorFromUnknown(cause)) - } - }) -} diff --git a/benchmarks/retrieval/router-job-worker.mjs b/benchmarks/retrieval/router-job-worker.mjs deleted file mode 100644 index f85144d..0000000 --- a/benchmarks/retrieval/router-job-worker.mjs +++ /dev/null @@ -1,181 +0,0 @@ -import { parentPort } from "node:worker_threads" - -import { - fitRecommendedEvidenceRouter, - fitRecommendedEvidenceRouterParallel, - optimizeEvidenceRouter, - optimizeEvidenceRouterParallel, -} from "./weight-search.ts" - -if (parentPort === null) throw new Error("Router job worker requires a parent port") - -let busy = false -let activeCandidateQueue - -const createRemoteCandidateQueue = (workerCount, batchSize) => { - const pending = new Map() - const snapshotIds = new WeakMap() - const sentSnapshots = new WeakSet() - let nextRequestId = 0 - let nextSnapshotId = 0 - let closed = false - - const queue = { - workerCount, - batchSize, - activeWorkerCount: () => 0, - evaluate(snapshot, candidates, signal, requestedSnapshotId) { - if (closed) return Promise.reject(new Error("Remote candidate queue is closed")) - if (signal?.aborted) - return Promise.reject(new Error("Remote candidate queue was interrupted")) - if (candidates.length === 0) return Promise.resolve([]) - let snapshotId = requestedSnapshotId ?? snapshotIds.get(snapshot) - if (snapshotId === undefined) { - snapshotId = String(nextSnapshotId++) - snapshotIds.set(snapshot, snapshotId) - } - const requestId = nextRequestId++ - return new Promise((resolve, reject) => { - const abort = () => { - pending.delete(requestId) - reject(new Error("Remote candidate queue was interrupted")) - } - if (signal !== undefined) { - signal.addEventListener("abort", abort, { once: true }) - if (signal.aborted) { - abort() - return - } - } - pending.set(requestId, { - resolve: (results) => { - signal?.removeEventListener("abort", abort) - resolve(results) - }, - reject: (error) => { - signal?.removeEventListener("abort", abort) - reject(error) - }, - }) - parentPort.postMessage({ - type: "candidate-evaluate", - requestId, - snapshotId, - snapshot: sentSnapshots.has(snapshot) ? undefined : snapshot, - candidates, - }) - sentSnapshots.add(snapshot) - }) - }, - handleResult(message) { - const request = pending.get(message.requestId) - if (request === undefined) return - pending.delete(message.requestId) - if (message.type === "candidate-error") request.reject(new Error(message.message)) - else request.resolve(message.results) - }, - close() { - closed = true - for (const request of pending.values()) - request.reject(new Error("Remote candidate queue closed")) - pending.clear() - }, - } - return queue -} - -const runJob = async (message) => { - const candidateQueue = message.useCandidateQueue - ? createRemoteCandidateQueue(message.candidateWorkerCount, message.candidateBatchSize) - : undefined - activeCandidateQueue = candidateQueue - const options = - candidateQueue === undefined - ? { workerCount: 0 } - : { - workerCount: candidateQueue.workerCount, - fallbackToSerial: false, - evaluationQueue: candidateQueue, - } - try { - if (message.job.kind === "holdout") { - const results = - candidateQueue === undefined - ? await optimizeEvidenceRouter( - message.job.model, - message.job.fusion, - message.job.strategy, - message.job.fold, - message.job.development, - message.job.validation, - message.profile, - ) - : await optimizeEvidenceRouterParallel( - message.job.model, - message.job.fusion, - message.job.strategy, - message.job.fold, - message.job.development, - message.job.validation, - message.profile, - options, - ) - parentPort.postMessage({ type: "result", jobId: message.jobId, kind: "holdout", results }) - } else if (message.job.kind === "fit-all") { - const results = - candidateQueue === undefined - ? await fitRecommendedEvidenceRouter( - message.job.model, - message.job.fusion, - message.job.samples, - message.profile, - ) - : await fitRecommendedEvidenceRouterParallel( - message.job.model, - message.job.fusion, - message.job.samples, - message.profile, - options, - ) - parentPort.postMessage({ type: "result", jobId: message.jobId, kind: "fit-all", results }) - } else { - throw new Error("Router job worker received an unknown job kind") - } - } finally { - candidateQueue?.close() - activeCandidateQueue = undefined - } -} - -parentPort.postMessage({ type: "ready" }) - -parentPort.on("message", async (message) => { - if (message?.type === "candidate-result" || message?.type === "candidate-error") { - activeCandidateQueue?.handleResult(message) - return - } - if (message?.type !== "run") { - parentPort.postMessage({ type: "error", message: "Router job worker received an invalid task" }) - return - } - if (busy) { - parentPort.postMessage({ - type: "error", - jobId: message.jobId, - message: "Router job worker is busy", - }) - return - } - busy = true - try { - await runJob(message) - } catch (error) { - parentPort.postMessage({ - type: "error", - jobId: message.jobId, - message: error instanceof Error ? error.message : String(error), - }) - } finally { - busy = false - } -}) diff --git a/benchmarks/retrieval/runner.ts b/benchmarks/retrieval/runner.ts index e95a96b..5cec8e6 100644 --- a/benchmarks/retrieval/runner.ts +++ b/benchmarks/retrieval/runner.ts @@ -1,94 +1,32 @@ import { mkdir, writeFile } from "node:fs/promises" import path from "node:path" -import { Effect, Stream } from "effect" +import { Effect } from "effect" -import type { Embedding } from "../../src/domain/chunk.js" import { DEFAULT_CONFIG } from "../../src/domain/config.js" -import type { EmbeddingDtype } from "../../src/domain/dtype.js" -import type { StoredChunk } from "../../src/domain/index-data.js" import { MODEL_REGISTRY } from "../../src/domain/models.js" -import type { BoundEmbedder } from "../../src/domain/ports.js" -import { IndexStore, SparseEmbedder } from "../../src/domain/ports.js" -import { FUSION_METHODS, type FusionMethod } from "../../src/domain/retrieval.js" -import type { SparseContract, SparseTerm, SparseVector } from "../../src/domain/sparse.js" -import { contentHash } from "../../src/lib/content-hash.js" -import { buildQueryTermCoverage } from "../../src/lib/retrieval/evidence-router.js" -import { createAutoBoundEmbedder } from "../../src/services/embedder.js" -import { loadCorpusManifests, prepareRepository } from "./corpus.js" -import { assignGroupedFolds, foldKey } from "./folds.js" +import { FUSION_METHODS } from "../../src/domain/retrieval.js" +import { loadCorpusManifests } from "./corpus/repository.js" +import { collectBenchmarkData } from "./evaluation/collect.js" +import { assignGroupedFolds } from "./evaluation/folds.js" import { - contextRecallAtBudget, - goldTargetRanks, - recallAt, - reciprocalRank, - resolveGoldTargets, - successAt, -} from "./metrics.js" -import { OPTIMIZATION_PROFILES, type OptimizationProfile } from "./optimization-profiles.js" -import { prepareCorpus, type PreparedCorpus } from "./prepare.js" -import { fuseVariant, rankLexicalChannels, RETRIEVAL_VARIANTS } from "./ranking.js" -import { renderMarkdownReport } from "./report.js" -import { - runEvidenceRouterJobs, - type EvidenceRouterFitAllJob, - type EvidenceRouterHoldoutJob, -} from "./router-job-pool.js" -import { withSqliteBenchmarkStore } from "./sqlite-index.js" + OPTIMIZATION_PROFILES, + type OptimizationProfile, +} from "./evaluation/optimization-profiles.js" +import { renderMarkdownReport } from "./evaluation/report.js" +import { runBenchmarkSearch, type BenchmarkSearchConfig } from "./evaluation/search.js" import { ROUTER_SEARCH_STRATEGY, type BenchmarkArtifact, type BenchmarkProfile, type CorpusManifest, - type QueryKind, - type QueryMeasurement, - type RecommendedEvidenceRouter, type ValidationStrategy, - type EvidenceRouterSearchResult, -} from "./types.js" -import { - fitRecommendedEvidenceRouterParallel, - fitRecommendedFusionWeightsParallel, - fitRecommendedWeightsParallel, - optimizeEvidenceRouterParallel, - optimizeFusionWeightsParallel, - optimizeWeightsParallel, - evaluateProductionRouter, - type ParallelSearchOptions, - type WeightSearchSample, -} from "./weight-search.js" -import { - createCandidateEvaluationQueue, - getDefaultWorkerCount, - resolveWorkerCount, - type CandidateEvaluationQueue, -} from "./worker-pool.js" +} from "./evaluation/types.js" +import { getDefaultWorkerCount, resolveWorkerCount } from "./execution/candidate-evaluation-pool.js" const CONTEXT_BUDGETS = [2_048, 4_096, 8_192, 16_384] as const -const EMBEDDING_BATCH_SIZE = 2 -const SINGLE_ITEM_ESTIMATED_TOKENS = 2_048 -const QUERY_KINDS: readonly QueryKind[] = [ - "identifier", - "searchPhrase", - "naturalQuestion", - "agentTask", -] - -/** Search and validation stages enabled by one benchmark profile. */ -interface BenchmarkProfileConfig { - /** Number of intent-grouped cross-validation folds. */ - readonly groupedFolds: number - /** Whether each selected repository is evaluated as an excluded holdout. */ - readonly repositoryHoldouts: boolean - /** Whether historical query-kind RRF grids and Shapley diagnostics run. */ - readonly legacyDiagnostics: boolean - /** Static fusion formulas evaluated by this profile. */ - readonly fusionMethods: readonly FusionMethod[] - /** Fusion formulas used when evaluating the evidence router. */ - readonly routerFusionMethods: readonly FusionMethod[] -} -const profileConfig = (profile: BenchmarkProfile): BenchmarkProfileConfig => { +const profileConfig = (profile: BenchmarkProfile): BenchmarkSearchConfig => { switch (profile) { case "smoke": return { @@ -125,101 +63,6 @@ const profileConfig = (profile: BenchmarkProfile): BenchmarkProfileConfig => { } } -const reportProgress = (message: string): void => { - process.stderr.write(`[retrieval benchmark] ${message}\n`) -} - -const runParallelSearch = ( - operation: (signal: AbortSignal) => Promise, -): Effect.Effect => - Effect.tryPromise({ - try: (signal) => operation(signal), - catch: (cause) => (cause instanceof Error ? cause : new Error(String(cause))), - }) - -const splitSamples = ( - samples: readonly WeightSearchSample[], - isValidation: (sample: WeightSearchSample) => boolean, -): { - readonly development: readonly WeightSearchSample[] - readonly validation: readonly WeightSearchSample[] -} => ({ - development: samples.filter((sample) => !isValidation(sample)), - validation: samples.filter(isValidation), -}) - -const planEvidenceRouterJobs = ( - samplesByModel: ReadonlyMap, - config: BenchmarkProfileConfig, - groupedStrategy: ValidationStrategy, -): readonly EvidenceRouterHoldoutJob[] => { - const jobs: EvidenceRouterHoldoutJob[] = [] - for (const [model, samples] of samplesByModel) { - const repositories = [...new Set(samples.map((sample) => sample.repository))] - for (const fusion of config.routerFusionMethods) { - for (let fold = 0; fold < config.groupedFolds; fold++) { - const split = splitSamples(samples, (sample) => sample.groupedFold === fold) - jobs.push({ - kind: "holdout", - model, - fusion, - strategy: groupedStrategy, - fold: String(fold + 1), - development: split.development, - validation: split.validation, - }) - } - if (config.repositoryHoldouts && repositories.length > 1) { - for (const repository of repositories) { - const split = splitSamples(samples, (sample) => sample.repository === repository) - jobs.push({ - kind: "holdout", - model, - fusion, - strategy: "leave-one-repository-out", - fold: repository, - development: split.development, - validation: split.validation, - }) - } - } - } - } - return jobs -} - -type RouterSearchJobResult = - | { readonly kind: "holdout"; readonly results: readonly EvidenceRouterSearchResult[] } - | { readonly kind: "fit-all"; readonly results: readonly RecommendedEvidenceRouter[] } - -const runRouterSearchJob = ( - job: EvidenceRouterHoldoutJob | EvidenceRouterFitAllJob, - profile: OptimizationProfile, - options: ParallelSearchOptions, -): Promise => { - if (job.kind === "holdout") - return optimizeEvidenceRouterParallel( - job.model, - job.fusion, - job.strategy, - job.fold, - job.development, - job.validation, - profile, - options, - ).then((results) => ({ kind: "holdout", results })) - return fitRecommendedEvidenceRouterParallel( - job.model, - job.fusion, - job.samples, - profile, - options, - ).then((results) => ({ kind: "fit-all", results })) -} - -const isLongInput = (text: string): boolean => - Buffer.byteLength(text, "utf8") / 4 > SINGLE_ITEM_ESTIMATED_TOKENS - const selectValues = (value: string | undefined): ReadonlySet | null => value ? new Set( @@ -271,80 +114,6 @@ const selectOptimizationProfile = (): Effect.Effect : Effect.succeed(selected) } -const embedTexts = ( - texts: readonly string[], - model: string, - embedder: BoundEmbedder, -): Effect.Effect => - Effect.gen(function* () { - const vectors: Float32Array[] = [] - let start = 0 - while (start < texts.length) { - const currentIsLong = isLongInput(texts[start]) - const next = texts[start + 1] - const nextIsLong = next !== undefined && isLongInput(next) - const batchSize = currentIsLong || nextIsLong ? 1 : EMBEDDING_BATCH_SIZE - const batch = texts.slice(start, start + batchSize) - const embeddings = yield* embedder.batch(batch) - vectors.push(...embeddings.map((embedding) => embedding.vector)) - start += batchSize - } - return vectors - }).pipe(Effect.mapError((cause) => new Error(`Embedding failed for ${model}`, { cause }))) - -const embedSparseTexts = ( - texts: readonly string[], - embedder: typeof SparseEmbedder.Service, -): Effect.Effect => - Effect.gen(function* () { - const vectors: SparseVector[] = [] - let start = 0 - while (start < texts.length) { - const batch = texts.slice(start, start + DEFAULT_CONFIG.sparseEmbedder.batchSize) - vectors.push(...(yield* embedder.batch(batch))) - start += DEFAULT_CONFIG.sparseEmbedder.batchSize - } - return vectors - }).pipe(Effect.mapError((cause) => new Error("Sparse document embedding failed", { cause }))) - -const toStoredChunk = (chunk: PreparedCorpus["chunks"][number]): StoredChunk => { - const { text, ...location } = chunk - return { ...location, contentHash: contentHash(text) } -} - -const persistBenchmarkCorpus = ( - store: typeof IndexStore.Service, - corpus: PreparedCorpus, - vectors: readonly Float32Array[], - sparseVectors: readonly SparseVector[], - dims: number, - dtype: EmbeddingDtype, - sparseContract: SparseContract, - sparseIdf: readonly SparseTerm[], -): Effect.Effect => - Effect.gen(function* () { - const pairs = corpus.chunks.map((chunk, index): readonly [StoredChunk, Embedding] => [ - toStoredChunk(chunk), - { vector: vectors[index]!, dims, dtype }, - ]) - yield* store.persistIndex({ - chunks: Stream.succeed( - pairs.map( - ([chunk, embedding], index) => [chunk, embedding, sparseVectors[index]!] as const, - ), - ), - identifierIndex: corpus.identifierIndex, - bm25Index: corpus.bm25Index, - files: [], - dims, - dtype, - embeddingCache: [], - sparseEmbeddingCache: [], - sparseContract, - sparseIdf, - }) - }) - const writeArtifact = (artifact: BenchmarkArtifact): Effect.Effect => Effect.gen(function* () { const outputDirectory = path.resolve("benchmarks/results") @@ -386,447 +155,25 @@ export const runRetrievalBenchmark = ( const manifests = yield* selectManifests(yield* loadCorpusManifests(), profile) const groupedFoldAssignments = assignGroupedFolds(manifests, config.groupedFolds) const models = yield* selectModels() - const repositories: BenchmarkArtifact["repositories"][number][] = [] - const embeddingRuns: BenchmarkArtifact["embeddingRuns"][number][] = [] - const sparseEmbeddingRuns: BenchmarkArtifact["sparseEmbeddingRuns"][number][] = [] - const measurements: QueryMeasurement[] = [] - const sampleGroups = new Map< - string, - { - model: string - queryKind: QueryKind - samples: WeightSearchSample[] - } - >() - const samplesByModel = new Map() - let retrievalDurationMs = 0 - let candidateQueueStartupDurationMs = 0 - let candidateQueueShutdownDurationMs = 0 - - for (const manifest of manifests) { - const repositoryPath = yield* prepareRepository(manifest) - const corpus = yield* prepareCorpus(repositoryPath, manifest) - repositories.push({ - id: manifest.id, - repository: manifest.repository, - revision: manifest.revision, - chunks: corpus.chunks.length, - preparationDurationMs: corpus.preparationDurationMs, - }) - - const targetsByQuestion: (readonly ReadonlySet[])[] = [] - for (const question of manifest.questions) { - const targets = resolveGoldTargets( - question.groundTruth, - corpus.chunks, - corpus.identifiersByChunk, - ) - const unresolved = question.groundTruth.filter((_, index) => targets[index].size === 0) - if (unresolved.length > 0) - return yield* Effect.fail( - new Error( - `${question.id} has unresolved gold targets: ${unresolved.map((target) => `${target.file}::${target.symbol}`).join(", ")}`, - ), - ) - targetsByQuestion.push(targets) - } - - const queries = manifest.questions.flatMap((question, questionIndex) => - QUERY_KINDS.map((queryKind) => ({ - questionIndex, - queryKind, - query: question.queries[queryKind], - })), - ) - for (const model of models) { - const info = MODEL_REGISTRY[model] - if (info === undefined) - return yield* Effect.fail(new Error(`Unknown embedding model ${model}`)) - const bound = yield* createAutoBoundEmbedder({ - model, - dtype: info.defaultDtype, - dims: info.dims, - }).pipe( - Effect.mapError( - (cause) => new Error(`Could not auto-select a device for ${model}`, { cause }), - ), - ) - const embeddingStartedAt = performance.now() - const chunkVectors = yield* embedTexts( - corpus.chunks.map((chunk) => chunk.text), - model, - bound.embedder, - ) - const chunkEmbeddingDurationMs = performance.now() - embeddingStartedAt - const queryEmbeddingStartedAt = performance.now() - const queryVectors = yield* embedTexts( - queries.map((entry) => entry.query), - model, - bound.embedder, - ) - embeddingRuns.push({ - repository: manifest.id, - model, - device: bound.device, - batchSize: EMBEDDING_BATCH_SIZE, - chunkEmbeddingDurationMs, - queryEmbeddingDurationMs: performance.now() - queryEmbeddingStartedAt, - }) - - const retrievalStartedAt = performance.now() - const modelRun = yield* withSqliteBenchmarkStore( - model, - info.defaultDtype, - Effect.gen(function* () { - const store = yield* IndexStore - const sparseEmbedder = yield* SparseEmbedder - const sparseStartedAt = performance.now() - const sparseVectors = yield* embedSparseTexts( - corpus.chunks.map((chunk) => chunk.text), - sparseEmbedder, - ) - const sparseChunkEmbeddingDurationMs = performance.now() - sparseStartedAt - const sparseIdf = yield* sparseEmbedder.loadIdf() - const sparseQueryStartedAt = performance.now() - const sparseQueries = yield* Effect.forEach(queries, ({ query }) => - sparseEmbedder.tokenizeQuery(query), - ) - const sparseQueryTokenizationDurationMs = performance.now() - sparseQueryStartedAt - yield* persistBenchmarkCorpus( - store, - corpus, - chunkVectors, - sparseVectors, - info.dims, - info.defaultDtype, - sparseEmbedder.contract, - sparseIdf, - ) - sparseEmbeddingRuns.push({ - repository: manifest.id, - model: sparseEmbedder.contract.model, - tokenizerModel: sparseEmbedder.contract.tokenizer, - batchSize: DEFAULT_CONFIG.sparseEmbedder.batchSize, - chunkEmbeddingDurationMs: sparseChunkEmbeddingDurationMs, - queryTokenizationDurationMs: sparseQueryTokenizationDurationMs, - }) - const searchData = yield* store.loadSearchData() - const modelMeasurements: QueryMeasurement[] = [] - const modelSamples: WeightSearchSample[] = [] - const samplesByQueryKind = new Map() - - for (let queryIndex = 0; queryIndex < queries.length; queryIndex++) { - const entry = queries[queryIndex] - const question = manifest.questions[entry.questionIndex] - const targets = targetsByQuestion[entry.questionIndex] - const groupedFold = groupedFoldAssignments.get(foldKey(manifest.id, question.id)) - if (groupedFold === undefined) - return yield* Effect.fail( - new Error(`No grouped fold assignment for ${manifest.id}/${question.id}`), - ) - const channelStartedAt = performance.now() - const lexicalRankings = rankLexicalChannels(entry.query, searchData) - const dense = yield* store.searchDense({ - vector: queryVectors[queryIndex]!, - dims: info.dims, - dtype: info.defaultDtype, - }) - const sparse = yield* store.searchSparse(sparseQueries[queryIndex]!) - const rankings = { ...lexicalRankings, dense, sparse } - const channelDurationMs = performance.now() - channelStartedAt - const sample: WeightSearchSample = { - repository: manifest.id, - intentId: question.id, - queryKind: entry.queryKind, - groupedFold, - query: entry.query, - rankings, - targets, - chunks: corpus.chunks, - termCoverage: buildQueryTermCoverage( - entry.query, - searchData.bm25Index, - searchData.identifierIndex, - ), - } - modelSamples.push(sample) - samplesByQueryKind.set(entry.queryKind, [ - ...(samplesByQueryKind.get(entry.queryKind) ?? []), - sample, - ]) - for (const variant of RETRIEVAL_VARIANTS) { - const variantStartedAt = performance.now() - const ranked = fuseVariant(variant, entry.query, rankings) - const queryDurationMs = channelDurationMs + performance.now() - variantStartedAt - modelMeasurements.push({ - repository: manifest.id, - language: manifest.language, - size: manifest.size, - revision: manifest.revision, - model, - variant, - questionId: question.id, - queryKind: entry.queryKind, - query: entry.query, - category: question.category, - difficulty: question.difficulty, - groupedFold, - recallAt5: recallAt(ranked, targets, 5), - recallAt10: recallAt(ranked, targets, 10), - recallAt20: recallAt(ranked, targets, 20), - recallAt50: recallAt(ranked, targets, 50), - successAt10: successAt(ranked, targets, 10), - successAt20: successAt(ranked, targets, 20), - reciprocalRank: reciprocalRank(ranked, targets), - goldRanks: goldTargetRanks(ranked, targets), - contextRecall: Object.fromEntries( - CONTEXT_BUDGETS.map((budget) => [ - String(budget), - contextRecallAtBudget(ranked, targets, corpus.chunks, budget), - ]), - ), - queryDurationMs, - }) - } - } - return { measurements: modelMeasurements, samples: modelSamples, samplesByQueryKind } - }), - ) - retrievalDurationMs += performance.now() - retrievalStartedAt - measurements.push(...modelRun.measurements) - samplesByModel.set(model, [...(samplesByModel.get(model) ?? []), ...modelRun.samples]) - for (const [queryKind, samples] of modelRun.samplesByQueryKind) { - const groupKey = `${model}\0${queryKind}` - const group = sampleGroups.get(groupKey) ?? { model, queryKind, samples: [] } - group.samples.push(...samples) - sampleGroups.set(groupKey, group) - } - } - } - - const weightSearchStartedAt = performance.now() - const weightSearch: BenchmarkArtifact["weightSearch"][number][] = [] - const recommendedWeights: BenchmarkArtifact["recommendedWeights"][number][] = [] - if (config.legacyDiagnostics) { - for (const group of sampleGroups.values()) { - for (let fold = 0; fold < config.groupedFolds; fold++) { - const split = splitSamples(group.samples, (sample) => sample.groupedFold === fold) - weightSearch.push( - yield* runParallelSearch((signal) => - optimizeWeightsParallel( - group.model, - group.queryKind, - groupedStrategy, - String(fold + 1), - split.development, - split.validation, - optimizationProfile, - { ...searchOptions, signal }, - ), - ), - ) - } - const repositories = [...new Set(group.samples.map((sample) => sample.repository))] - if (config.repositoryHoldouts && repositories.length > 1) { - for (const repository of repositories) { - const split = splitSamples(group.samples, (sample) => sample.repository === repository) - weightSearch.push( - yield* runParallelSearch((signal) => - optimizeWeightsParallel( - group.model, - group.queryKind, - "leave-one-repository-out", - repository, - split.development, - split.validation, - optimizationProfile, - { ...searchOptions, signal }, - ), - ), - ) - } - } - recommendedWeights.push( - yield* runParallelSearch((signal) => - fitRecommendedWeightsParallel( - group.model, - group.queryKind, - group.samples, - optimizationProfile, - { ...searchOptions, signal }, - ), - ), - ) - } - } - const weightSearchDurationMs = performance.now() - weightSearchStartedAt - - const fusionSearchStartedAt = performance.now() - const productionRouterSearch: BenchmarkArtifact["productionRouterSearch"][number][] = [] - for (const [model, samples] of samplesByModel) { - const repositories = [...new Set(samples.map((sample) => sample.repository))] - for (let fold = 0; fold < config.groupedFolds; fold++) { - productionRouterSearch.push( - evaluateProductionRouter( - model, - groupedStrategy, - String(fold + 1), - samples.filter((sample) => sample.groupedFold !== fold), - samples.filter((sample) => sample.groupedFold === fold), - optimizationProfile, - ), - ) - } - if (config.repositoryHoldouts && repositories.length > 1) { - for (const repository of repositories) { - productionRouterSearch.push( - evaluateProductionRouter( - model, - "leave-one-repository-out", - repository, - samples.filter((sample) => sample.repository !== repository), - samples.filter((sample) => sample.repository === repository), - optimizationProfile, - ), - ) - } - } - } - const fusionSearch: BenchmarkArtifact["fusionSearch"][number][] = [] - for (const [model, samples] of samplesByModel) { - const repositories = [...new Set(samples.map((sample) => sample.repository))] - for (const fusion of config.fusionMethods) { - reportProgress(`${model}: selecting static ${fusion} fusion weights`) - for (let fold = 0; fold < config.groupedFolds; fold++) { - const split = splitSamples(samples, (sample) => sample.groupedFold === fold) - fusionSearch.push( - yield* runParallelSearch((signal) => - optimizeFusionWeightsParallel( - model, - fusion, - groupedStrategy, - String(fold + 1), - split.development, - split.validation, - optimizationProfile, - { ...searchOptions, signal }, - ), - ), - ) - } - if (config.repositoryHoldouts && repositories.length > 1) { - for (const repository of repositories) { - const split = splitSamples(samples, (sample) => sample.repository === repository) - fusionSearch.push( - yield* runParallelSearch((signal) => - optimizeFusionWeightsParallel( - model, - fusion, - "leave-one-repository-out", - repository, - split.development, - split.validation, - optimizationProfile, - { ...searchOptions, signal }, - ), - ), - ) - } - } - } - } - const recommendedFusionWeights: BenchmarkArtifact["recommendedFusionWeights"][number][] = [] - for (const [model, samples] of samplesByModel) - for (const fusion of config.fusionMethods) - recommendedFusionWeights.push( - yield* runParallelSearch((signal) => - fitRecommendedFusionWeightsParallel(model, fusion, samples, optimizationProfile, { - ...searchOptions, - signal, - }), - ), - ) - const fusionSearchDurationMs = performance.now() - fusionSearchStartedAt - - const evidenceRouterSearchStartedAt = performance.now() - const routerJobs = planEvidenceRouterJobs(samplesByModel, config, groupedStrategy) - const routerWorkerBudget = Math.min( - resolveWorkerCount(searchOptions.workerCount), - getDefaultWorkerCount(), - ) - const recommendedJobs: EvidenceRouterFitAllJob[] = [] - for (const [model, samples] of samplesByModel) - for (const fusion of config.routerFusionMethods) - recommendedJobs.push({ kind: "fit-all", model, fusion, samples }) - const allRouterJobs: readonly (EvidenceRouterHoldoutJob | EvidenceRouterFitAllJob)[] = [ - ...routerJobs, - ...recommendedJobs, - ] - const canParallelizeRouterJobs = - !serialSearch && routerWorkerBudget >= 3 && allRouterJobs.length > 0 - const routerControllerCount = canParallelizeRouterJobs - ? Math.min(allRouterJobs.length, Math.max(1, Math.floor(routerWorkerBudget / 2))) - : 0 - const candidateWorkerCount = canParallelizeRouterJobs - ? Math.max(1, routerWorkerBudget - routerControllerCount) - : 0 - let candidateQueue: CandidateEvaluationQueue | undefined - if (canParallelizeRouterJobs) { - const candidateQueueStartedAt = performance.now() - candidateQueue = yield* runParallelSearch(() => - createCandidateEvaluationQueue({ workerCount: candidateWorkerCount }), - ) - candidateQueueStartupDurationMs = performance.now() - candidateQueueStartedAt - } - reportProgress( - `running ${allRouterJobs.length} evidence-router jobs with ` + - `${routerControllerCount} native controllers and ` + - `${candidateQueue?.workerCount ?? 0} shared candidate workers`, - ) - const closeCandidateQueue = - candidateQueue === undefined - ? Effect.void - : Effect.gen(function* () { - const candidateQueueStartedAt = performance.now() - yield* runParallelSearch(() => candidateQueue!.close()) - candidateQueueShutdownDurationMs = performance.now() - candidateQueueStartedAt - }) - const routerResults = yield* Effect.ensuring( - canParallelizeRouterJobs - ? runParallelSearch((signal) => - runEvidenceRouterJobs(allRouterJobs, optimizationProfile, { - workerCount: routerControllerCount, - candidateQueue, - signal, - }).then((results) => - results.map((result): RouterSearchJobResult => { - if (result.kind === "holdout") return { kind: "holdout", results: result.results } - return { kind: "fit-all", results: result.results } - }), - ), - ) - : Effect.forEach( - allRouterJobs, - (job) => - runParallelSearch((signal) => - runRouterSearchJob(job, optimizationProfile, { - ...searchOptions, - workerCount: 0, - signal, - }), - ), - { concurrency: 1 }, - ), - Effect.orDie(closeCandidateQueue), + const collected = yield* collectBenchmarkData(manifests, models, groupedFoldAssignments) + const { + repositories, + embeddingRuns, + sparseEmbeddingRuns, + measurements, + sampleGroups, + samplesByModel, + retrievalDurationMs, + } = collected + const search = yield* runBenchmarkSearch( + config, + sampleGroups, + samplesByModel, + groupedStrategy, + optimizationProfile, + serialSearch, + searchOptions, ) - let evidenceRouterSearch: readonly EvidenceRouterSearchResult[] = [] - let recommendedEvidenceRouters: readonly RecommendedEvidenceRouter[] = [] - for (const result of routerResults) { - if (result.kind === "holdout") - evidenceRouterSearch = [...evidenceRouterSearch, ...result.results] - else recommendedEvidenceRouters = [...recommendedEvidenceRouters, ...result.results] - } - const evidenceRouterSearchDurationMs = performance.now() - evidenceRouterSearchStartedAt const embeddingDurationMs = embeddingRuns.reduce( (sum, run) => sum + run.chunkEmbeddingDurationMs + run.queryEmbeddingDurationMs, @@ -858,11 +205,11 @@ export const runRetrievalBenchmark = ( corpusPreparationDurationMs, embeddingDurationMs, retrievalDurationMs, - weightSearchDurationMs, - fusionSearchDurationMs, - evidenceRouterSearchDurationMs, - candidateQueueStartupDurationMs, - candidateQueueShutdownDurationMs, + weightSearchDurationMs: search.weightSearchDurationMs, + fusionSearchDurationMs: search.fusionSearchDurationMs, + evidenceRouterSearchDurationMs: search.evidenceRouterSearchDurationMs, + candidateQueueStartupDurationMs: search.candidateQueueStartupDurationMs, + candidateQueueShutdownDurationMs: search.candidateQueueShutdownDurationMs, }, chunkConfig: { chunkLines: DEFAULT_CONFIG.chunkLines, @@ -884,13 +231,13 @@ export const runRetrievalBenchmark = ( embeddingRuns, sparseEmbeddingRuns, measurements, - weightSearch, - recommendedWeights, - productionRouterSearch, - fusionSearch, - recommendedFusionWeights, - evidenceRouterSearch, - recommendedEvidenceRouters, + weightSearch: search.weightSearch, + recommendedWeights: search.recommendedWeights, + productionRouterSearch: search.productionRouterSearch, + fusionSearch: search.fusionSearch, + recommendedFusionWeights: search.recommendedFusionWeights, + evidenceRouterSearch: search.evidenceRouterSearch, + recommendedEvidenceRouters: search.recommendedEvidenceRouters, } const outputPath = yield* writeArtifact(artifact) return { artifact, outputPath } diff --git a/benchmarks/retrieval/ts-loader.mjs b/benchmarks/retrieval/ts-loader.mjs deleted file mode 100644 index 73ed509..0000000 --- a/benchmarks/retrieval/ts-loader.mjs +++ /dev/null @@ -1,10 +0,0 @@ -export async function resolve(specifier, context, defaultResolve) { - if (specifier.startsWith(".") && specifier.endsWith(".js")) { - try { - return await defaultResolve(`${specifier.slice(0, -3)}.ts`, context, defaultResolve) - } catch { - // Keep normal JavaScript resolution for actual .js modules. - } - } - return defaultResolve(specifier, context, defaultResolve) -} diff --git a/benchmarks/retrieval/worker-message.ts b/benchmarks/retrieval/worker-message.ts deleted file mode 100644 index dd52042..0000000 --- a/benchmarks/retrieval/worker-message.ts +++ /dev/null @@ -1,8 +0,0 @@ -/** Check whether an unknown worker message has one of the expected protocol types. */ -export const hasWorkerMessageType = ( - message: unknown, - types: readonly string[], -): message is { readonly type: string } => { - if (typeof message !== "object" || message === null || !("type" in message)) return false - return typeof message.type === "string" && types.includes(message.type) -} diff --git a/benchmarks/tests/channels.test.ts b/benchmarks/tests/channels.test.ts index aca5ca5..d4e2891 100644 --- a/benchmarks/tests/channels.test.ts +++ b/benchmarks/tests/channels.test.ts @@ -9,22 +9,25 @@ import { } from "../../src/lib/retrieval/evidence-router.js" import { fuseRankings } from "../../src/lib/retrieval/fusion.js" import { buildIdentifierIndex } from "../../src/lib/retrieval/identifier-index.js" -import { prepareFusion } from "../retrieval/fusion.js" +import type { PreparedCorpus } from "../retrieval/corpus/prepare.js" import { contextRecallAtBudget, recallAt, reciprocalRank, resolveGoldTargets, -} from "../retrieval/metrics.js" -import type { PreparedCorpus } from "../retrieval/prepare.js" -import { fuseVariant, rankLexicalChannels, RETRIEVAL_VARIANTS } from "../retrieval/ranking.js" -import { ROUTER_OBJECTIVES } from "../retrieval/types.js" +} from "../retrieval/evaluation/metrics.js" +import { prepareFusion } from "../retrieval/evaluation/prepared-fusion.js" +import { + fuseVariant, + rankLexicalChannels, + RETRIEVAL_VARIANTS, +} from "../retrieval/evaluation/ranking.js" +import { ROUTER_OBJECTIVES } from "../retrieval/evaluation/types.js" import { optimizeEvidenceRouter, - optimizeEvidenceRouterParallel, optimizeWeights, selectEligibleCandidate, -} from "../retrieval/weight-search.js" +} from "../retrieval/evaluation/weight-search.js" const texts = [ "export function loadProjectConfiguration() { return config }", @@ -118,6 +121,15 @@ const selectTop20Router = (results: re return result } +const withoutSearchTimings = < + T extends { readonly searchDiagnostics: { readonly timings: object } }, +>( + result: T, +) => { + const { timings: _timings, ...searchDiagnostics } = result.searchDiagnostics + return { ...result, searchDiagnostics } +} + describe("retrieval benchmark fixture", () => { it("covers the benchmark's five-channel retrieval variants", () => { expect(new Set(RETRIEVAL_VARIANTS).size).toBe(21) @@ -591,7 +603,7 @@ describe("retrieval benchmark fixture", () => { await optimizeEvidenceRouter("fixture", "dbsf", "grouped-5-fold", "1", samples, samples), ) const parallel = selectTop20Router( - await optimizeEvidenceRouterParallel( + await optimizeEvidenceRouter( "fixture", "dbsf", "grouped-5-fold", @@ -603,7 +615,7 @@ describe("retrieval benchmark fixture", () => { ), ) - expect(parallel).toEqual(serial) + expect(withoutSearchTimings(parallel)).toEqual(withoutSearchTimings(serial)) }) it("does not treat a guardrail-failing fallback as promotable", () => { diff --git a/benchmarks/tests/corpus.test.ts b/benchmarks/tests/corpus.test.ts index e8153a5..fad643a 100644 --- a/benchmarks/tests/corpus.test.ts +++ b/benchmarks/tests/corpus.test.ts @@ -1,9 +1,9 @@ import { expect, it } from "@effect/vitest" import { Effect } from "effect" -import { loadCorpusManifests, prepareRepository } from "../retrieval/corpus.js" -import { resolveGoldTargets } from "../retrieval/metrics.js" -import { prepareCorpus } from "../retrieval/prepare.js" +import { prepareCorpus } from "../retrieval/corpus/prepare.js" +import { loadCorpusManifests, prepareRepository } from "../retrieval/corpus/repository.js" +import { resolveGoldTargets } from "../retrieval/evaluation/metrics.js" // This validation intentionally reads real pinned checkouts; memfs is used by other adapter tests. it.effect("resolves every authored gold symbol in each pinned corpus", () => diff --git a/benchmarks/tests/optimization-profiles.test.ts b/benchmarks/tests/optimization-profiles.test.ts index 977958f..34d5cb5 100644 --- a/benchmarks/tests/optimization-profiles.test.ts +++ b/benchmarks/tests/optimization-profiles.test.ts @@ -4,7 +4,7 @@ import { OPTIMIZATION_PROFILES, SEARCH_PRIORITY_PROFILE, decodeOptimizationProfile, -} from "../retrieval/optimization-profiles.js" +} from "../retrieval/evaluation/optimization-profiles.js" it("records the authored weighted search-priority objective without a profile schema version", () => { expect(SEARCH_PRIORITY_PROFILE.queryFormWeights).toEqual({ diff --git a/benchmarks/tests/retrieval.test.ts b/benchmarks/tests/retrieval.test.ts index 3d723ad..c815fef 100644 --- a/benchmarks/tests/retrieval.test.ts +++ b/benchmarks/tests/retrieval.test.ts @@ -1,13 +1,13 @@ import { expect, it } from "@effect/vitest" import { Effect } from "effect" -import { assignGroupedFolds } from "../retrieval/folds.js" -import { runRetrievalBenchmark } from "../retrieval/runner.js" +import { assignGroupedFolds } from "../retrieval/evaluation/folds.js" import { ROUTER_OBJECTIVES, ROUTER_SEARCH_STRATEGY, type BenchmarkProfile, -} from "../retrieval/types.js" +} from "../retrieval/evaluation/types.js" +import { runRetrievalBenchmark } from "../retrieval/runner.js" const foldQuestions = (prefix: string) => Array.from({ length: 12 }, (_, index) => ({ diff --git a/benchmarks/tests/worker-pool.test.ts b/benchmarks/tests/worker-pool.test.ts index 5592d02..272f0ff 100644 --- a/benchmarks/tests/worker-pool.test.ts +++ b/benchmarks/tests/worker-pool.test.ts @@ -3,23 +3,16 @@ import { availableParallelism } from "node:os" import { describe, expect, it } from "@effect/vitest" import type { Chunk } from "../../src/domain/chunk.js" -import { prepareFusion } from "../retrieval/fusion.js" -import { SEARCH_PRIORITY_PROFILE } from "../retrieval/optimization-profiles.js" -import { - runEvidenceRouterJobs, - type EvidenceRouterFitAllJob, - type EvidenceRouterHoldoutJob, -} from "../retrieval/router-job-pool.js" +import { SEARCH_PRIORITY_PROFILE } from "../retrieval/evaluation/optimization-profiles.js" +import { prepareFusion } from "../retrieval/evaluation/prepared-fusion.js" import { fitRecommendedEvidenceRouter, fitRecommendedFusionWeights, - fitRecommendedFusionWeightsParallel, fitRecommendedWeights, - fitRecommendedWeightsParallel, optimizeEvidenceRouter, optimizeFusionWeights, summarize, -} from "../retrieval/weight-search.js" +} from "../retrieval/evaluation/weight-search.js" import { createCandidateEvaluationPool, createCandidateEvaluationPoolOnQueue, @@ -29,7 +22,7 @@ import { getDefaultWorkerCount, resolveWorkerCount, type EvaluationCandidate, -} from "../retrieval/worker-pool.js" +} from "../retrieval/execution/candidate-evaluation-pool.js" const chunks: readonly Chunk[] = [0, 1, 2, 3].map((index) => ({ id: String(index), @@ -106,60 +99,47 @@ const withoutSearchTimings = < describe("benchmark candidate evaluation pool", () => { it("runs holdout and fit-all jobs through one native queue", async () => { - const holdoutJob: EvidenceRouterHoldoutJob = { - kind: "holdout", - model: "fixture", - fusion: "dbsf", - strategy: "grouped-5-fold", - fold: "1", - development: [searchSample], - validation: [searchSample], - } - const fitAllJob: EvidenceRouterFitAllJob = { - kind: "fit-all", - model: "fixture", - fusion: "dbsf", - samples: [searchSample], - } const candidateQueue = await createCandidateEvaluationQueue({ workerCount: 2 }) try { - const [parallelHoldout, parallelFitAll] = await runEvidenceRouterJobs( - [holdoutJob, fitAllJob], - SEARCH_PRIORITY_PROFILE, - { - workerCount: 2, - candidateQueue, - }, - ) + const [parallelHoldout, parallelFitAll] = await Promise.all([ + optimizeEvidenceRouter( + "fixture", + "dbsf", + "grouped-5-fold", + "1", + [searchSample], + [searchSample], + SEARCH_PRIORITY_PROFILE, + { workerCount: 0, evaluationQueue: candidateQueue }, + ), + fitRecommendedEvidenceRouter("fixture", "dbsf", [searchSample], SEARCH_PRIORITY_PROFILE, { + workerCount: 0, + evaluationQueue: candidateQueue, + }), + ]) const serialHoldout = await optimizeEvidenceRouter( - holdoutJob.model, - holdoutJob.fusion, - holdoutJob.strategy, - holdoutJob.fold, - holdoutJob.development, - holdoutJob.validation, + "fixture", + "dbsf", + "grouped-5-fold", + "1", + [searchSample], + [searchSample], SEARCH_PRIORITY_PROFILE, ) const serialFitAll = await fitRecommendedEvidenceRouter( - fitAllJob.model, - fitAllJob.fusion, - fitAllJob.samples, + "fixture", + "dbsf", + [searchSample], SEARCH_PRIORITY_PROFILE, ) - expect(parallelHoldout.jobId).toBe(0) - expect(parallelHoldout.kind).toBe("holdout") - expect(parallelHoldout.results.map(withoutSearchTimings)).toEqual( + expect(parallelHoldout.map(withoutSearchTimings)).toEqual( serialHoldout.map(withoutSearchTimings), ) - expect(parallelFitAll.jobId).toBe(1) - expect(parallelFitAll.kind).toBe("fit-all") - expect(parallelFitAll.results.map(withoutSearchTimings)).toEqual( + expect(parallelFitAll.map(withoutSearchTimings)).toEqual( serialFitAll.map(withoutSearchTimings), ) - expect( - parallelHoldout.results[0]?.searchDiagnostics.timings.candidateEvaluationMs, - ).toBeGreaterThan(0) + expect(parallelHoldout[0]?.searchDiagnostics.timings.candidateEvaluationMs).toBeGreaterThan(0) } finally { await candidateQueue.close() } @@ -276,7 +256,7 @@ describe("benchmark candidate evaluation pool", () => { it("keeps the explicit parallel search result equal to the serial search", async () => { const serial = await fitRecommendedWeights("fixture", "identifier", [searchSample]) - const parallel = await fitRecommendedWeightsParallel( + const parallel = await fitRecommendedWeights( "fixture", "identifier", [searchSample], @@ -289,7 +269,7 @@ describe("benchmark candidate evaluation pool", () => { it("keeps static fusion fitting serial and parallel paths equivalent", async () => { const serial = await fitRecommendedFusionWeights("fixture", "dbsf", [searchSample]) - const parallel = await fitRecommendedFusionWeightsParallel( + const parallel = await fitRecommendedFusionWeights( "fixture", "dbsf", [searchSample], @@ -333,7 +313,10 @@ describe("benchmark candidate evaluation pool", () => { }) it("falls back to serial evaluation when worker startup is unavailable", async () => { - const unavailableWorkerUrl = new URL("../retrieval/fusion-worker.mjs", import.meta.url) + const unavailableWorkerUrl = new URL( + "../retrieval/execution/candidate-evaluation-worker.mjs", + import.meta.url, + ) unavailableWorkerUrl.pathname += ".missing" const pool = await createCandidateEvaluationPool(snapshot, { workerCount: 2, From 4877bee97dc9d44aa1e475d422989255a7e8496c Mon Sep 17 00:00:00 2001 From: Lucas Burmeister Date: Wed, 5 Aug 2026 02:53:41 +0200 Subject: [PATCH 7/8] perf: persist retrieval benchmark cache --- CONTEXT.md | 19 +- benchmarks/BASELINE.md | 14 + benchmarks/retrieval/evaluation/collect.ts | 268 +++++++++++++----- .../retrieval/execution/benchmark-cache.ts | 209 ++++++++++++++ .../retrieval/execution/sqlite-index.ts | 23 +- benchmarks/tests/benchmark-cache.test.ts | 85 ++++++ ...-based-fusion-and-optimization-profiles.md | 97 ++++++- 7 files changed, 629 insertions(+), 86 deletions(-) create mode 100644 benchmarks/retrieval/execution/benchmark-cache.ts create mode 100644 benchmarks/tests/benchmark-cache.test.ts diff --git a/CONTEXT.md b/CONTEXT.md index 384f3b5..116bf65 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -214,6 +214,11 @@ Short-profile fusion runs use Relative Score and DBSF; the full profile also run comparisons. Schema 14 records the router search strategy and compute-time breakdown in each artifact, including corpus preparation, embedding, retrieval, static fusion search, and evidence-router search duration. +For runtime planning, `timings.evidenceRouterSearchDurationMs` is the relevant embedding-free wall-clock +measurement. The current `develop` calibration uses one DBSF fusion, grouped 3-fold, no repository +holdouts, one fit-all job, and three objective selections from each shared dynamic search. Its empirical +estimate is `T ~= 32.15 + 0.14 * chunks` seconds for 60 query samples; the complete job-count and +sample-scaling model is documented in ADR-0019. Schema 15 adds deterministic one-stage proxy promotion to the evidence-router search. The benchmark evaluates the current production router as an explicit holdout baseline, adds Recall@50, and uses one shared Pareto search to select objective-specific candidates for direct retrieval, reranker top-20 @@ -229,10 +234,11 @@ in schema-17 artifacts. ADR-0020 promotes the validated Sparse contract to the p fusion path and persists its IDF and postings in `.pix/index.db`. Schema 19 removes the benchmark-owned Sparse encoder, in-process postings implementation, and separate embedding caches. Benchmark profile fitting and optimizer search remain benchmark-owned, while the -fusion adapters and evidence signals are shared with production. Benchmarks -compose the production SparseEmbedder and IndexStore around a migrated in-memory SQLite database; -Dense and Sparse ranking therefore execute through the same adapters as product queries. Experimental -profile fitting remains benchmark-owned. Every artifact includes the authored file-qualified ground truth and +fusion adapters and evidence signals are shared with production. Benchmarks compose the production +SparseEmbedder and IndexStore around a migrated SQLite database; current schema-23 runs persist that +benchmark database and channel rankings under `benchmarks/.cache/retrieval/v1/` for warm reuse. Dense +and Sparse ranking therefore execute through the same adapters as product queries. Experimental profile +fitting remains benchmark-owned. Every artifact includes the authored file-qualified ground truth and both the current Production router and a fixed five-channel `1/1/1/1/1` historical RRF baseline. Channel combinations and leave-one-channel-out variants use equal weights so channel contribution is not confounded by routing. New repositories are represented by JSON manifests in `benchmarks/corpus/`, selected with @@ -268,8 +274,9 @@ once per worker and defaults to one candidate per task because the benchmark's c Serial mode keeps the same algorithm on the main thread for comparison. Native worker startup or task failures are surfaced and all workers are terminated before the benchmark fails. Repository checkouts live under ignored `benchmarks/.cache/repos/`; generated artifacts live under -ignored `benchmarks/results/`. Dense and Sparse vectors are held only in the production in-memory -SQLite adapter during a benchmark run. See `benchmarks/README.md` and `benchmarks/BASELINE.md`. +ignored `benchmarks/results/`. Benchmark Dense and Sparse vectors and channel rankings live under the +ignored `benchmarks/.cache/retrieval/v1/` cache; production indexes remain separate. See +`benchmarks/README.md` and `benchmarks/BASELINE.md`. ### Scorer diff --git a/benchmarks/BASELINE.md b/benchmarks/BASELINE.md index ea8b89b..ee5c518 100644 --- a/benchmarks/BASELINE.md +++ b/benchmarks/BASELINE.md @@ -200,6 +200,20 @@ and SQLite retrieval `2.50 s`. The fit-all router evaluated `3,693` proxy candid candidates. The small hold-out movements are recorded rather than treated as a universal quality gain; future optimizer changes must compare the same pinned corpora, model, folds, and static controls. +### Strategy Equivalence Check + +There is a matched Schema-19 `fd` smoke comparison between +`retrieval-2026-08-02T01-31-54.105Z.json` (`successive-halving-pareto`) and +`retrieval-2026-08-03T15-09-38.902Z.json` (`proxy-promotion`). Both use the same pinned corpus, model, +folds, DBSF router, objectives, and guardrail settings. The weighted evidence-router holdout summaries +are identical at the reported precision for all three objectives. Successive Halving took `97.01 s` +of router time; proxy promotion took `111.90 s`. + +This supports using Successive Halving as a faster benchmark-search mode with no observed quality loss +on this control. It is an artifact comparison, not an automated strategy A/B test, and it does not yet +establish equivalence on the larger FastAPI and Effect-TS corpora. The older Schema-14/15 comparison +also found only rounding-level holdout movement while successive halving reduced router time by `9.3%`. + ## Schema 14: Router Search Strategy And Runtime Telemetry Schema 14 records the deterministic `halton-global-scout-elitist-beam` router search strategy and diff --git a/benchmarks/retrieval/evaluation/collect.ts b/benchmarks/retrieval/evaluation/collect.ts index 458254f..5df53c2 100644 --- a/benchmarks/retrieval/evaluation/collect.ts +++ b/benchmarks/retrieval/evaluation/collect.ts @@ -1,4 +1,5 @@ import { Effect, Stream } from "effect" +import { SqlClient } from "effect/unstable/sql" import type { Embedding } from "../../../src/domain/chunk.js" import { DEFAULT_CONFIG } from "../../../src/domain/config.js" @@ -7,6 +8,7 @@ import type { StoredChunk } from "../../../src/domain/index-data.js" import { MODEL_REGISTRY } from "../../../src/domain/models.js" import type { BoundEmbedder, SearchData } from "../../../src/domain/ports.js" import { IndexStore, SparseEmbedder } from "../../../src/domain/ports.js" +import type { ChannelRankings } from "../../../src/domain/retrieval.js" import type { SparseContract, SparseQuery, @@ -18,6 +20,13 @@ import { buildQueryTermCoverage } from "../../../src/lib/retrieval/evidence-rout import { createAutoBoundEmbedder } from "../../../src/services/embedder.js" import { prepareCorpus, type PreparedCorpus } from "../corpus/prepare.js" import { prepareRepository } from "../corpus/repository.js" +import { + benchmarkCachePaths, + loadCachedRankings, + saveCachedRankings, + type BenchmarkCachePaths, + type CachedRankingQuery, +} from "../execution/benchmark-cache.js" import { withSqliteBenchmarkStore } from "../execution/sqlite-index.js" import { foldKey } from "./folds.js" import { @@ -53,6 +62,7 @@ interface ModelMeasurements { readonly measurements: readonly QueryMeasurement[] readonly samples: readonly WeightSearchSample[] readonly samplesByQueryKind: ReadonlyMap + readonly rankings: readonly ChannelRankings[] readonly retrievalDurationMs: number } @@ -168,27 +178,22 @@ const persistBenchmarkCorpus = ( }) }) -const collectModelSamples = ( +interface BuiltModelSamples { + readonly measurements: readonly QueryMeasurement[] + readonly samples: readonly WeightSearchSample[] + readonly samplesByQueryKind: ReadonlyMap +} + +const buildModelSamples = ( manifest: CorpusManifest, corpus: PreparedCorpus, queries: readonly BenchmarkQuery[], targetsByQuestion: readonly (readonly ReadonlySet[])[], groupedFoldAssignments: ReadonlyMap, model: string, - queryVectors: readonly Float32Array[], - sparseQueries: readonly SparseQuery[], - searchData: SearchData, - store: typeof IndexStore.Service, - dims: number, - dtype: EmbeddingDtype, -): Effect.Effect< - { - readonly measurements: readonly QueryMeasurement[] - readonly samples: readonly WeightSearchSample[] - readonly samplesByQueryKind: ReadonlyMap - }, - Error -> => + rankingsByQuery: readonly ChannelRankings[], + channelDurations: readonly number[], +): Effect.Effect => Effect.gen(function* () { const modelMeasurements: QueryMeasurement[] = [] const modelSamples: WeightSearchSample[] = [] @@ -203,16 +208,9 @@ const collectModelSamples = ( return yield* Effect.fail( new Error(`No grouped fold assignment for ${manifest.id}/${question.id}`), ) - const channelStartedAt = performance.now() - const lexicalRankings = rankLexicalChannels(entry.query, searchData) - const dense = yield* store.searchDense({ - vector: queryVectors[queryIndex]!, - dims, - dtype, - }) - const sparse = yield* store.searchSparse(sparseQueries[queryIndex]!) - const rankings = { ...lexicalRankings, dense, sparse } - const channelDurationMs = performance.now() - channelStartedAt + const rankings = rankingsByQuery[queryIndex] + if (rankings === undefined) + return yield* Effect.fail(new Error(`No cached rankings for query ${queryIndex}`)) const sample: WeightSearchSample = { repository: manifest.id, intentId: question.id, @@ -222,11 +220,7 @@ const collectModelSamples = ( rankings, targets, chunks: corpus.chunks, - termCoverage: buildQueryTermCoverage( - entry.query, - searchData.bm25Index, - searchData.identifierIndex, - ), + termCoverage: buildQueryTermCoverage(entry.query, corpus.bm25Index, corpus.identifierIndex), } modelSamples.push(sample) samplesByQueryKind.set(entry.queryKind, [ @@ -236,7 +230,8 @@ const collectModelSamples = ( for (const variant of RETRIEVAL_VARIANTS) { const variantStartedAt = performance.now() const ranked = fuseVariant(variant, entry.query, rankings) - const queryDurationMs = channelDurationMs + performance.now() - variantStartedAt + const queryDurationMs = + (channelDurations[queryIndex] ?? 0) + performance.now() - variantStartedAt modelMeasurements.push({ repository: manifest.id, language: manifest.language, @@ -272,6 +267,53 @@ const collectModelSamples = ( return { measurements: modelMeasurements, samples: modelSamples, samplesByQueryKind } }) +const collectModelSamples = ( + manifest: CorpusManifest, + corpus: PreparedCorpus, + queries: readonly BenchmarkQuery[], + targetsByQuestion: readonly (readonly ReadonlySet[])[], + groupedFoldAssignments: ReadonlyMap, + model: string, + queryVectors: readonly Float32Array[], + sparseQueries: readonly SparseQuery[], + searchData: SearchData, + store: typeof IndexStore.Service, + dims: number, + dtype: EmbeddingDtype, +): Effect.Effect => + Effect.gen(function* () { + const rankings: ChannelRankings[] = [] + const channelDurations: number[] = [] + for (let queryIndex = 0; queryIndex < queries.length; queryIndex++) { + const entry = queries[queryIndex] + if (entry === undefined) + return yield* Effect.fail(new Error(`Missing benchmark query ${queryIndex}`)) + const channelStartedAt = performance.now() + const lexicalRankings = rankLexicalChannels(entry.query, searchData) + const dense = yield* store.searchDense({ + vector: queryVectors[queryIndex]!, + dims, + dtype, + }) + const sparse = yield* store.searchSparse(sparseQueries[queryIndex]!) + rankings.push({ ...lexicalRankings, dense, sparse }) + channelDurations.push(performance.now() - channelStartedAt) + } + return { + ...(yield* buildModelSamples( + manifest, + corpus, + queries, + targetsByQuestion, + groupedFoldAssignments, + model, + rankings, + channelDurations, + )), + rankings, + } + }) + const collectModelMeasurements = ( manifest: CorpusManifest, corpus: PreparedCorpus, @@ -282,37 +324,68 @@ const collectModelMeasurements = ( info: (typeof MODEL_REGISTRY)[string] & object, chunkVectors: readonly Float32Array[], queryVectors: readonly Float32Array[], + cachePaths: BenchmarkCachePaths, + hasPersistedIndex: boolean, + cachedRankings?: readonly ChannelRankings[], ): Effect.Effect => Effect.gen(function* () { const retrievalStartedAt = performance.now() + if (cachedRankings !== undefined) { + const built = yield* buildModelSamples( + manifest, + corpus, + queries, + targetsByQuestion, + groupedFoldAssignments, + model, + cachedRankings, + queries.map(() => 0), + ) + return { + ...built, + rankings: cachedRankings, + sparseEmbeddingRun: { + repository: manifest.id, + model: DEFAULT_CONFIG.sparseEmbedder.model, + tokenizerModel: DEFAULT_CONFIG.sparseEmbedder.queryModel, + batchSize: DEFAULT_CONFIG.sparseEmbedder.batchSize, + chunkEmbeddingDurationMs: 0, + queryTokenizationDurationMs: 0, + }, + retrievalDurationMs: performance.now() - retrievalStartedAt, + } + } const modelRun = yield* withSqliteBenchmarkStore( model, info.defaultDtype, Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient const store = yield* IndexStore const sparseEmbedder = yield* SparseEmbedder - const sparseStartedAt = performance.now() - const sparseVectors = yield* embedSparseTexts( - corpus.chunks.map((chunk) => chunk.text), - sparseEmbedder, - ) - const sparseChunkEmbeddingDurationMs = performance.now() - sparseStartedAt - const sparseIdf = yield* sparseEmbedder.loadIdf() + let sparseChunkEmbeddingDurationMs = 0 + if (!hasPersistedIndex) { + const sparseStartedAt = performance.now() + const sparseVectors = yield* embedSparseTexts( + corpus.chunks.map((chunk) => chunk.text), + sparseEmbedder, + ) + sparseChunkEmbeddingDurationMs = performance.now() - sparseStartedAt + yield* persistBenchmarkCorpus( + store, + corpus, + chunkVectors, + sparseVectors, + info.dims, + info.defaultDtype, + sparseEmbedder.contract, + yield* sparseEmbedder.loadIdf(), + ) + } const sparseQueryStartedAt = performance.now() const sparseQueries = yield* Effect.forEach(queries, ({ query }) => sparseEmbedder.tokenizeQuery(query), ) const sparseQueryTokenizationDurationMs = performance.now() - sparseQueryStartedAt - yield* persistBenchmarkCorpus( - store, - corpus, - chunkVectors, - sparseVectors, - info.dims, - info.defaultDtype, - sparseEmbedder.contract, - sparseIdf, - ) const searchData = yield* store.loadSearchData() const collected = yield* collectModelSamples( manifest, @@ -328,6 +401,7 @@ const collectModelMeasurements = ( info.dims, info.defaultDtype, ) + yield* saveCachedRankings(sql, cachePaths.cacheKey, queries, collected.rankings) return { sparseEmbeddingRun: { repository: manifest.id, @@ -340,16 +414,48 @@ const collectModelMeasurements = ( ...collected, } }), + cachePaths.databasePath, ) return { sparseEmbeddingRun: modelRun.sparseEmbeddingRun, measurements: modelRun.measurements, samples: modelRun.samples, samplesByQueryKind: modelRun.samplesByQueryKind, + rankings: modelRun.rankings, retrievalDurationMs: performance.now() - retrievalStartedAt, } }) +interface BenchmarkCacheState { + readonly rankings: readonly ChannelRankings[] | undefined + readonly hasPersistedIndex: boolean +} + +const inspectBenchmarkCache = ( + model: string, + dtype: EmbeddingDtype, + cachePaths: BenchmarkCachePaths, + queries: readonly CachedRankingQuery[], + chunkCount: number, +): Effect.Effect => + withSqliteBenchmarkStore( + model, + dtype, + Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient + const store = yield* IndexStore + const [rankings, status] = yield* Effect.all([ + loadCachedRankings(sql, cachePaths.cacheKey, queries), + store.getStatus(), + ]) + return { + rankings, + hasPersistedIndex: status.chunks === chunkCount && status.files > 0, + } + }), + cachePaths.databasePath, + ) + const collectRepositoryMeasurements = ( manifest: CorpusManifest, models: readonly string[], @@ -400,28 +506,47 @@ const collectRepositoryMeasurements = ( const info = MODEL_REGISTRY[model] if (info === undefined) return yield* Effect.fail(new Error(`Unknown embedding model ${model}`)) - const bound = yield* createAutoBoundEmbedder({ + const cachePaths = benchmarkCachePaths(manifest, model, info.dims, info.defaultDtype) + const cacheState = yield* inspectBenchmarkCache( model, - dtype: info.defaultDtype, - dims: info.dims, - }).pipe( - Effect.mapError( - (cause) => new Error(`Could not auto-select a device for ${model}`, { cause }), - ), - ) - const embeddingStartedAt = performance.now() - const chunkVectors = yield* embedTexts( - corpus.chunks.map((chunk) => chunk.text), - model, - bound.embedder, - ) - const chunkEmbeddingDurationMs = performance.now() - embeddingStartedAt - const queryEmbeddingStartedAt = performance.now() - const queryVectors = yield* embedTexts( - queries.map((entry) => entry.query), - model, - bound.embedder, + info.defaultDtype, + cachePaths, + queries.map(({ queryKind, query }) => ({ queryKind, query })), + corpus.chunks.length, ) + let device = "cache" + let chunkEmbeddingDurationMs = 0 + let queryEmbeddingDurationMs = 0 + let chunkVectors: readonly Float32Array[] = [] + let queryVectors: readonly Float32Array[] = [] + if (cacheState.rankings === undefined) { + const bound = yield* createAutoBoundEmbedder({ + model, + dtype: info.defaultDtype, + dims: info.dims, + }).pipe( + Effect.mapError( + (cause) => new Error(`Could not auto-select a device for ${model}`, { cause }), + ), + ) + device = bound.device + if (!cacheState.hasPersistedIndex) { + const embeddingStartedAt = performance.now() + chunkVectors = yield* embedTexts( + corpus.chunks.map((chunk) => chunk.text), + model, + bound.embedder, + ) + chunkEmbeddingDurationMs = performance.now() - embeddingStartedAt + } + const queryEmbeddingStartedAt = performance.now() + queryVectors = yield* embedTexts( + queries.map((entry) => entry.query), + model, + bound.embedder, + ) + queryEmbeddingDurationMs = performance.now() - queryEmbeddingStartedAt + } const modelData = yield* collectModelMeasurements( manifest, corpus, @@ -432,14 +557,17 @@ const collectRepositoryMeasurements = ( info, chunkVectors, queryVectors, + cachePaths, + cacheState.hasPersistedIndex, + cacheState.rankings, ) embeddingRuns.push({ repository: manifest.id, model, - device: bound.device, + device, batchSize: EMBEDDING_BATCH_SIZE, chunkEmbeddingDurationMs, - queryEmbeddingDurationMs: performance.now() - queryEmbeddingStartedAt, + queryEmbeddingDurationMs, }) sparseEmbeddingRuns.push(modelData.sparseEmbeddingRun) measurements.push(...modelData.measurements) diff --git a/benchmarks/retrieval/execution/benchmark-cache.ts b/benchmarks/retrieval/execution/benchmark-cache.ts new file mode 100644 index 0000000..fbd8169 --- /dev/null +++ b/benchmarks/retrieval/execution/benchmark-cache.ts @@ -0,0 +1,209 @@ +import path from "node:path" + +import { Effect, Schema } from "effect" +import { SqlClient, SqlSchema } from "effect/unstable/sql" + +import { DEFAULT_CONFIG } from "../../../src/domain/config.js" +import type { ChannelRankings } from "../../../src/domain/retrieval.js" +import { contentHash } from "../../../src/lib/content-hash.js" +import type { CorpusManifest, QueryKind } from "../evaluation/types.js" + +const CACHE_ROOT = path.resolve("benchmarks/.cache/retrieval/v1") +const CACHE_VERSION = 1 +const RANKING_IMPLEMENTATION_VERSION = 1 +const CACHE_BATCH_SIZE = 100 + +const RankedChunkSchema = Schema.Struct({ + chunkIndex: Schema.Number, + score: Schema.Number, +}) + +const ChannelRankingsSchema = Schema.Struct({ + identity: Schema.Array(RankedChunkSchema), + camelcase: Schema.Array(RankedChunkSchema), + bm25: Schema.Array(RankedChunkSchema), + dense: Schema.Array(RankedChunkSchema), + sparse: Schema.Array(RankedChunkSchema), +}) + +const RankingRowSchema = Schema.Struct({ + cacheKey: Schema.String, + queryIndex: Schema.Number, + queryKind: Schema.String, + query: Schema.String, + rankingsJson: Schema.String, +}) + +const asCacheError = + (message: string) => + (cause: unknown): Error => + cause instanceof Error ? cause : new Error(message, { cause }) + +const selectRankingRows = (sql: SqlClient.SqlClient) => { + const select = SqlSchema.findAll({ + Request: Schema.String, + Result: RankingRowSchema, + execute: (cacheKey) => sql` + SELECT cache_key, query_index, query_kind, query, rankings_json + FROM benchmark_channel_rankings + WHERE cache_key = ${cacheKey} + ORDER BY query_index + `, + }) + return (cacheKey: string) => + select(cacheKey).pipe( + Effect.mapError(asCacheError("Could not load benchmark channel rankings")), + ) +} + +export interface BenchmarkCachePaths { + readonly cacheKey: string + readonly databasePath: string +} + +/** Stable identity for one corpus/model/index configuration. */ +export const benchmarkCachePaths = ( + manifest: CorpusManifest, + model: string, + dims: number, + dtype: string, +): BenchmarkCachePaths => { + const identity = JSON.stringify({ + cacheVersion: CACHE_VERSION, + rankingImplementationVersion: RANKING_IMPLEMENTATION_VERSION, + manifest: { + id: manifest.id, + repository: manifest.repository, + revision: manifest.revision, + includeRoots: manifest.includeRoots, + excludePaths: manifest.excludePaths, + extensions: manifest.extensions, + }, + model, + dims, + dtype, + chunk: { + lines: DEFAULT_CONFIG.chunkLines, + overlap: DEFAULT_CONFIG.overlapLines, + minimumCharacters: DEFAULT_CONFIG.minChunkChars, + }, + sparse: { + model: DEFAULT_CONFIG.sparseEmbedder.model, + modelRevision: DEFAULT_CONFIG.sparseEmbedder.modelRevision, + queryModel: DEFAULT_CONFIG.sparseEmbedder.queryModel, + queryRevision: DEFAULT_CONFIG.sparseEmbedder.queryRevision, + idfContentHash: DEFAULT_CONFIG.sparseEmbedder.idfContentHash, + }, + }) + const cacheKey = contentHash(identity) + return { + cacheKey, + databasePath: path.join(CACHE_ROOT, `${cacheKey}.db`), + } +} + +/** Create the benchmark-only table without changing production migrations. */ +export const ensureBenchmarkCacheTable = (sql: SqlClient.SqlClient): Effect.Effect => + sql` + CREATE TABLE IF NOT EXISTS benchmark_channel_rankings ( + cache_key TEXT NOT NULL, + query_index INTEGER NOT NULL, + query_kind TEXT NOT NULL, + query TEXT NOT NULL, + rankings_json TEXT NOT NULL CHECK (json_valid(rankings_json)), + PRIMARY KEY (cache_key, query_index) + ) STRICT + `.pipe( + Effect.asVoid, + Effect.mapError((cause) => + cause instanceof Error + ? cause + : new Error("Could not create benchmark cache table", { cause }), + ), + ) + +const decodeRankings = (value: string): ChannelRankings | undefined => { + try { + return Schema.decodeUnknownSync(ChannelRankingsSchema)(JSON.parse(value)) + } catch { + return undefined + } +} + +export interface CachedRankingQuery { + readonly queryKind: QueryKind + readonly query: string +} + +/** Load a complete ranking set, rejecting partial or stale query payloads. */ +export const loadCachedRankings = ( + sql: SqlClient.SqlClient, + cacheKey: string, + queries: readonly CachedRankingQuery[], +): Effect.Effect => + Effect.gen(function* () { + const rows = yield* selectRankingRows(sql)(cacheKey) + if (rows.length !== queries.length) return undefined + + const rankings: ChannelRankings[] = [] + for (let index = 0; index < queries.length; index++) { + const row = rows[index] + const query = queries[index] + if ( + row === undefined || + query === undefined || + row.cacheKey !== cacheKey || + row.queryIndex !== index || + row.queryKind !== query.queryKind || + row.query !== query.query + ) + return undefined + const decoded = decodeRankings(row.rankingsJson) + if (decoded === undefined) return undefined + rankings.push(decoded) + } + return rankings + }) + +/** Replace one complete ranking set transactionally after a successful collection. */ +export const saveCachedRankings = ( + sql: SqlClient.SqlClient, + cacheKey: string, + queries: readonly CachedRankingQuery[], + rankings: readonly ChannelRankings[], +): Effect.Effect => + Effect.gen(function* () { + if (queries.length !== rankings.length) + return yield* Effect.fail(new Error("Benchmark ranking cache length mismatch")) + + yield* sql + .withTransaction( + Effect.gen(function* () { + yield* sql`DELETE FROM benchmark_channel_rankings WHERE cache_key = ${cacheKey}` + for (let start = 0; start < queries.length; start += CACHE_BATCH_SIZE) { + const values = sql.join( + ", ", + false, + )( + queries.slice(start, start + CACHE_BATCH_SIZE).map((query, offset) => { + const ranking = rankings[start + offset] + if (ranking === undefined) throw new Error("Missing benchmark ranking cache row") + return sql`( + ${cacheKey}, + ${start + offset}, + ${query.queryKind}, + ${query.query}, + ${JSON.stringify(ranking)} + )` + }), + ) + yield* sql` + INSERT INTO benchmark_channel_rankings ( + cache_key, query_index, query_kind, query, rankings_json + ) VALUES ${values} + ` + } + }), + ) + .pipe(Effect.mapError(asCacheError("Could not save benchmark channel rankings"))) + }) diff --git a/benchmarks/retrieval/execution/sqlite-index.ts b/benchmarks/retrieval/execution/sqlite-index.ts index 15408ec..e20bb5f 100644 --- a/benchmarks/retrieval/execution/sqlite-index.ts +++ b/benchmarks/retrieval/execution/sqlite-index.ts @@ -1,5 +1,6 @@ +import { NodeServices } from "@effect/platform-node" import { Effect, Layer } from "effect" -import { layerNoop } from "effect/FileSystem" +import { SqlClient } from "effect/unstable/sql" import { DEFAULT_CONFIG, type Config } from "../../../src/domain/config.js" import type { EmbeddingDtype } from "../../../src/domain/dtype.js" @@ -7,6 +8,7 @@ import { ConfigStore, IndexStore, SparseEmbedder } from "../../../src/domain/por import { SparseEmbedderBase } from "../../../src/services/sparse-embedder.js" import { SqliteIndexStoreBase } from "../../../src/services/sqlite-index-store.js" import { sqliteIndexDatabaseLayer } from "../../../src/services/sqlite-index-store/client.js" +import { ensureBenchmarkCacheTable } from "./benchmark-cache.js" const benchmarkConfig = (model: string, dtype: EmbeddingDtype): Config => ({ ...DEFAULT_CONFIG, @@ -28,21 +30,28 @@ const benchmarkConfigStore = (config: Config): typeof ConfigStore.Service => ({ configExists: () => Effect.succeed(true), }) -const sqliteBenchmarkLayer = (model: string, dtype: EmbeddingDtype) => +const sqliteBenchmarkLayer = (model: string, dtype: EmbeddingDtype, databasePath: string) => Layer.merge(SqliteIndexStoreBase, SparseEmbedderBase).pipe( Layer.provideMerge( Layer.merge( Layer.succeed(ConfigStore, benchmarkConfigStore(benchmarkConfig(model, dtype))), - sqliteIndexDatabaseLayer(":memory:"), + sqliteIndexDatabaseLayer(databasePath), ), ), - Layer.provideMerge(layerNoop({})), + Layer.provideMerge(NodeServices.layer), ) -/** Run one benchmark operation against production retrieval adapters and an in-memory index. */ +/** Run one benchmark operation against production retrieval adapters and a benchmark SQLite index. */ export const withSqliteBenchmarkStore = ( model: string, dtype: EmbeddingDtype, - effect: Effect.Effect, + effect: Effect.Effect, + databasePath = ":memory:", ): Effect.Effect => - Effect.scoped(effect.pipe(Effect.provide(sqliteBenchmarkLayer(model, dtype)))) + Effect.scoped( + Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient + yield* ensureBenchmarkCacheTable(sql) + return yield* effect + }).pipe(Effect.provide(sqliteBenchmarkLayer(model, dtype, databasePath))), + ) diff --git a/benchmarks/tests/benchmark-cache.test.ts b/benchmarks/tests/benchmark-cache.test.ts new file mode 100644 index 0000000..505f669 --- /dev/null +++ b/benchmarks/tests/benchmark-cache.test.ts @@ -0,0 +1,85 @@ +import { randomUUID } from "node:crypto" +import { tmpdir } from "node:os" +import { join } from "node:path" + +import { NodeServices } from "@effect/platform-node" +import { expect, it } from "@effect/vitest" +import { Effect } from "effect" +import { FileSystem } from "effect/FileSystem" +import { SqlClient } from "effect/unstable/sql" + +import type { ChannelRankings } from "../../src/domain/retrieval.js" +import { + loadCachedRankings, + saveCachedRankings, + type CachedRankingQuery, +} from "../retrieval/execution/benchmark-cache.js" +import { withSqliteBenchmarkStore } from "../retrieval/execution/sqlite-index.js" + +const queries = [ + { queryKind: "searchPhrase", query: "find the cache" }, + { queryKind: "naturalQuestion", query: "where is the cache stored?" }, +] as const satisfies readonly CachedRankingQuery[] + +const rankings: readonly ChannelRankings[] = [ + { + identity: [{ chunkIndex: 2, score: 1 }], + camelcase: [], + bm25: [{ chunkIndex: 1, score: 0.8 }], + dense: [{ chunkIndex: 0, score: 0.7 }], + sparse: [], + }, + { + identity: [], + camelcase: [{ chunkIndex: 3, score: 1 }], + bm25: [], + dense: [{ chunkIndex: 4, score: 0.6 }], + sparse: [{ chunkIndex: 5, score: 0.5 }], + }, +] + +it.effect("benchmark ranking cache survives a physical SQLite reopen", () => { + const databasePath = join(tmpdir(), `pix-benchmark-cache-${randomUUID()}.db`) + const run = (effect: Effect.Effect) => + withSqliteBenchmarkStore("cache-test", "fp32", effect, databasePath) + + return Effect.gen(function* () { + const fs = yield* FileSystem + yield* Effect.addFinalizer(() => + Effect.forEach( + [databasePath, `${databasePath}-shm`, `${databasePath}-wal`], + (candidate) => + Effect.gen(function* () { + if (yield* fs.exists(candidate)) yield* fs.remove(candidate) + }), + { discard: true }, + ).pipe(Effect.orDie), + ) + + yield* run( + Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient + yield* saveCachedRankings(sql, "cache-key", queries, rankings) + }), + ) + + const loaded = yield* run( + Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient + return yield* loadCachedRankings(sql, "cache-key", queries) + }), + ) + expect(loaded).toEqual(rankings) + + const stale = yield* run( + Effect.gen(function* () { + const sql = yield* SqlClient.SqlClient + return yield* loadCachedRankings(sql, "cache-key", [ + { ...queries[0]! }, + { queryKind: queries[1]!.queryKind, query: "changed query" }, + ]) + }), + ) + expect(stale).toBeUndefined() + }).pipe(Effect.provide(NodeServices.layer), Effect.scoped) +}) diff --git a/docs/adr/0019-evidence-based-fusion-and-optimization-profiles.md b/docs/adr/0019-evidence-based-fusion-and-optimization-profiles.md index 22bb5f9..4f6a60c 100644 --- a/docs/adr/0019-evidence-based-fusion-and-optimization-profiles.md +++ b/docs/adr/0019-evidence-based-fusion-and-optimization-profiles.md @@ -26,9 +26,11 @@ versus `73.3%`. Use the current Production router as the benchmark guardrail. Re historical diagnostic and rollback baseline while the broader matrix validation continues in issue #166. -Benchmarks compose production embedders and the production IndexStore with an in-memory SQLite database. -They may add fusion candidates at the `RankedChunk[]` seam, but must not reimplement production encoding, -persistence, or scoring. A fixed equal-weight RRF is always reported separately from production routing. +Benchmarks compose production embedders and the production IndexStore with a benchmark-owned SQLite +database. Cold runs persist the physical benchmark index and channel rankings under +`benchmarks/.cache/retrieval/v1/`; warm runs reuse them. They may add fusion candidates at the +`RankedChunk[]` seam, but must not reimplement production encoding, persistence, or scoring. A fixed +equal-weight RRF is always reported separately from production routing. The router configuration will contain: @@ -78,6 +80,95 @@ metadata, not reliably observable production inputs. Observable evidence can sti identifier coverage, query length, score geometry, and channel agreement. Explicit user-selected profiles may be added later. +## Runtime Estimation + +The benchmark records `timings.evidenceRouterSearchDurationMs` separately from embedding, physical +retrieval, and static fusion. It is the wall-clock duration of the evidence-router stage, including +candidate preparation/evaluation, all router holdout jobs, fit-all jobs, and shared candidate-worker +queue overhead. + +The number of router jobs is: + +```text +J = M * F * (K + H + 1) +``` + +where `M` is the number of embedding models, `F` the number of router fusion methods, `K` the number +of grouped folds, and `H` the number of repository holdout jobs: one per selected repository when +repository holdouts are enabled and more than one repository is selected, otherwise `0`. The `+1` is +the fit-all recommendation job for each model/fusion pair. The current three objectives (`direct`, +`reranker-top20`, and +`reranker-top50`) are selected from one shared candidate search per job; they do not multiply the main +dynamic search by three. They add result selection and validation work only. + +The current job counts are therefore: + +| Profile and corpus selection | `M` | `F` | `K` | `H` | Router jobs `J` | +| ----------------------------------- | --: | --: | --: | --: | --------------: | +| `develop`, one or more repositories | 1 | 1 | 3 | 0 | 4 | +| `validate`, all three repositories | 1 | 1 | 5 | 3 | 9 | +| `full`, all three repositories | 1 | 3 | 5 | 3 | 27 | + +The implementation runs all planned jobs through one shared eleven-worker candidate queue. Therefore +`J` scales candidate work, but not necessarily wall-clock time linearly: jobs overlap and compete for the +same workers. The implementation also holds out all `K` folds and all selected repositories; the +factors are not `K - 1` and `H - 1`. + +For the current `develop` calibration, each corpus has 15 authored questions and four query forms, +so the router sees 60 query samples. With one MiniLM model, DBSF only, grouped 3-fold, no repository +holdouts, the measured points are: + +| Chunks `N` | Router time `T` | +| ---------: | --------------: | +| 91 | 8.98 s | +| 411 | 126.94 s | +| 6,386 | 954.05 s | + +Each of the four jobs searches the same 40 router parameters with 64 global scouts, beam width 6, +and two coordinate passes. A job currently evaluates roughly 5.3k-5.9k proxy candidates and +3.2k-3.5k full candidates. The diagnostics are copied into one result row per objective, so these +counts must be deduplicated per job; they must not be summed across the three objective rows. + +A provisional line for this exact workload is: + +```text +T_develop(N) ~= 32.15 + 0.14 * N seconds +``` + +It is an empirical calibration from only three points (`R^2 ~= 0.9935`), not a universal complexity +law. For a first-order estimate of a different sample count `Q` and router-job count `J`, use the +variable work term as: + +```text +T_rough(N, Q, J) ~= 32.15 + 0.14 * N * (Q / 60) * (J / 4) seconds +``` + +This intentionally excludes embedding time and should be treated as a planning estimate until runs +with controlled `K`, `H`, and `F` variations provide separate calibration for fixed overhead, worker +contention, and query-sample scaling. Full-run estimates must add embedding, physical retrieval, and +static fusion timings separately. + +Using the current three corpora as one combined corpus (`N = 6,888` chunks and `Q = 180` query +samples), this provisional model predicts approximately 49 minutes for all-repository `develop` +with DBSF only (`J = 4`), 1 hour 49 minutes for all-repository `validate` with DBSF only (`J = 9`), +and 5 hours 26 minutes for the current all-repository `full` router stage (`J = 27`). These values +are deliberately estimates, not acceptance thresholds; a controlled multi-repository develop run is +still required to calibrate the query-sample and shared-worker effects. + +The historical full artifact `retrieval-2026-08-01T10-12-09.945Z.json` completed in about 49 minutes +(43.8 minutes in the router stage), but it is schema 17 and uses +`halton-global-scout-elitist-beam-successive-halving-pareto`. Current schema 23 uses +`halton-global-scout-elitist-beam-proxy-promotion`. Its individual develop measurements are therefore +not directly comparable to that historical full run; the 5 hour 26 minute value is a projection for +the current strategy, not a claim about the older artifact. + +A matched Schema-19 `fd` smoke comparison provides the current equivalence signal: the successive- +halving-pareto artifact reports `97.01 s` router time and the proxy-promotion artifact reports +`111.90 s`; their weighted evidence-router holdout summaries are identical at reported precision for +`direct`, `reranker-top20`, and `reranker-top50`. This is encouraging evidence for a faster successive- +halving mode, but it is not an automated A/B test and has not yet been repeated on FastAPI or Effect-TS. +The existing benchmark tests assert artifact structure and guardrails, not cross-strategy quality. + ## Consequences - DBSF is the active compatibility fusion; the current Production router is the benchmark guardrail, while From 6be00baac82f0adcae5e74841668862ffdbd17d7 Mon Sep 17 00:00:00 2001 From: Lucas Burmeister Date: Wed, 5 Aug 2026 04:04:15 +0200 Subject: [PATCH 8/8] perf: restore selectable successive halving search --- CONTEXT.md | 7 +- benchmarks/BASELINE.md | 33 +- benchmarks/README.md | 12 +- benchmarks/retrieval/evaluation/report.ts | 15 +- benchmarks/retrieval/evaluation/types.ts | 54 ++-- .../retrieval/evaluation/weight-search.ts | 302 +++++++++++++----- benchmarks/retrieval/runner.ts | 35 +- benchmarks/tests/retrieval.test.ts | 23 +- benchmarks/tests/worker-pool.test.ts | 51 ++- 9 files changed, 424 insertions(+), 108 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index 116bf65..b45d35b 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -235,7 +235,7 @@ fusion path and persists its IDF and postings in `.pix/index.db`. Schema 19 removes the benchmark-owned Sparse encoder, in-process postings implementation, and separate embedding caches. Benchmark profile fitting and optimizer search remain benchmark-owned, while the fusion adapters and evidence signals are shared with production. Benchmarks compose the production -SparseEmbedder and IndexStore around a migrated SQLite database; current schema-23 runs persist that +SparseEmbedder and IndexStore around a migrated SQLite database; current schema-24 runs persist that benchmark database and channel rankings under `benchmarks/.cache/retrieval/v1/` for warm reuse. Dense and Sparse ranking therefore execute through the same adapters as product queries. Experimental profile fitting remains benchmark-owned. Every artifact includes the authored file-qualified ground truth and @@ -253,6 +253,11 @@ current beam elites, so a later coordinate cannot regress the best development c development folds and evaluated unchanged against static weights on excluded intent folds and repositories; authored query-form labels remain informed reference strata and are not router inputs. This router remains benchmark-only until holdouts justify a production change. +Schema 24 exposes the historical `successive-halving` router search beside the default `proxy-promotion` +mode through `PIX_BENCH_ROUTER_STRATEGY`. Successive Halving uses the original lexicographic quality +comparator and `halvingKeepFactor`; both modes share the proxy/full evaluator pools and native worker +queue. Artifacts record the selected algorithm so FastAPI and other corpus comparisons remain +reproducible. Benchmark weight and router searches use a benchmark-only prepared evaluator: each sample and fusion method materializes per-chunk normalized or RRF contributions once, then candidate weights reuse that data. Public `fuseRankings` semantics and its existing ranking, normalization, and typed-array caches diff --git a/benchmarks/BASELINE.md b/benchmarks/BASELINE.md index ee5c518..4eb97f6 100644 --- a/benchmarks/BASELINE.md +++ b/benchmarks/BASELINE.md @@ -1,7 +1,38 @@ # Preliminary Retrieval Baseline The schema-17 entries below are historical artifacts from the benchmark-owned Sparse implementation. -Current schema-23 runs use the production SparseEmbedder and IndexStore without benchmark vector caches. +Current schema-24 runs use the production SparseEmbedder and IndexStore without benchmark vector caches. + +## Schema 24: Selectable Router Search + +Schema 24 exposes the current `proxy-promotion` search and the historical `successive-halving` +algorithm from commits `1754725` and `2f92428` through `PIX_BENCH_ROUTER_STRATEGY`. The legacy path +uses the original `halvingKeepFactor: 8`, lexicographic `R@20/R@10/Context@4k/MRR` comparator, and +does not run the later random-scout baseline. Both runs below used the same MiniLM model, pinned corpus, +warm retrieval cache, DBSF fusion, and 11 native worker threads. + +### FastAPI Develop Sanity Check + +| Strategy | Artifact | Router time | Total time | Dynamic holdout summary (R@5/R@10/R@20/R@50/Context@4k) | +| -------------------- | ----------------------------------------- | ----------: | ---------: | ------------------------------------------------------------ | +| `successive-halving` | `retrieval-2026-08-05T01-57-19.822Z.json` | 18.56 s | 23.59 s | 51.3%/64.3%/77.3%/88.7%/58.7% | +| `proxy-promotion` | `retrieval-2026-08-05T01-59-43.049Z.json` | 114.59 s | 119.64 s | objective-specific; 54.0%/65.0%/72.3%/85.0%/58.7% for direct | + +Successive Halving was 83.8% faster on this warm-cache control. Its objective-specific output is one +historical candidate repeated across the current three-objective artifact rows; Proxy Promotion selects +objective-specific candidates. This is a sanity check of behavior and queue integration, not a quality +promotion decision. + +### Effect v4 Develop Sanity Check + +| Strategy | Artifact | Router time | Total time | Dynamic holdout summary (R@5/R@10/R@20/R@50/Context@4k) | +| -------------------- | ----------------------------------------- | ----------: | ---------: | ------------------------------------------------------------ | +| `successive-halving` | `retrieval-2026-08-05T02-00-45.058Z.json` | 29.62 s | 46.62 s | 40.0%/52.0%/59.3%/82.0%/48.0% | +| `proxy-promotion` | `retrieval-2026-08-05T01-52-54.327Z.json` | 938.61 s | 954.79 s | objective-specific; 32.7%/42.0%/52.0%/80.0%/42.0% for direct | + +Successive Halving was 96.8% faster on the larger Effect v4 control. The worker queue completed both +strategies successfully; the large runtime gap comes from candidate selection/evaluation volume, not +embedding, which was reused from the persistent cache. ## Schema 20: Search-Priority DBSF Selection diff --git a/benchmarks/README.md b/benchmarks/README.md index 5c0eb81..25c3fa7 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -95,7 +95,7 @@ vp run bench:retrieval:full `bench:retrieval` aliases `bench:retrieval:validate`. Every profile measures the same physical rankings and retrieval variants; profiles only control matrix size, holdout coverage, and expensive -diagnostics. The selected profile is recorded in schema-23 artifacts without changing retrieval +diagnostics. The selected profile is recorded in schema-24 artifacts without changing retrieval semantics. The full profile includes all three fusion methods; short profiles intentionally omit RRF to keep development runs fast. @@ -110,6 +110,7 @@ Limit an exploratory run with comma-separated environment variables: $env:PIX_BENCH_REPOS = "fd" $env:PIX_BENCH_MODELS = "Xenova/all-MiniLM-L6-v2" $env:PIX_BENCH_OPTIMIZATION_PROFILE = "search-priority" +$env:PIX_BENCH_ROUTER_STRATEGY = "proxy-promotion" vp run bench:retrieval:validate ``` @@ -154,6 +155,12 @@ defaulting to MiniLM. Select another with `PIX_BENCH_MODELS`. Supported values a - `Xenova/bge-small-en-v1.5` - `jinaai/jina-embeddings-v2-base-code` +The router search defaults to `proxy-promotion`. Set `PIX_BENCH_ROUTER_STRATEGY` to +`successive-halving` to select the historical Successive-Halving variant. It uses the original +lexicographic `R@20`, `R@10`, `Context@4k`, and MRR comparator plus its `halvingKeepFactor`. +Both strategies use the same candidate evaluator and native worker queue, so their artifacts can be +compared directly. + The Jina code model cannot embed Effect's longest 7,103-token AST chunk on the tested DML GPU even as a single-item batch. Do not silently truncate, re-chunk only one model, or mix CPU and GPU vectors to complete that cell: any of those choices changes the comparison. Treat the cell as unsupported until @@ -343,7 +350,8 @@ output size without introducing an LLM or provider-specific tokenizer. Each run writes ignored JSON and Markdown artifacts under `benchmarks/results`. JSON rows retain the repository, revision, language, size, category, difficulty, query form, grouped fold, model, variant, -individual gold ranks, timing, and every metric. Schema 23 adds shared candidate-queue lifecycle and +individual gold ranks, timing, and every metric. Schema 24 adds selectable router strategies while +retaining shared candidate-queue lifecycle and per-router candidate-pool initialization timings. Each artifact stores each authored query and its exact file-qualified ground truth once, records productive Sparse timings, and adds the fixed equal-weight RRF baseline. The Markdown report includes quality by query form, diff --git a/benchmarks/retrieval/evaluation/report.ts b/benchmarks/retrieval/evaluation/report.ts index 0177ae5..6da5715 100644 --- a/benchmarks/retrieval/evaluation/report.ts +++ b/benchmarks/retrieval/evaluation/report.ts @@ -71,6 +71,13 @@ export const renderMarkdownReport = (artifact: BenchmarkArtifact): string => { group.push(row) groups.set(key, group) } + const strategyFactorLabel = artifact.searchStrategy.algorithm.includes("successive-halving") + ? "keep" + : "promotion" + const strategyFactor = + "halvingKeepFactor" in artifact.searchStrategy + ? artifact.searchStrategy.halvingKeepFactor + : artifact.searchStrategy.proxyPromotionFactor const lines = [ "# Retrieval Quality Benchmark", @@ -87,7 +94,7 @@ export const renderMarkdownReport = (artifact: BenchmarkArtifact): string => { "", `Validation protocol: candidates use ${artifact.validationProtocol.selection}; holdouts are ${artifact.validationProtocol.holdouts.join(" and ")}; final promotion requires the recorded ${artifact.validationProtocol.finalTest}.`, "", - `Search strategy: \`${artifact.searchStrategy.algorithm}\` (${artifact.searchStrategy.globalScouts} global scouts, beam ${artifact.searchStrategy.beamWidth}, ${artifact.searchStrategy.coordinatePasses} coordinate passes, ${artifact.searchStrategy.proxySampleFraction * 100}% proxy with minimum ${artifact.searchStrategy.proxyMinimumSamples}, promotion factor ${artifact.searchStrategy.proxyPromotionFactor}x).`, + `Search strategy: \`${artifact.searchStrategy.algorithm}\` (${artifact.searchStrategy.globalScouts} global scouts, beam ${artifact.searchStrategy.beamWidth}, ${artifact.searchStrategy.coordinatePasses} coordinate passes, ${artifact.searchStrategy.proxySampleFraction * 100}% proxy with minimum ${artifact.searchStrategy.proxyMinimumSamples}, ${strategyFactorLabel} factor ${strategyFactor}x).`, "", `Context budgets use the documented \`${artifact.contextTokenEstimator}\` estimator.`, "", @@ -268,13 +275,13 @@ export const renderMarkdownReport = (artifact: BenchmarkArtifact): string => { "", "Validation metrics are weighted by each excluded fold's query count. The current Production router is shown beside the selected dynamic router.", "", - "| Model | Fusion | Objective | Strategy | Promotion | Production R@5 | Dynamic R@5 | Production R@10 | Dynamic R@10 | Production R@20 | Dynamic R@20 | Production R@50 | Dynamic R@50 | Production Ctx@4k | Dynamic Ctx@4k | Random R@20 | Random Ctx@4k |", - "| --- | --- | --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", + "| Model | Fusion | Objective | Strategy | Promotion | Search baseline | Production R@5 | Dynamic R@5 | Production R@10 | Dynamic R@10 | Production R@20 | Dynamic R@20 | Production R@50 | Dynamic R@50 | Production Ctx@4k | Dynamic Ctx@4k | Random R@20 | Random Ctx@4k |", + "| --- | --- | --- | --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", ) for (const [key, rows] of routerGroups) { const [model, fusion, objective, strategy] = key.split("\0") lines.push( - `| ${model} | ${fusion} | ${objective} | ${strategy} | ${promotionLabel(rows.every((row) => row.promotionStatus === "eligible") ? "eligible" : "no-eligible-candidate")} | ${percent(weightedAverage(rows, (row) => row.productionValidation.recallAt5))} | ${percent(weightedAverage(rows, (row) => row.validation.recallAt5))} | ${percent(weightedAverage(rows, (row) => row.productionValidation.recallAt10))} | ${percent(weightedAverage(rows, (row) => row.validation.recallAt10))} | ${percent(weightedAverage(rows, (row) => row.productionValidation.recallAt20))} | ${percent(weightedAverage(rows, (row) => row.validation.recallAt20))} | ${percent(weightedAverage(rows, (row) => row.productionValidation.recallAt50))} | ${percent(weightedAverage(rows, (row) => row.validation.recallAt50))} | ${percent(weightedAverage(rows, (row) => row.productionValidation.contextRecallAt4096))} | ${percent(weightedAverage(rows, (row) => row.validation.contextRecallAt4096))} | ${percent(weightedAverage(rows, (row) => row.searchBaseline.validation.recallAt20))} | ${percent(weightedAverage(rows, (row) => row.searchBaseline.validation.contextRecallAt4096))} |`, + `| ${model} | ${fusion} | ${objective} | ${strategy} | ${promotionLabel(rows.every((row) => row.promotionStatus === "eligible") ? "eligible" : "no-eligible-candidate")} | ${rows[0]?.searchBaseline.algorithm ?? "unknown"} | ${percent(weightedAverage(rows, (row) => row.productionValidation.recallAt5))} | ${percent(weightedAverage(rows, (row) => row.validation.recallAt5))} | ${percent(weightedAverage(rows, (row) => row.productionValidation.recallAt10))} | ${percent(weightedAverage(rows, (row) => row.validation.recallAt10))} | ${percent(weightedAverage(rows, (row) => row.productionValidation.recallAt20))} | ${percent(weightedAverage(rows, (row) => row.validation.recallAt20))} | ${percent(weightedAverage(rows, (row) => row.productionValidation.recallAt50))} | ${percent(weightedAverage(rows, (row) => row.validation.recallAt50))} | ${percent(weightedAverage(rows, (row) => row.productionValidation.contextRecallAt4096))} | ${percent(weightedAverage(rows, (row) => row.validation.contextRecallAt4096))} | ${percent(weightedAverage(rows, (row) => row.searchBaseline.validation.recallAt20))} | ${percent(weightedAverage(rows, (row) => row.searchBaseline.validation.contextRecallAt4096))} |`, ) } diff --git a/benchmarks/retrieval/evaluation/types.ts b/benchmarks/retrieval/evaluation/types.ts index 5f56cff..d08db1b 100644 --- a/benchmarks/retrieval/evaluation/types.ts +++ b/benchmarks/retrieval/evaluation/types.ts @@ -62,23 +62,41 @@ type FusionMethod = ProductionFusionMethod export const ROUTER_OBJECTIVES = ["direct", "reranker-top20", "reranker-top50"] as const export type RouterObjective = (typeof ROUTER_OBJECTIVES)[number] -/** Versioned evidence-router search strategy recorded in every benchmark artifact. */ -export const ROUTER_SEARCH_STRATEGY = { - algorithm: "halton-global-scout-elitist-beam-proxy-promotion", - globalScouts: 64, - beamWidth: 6, - coordinatePasses: 2, - candidateDepth: 200, - proxySampleFraction: 0.25, - proxyMinimumSamples: 32, - proxyPromotionFactor: 8, - objectives: ROUTER_OBJECTIVES, - guardrailTolerance: 0.01, - seed: 0, - normalization: "per-channel-max-weight", - tieBreaking: "guardrails>objective>complexity>stable-key", +/** Versioned evidence-router search strategies recorded in benchmark artifacts. */ +export const ROUTER_SEARCH_STRATEGIES = { + "proxy-promotion": { + algorithm: "halton-global-scout-elitist-beam-proxy-promotion", + globalScouts: 64, + beamWidth: 6, + coordinatePasses: 2, + candidateDepth: 200, + proxySampleFraction: 0.25, + proxyMinimumSamples: 32, + proxyPromotionFactor: 8, + objectives: ROUTER_OBJECTIVES, + guardrailTolerance: 0.01, + seed: 0, + normalization: "per-channel-max-weight", + tieBreaking: "guardrails>objective>complexity>stable-key", + }, + "successive-halving": { + algorithm: "halton-global-scout-elitist-beam-successive-halving", + globalScouts: 64, + beamWidth: 6, + coordinatePasses: 2, + candidateDepth: 200, + proxySampleFraction: 0.25, + proxyMinimumSamples: 32, + halvingKeepFactor: 8, + }, } as const +export type RouterSearchStrategyName = keyof typeof ROUTER_SEARCH_STRATEGIES +export type RouterSearchStrategy = (typeof ROUTER_SEARCH_STRATEGIES)[RouterSearchStrategyName] + +export const DEFAULT_ROUTER_SEARCH_STRATEGY: RouterSearchStrategyName = "proxy-promotion" +export const ROUTER_SEARCH_STRATEGY = ROUTER_SEARCH_STRATEGIES[DEFAULT_ROUTER_SEARCH_STRATEGY] + /** Runtime/coverage trade-off selected for one benchmark invocation. */ export type BenchmarkProfile = "smoke" | "develop" | "validate" | "full" @@ -266,7 +284,7 @@ export interface RouterSearchDiagnostics { /** Holdout comparison against a deterministic random-search baseline. */ export interface SearchBaselineComparison { - readonly algorithm: "random-scout" + readonly algorithm: "random-scout" | "not-run" readonly seed: number readonly candidates: number readonly development: QualitySummary @@ -343,14 +361,14 @@ export interface BenchmarkTimings { /** Reproducible machine-readable output of one complete benchmark run. */ export interface BenchmarkArtifact { - readonly schemaVersion: 23 + readonly schemaVersion: 24 /** Profile controlling benchmark coverage without changing retrieval behavior. */ readonly benchmarkProfile: BenchmarkProfile /** Versioned objective profile used for candidate selection and aggregate metrics. */ readonly optimizationProfile: OptimizationProfile readonly validationProtocol: ValidationProtocol readonly generatedAt: string - readonly searchStrategy: typeof ROUTER_SEARCH_STRATEGY + readonly searchStrategy: RouterSearchStrategy readonly timings: BenchmarkTimings readonly chunkConfig: { readonly chunkLines: number diff --git a/benchmarks/retrieval/evaluation/weight-search.ts b/benchmarks/retrieval/evaluation/weight-search.ts index 19ca1c3..6e693d5 100644 --- a/benchmarks/retrieval/evaluation/weight-search.ts +++ b/benchmarks/retrieval/evaluation/weight-search.ts @@ -34,6 +34,8 @@ import { prepareFusion, type PreparedFusionEvaluator } from "./prepared-fusion.j import { ROUTER_OBJECTIVES, ROUTER_SEARCH_STRATEGY, + ROUTER_SEARCH_STRATEGIES, + type RouterSearchStrategyName, type EvidenceRouterSearchResult, type FusionSearchResult, type HoldoutQuality, @@ -63,6 +65,7 @@ const SEARCH_GLOBAL_SCOUTS = ROUTER_SEARCH_STRATEGY.globalScouts const SEARCH_PROXY_SAMPLE_FRACTION = ROUTER_SEARCH_STRATEGY.proxySampleFraction const SEARCH_PROXY_MINIMUM_SAMPLES = ROUTER_SEARCH_STRATEGY.proxyMinimumSamples const SEARCH_PROXY_PROMOTION_FACTOR = ROUTER_SEARCH_STRATEGY.proxyPromotionFactor +const SEARCH_HALVING_KEEP_FACTOR = ROUTER_SEARCH_STRATEGIES["successive-halving"].halvingKeepFactor const HALTON_PRIMES = [ 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127, 131, 137, 139, 149, @@ -93,6 +96,8 @@ interface SearchOptions extends CandidateEvaluationPoolOptions { readonly signal?: AbortSignal /** Shared candidate scheduler used to interleave multiple router searches. */ readonly evaluationQueue?: CandidateEvaluationQueue + /** Benchmark-only router search algorithm; defaults to the current proxy promotion mode. */ + readonly routerSearchStrategy?: RouterSearchStrategyName } /** Options for benchmark search APIs without colliding with the product query options type. */ @@ -493,6 +498,26 @@ const compareObjectiveQuality = ( return 0 } +const compareSuccessiveHalvingQuality = (left: QualitySummary, right: QualitySummary): number => { + const leftValues = [ + left.recallAt20, + left.recallAt10, + left.contextRecallAt4096, + left.meanReciprocalRank, + ] + const rightValues = [ + right.recallAt20, + right.recallAt10, + right.contextRecallAt4096, + right.meanReciprocalRank, + ] + for (let index = 0; index < leftValues.length; index++) { + if (leftValues[index] > rightValues[index]) return -1 + if (leftValues[index] < rightValues[index]) return 1 + } + return 0 +} + const weightCandidates = (): readonly ChannelWeights[] => { const candidates: ChannelWeights[] = [] for (const identity of WEIGHT_LEVELS) @@ -604,6 +629,7 @@ const sortWeightCandidates = ( objective: RouterObjective, baseline: QualitySummary | undefined, profile: OptimizationProfile, + successiveHalving: boolean, ): readonly WeightCandidate[] => entries .map(([key, weights]) => { @@ -613,8 +639,12 @@ const sortWeightCandidates = ( }) .sort( (left, right) => - compareObjectiveQuality(left.quality, right.quality, objective, baseline, profile) || - activeChannelsKey(left.weights).localeCompare(activeChannelsKey(right.weights)), + (successiveHalving + ? compareSuccessiveHalvingQuality(left.quality, right.quality) + : compareObjectiveQuality(left.quality, right.quality, objective, baseline, profile)) || + (successiveHalving + ? 0 + : activeChannelsKey(left.weights).localeCompare(activeChannelsKey(right.weights))), ) .slice(0, limit) @@ -640,6 +670,7 @@ const rankWeightCandidates = async ( objective: RouterObjective = "reranker-top20", baseline?: QualitySummary, profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + successiveHalving = false, ): Promise => { const unique = uniqueWeightCandidates(candidates) const pending = unique.filter(([key]) => qualityCache.get(key) === undefined) @@ -650,7 +681,15 @@ const rankWeightCandidates = async ( qualityCache, "Candidate evaluation returned an incomplete weight result", ) - return sortWeightCandidates(unique, qualityCache, limit, objective, baseline, profile) + return sortWeightCandidates( + unique, + qualityCache, + limit, + objective, + baseline, + profile, + successiveHalving, + ) } const withWeight = (weights: ChannelWeights, channel: ChannelName, value: number): ChannelWeights => @@ -675,6 +714,7 @@ const selectBestWeights = async ( objective: RouterObjective = "reranker-top20", baseline?: QualitySummary, profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + successiveHalving = false, ): Promise<{ readonly weights: ChannelWeights; readonly quality: QualitySummary }> => { const qualityCache = new Map() let beam = await rankWeightCandidates( @@ -685,6 +725,7 @@ const selectBestWeights = async ( objective, baseline, profile, + successiveHalving, ) for (let pass = 0; pass < SEARCH_PASSES; pass++) { for (const channel of CHANNELS) { @@ -696,6 +737,7 @@ const selectBestWeights = async ( objective, baseline, profile, + successiveHalving, ) } } @@ -709,6 +751,7 @@ const selectWeightSubsetCandidates = ( candidates: readonly ChannelWeights[], qualities: readonly QualitySummary[], profile: OptimizationProfile, + successiveHalving: boolean, ): readonly ChannelWeights[] => { const selected = new Map< string, @@ -723,7 +766,10 @@ const selectWeightSubsetCandidates = ( const current = selected.get(key) if ( current === undefined || - compareObjectiveQuality(quality, current.quality, "reranker-top20", undefined, profile) < 0 + (successiveHalving + ? compareSuccessiveHalvingQuality(quality, current.quality) + : compareObjectiveQuality(quality, current.quality, "reranker-top20", undefined, profile)) < + 0 ) selected.set(key, { weights, quality }) } @@ -733,10 +779,11 @@ const selectWeightSubsetCandidates = ( const selectBestWeightsPerSubset = async ( pool: CandidateEvaluationPool, profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, + successiveHalving = false, ): Promise => { const candidates = weightCandidates() const qualities = await pool.evaluate(candidates.map((weights) => ({ weights }))) - return selectWeightSubsetCandidates(candidates, qualities, profile) + return selectWeightSubsetCandidates(candidates, qualities, profile, successiveHalving) } const selectStaticWeights = ( @@ -1017,6 +1064,7 @@ interface RouterSearchContext { readonly stats: SearchEvaluationStats readonly fullPool: CandidateEvaluationPool readonly proxyPool: CandidateEvaluationPool + readonly routerSearchStrategy: RouterSearchStrategyName } const compareRouterCandidates = ( @@ -1029,6 +1077,19 @@ const compareRouterCandidates = ( compareObjectiveQuality(left.quality, right.quality, objective, baseline, profile) || routerComplexity(left.config) - routerComplexity(right.config) +const compareSuccessiveHalvingCandidates = ( + left: RouterCandidate, + right: RouterCandidate, +): number => + compareSuccessiveHalvingQuality(left.quality, right.quality) || + routerComplexity(left.config) - routerComplexity(right.config) + +const selectSuccessiveHalvingCandidates = ( + candidates: readonly RouterCandidate[], + limit: number, +): readonly RouterCandidate[] => + [...candidates].sort(compareSuccessiveHalvingCandidates).slice(0, limit) + const selectRandomRouter = async ( samples: readonly EvidenceSearchSample[], baseSeed: EvidenceRouterConfig, @@ -1116,16 +1177,28 @@ const finalizeRouterCandidates = ( } } } - for (const candidate of rankedFull) context.archive.set(routerKey(candidate.config), candidate) - const rankedElites = selectObjectiveCandidates( - [...context.elites.values(), ...rankedFull], - SEARCH_BEAM_WIDTH, - context.baseline, - context.profile, - ) + const orderedFull = + context.routerSearchStrategy === "successive-halving" + ? [...rankedFull].sort(compareSuccessiveHalvingCandidates) + : rankedFull + for (const candidate of orderedFull) context.archive.set(routerKey(candidate.config), candidate) + const rankedElites = + context.routerSearchStrategy === "successive-halving" + ? selectSuccessiveHalvingCandidates( + [...context.elites.values(), ...orderedFull], + SEARCH_BEAM_WIDTH, + ) + : selectObjectiveCandidates( + [...context.elites.values(), ...orderedFull], + SEARCH_BEAM_WIDTH, + context.baseline, + context.profile, + ) context.elites.clear() for (const candidate of rankedElites) context.elites.set(routerKey(candidate.config), candidate) - return selectObjectiveCandidates(rankedFull, limit, context.baseline, context.profile) + return context.routerSearchStrategy === "successive-halving" + ? orderedFull.slice(0, limit) + : selectObjectiveCandidates(rankedFull, limit, context.baseline, context.profile) } const evaluateRouterConfigs = async ( @@ -1188,17 +1261,28 @@ const rankRouterCandidates = async ( context.stats.timings.candidatePreparationMs += rankedProxy.candidatePreparationMs context.stats.timings.candidateEvaluationMs += rankedProxy.candidateEvaluationMs const proxySelectionStartedAt = performance.now() - const selectedProxy = selectObjectiveCandidates( - rankedProxy.candidates, - limit * SEARCH_PROXY_PROMOTION_FACTOR, - context.proxyBaseline, - context.profile, - ) + const selectedProxy = + context.routerSearchStrategy === "successive-halving" + ? selectSuccessiveHalvingCandidates( + rankedProxy.candidates, + limit * SEARCH_HALVING_KEEP_FACTOR, + ) + : selectObjectiveCandidates( + rankedProxy.candidates, + limit * SEARCH_PROXY_PROMOTION_FACTOR, + context.proxyBaseline, + context.profile, + ) const selected = new Map( selectedProxy.map((candidate) => [routerKey(candidate.config), candidate.config]), ) - proxyKeys = selectedProxy.map((candidate) => routerKey(candidate.config)) - context.stats.proxyPromotions += proxyKeys.length + // The historical mode uses the original global lexicographic halving set and skips + // newer objective-diverse promotion and proxy/full agreement diagnostics. + proxyKeys = + context.routerSearchStrategy === "successive-halving" + ? [] + : selectedProxy.map((candidate) => routerKey(candidate.config)) + context.stats.proxyPromotions += selectedProxy.length const protectedKeys = new Set([...protectedConfigs.map(routerKey), ...context.elites.keys()]) for (const [key, config] of fullCandidates) { if (protectedKeys.has(key)) selected.set(key, config) @@ -1313,6 +1397,7 @@ const selectBestEvidenceRouter = async ( profile: OptimizationProfile, fullPool: CandidateEvaluationPool, proxyPool: CandidateEvaluationPool, + routerSearchStrategy: RouterSearchStrategyName, ): Promise<{ readonly selections: readonly EvidenceRouterSelection[] readonly productionQuality: QualitySummary @@ -1331,6 +1416,7 @@ const selectBestEvidenceRouter = async ( proxyBaseline, stats, } = preparation + const useSuccessiveHalving = routerSearchStrategy === "successive-halving" const searchContext: RouterSearchContext = { samples: evidenceSamples, proxySamples, @@ -1345,6 +1431,7 @@ const selectBestEvidenceRouter = async ( stats, fullPool, proxyPool, + routerSearchStrategy, } const profileSeed = profileSeedFor(profile) const baseSearchStartedAt = performance.now() @@ -1353,25 +1440,31 @@ const selectBestEvidenceRouter = async ( "reranker-top20", productionQuality, profile, + useSuccessiveHalving, ) - const subsetWeights = await selectBestWeightsPerSubset(fullPool, profile) + const subsetWeights = await selectBestWeightsPerSubset(fullPool, profile, useSuccessiveHalving) const baseSeeds = [baseWeights.weights, ...subsetWeights].map(emptyRouterConfig) - baseSeeds.unshift(profileSeed) + if (!useSuccessiveHalving) baseSeeds.unshift(profileSeed) stats.timings.baseWeightSearchMs += performance.now() - baseSearchStartedAt const parameters = routerParameters() const randomBaseSeed = baseSeeds[0] if (randomBaseSeed === undefined) throw new Error("Evidence router search has no base seed") - const randomSearchStartedAt = performance.now() - const randomSearch = await selectRandomRouter( - evidenceSamples, - randomBaseSeed, - parameters, - fullPool, - productionQuality, - profile, - stats, - ) - stats.timings.randomSearchMs += performance.now() - randomSearchStartedAt + const randomSearch = useSuccessiveHalving + ? undefined + : await (async () => { + const randomSearchStartedAt = performance.now() + const result = await selectRandomRouter( + evidenceSamples, + randomBaseSeed, + parameters, + fullPool, + productionQuality, + profile, + stats, + ) + stats.timings.randomSearchMs += performance.now() - randomSearchStartedAt + return result + })() const beamSearchStartedAt = performance.now() let beam = await rankRouterCandidates( searchContext, @@ -1398,32 +1491,52 @@ const selectBestEvidenceRouter = async ( } stats.timings.beamSearchMs += performance.now() - beamSearchStartedAt const fallback = beam[0] + if (fallback === undefined) throw new Error("Evidence router search produced no candidate") const candidates = [...searchContext.archive.values()] - const selections = ROUTER_OBJECTIVES.map((objective) => { - const rankedCandidates = [...candidates].sort((left, right) => - compareRouterCandidates(left, right, objective, productionQuality, profile), - ) - const { candidate: eligibleCandidate, promotionStatus } = selectEligibleCandidate( - rankedCandidates, - (entry) => - evidenceRouterGuardrailsMet(samples, entry.config, fusion, guardrailBaselines, profile), - ) - const candidate = eligibleCandidate ?? rankedCandidates[0] ?? fallback - if (candidate === undefined) throw new Error("Evidence router search produced no candidate") - return { - objective, - config: candidate.config, - quality: candidate.quality, - guardrailsMet: eligibleCandidate !== undefined, - promotionStatus, - } - }) + const selections = useSuccessiveHalving + ? (() => { + const guardrailsMet = evidenceRouterGuardrailsMet( + samples, + fallback.config, + fusion, + guardrailBaselines, + profile, + ) + return ROUTER_OBJECTIVES.map((objective) => ({ + objective, + config: fallback.config, + quality: fallback.quality, + guardrailsMet, + promotionStatus: (guardrailsMet + ? "eligible" + : "no-eligible-candidate") as PromotionStatus, + })) + })() + : ROUTER_OBJECTIVES.map((objective) => { + const rankedCandidates = [...candidates].sort((left, right) => + compareRouterCandidates(left, right, objective, productionQuality, profile), + ) + const { candidate: eligibleCandidate, promotionStatus } = selectEligibleCandidate( + rankedCandidates, + (entry) => + evidenceRouterGuardrailsMet(samples, entry.config, fusion, guardrailBaselines, profile), + ) + const candidate = eligibleCandidate ?? rankedCandidates[0] ?? fallback + return { + objective, + config: candidate.config, + quality: candidate.quality, + guardrailsMet: eligibleCandidate !== undefined, + promotionStatus, + } + }) + const randomCandidate = randomSearch?.candidate ?? fallback return { selections, productionQuality, searchDiagnostics: buildSearchDiagnostics(parameters, stats), - randomCandidate: randomSearch.candidate, - randomCandidates: randomSearch.candidates, + randomCandidate, + randomCandidates: randomSearch?.candidates ?? 0, proxyEvaluations: stats.proxyEvaluations, fullEvaluations: stats.fullEvaluations, } @@ -1677,6 +1790,7 @@ const withEvidencePools = async ( profile, fullPool, proxyPool, + options.routerSearchStrategy ?? "proxy-promotion", ) return await operation(selection, fullPool) } finally { @@ -1698,7 +1812,18 @@ const selectStaticWeightsForSelections = async ( fullPool: CandidateEvaluationPool, productionQuality: QualitySummary, profile: OptimizationProfile, + routerSearchStrategy: RouterSearchStrategyName, ): Promise => { + if (routerSearchStrategy === "successive-halving") { + const staticSelection = await selectBestWeights( + fullPool, + "reranker-top20", + undefined, + profile, + true, + ) + return selections.map((selection) => ({ selection, staticSelection })) + } const selected: StaticRouterSelection[] = [] for (const selection of selections) { selected.push({ @@ -1714,6 +1839,20 @@ const selectStaticWeightsForSelections = async ( return selected } +const selectStaticWeightsForSearch = ( + dynamicSelection: Awaited>, + fullPool: CandidateEvaluationPool, + profile: OptimizationProfile, + options: SearchOptions, +): Promise => + selectStaticWeightsForSelections( + dynamicSelection.selections, + fullPool, + dynamicSelection.productionQuality, + profile, + options.routerSearchStrategy ?? "proxy-promotion", + ) + export const optimizeEvidenceRouter = async ( model: string, fusion: FusionMethod, @@ -1728,23 +1867,40 @@ export const optimizeEvidenceRouter = async ( const productionValidation = summarizeProductionRouter(validation, profile) const validationEvidence = prepareEvidenceSamples(validation) const holdoutProfile = unweightedProfile(profile) - const randomBaseline: SearchBaselineComparison = { - algorithm: "random-scout", - seed: RANDOM_SEARCH_SEED, - candidates: dynamicSelection.randomCandidates, - development: dynamicSelection.randomCandidate.quality, - validation: summarizeEvidenceRouter( - validationEvidence, - dynamicSelection.randomCandidate.config, - fusion, - profile, - ), - } - const staticSelections = await selectStaticWeightsForSelections( - dynamicSelection.selections, + const randomBaseline: SearchBaselineComparison = + dynamicSelection.randomCandidates === 0 + ? { + algorithm: "not-run", + seed: RANDOM_SEARCH_SEED, + candidates: 0, + development: summarizeEvidenceRouter( + [], + dynamicSelection.randomCandidate.config, + fusion, + ), + validation: summarizeEvidenceRouter( + [], + dynamicSelection.randomCandidate.config, + fusion, + ), + } + : { + algorithm: "random-scout", + seed: RANDOM_SEARCH_SEED, + candidates: dynamicSelection.randomCandidates, + development: dynamicSelection.randomCandidate.quality, + validation: summarizeEvidenceRouter( + validationEvidence, + dynamicSelection.randomCandidate.config, + fusion, + profile, + ), + } + const staticSelections = await selectStaticWeightsForSearch( + dynamicSelection, fullPool, - dynamicSelection.productionQuality, profile, + options, ) const results: EvidenceRouterSearchResult[] = staticSelections.map( ({ selection, staticSelection }) => ({ @@ -1883,11 +2039,11 @@ const fitRecommendedEvidenceRouterWithOptions = ( options: SearchOptions, ): Promise => withEvidencePools(samples, fusion, profile, options, async (dynamicSelection, fullPool) => { - const staticSelections = await selectStaticWeightsForSelections( - dynamicSelection.selections, + const staticSelections = await selectStaticWeightsForSearch( + dynamicSelection, fullPool, - dynamicSelection.productionQuality, profile, + options, ) const results: RecommendedEvidenceRouter[] = staticSelections.map( ({ selection, staticSelection }) => ({ diff --git a/benchmarks/retrieval/runner.ts b/benchmarks/retrieval/runner.ts index 5cec8e6..51ef65f 100644 --- a/benchmarks/retrieval/runner.ts +++ b/benchmarks/retrieval/runner.ts @@ -16,10 +16,13 @@ import { import { renderMarkdownReport } from "./evaluation/report.js" import { runBenchmarkSearch, type BenchmarkSearchConfig } from "./evaluation/search.js" import { - ROUTER_SEARCH_STRATEGY, + DEFAULT_ROUTER_SEARCH_STRATEGY, + ROUTER_SEARCH_STRATEGIES, type BenchmarkArtifact, type BenchmarkProfile, type CorpusManifest, + type RouterSearchStrategy, + type RouterSearchStrategyName, type ValidationStrategy, } from "./evaluation/types.js" import { getDefaultWorkerCount, resolveWorkerCount } from "./execution/candidate-evaluation-pool.js" @@ -114,6 +117,29 @@ const selectOptimizationProfile = (): Effect.Effect : Effect.succeed(selected) } +export const resolveRouterSearchStrategy = ( + requested: string | undefined, +): RouterSearchStrategyName => { + if (requested === undefined) return DEFAULT_ROUTER_SEARCH_STRATEGY + if (!Object.hasOwn(ROUTER_SEARCH_STRATEGIES, requested)) { + throw new Error( + `Unknown PIX_BENCH_ROUTER_STRATEGY value: ${requested}; expected one of ${Object.keys(ROUTER_SEARCH_STRATEGIES).join(", ")}`, + ) + } + const selected = (ROUTER_SEARCH_STRATEGIES as Record)[ + requested + ] + if (selected === undefined) + throw new Error(`Router search strategy is not configured: ${requested}`) + return requested as RouterSearchStrategyName +} + +const selectRouterSearchStrategy = (): Effect.Effect => + Effect.try({ + try: () => resolveRouterSearchStrategy(process.env.PIX_BENCH_ROUTER_STRATEGY), + catch: (cause) => (cause instanceof Error ? cause : new Error(String(cause))), + }) + const writeArtifact = (artifact: BenchmarkArtifact): Effect.Effect => Effect.gen(function* () { const outputDirectory = path.resolve("benchmarks/results") @@ -150,6 +176,7 @@ export const runRetrievalBenchmark = ( fallbackToSerial: false, } const optimizationProfile = yield* selectOptimizationProfile() + const routerSearchStrategy = yield* selectRouterSearchStrategy() const groupedStrategy: ValidationStrategy = config.groupedFolds === 3 ? "grouped-3-fold" : "grouped-5-fold" const manifests = yield* selectManifests(yield* loadCorpusManifests(), profile) @@ -172,7 +199,7 @@ export const runRetrievalBenchmark = ( groupedStrategy, optimizationProfile, serialSearch, - searchOptions, + { ...searchOptions, routerSearchStrategy }, ) const embeddingDurationMs = embeddingRuns.reduce( @@ -185,7 +212,7 @@ export const runRetrievalBenchmark = ( ) const artifact: BenchmarkArtifact = { - schemaVersion: 23, + schemaVersion: 24, benchmarkProfile: profile, optimizationProfile, validationProtocol: { @@ -199,7 +226,7 @@ export const runRetrievalBenchmark = ( nestedInnerFolds: Math.max(3, config.groupedFolds - 2), }, generatedAt: new Date().toISOString(), - searchStrategy: ROUTER_SEARCH_STRATEGY, + searchStrategy: ROUTER_SEARCH_STRATEGIES[routerSearchStrategy], timings: { totalDurationMs: performance.now() - benchmarkStartedAt, corpusPreparationDurationMs, diff --git a/benchmarks/tests/retrieval.test.ts b/benchmarks/tests/retrieval.test.ts index c815fef..c7843a4 100644 --- a/benchmarks/tests/retrieval.test.ts +++ b/benchmarks/tests/retrieval.test.ts @@ -4,10 +4,10 @@ import { Effect } from "effect" import { assignGroupedFolds } from "../retrieval/evaluation/folds.js" import { ROUTER_OBJECTIVES, - ROUTER_SEARCH_STRATEGY, + ROUTER_SEARCH_STRATEGIES, type BenchmarkProfile, } from "../retrieval/evaluation/types.js" -import { runRetrievalBenchmark } from "../retrieval/runner.js" +import { resolveRouterSearchStrategy, runRetrievalBenchmark } from "../retrieval/runner.js" const foldQuestions = (prefix: string) => Array.from({ length: 12 }, (_, index) => ({ @@ -99,8 +99,10 @@ const runProfile = (profile: BenchmarkProfile, groupedFolds: number, fusionMetho expect(artifact.evaluationCases.every(({ groundTruth }) => groundTruth.length > 0)).toBe(true) expect(artifact.models.length).toBeGreaterThan(0) expect(artifact.measurements.length).toBeGreaterThan(0) - expect(artifact.schemaVersion).toBe(23) - expect(artifact.searchStrategy).toEqual(ROUTER_SEARCH_STRATEGY) + expect(artifact.schemaVersion).toBe(24) + expect(artifact.searchStrategy).toEqual( + ROUTER_SEARCH_STRATEGIES[resolveRouterSearchStrategy(process.env.PIX_BENCH_ROUTER_STRATEGY)], + ) expect(artifact.timings.totalDurationMs).toBeGreaterThan(0) expect(Object.values(artifact.timings).every((duration) => duration >= 0)).toBe(true) expect( @@ -178,3 +180,16 @@ it("shuffles intent groups deterministically before assigning folds", () => { expectStratifiedClasses(manifests, assignments) }) + +it("resolves the selectable router search strategies", () => { + expect(resolveRouterSearchStrategy(undefined)).toBe("proxy-promotion") + expect(resolveRouterSearchStrategy("successive-halving")).toBe("successive-halving") + expect(ROUTER_SEARCH_STRATEGIES["successive-halving"]).toMatchObject({ + algorithm: "halton-global-scout-elitist-beam-successive-halving", + halvingKeepFactor: 8, + }) + expect("proxyPromotionFactor" in ROUTER_SEARCH_STRATEGIES["successive-halving"]).toBe(false) + expect(() => resolveRouterSearchStrategy("unknown")).toThrow( + "Unknown PIX_BENCH_ROUTER_STRATEGY value: unknown", + ) +}) diff --git a/benchmarks/tests/worker-pool.test.ts b/benchmarks/tests/worker-pool.test.ts index 272f0ff..9bfbd7a 100644 --- a/benchmarks/tests/worker-pool.test.ts +++ b/benchmarks/tests/worker-pool.test.ts @@ -88,6 +88,16 @@ const searchSample = { chunks, } +const halvingSamples = Array.from({ length: 40 }, (_, index) => ({ + ...searchSample, + intentId: `fixture-${String(index + 1).padStart(3, "0")}`, + queryKind: ["identifier", "searchPhrase", "naturalQuestion", "agentTask"][index % 4] as + | "identifier" + | "searchPhrase" + | "naturalQuestion" + | "agentTask", +})) + const withoutSearchTimings = < T extends { readonly searchDiagnostics: { readonly timings: object } }, >( @@ -110,11 +120,16 @@ describe("benchmark candidate evaluation pool", () => { [searchSample], [searchSample], SEARCH_PRIORITY_PROFILE, - { workerCount: 0, evaluationQueue: candidateQueue }, + { + workerCount: 0, + evaluationQueue: candidateQueue, + routerSearchStrategy: "successive-halving", + }, ), fitRecommendedEvidenceRouter("fixture", "dbsf", [searchSample], SEARCH_PRIORITY_PROFILE, { workerCount: 0, evaluationQueue: candidateQueue, + routerSearchStrategy: "successive-halving", }), ]) const serialHoldout = await optimizeEvidenceRouter( @@ -125,12 +140,14 @@ describe("benchmark candidate evaluation pool", () => { [searchSample], [searchSample], SEARCH_PRIORITY_PROFILE, + { workerCount: 0, routerSearchStrategy: "successive-halving" }, ) const serialFitAll = await fitRecommendedEvidenceRouter( "fixture", "dbsf", [searchSample], SEARCH_PRIORITY_PROFILE, + { workerCount: 0, routerSearchStrategy: "successive-halving" }, ) expect(parallelHoldout.map(withoutSearchTimings)).toEqual( @@ -145,6 +162,38 @@ describe("benchmark candidate evaluation pool", () => { } }) + it("runs the historical halving stage through the worker queue", async () => { + const candidateQueue = await createCandidateEvaluationQueue({ workerCount: 2 }) + try { + const parallel = await fitRecommendedEvidenceRouter( + "fixture", + "dbsf", + halvingSamples, + SEARCH_PRIORITY_PROFILE, + { + workerCount: 0, + evaluationQueue: candidateQueue, + routerSearchStrategy: "successive-halving", + }, + ) + const serial = await fitRecommendedEvidenceRouter( + "fixture", + "dbsf", + halvingSamples, + SEARCH_PRIORITY_PROFILE, + { workerCount: 0, routerSearchStrategy: "successive-halving" }, + ) + expect(parallel.map(withoutSearchTimings)).toEqual(serial.map(withoutSearchTimings)) + const result = parallel[0] + if (result === undefined) throw new Error("Missing halving router result") + expect(result.searchDiagnostics.proxyEvaluations).toBeGreaterThan(0) + expect(result.searchDiagnostics.proxyPromotions).toBeGreaterThan(0) + expect(result.searchDiagnostics.timings.randomSearchMs).toBe(0) + } finally { + await candidateQueue.close() + } + }) + it("derives a bounded default and honors explicit sizing", () => { expect(getDefaultWorkerCount()).toBeGreaterThanOrEqual(1) expect(getDefaultWorkerCount()).toBeLessThanOrEqual(availableParallelism())