Fix OOD gate embedding-scale mismatch that flags all in-distribution traffic

[aithal007]

Bug

fit_ood_gate() (src/tracer/fit/ood.py) calibrates its kNN-distance
thresholds on the fit-time embeddings. Those same embeddings are then
passed to EmbeddingIndex.build() (src/tracer/embeddings/index.py), which
-- when faiss is installed, the default metric="cosine" path -- calls
faiss.normalize_L2(X) on them in place. Because of a zero-copy
astype(..., copy=False) alias earlier in the same array, this silently
rescales the exact array fit_ood_gate had just calibrated on
(src/tracer/api.py, lines 159 vs. 176 on main).

Router.load() (src/tracer/runtime/router.py) then loads
EmbeddingIndex.load(...).embeddings as train_embeddings for the OOD
gate -- i.e. the normalized copy -- and applies it against thresholds
that were computed on the raw scale.

Impact: with faiss installed (the all extra / most production
setups), every training point -- the exact data the gate was calibrated on
-- gets flagged out-of-distribution at inference. A router in this state
silently defers essentially all traffic to the teacher LLM instead of ever
letting the surrogate handle it, defeating the library's entire
cost-saving purpose, with no error or warning surfaced anywhere.

Reproduction (measured, not just described)

Fitted a router on 240 synthetic training points, then checked how many of
those exact same points the OOD gate flags as out-of-distribution --
should be ~0, since the gate calibrated directly on them:

$ python repro_ood_bug.py   # on main (before this fix), faiss-cpu installed
training points flagged as out-of-distribution: 240/240
(these are the exact rows the gate calibrated on -- should be 0/240)

$ python repro_ood_bug.py   # on this branch (after the fix)
training points flagged as out-of-distribution: 0/240
(these are the exact rows the gate calibrated on -- should be 0/240)

Same result is captured as an automated regression test -- see Testing
below.

Fix

• src/tracer/embeddings/index.py: EmbeddingIndex.build() no longer
  mutates its caller's array; it now always copies before normalizing, and
  .embeddings is kept consistent with whatever was actually indexed.
• src/tracer/api.py: the OOD gate's calibration embeddings are now saved
  to their own dedicated artifact (ood_train_embeddings.npy), decoupled
  from the FAISS index's array entirely. The two can no longer drift out of
  scale with each other, regardless of call order or whether faiss is
  installed.
• src/tracer/runtime/router.py: Router.load() reads from the dedicated
  file, falling back to the old EmbeddingIndex path for artifacts fitted
  before this change (no forced re-fit needed for existing users).

Also: avoid rebuilding the k-NN index on every prediction

While in this code path: ood_mask() was rebuilding a fresh
sklearn.neighbors.NearestNeighbors index from scratch on every single
predict() call, via Router._ood_flags(). Added build_ood_index()
(fit/ood.py) so Router.load() builds it once and every prediction reuses
it -- same algorithm, same results (test_prebuilt_index_matches_fresh_index),
just not re-fit per request.

What I deliberately did not do

An earlier version of this fix L2-normalized embeddings unconditionally
inside fit_ood_gate/ood_mask as a more "robust-by-construction" fix.
I reverted that: it silently changes the gate's distance metric from
magnitude-sensitive Euclidean to purely angular, which broke a real
detection case (a synthetic magnitude-only outlier stopped being flagged,
tests/test_ood.py::test_far_points_are_flagged_in_distribution_passes).
That's a behavioral/metric change worth its own discussion and ablation,
not something to bundle silently into a bug fix. The fix here is scoped to
the actual bug: making calibration and inference agree on scale, without
changing what "far" means.

Testing

• tests/test_embedding_index.py (new): EmbeddingIndex.build() doesn't
  mutate its input; .embeddings matches what's indexed; the no-faiss
  fallback path is unaffected.
• tests/test_ood.py: build_ood_index() + cached-index path give
  identical results to the fresh-fit path.
• tests/test_router.py (new): test_ood_gate_correct_after_fit_and_load_with_faiss
  -- full fit() -> load_router() round trip with faiss installed, the
  same scenario as the reproduction above, now automated.
• Full suite: pytest tests/ -v -- 135 passed, 2 skipped (pre-existing,
  unrelated: "no deployable pipeline" on an unlucky synthetic split), 0
  failed.

To reproduce the bug directly: check out src/tracer/api.py and
src/tracer/runtime/router.py from main on top of this branch,
pip install faiss-cpu, run
pytest tests/test_router.py::test_ood_gate_correct_after_fit_and_load_with_faiss -v
-- it fails with every one of 240 training points flagged OOD.

[Copilot]

🟢 Approval recommended

The changes directly address the described scale-mismatch root cause, preserve backward compatibility for existing artifacts, and add targeted regression tests for both FAISS and cached OOD-index behavior.

Pull request overview

Fixes a critical OOD-gate calibration/inference scale mismatch that can cause routers (especially with FAISS + cosine indexing) to incorrectly flag all in-distribution traffic as out-of-distribution, leading to near-total deferral to the teacher and defeating TRACER’s cost-saving routing.

Changes:

• Prevent EmbeddingIndex.build() from mutating caller-provided embeddings during FAISS cosine normalization, and ensure EmbeddingIndex.embeddings matches the actually-indexed representation.
• Persist OOD calibration embeddings to a dedicated artifact (ood_train_embeddings.npy) and load from it at runtime (with backward-compatible fallback to the prior index embeddings file).
• Avoid per-request k-NN refits by prebuilding and caching the OOD k-NN index at router load time, with regression tests validating equivalence.

File summaries

File	Description
src/tracer/embeddings/index.py	Copy-before-normalize to avoid in-place mutation and keep .embeddings consistent with indexed vectors.
src/tracer/api.py	Writes ood_train_embeddings.npy alongside ood.json to decouple OOD calibration embeddings from the FAISS index representation.
src/tracer/runtime/router.py	Loads dedicated OOD train embeddings (with fallback) and prebuilds a reusable k-NN index used for OOD checks.
src/tracer/fit/ood.py	Adds build_ood_index() and updates ood_mask() to optionally reuse a prefit index, avoiding repeated index construction.
tests/test_embedding_index.py	New regression tests asserting EmbeddingIndex.build() doesn’t mutate inputs and that indexed embeddings match the stored representation.
tests/test_router.py	New end-to-end regression covering fit→load with FAISS installed and validating OOD behavior on train vs. far points.
tests/test_ood.py	Adds tests verifying cached-index OOD results match the previous fresh-fit-per-call behavior.
tests/test_gate.py	Clarifies the intent/scope of the “extreme OOD inputs” routing test (acceptor behavior, not the distance-based OOD gate).

Review details

• Files reviewed: 8/8 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.