diff --git a/CONTEXT.md b/CONTEXT.md index 72b91c3..b45d35b 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 @@ -212,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 @@ -227,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-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 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 @@ -245,9 +253,35 @@ 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 +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, 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`. +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 eedf1f7..4eb97f6 100644 --- a/benchmarks/BASELINE.md +++ b/benchmarks/BASELINE.md @@ -1,14 +1,45 @@ # 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-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 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. @@ -200,6 +231,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/README.md b/benchmarks/README.md index 7c1f1d8..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-22 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 @@ -310,9 +317,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 @@ -339,7 +350,9 @@ 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 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, marginal leave-one-channel-out contribution, cross-validation folds, Shapley values, and final fitted 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..5df53c2 --- /dev/null +++ b/benchmarks/retrieval/evaluation/collect.ts @@ -0,0 +1,653 @@ +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" +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 { ChannelRankings } from "../../../src/domain/retrieval.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 { + 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 { + 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 rankings: readonly ChannelRankings[] + 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, + }) + }) + +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, + rankingsByQuery: readonly ChannelRankings[], + channelDurations: readonly number[], +): Effect.Effect => + 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 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, + queryKind: entry.queryKind, + groupedFold, + query: entry.query, + rankings, + targets, + chunks: corpus.chunks, + termCoverage: buildQueryTermCoverage(entry.query, corpus.bm25Index, corpus.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 = + (channelDurations[queryIndex] ?? 0) + 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 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, + queries: readonly BenchmarkQuery[], + targetsByQuestion: readonly (readonly ReadonlySet[])[], + groupedFoldAssignments: ReadonlyMap, + model: string, + 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 + 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 + const searchData = yield* store.loadSearchData() + const collected = yield* collectModelSamples( + manifest, + corpus, + queries, + targetsByQuestion, + groupedFoldAssignments, + model, + queryVectors, + sparseQueries, + searchData, + store, + info.dims, + info.defaultDtype, + ) + yield* saveCachedRankings(sql, cachePaths.cacheKey, queries, collected.rankings) + return { + sparseEmbeddingRun: { + repository: manifest.id, + model: sparseEmbedder.contract.model, + tokenizerModel: sparseEmbedder.contract.tokenizer, + batchSize: DEFAULT_CONFIG.sparseEmbedder.batchSize, + chunkEmbeddingDurationMs: sparseChunkEmbeddingDurationMs, + queryTokenizationDurationMs: sparseQueryTokenizationDurationMs, + }, + ...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[], + 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 cachePaths = benchmarkCachePaths(manifest, model, info.dims, info.defaultDtype) + const cacheState = yield* inspectBenchmarkCache( + model, + 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, + queries, + targetsByQuestion, + groupedFoldAssignments, + model, + info, + chunkVectors, + queryVectors, + cachePaths, + cacheState.hasPersistedIndex, + cacheState.rankings, + ) + embeddingRuns.push({ + repository: manifest.id, + model, + device, + batchSize: EMBEDDING_BATCH_SIZE, + chunkEmbeddingDurationMs, + queryEmbeddingDurationMs, + }) + 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/evaluation/prepared-fusion-core.d.mts b/benchmarks/retrieval/evaluation/prepared-fusion-core.d.mts new file mode 100644 index 0000000..591ec9e --- /dev/null +++ b/benchmarks/retrieval/evaluation/prepared-fusion-core.d.mts @@ -0,0 +1,20 @@ +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" + +export function evaluatePreparedContributions( + matrix: PreparedFusionSnapshot, + weights: ChannelWeights, +): { + readonly chunkIndex: number + readonly score: number +}[] + +export function evaluateCandidate( + snapshot: EvaluationSnapshot, + candidate: EvaluationCandidate, +): QualitySummary diff --git a/benchmarks/retrieval/evaluation/prepared-fusion-core.mjs b/benchmarks/retrieval/evaluation/prepared-fusion-core.mjs new file mode 100644 index 0000000..d8a9f64 --- /dev/null +++ b/benchmarks/retrieval/evaluation/prepared-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/evaluation/prepared-fusion.ts b/benchmarks/retrieval/evaluation/prepared-fusion.ts new file mode 100644 index 0000000..ce88150 --- /dev/null +++ b/benchmarks/retrieval/evaluation/prepared-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 "./prepared-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/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 87% rename from benchmarks/retrieval/report.ts rename to benchmarks/retrieval/evaluation/report.ts index 02c994a..6da5715 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() @@ -62,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", @@ -78,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.`, "", @@ -86,9 +102,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", "", @@ -242,9 +258,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)} |`, ) @@ -261,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))} |`, ) } @@ -319,9 +333,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 82% rename from benchmarks/retrieval/types.ts rename to benchmarks/retrieval/evaluation/types.ts index 69efefc..d08db1b 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. */ @@ -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" @@ -229,6 +247,26 @@ 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 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. */ + 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,11 +279,12 @@ export interface RouterSearchDiagnostics { readonly proxyPromotions: number readonly proxyFullAgreement: number readonly protectedEliteCount: number + readonly timings: RouterSearchTimings } /** 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 @@ -314,18 +353,22 @@ 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: 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/weight-search.ts b/benchmarks/retrieval/evaluation/weight-search.ts similarity index 54% rename from benchmarks/retrieval/weight-search.ts rename to benchmarks/retrieval/evaluation/weight-search.ts index 6d9d500..6e693d5 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,19 +12,30 @@ 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 { fuseRankings } from "../../src/lib/retrieval/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, + ROUTER_SEARCH_STRATEGIES, + type RouterSearchStrategyName, type EvidenceRouterSearchResult, type FusionSearchResult, type HoldoutQuality, @@ -54,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, @@ -79,11 +91,76 @@ interface EvidenceSearchSample { readonly evidence: RoutingEvidence } +/** Native pool controls plus the cancellation signal owned by the benchmark Effect. */ +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. */ +export type BenchmarkSearchOptions = SearchOptions + +const preparedFusionCache = new WeakMap< + WeightSearchSample, + Map +>() +const candidatePoolInitializationMs = new WeakMap() + +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: SearchOptions, +): Promise => { + const snapshotPreparationStartedAt = performance.now() + const snapshot = createEvaluationSnapshot( + samples.map((sample) => ({ + evaluator: preparedFusionEvaluatorFor(sample, fusion), + targets: sample.targets, + chunks: sample.chunks, + sampleWeight: profile.queryFormWeights[sample.queryKind], + })), + ) + 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 = ( + samples: readonly EvidenceSearchSample[], + config: EvidenceRouterConfig, +): EvaluationCandidate => ({ + weights: samples.map(({ evidence }) => routeWithEvidence(evidence, config)), +}) const summarizeRanked = ( samples: readonly T[], @@ -142,7 +219,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", @@ -150,7 +228,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 +242,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 +257,7 @@ const summarizeProductionRouter = ( samples, (sample) => fuseWithWeights( - sample.rankings, + sample, routeWithEvidence( buildRoutingEvidence(sample.query, sample.rankings, sample.termCoverage), PRODUCTION_COMPATIBILITY_CONFIG, @@ -421,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) @@ -514,35 +611,85 @@ 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, + successiveHalving: boolean, +): 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( (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) + +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 ( + candidates: readonly ChannelWeights[], + limit: number, + qualityCache: Map, + pool: CandidateEvaluationPool, + 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) + const qualities = await pool.evaluate(pending.map(([, weights]) => ({ weights }))) + storeQualityResults( + pending, + qualities, + qualityCache, + "Candidate evaluation returned an incomplete weight result", + ) + return sortWeightCandidates( + unique, + qualityCache, + limit, + objective, + baseline, + profile, + successiveHalving, + ) } const withWeight = (weights: ChannelWeights, channel: ChannelName, value: number): ChannelWeights => @@ -551,41 +698,46 @@ const withWeight = (weights: ChannelWeights, channel: ChannelName, value: number [channel]: value, }) -const selectBestWeights = ( - samples: readonly WeightSearchSample[], - fusion: FusionMethod = "rrf", +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 ( + pool: CandidateEvaluationPool, objective: RouterObjective = "reranker-top20", baseline?: QualitySummary, profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): { readonly weights: ChannelWeights; readonly quality: QualitySummary } => { + successiveHalving = false, +): Promise<{ readonly weights: ChannelWeights; readonly quality: QualitySummary }> => { const qualityCache = new Map() - let beam = rankWeightCandidates( - samples, + let beam = await rankWeightCandidates( weightCandidates(), SEARCH_BEAM_WIDTH, qualityCache, - fusion, + pool, objective, baseline, profile, + successiveHalving, ) for (let pass = 0; pass < SEARCH_PASSES; pass++) { for (const channel of CHANNELS) { - beam = 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)), - ), - ], + beam = await rankWeightCandidates( + coordinateWeightCandidates(beam, channel), SEARCH_BEAM_WIDTH, qualityCache, - fusion, + pool, objective, baseline, profile, + successiveHalving, ) } } @@ -595,28 +747,52 @@ 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, + successiveHalving: boolean, ): 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 || - 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 }) } return [...selected.values()].map((entry) => entry.weights) } +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, successiveHalving) +} + +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" @@ -748,6 +924,7 @@ const buildSearchDiagnostics = ( ? 1 : stats.proxyAgreementMatches / stats.proxyAgreementComparisons, protectedEliteCount: stats.protectedEliteCount, + timings: { ...stats.timings }, }) const radicalInverse = (index: number, base: number): number => { @@ -859,6 +1036,18 @@ interface SearchEvaluationStats { proxyAgreementMatches: number proxyAgreementComparisons: number protectedEliteCount: number + timings: MutableRouterSearchTimings +} + +interface MutableRouterSearchTimings { + preparationMs: number + candidatePoolInitializationMs: number + baseWeightSearchMs: number + randomSearchMs: number + beamSearchMs: number + candidatePreparationMs: number + candidateEvaluationMs: number + candidateSelectionMs: number } interface RouterSearchContext { @@ -873,6 +1062,9 @@ interface RouterSearchContext { readonly fusion: FusionMethod readonly profile: OptimizationProfile readonly stats: SearchEvaluationStats + readonly fullPool: CandidateEvaluationPool + readonly proxyPool: CandidateEvaluationPool + readonly routerSearchStrategy: RouterSearchStrategyName } const compareRouterCandidates = ( @@ -885,21 +1077,48 @@ const compareRouterCandidates = ( compareObjectiveQuality(left.quality, right.quality, objective, baseline, profile) || routerComplexity(left.config) - routerComplexity(right.config) -const selectRandomRouter = ( +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, 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), - })) + stats: SearchEvaluationStats, +): Promise<{ readonly candidate: RouterCandidate; readonly candidates: number }> => { + const configs = buildRandomRouterSeeds(baseSeed, parameters) + 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] + const quality = qualities[index] + if (config === undefined || quality === undefined) + 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] ?? { config: baseSeed, quality: summarizeEvidenceRouter(samples, baseSeed, fusion, profile) } + )[0] + if (candidate === undefined) + throw new Error("Candidate evaluation produced no random router candidate") return { candidate, candidates: candidates.length } } @@ -929,13 +1148,100 @@ const selectObjectiveCandidates = ( return [...selected.values()].slice(0, limit) } -const rankRouterCandidates = ( +interface RouterEvaluationResult { + readonly candidates: readonly RouterCandidate[] + readonly cacheHits: number + readonly evaluations: number + readonly candidatePreparationMs: number + readonly candidateEvaluationMs: 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++ + } + } + } + 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 context.routerSearchStrategy === "successive-halving" + ? orderedFull.slice(0, limit) + : 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 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, + qualityCache, + "Candidate evaluation returned an incomplete router 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, + candidatePreparationMs, + candidateEvaluationMs, + } +} + +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,80 +1250,60 @@ 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 selectedProxy = selectObjectiveCandidates( - rankedProxy.map((candidate) => ({ config: candidate.config, quality: candidate.quality })), - limit * SEARCH_PROXY_PROMOTION_FACTOR, - context.proxyBaseline, - context.profile, + const rankedProxy = await evaluateRouterConfigs( + fullCandidates, + context.proxySamples, + context.proxyPool, + context.proxyQualityCache, ) + 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 = + 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) } fullCandidates = [...selected] + context.stats.timings.candidateSelectionMs += performance.now() - proxySelectionStartedAt } - 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 + 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 { @@ -1040,39 +1326,97 @@ export const selectEligibleCandidate = ( } } -const selectBestEvidenceRouter = ( - samples: readonly WeightSearchSample[], - fusion: FusionMethod, - profile: OptimizationProfile = SEARCH_PRIORITY_PROFILE, -): { - readonly selections: readonly EvidenceRouterSelection[] +interface RouterSearchPreparation { + readonly samples: readonly WeightSearchSample[] + 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 preparationStartedAt = performance.now() 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 { + samples, + 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, + timings: { + preparationMs: performance.now() - preparationStartedAt, + candidatePoolInitializationMs: 0, + baseWeightSearchMs: 0, + randomSearchMs: 0, + beamSearchMs: 0, + candidatePreparationMs: 0, + candidateEvaluationMs: 0, + candidateSelectionMs: 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 ( + preparation: RouterSearchPreparation, + fusion: FusionMethod, + profile: OptimizationProfile, + fullPool: CandidateEvaluationPool, + proxyPool: CandidateEvaluationPool, + routerSearchStrategy: RouterSearchStrategyName, +): Promise<{ + readonly selections: readonly EvidenceRouterSelection[] + readonly productionQuality: QualitySummary + readonly proxyEvaluations: number + readonly fullEvaluations: number + readonly searchDiagnostics: RouterSearchDiagnostics + readonly randomCandidate: RouterCandidate + readonly randomCandidates: number +}> => { + const { + samples, + evidenceSamples, + proxySamples, + guardrailBaselines, + productionQuality, + proxyBaseline, + stats, + } = preparation + const useSuccessiveHalving = routerSearchStrategy === "successive-halving" const searchContext: RouterSearchContext = { samples: evidenceSamples, proxySamples, @@ -1081,38 +1425,48 @@ const selectBestEvidenceRouter = ( elites: new Map(), archive: new Map(), baseline: productionQuality, - proxyBaseline: proxyBaselines.overall, + proxyBaseline, fusion, profile, stats, + fullPool, + proxyPool, + routerSearchStrategy, } - 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) - 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( - evidenceSamples, - randomBaseSeed, - parameters, - fusion, + const profileSeed = profileSeedFor(profile) + const baseSearchStartedAt = performance.now() + const baseWeights = await selectBestWeights( + fullPool, + "reranker-top20", productionQuality, profile, + useSuccessiveHalving, ) - let beam = rankRouterCandidates( + const subsetWeights = await selectBestWeightsPerSubset(fullPool, profile, useSuccessiveHalving) + const baseSeeds = [baseWeights.weights, ...subsetWeights].map(emptyRouterConfig) + 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 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, [...baseSeeds, ...buildGlobalRouterSeeds(baseSeeds, parameters)], SEARCH_BEAM_WIDTH, @@ -1121,7 +1475,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. @@ -1135,40 +1489,82 @@ const selectBestEvidenceRouter = ( ) } } + 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, } } +const withCandidatePool = async ( + samples: readonly WeightSearchSample[], + fusion: FusionMethod, + profile: OptimizationProfile, + options: SearchOptions, + 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() + } +} + /** 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 +1572,9 @@ 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(pool, "reranker-top20", undefined, profile) return { model, queryKind, @@ -1192,8 +1589,52 @@ 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: SearchOptions, +): 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, + options: SearchOptions = { workerCount: 0 }, +): 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 +1642,9 @@ 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(pool, "reranker-top20", undefined, profile) const guardrailBaselines = buildGuardrailBaselines(development, profile) const holdoutProfile = unweightedProfile(profile) const validationQuality = summarize(validationSamples, selected.weights, fusion, profile) @@ -1233,6 +1675,50 @@ 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: SearchOptions, +): 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, + options: SearchOptions = { workerCount: 0 }, +): 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 +1737,123 @@ 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 withEvidencePools = async ( + samples: readonly WeightSearchSample[], + fusion: FusionMethod, + profile: OptimizationProfile, + options: SearchOptions, + operation: ( + selection: Awaited>, + fullPool: CandidateEvaluationPool, + ) => Promise, +): Promise => { + const preparation = prepareRouterSearch(samples, profile) + const { evidenceSamples, proxySamples } = preparation + const fullPool = await createEvaluationPoolForSamples( + evidenceSamples.map(({ sample }) => sample), + fusion, + profile, + options, + ) + preparation.stats.timings.candidatePoolInitializationMs += + candidatePoolInitializationMs.get(fullPool) ?? 0 + 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 + ? fullPool + : await createEvaluationPoolForSamples( + proxySamples.map(({ sample }) => sample), + fusion, + profile, + options, + ) + if (proxyPool !== undefined && proxyPool !== fullPool) + preparation.stats.timings.candidatePoolInitializationMs += + candidatePoolInitializationMs.get(proxyPool) ?? 0 + 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, + options.routerSearchStrategy ?? "proxy-promotion", + ) + 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, + 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({ + selection, + staticSelection: await selectStaticWeights( + fullPool, + selection.objective, + productionQuality, + profile, + ), + }) + } + 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, strategy: EvidenceRouterSearchResult["strategy"], @@ -1260,76 +1861,95 @@ 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, + options: SearchOptions = { workerCount: 0 }, +): Promise => + withEvidencePools(development, fusion, profile, options, async (dynamicSelection, fullPool) => { + const productionValidation = summarizeProductionRouter(validation, profile) + const validationEvidence = prepareEvidenceSamples(validation) + const holdoutProfile = unweightedProfile(profile) + 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, profile, + options, ) - 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), - selection.config, - fusion, - holdoutProfile, - ), - 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 }) -} /** 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(pool, "reranker-top20", undefined, profile) return { model, queryKind, @@ -1339,14 +1959,35 @@ export const fitRecommendedWeights = ( } } +const fitRecommendedWeightsWithOptions = ( + model: string, + queryKind: QueryKind, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile, + options: SearchOptions, +): 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, + options: SearchOptions = { workerCount: 0 }, +): 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(pool, "reranker-top20", undefined, profile) const guardrailBaselines = buildGuardrailBaselines(samples, profile) const guardrailsMet = fusionGuardrailsMet( samples, @@ -1366,40 +2007,70 @@ export const fitRecommendedFusionWeights = ( } } +const fitRecommendedFusionWeightsWithOptions = ( + model: string, + fusion: FusionMethod, + samples: readonly WeightSearchSample[], + profile: OptimizationProfile, + options: SearchOptions, +): 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, + options: SearchOptions = { workerCount: 0 }, +): Promise => + fitRecommendedFusionWeightsWithOptions(model, fusion, samples, profile, options) + /** * Fit one evidence-router candidate on all samples after cross-validation has measured * generalization. */ -export const fitRecommendedEvidenceRouter = ( +const fitRecommendedEvidenceRouterWithOptions = ( 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, + options: SearchOptions, +): Promise => + withEvidencePools(samples, fusion, profile, options, async (dynamicSelection, fullPool) => { + const staticSelections = await selectStaticWeightsForSearch( + dynamicSelection, + fullPool, profile, + options, ) - 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, - } + 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, + options: SearchOptions = { workerCount: 0 }, +): Promise => + fitRecommendedEvidenceRouterWithOptions(model, fusion, samples, profile, options) 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/candidate-evaluation-pool.ts b/benchmarks/retrieval/execution/candidate-evaluation-pool.ts new file mode 100644 index 0000000..f78039e --- /dev/null +++ b/benchmarks/retrieval/execution/candidate-evaluation-pool.ts @@ -0,0 +1,721 @@ +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 "../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("./candidate-evaluation-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 +} + +/** 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" +} + +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 ReadyWorker { + readonly worker: Worker + readonly ready: Promise + readonly resolveReady: () => void + readonly rejectReady: (error: Error) => void +} + +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 attachFusionWorkerLifecycle = ( + slot: ReadyWorker, + onMessage: (message: unknown) => void, + onError: (cause: Error) => void, + isClosed: () => boolean, +): void => + attachWorkerLifecycle(slot.worker, onMessage, onError, (code) => { + if (!isClosed()) onError(new Error(`Candidate evaluation worker exited with code ${code}`)) + }) + +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) + 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 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") / + 4, + ) + +/** Build the compact, cloneable snapshot used by serial and worker evaluators. */ +export const createEvaluationSnapshot = ( + inputs: readonly EvaluationSampleInput[], +): 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 = ( + snapshot: EvaluationSnapshot, + candidates: readonly EvaluationCandidate[], +): readonly QualitySummary[] => + candidates.map((candidate) => evaluateCandidate(snapshot, candidate)) + +interface QueueWorkerSlot extends ReadyWorker { + readonly knownSnapshots: Set + busy: boolean +} + +interface CandidateQueueRequest { + 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 WorkerMessage => { + if (!isQueueRecord(message) || typeof message.type !== "string") return false + if (message.type === "ready") return true + if (message.type === "error") return typeof message.message === "string" + return ( + message.type === "result" && + 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 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) => + 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) + } + + /** 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(`Candidate evaluation 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("Candidate evaluation 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(`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("Candidate evaluation 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 = { + 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: 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, + 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 + } + + 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 { + if (this.closed) return + this.closed = true + if (this.ownsQueue) await this.queue.close() + } + + 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) + +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 ( + snapshot: EvaluationSnapshot, + options: CandidateEvaluationPoolOptions = {}, +): Promise => { + const requestedWorkerCount = resolveWorkerCount(options.workerCount) + const batchSize = resolveBatchSize(options.batchSize) + if (requestedWorkerCount <= 1 || snapshot.samples.length === 0) + return createOwnedSerialPool(snapshot, batchSize) + + try { + const queue = await createCandidateEvaluationQueue({ + workerCount: requestedWorkerCount, + batchSize, + workerUrl: options.workerUrl ?? DEFAULT_WORKER_URL, + }) + return new QueuedCandidateEvaluationPool(snapshot, queue, undefined, true) + } catch (cause) { + if (options.fallbackToSerial !== false) return createOwnedSerialPool(snapshot, batchSize) + throw errorFromUnknown(cause) + } +} diff --git a/benchmarks/retrieval/execution/candidate-evaluation-worker.mjs b/benchmarks/retrieval/execution/candidate-evaluation-worker.mjs new file mode 100644 index 0000000..5c401e5 --- /dev/null +++ b/benchmarks/retrieval/execution/candidate-evaluation-worker.mjs @@ -0,0 +1,32 @@ +import { parentPort } from "node:worker_threads" + +import { evaluateCandidate } from "../evaluation/prepared-fusion-core.mjs" + +if (parentPort === null) throw new Error("Fusion worker requires a parent port") + +const snapshots = new Map() + +parentPort.postMessage({ type: "ready" }) + +parentPort.on("message", (message) => { + try { + 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(`Candidate evaluation worker has no snapshot ${snapshotId}`) + 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/execution/sqlite-index.ts b/benchmarks/retrieval/execution/sqlite-index.ts new file mode 100644 index 0000000..e20bb5f --- /dev/null +++ b/benchmarks/retrieval/execution/sqlite-index.ts @@ -0,0 +1,57 @@ +import { NodeServices } from "@effect/platform-node" +import { Effect, Layer } from "effect" +import { SqlClient } from "effect/unstable/sql" + +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 { ensureBenchmarkCacheTable } from "./benchmark-cache.js" + +const benchmarkConfig = (model: string, dtype: EmbeddingDtype): Config => ({ + ...DEFAULT_CONFIG, + embedder: { + ...DEFAULT_CONFIG.embedder, + model, + dtype, + }, + vectorSearch: { + ...DEFAULT_CONFIG.vectorSearch, + mode: "exact", + }, +}) + +const benchmarkConfigStore = (config: Config): typeof ConfigStore.Service => ({ + readConfig: () => Effect.succeed(config), + healConfig: () => Effect.succeed({ config, conflicts: [] }), + writeConfig: () => Effect.void, + configExists: () => Effect.succeed(true), +}) + +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(databasePath), + ), + ), + Layer.provideMerge(NodeServices.layer), + ) + +/** Run one benchmark operation against production retrieval adapters and a benchmark SQLite index. */ +export const withSqliteBenchmarkStore = ( + model: string, + dtype: EmbeddingDtype, + effect: Effect.Effect, + databasePath = ":memory:", +): Effect.Effect => + 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/retrieval/runner.ts b/benchmarks/retrieval/runner.ts index 2709eba..51ef65f 100644 --- a/benchmarks/retrieval/runner.ts +++ b/benchmarks/retrieval/runner.ts @@ -1,80 +1,35 @@ 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 { 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, + DEFAULT_ROUTER_SEARCH_STRATEGY, + ROUTER_SEARCH_STRATEGIES, type BenchmarkArtifact, type BenchmarkProfile, type CorpusManifest, - type QueryKind, - type QueryMeasurement, + type RouterSearchStrategy, + type RouterSearchStrategyName, type ValidationStrategy, -} from "./types.js" -import { - fitRecommendedEvidenceRouter, - fitRecommendedFusionWeights, - fitRecommendedWeights, - optimizeEvidenceRouter, - optimizeFusionWeights, - optimizeWeights, - evaluateProductionRouter, - type WeightSearchSample, -} from "./weight-search.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 { @@ -111,13 +66,6 @@ const profileConfig = (profile: BenchmarkProfile): BenchmarkProfileConfig => { } } -const reportProgress = (message: string): void => { - process.stderr.write(`[retrieval benchmark] ${message}\n`) -} - -const isLongInput = (text: string): boolean => - Buffer.byteLength(text, "utf8") / 4 > SINGLE_ITEM_ESTIMATED_TOKENS - const selectValues = (value: string | undefined): ReadonlySet | null => value ? new Set( @@ -169,78 +117,27 @@ 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) } +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 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 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 => @@ -271,388 +168,39 @@ 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 } + : { + workerCount: Math.min(resolveWorkerCount(), getDefaultWorkerCount()), + 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) 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 - - 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 = config.legacyDiagnostics - ? [...sampleGroups.values()].flatMap((group) => { - const groupedFolds = Array.from({ length: config.groupedFolds }, (_, fold) => - optimizeWeights( - group.model, - group.queryKind, - groupedStrategy, - String(fold + 1), - group.samples.filter((sample) => sample.groupedFold !== fold), - group.samples.filter((sample) => sample.groupedFold === fold), - optimizationProfile, - ), - ) - 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() - 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++) { - fusionSearch.push( - optimizeFusionWeights( - model, - fusion, - 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) { - fusionSearch.push( - optimizeFusionWeights( - model, - fusion, - "leave-one-repository-out", - repository, - samples.filter((sample) => sample.repository !== repository), - samples.filter((sample) => sample.repository === repository), - optimizationProfile, - ), - ) - } - } - } - } - const recommendedFusionWeights = [...samplesByModel].flatMap(([model, samples]) => - config.fusionMethods.map((fusion) => - fitRecommendedFusionWeights(model, fusion, samples, optimizationProfile), - ), - ) - 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++) { - 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, - ), - ) - } - const repositories = [...new Set(samples.map((sample) => sample.repository))] - if (config.repositoryHoldouts && repositories.length > 1) { - for (const repository of repositories) { - 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, - ), - ) - } - } - } - } - reportProgress("fitting final evidence router on all samples") - const recommendedEvidenceRouters = [...samplesByModel].flatMap(([model, samples]) => - config.routerFusionMethods.flatMap((fusion) => - fitRecommendedEvidenceRouter(model, fusion, samples, optimizationProfile), - ), + 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, routerSearchStrategy }, ) - const evidenceRouterSearchDurationMs = performance.now() - evidenceRouterSearchStartedAt const embeddingDurationMs = embeddingRuns.reduce( (sum, run) => sum + run.chunkEmbeddingDurationMs + run.queryEmbeddingDurationMs, @@ -664,7 +212,7 @@ export const runRetrievalBenchmark = ( ) const artifact: BenchmarkArtifact = { - schemaVersion: 22, + schemaVersion: 24, benchmarkProfile: profile, optimizationProfile, validationProtocol: { @@ -678,15 +226,17 @@ 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, embeddingDurationMs, retrievalDurationMs, - weightSearchDurationMs, - fusionSearchDurationMs, - evidenceRouterSearchDurationMs, + weightSearchDurationMs: search.weightSearchDurationMs, + fusionSearchDurationMs: search.fusionSearchDurationMs, + evidenceRouterSearchDurationMs: search.evidenceRouterSearchDurationMs, + candidateQueueStartupDurationMs: search.candidateQueueStartupDurationMs, + candidateQueueShutdownDurationMs: search.candidateQueueShutdownDurationMs, }, chunkConfig: { chunkLines: DEFAULT_CONFIG.chunkLines, @@ -708,13 +258,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/sqlite-index.ts b/benchmarks/retrieval/sqlite-index.ts deleted file mode 100644 index 599c67c..0000000 --- a/benchmarks/retrieval/sqlite-index.ts +++ /dev/null @@ -1,48 +0,0 @@ -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" - -const benchmarkConfig = (model: string, dtype: EmbeddingDtype): Config => ({ - ...DEFAULT_CONFIG, - embedder: { - ...DEFAULT_CONFIG.embedder, - model, - dtype, - }, - vectorSearch: { - ...DEFAULT_CONFIG.vectorSearch, - mode: "exact", - }, -}) - -const benchmarkConfigStore = (config: Config): typeof ConfigStore.Service => ({ - readConfig: () => Effect.succeed(config), - healConfig: () => Effect.succeed({ config, conflicts: [] }), - writeConfig: () => Effect.void, - configExists: () => Effect.succeed(true), -}) - -const sqliteBenchmarkLayer = (model: string, dtype: EmbeddingDtype) => - Layer.merge(SqliteIndexStoreBase, SparseEmbedderBase).pipe( - Layer.provideMerge( - Layer.merge( - Layer.succeed(ConfigStore, benchmarkConfigStore(benchmarkConfig(model, dtype))), - sqliteIndexDatabaseLayer(":memory:"), - ), - ), - Layer.provideMerge(layerNoop({})), - ) - -/** Run one benchmark operation against production retrieval adapters and an in-memory index. */ -export const withSqliteBenchmarkStore = ( - model: string, - dtype: EmbeddingDtype, - effect: Effect.Effect, -): Effect.Effect => - Effect.scoped(effect.pipe(Effect.provide(sqliteBenchmarkLayer(model, dtype)))) 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/benchmarks/tests/channels.test.ts b/benchmarks/tests/channels.test.ts index d2977f5..d4e2891 100644 --- a/benchmarks/tests/channels.test.ts +++ b/benchmarks/tests/channels.test.ts @@ -9,15 +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 { recallAt, 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 type { PreparedCorpus } from "../retrieval/corpus/prepare.js" +import { + contextRecallAtBudget, + recallAt, + reciprocalRank, + resolveGoldTargets, +} 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, optimizeWeights, selectEligibleCandidate, -} from "../retrieval/weight-search.js" +} from "../retrieval/evaluation/weight-search.js" const texts = [ "export function loadProjectConfiguration() { return config }", @@ -111,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) @@ -368,6 +387,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 +509,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 +526,7 @@ describe("retrieval benchmark fixture", () => { targets: [new Set([0])], chunks, } - const result = optimizeWeights( + const result = await optimizeWeights( "fixture", "identifier", "grouped-5-fold", @@ -464,10 +542,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 +597,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 optimizeEvidenceRouter( + "fixture", + "dbsf", + "grouped-5-fold", + "1", + samples, + samples, + undefined, + { workerCount: 2, batchSize: 8, fallbackToSerial: false }, + ), + ) + + expect(withoutSearchTimings(parallel)).toEqual(withoutSearchTimings(serial)) + }) + it("does not treat a guardrail-failing fallback as promotable", () => { const selection = selectEligibleCandidate([{ name: "fallback" }], () => false) @@ -528,20 +627,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/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 269ffb9..c7843a4 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, + ROUTER_SEARCH_STRATEGIES, type BenchmarkProfile, -} from "../retrieval/types.js" +} from "../retrieval/evaluation/types.js" +import { resolveRouterSearchStrategy, runRetrievalBenchmark } from "../retrieval/runner.js" const foldQuestions = (prefix: string) => Array.from({ length: 12 }, (_, index) => ({ @@ -99,10 +99,17 @@ 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.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( + 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( @@ -173,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 new file mode 100644 index 0000000..9bfbd7a --- /dev/null +++ b/benchmarks/tests/worker-pool.test.ts @@ -0,0 +1,395 @@ +import { availableParallelism } from "node:os" + +import { describe, expect, it } from "@effect/vitest" + +import type { Chunk } from "../../src/domain/chunk.js" +import { SEARCH_PRIORITY_PROFILE } from "../retrieval/evaluation/optimization-profiles.js" +import { prepareFusion } from "../retrieval/evaluation/prepared-fusion.js" +import { + fitRecommendedEvidenceRouter, + fitRecommendedFusionWeights, + fitRecommendedWeights, + optimizeEvidenceRouter, + optimizeFusionWeights, + summarize, +} from "../retrieval/evaluation/weight-search.js" +import { + createCandidateEvaluationPool, + createCandidateEvaluationPoolOnQueue, + createCandidateEvaluationQueue, + createEvaluationSnapshot, + evaluateCandidatesSerial, + getDefaultWorkerCount, + resolveWorkerCount, + type EvaluationCandidate, +} from "../retrieval/execution/candidate-evaluation-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, +} + +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 } }, +>( + result: T, +) => { + const { timings: _timings, ...searchDiagnostics } = result.searchDiagnostics + return { ...result, searchDiagnostics } +} + +describe("benchmark candidate evaluation pool", () => { + it("runs holdout and fit-all jobs through one native queue", async () => { + const candidateQueue = await createCandidateEvaluationQueue({ workerCount: 2 }) + try { + const [parallelHoldout, parallelFitAll] = await Promise.all([ + optimizeEvidenceRouter( + "fixture", + "dbsf", + "grouped-5-fold", + "1", + [searchSample], + [searchSample], + SEARCH_PRIORITY_PROFILE, + { + 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( + "fixture", + "dbsf", + "grouped-5-fold", + "1", + [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( + serialHoldout.map(withoutSearchTimings), + ) + expect(parallelFitAll.map(withoutSearchTimings)).toEqual( + serialFitAll.map(withoutSearchTimings), + ) + expect(parallelHoldout[0]?.searchDiagnostics.timings.candidateEvaluationMs).toBeGreaterThan(0) + } finally { + await candidateQueue.close() + } + }) + + 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()) + 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("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, + 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 fitRecommendedWeights( + "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 fitRecommendedFusionWeights( + "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/execution/candidate-evaluation-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") + }) +}) 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