Fix OOD gate embedding-scale mismatch between calibration and inference

fit_ood_gate() calibrated its kNN-distance thresholds on the raw fit-time
embeddings, but EmbeddingIndex.build() then L2-normalized that same array
in place for its cosine FAISS index -- so Router loaded a rescaled copy as
train_embeddings, out of sync with the thresholds. With faiss installed,
this made every training point (the exact data the gate was calibrated on)
register as out-of-distribution, verified by reverting the fix and
re-running the new regression test: 240/240 false positives.

Fix: EmbeddingIndex.build() no longer mutates its input, and the OOD
gate's training embeddings are now saved as their own dedicated artifact
(ood_train_embeddings.npy), decoupled from whatever the FAISS index does
to its own copy, so the two can no longer drift out of scale with each
other regardless of call order or faiss availability.

Also: Router now builds the OOD gate's k-NN index once at load time
(build_ood_index) instead of re-fitting it from scratch on every single
predict() call.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: aithal007

src/tracer/api.py

@@ -153,12 +153,22 @@ def log_fn(msg: str) -> None:  # noqa: E301, local helper
         # Distance-based OOD gate: calibrate kNN-distance thresholds (global +
         # per predicted label) so the runtime can defer off-distribution inputs
         # the parity gate never saw. Safety net, not the partition.
+        #
+        # The gate's training embeddings are saved to their own dedicated
+        # array (ood_train_embeddings.npy) rather than reusing the FAISS
+        # index's embeddings.npy: the index build below normalizes its own
+        # copy for cosine similarity, which is a different representation
+        # than whatever fit_ood_gate calibrated its distance thresholds on.
+        # Keeping them as two independent artifacts means the two can never
+        # silently drift out of scale with each other, regardless of what
+        # normalization the FAISS index does or whether faiss is installed.
         try:
             from tracer.fit.ood import fit_ood_gate
             pred_label_strs = [idx_to_label.get(int(p), "?") for p in preds]
             ood_gate = fit_ood_gate(X, pred_label_strs)
             if ood_gate is not None:
                 (artifact_dir / "ood.json").write_text(json.dumps(ood_gate))
+                np.save(artifact_dir / "ood_train_embeddings.npy", X)
                 notes.append(
                     f"OOD gate calibrated (kNN dist, global thr={ood_gate['global_thr']:.3f}, "
                     f"{len(ood_gate['per_label_thr'])} per-label)")

src/tracer/embeddings/index.py

@@ -22,14 +22,20 @@ def build(cls, embeddings: np.ndarray, metric: str = "cosine") -> "EmbeddingInde
         except ImportError:
             return cls(embeddings, index=None)
 
-        X = embeddings.astype(np.float32, copy=False)
+        # Always copy: normalize_L2 mutates in place, and callers (e.g.
+        # api.py, which passes the same array it just used to calibrate the
+        # OOD gate and train the surrogate) must not have their embeddings
+        # silently rescaled out from under them as a side effect of building
+        # this index. Store the exact array that was indexed on `.embeddings`
+        # so the two stay consistent with each other.
+        X = np.array(embeddings, dtype=np.float32, copy=True)
         if metric == "cosine":
             faiss.normalize_L2(X)
             idx = faiss.IndexFlatIP(X.shape[1])
         else:
             idx = faiss.IndexFlatL2(X.shape[1])
         idx.add(X)
-        return cls(embeddings, index=idx)
+        return cls(X, index=idx)
 
     def search(self, query: np.ndarray, k: int = 5):
         if self._index is None:

src/tracer/fit/ood.py

@@ -15,6 +15,16 @@
 detection. It is intentionally NOT keyed on the partition cells, the cell
 construction is out of scope here; this gate only needs the input embeddings and
 the surrogate's predicted label.
+
+The caller is responsible for passing fit_ood_gate() and ood_mask() embeddings
+on the same scale (whatever that is -- raw, L2-normalized, or otherwise). This
+module does not rescale its input: doing so would silently change what "far"
+means (angular vs. magnitude-sensitive distance), which is a real behavioural
+choice, not something to bake in unconditionally. See api.py/router.py: the
+training embeddings used to calibrate this gate are saved and reloaded as their
+own dedicated artifact, kept independent of whatever the (possibly
+cosine-normalizing) FAISS index build does to its own copy, specifically so the
+two can never silently drift out of scale with each other.
 """
 
 from __future__ import annotations
@@ -52,20 +62,42 @@ def fit_ood_gate(X_train: np.ndarray, pred_labels, k: int = 10,
             "global_thr": global_thr, "per_label_thr": per_label}
 
 
-def ood_mask(X_query: np.ndarray, X_train: np.ndarray, query_labels, gate) -> np.ndarray:
-    """Boolean array, True where a query is out-of-distribution (should defer).
+def build_ood_index(X_train: np.ndarray, gate: Optional[dict]):
+    """Fit a reusable k-NN index over the training embeddings.
 
-    Uses the same mean-k-NN-distance rule the gate was calibrated with, comparing
-    against the per-predicted-label threshold (global fallback)."""
-    X_query = np.asarray(X_query, dtype=np.float32)
-    if gate is None or X_train is None or len(X_train) == 0 or len(X_query) == 0:
-        return np.zeros(len(X_query), dtype=bool)
+    A caller that invokes ``ood_mask`` repeatedly against the same X_train --
+    e.g. a live Router calling it once per prediction -- should build this
+    once (at load time) and pass it back in via ``ood_mask(..., index=...)``,
+    instead of paying the index-fit cost fresh on every single call.
+    """
+    if gate is None or X_train is None or len(X_train) == 0:
+        return None
     from sklearn.neighbors import NearestNeighbors
     k = int(min(gate.get("k", 10), len(X_train)))
     if k < 1:
+        return None
+    return NearestNeighbors(n_neighbors=k).fit(np.asarray(X_train, dtype=np.float32))
+
+
+def ood_mask(X_query: np.ndarray, X_train: np.ndarray, query_labels, gate,
+             index=None) -> np.ndarray:
+    """Boolean array, True where a query is out-of-distribution (should defer).
+
+    Uses the same mean-k-NN-distance rule the gate was calibrated with, comparing
+    against the per-predicted-label threshold (global fallback).
+
+    ``index``: an optional pre-fit index from ``build_ood_index()``. If omitted,
+    one is built fresh from X_train on every call (fine for one-off/batch use,
+    wasteful for a live router serving one prediction at a time).
+    """
+    X_query = np.asarray(X_query, dtype=np.float32)
+    if gate is None or len(X_query) == 0 or (index is None and (X_train is None or len(X_train) == 0)):
         return np.zeros(len(X_query), dtype=bool)
-    nn = NearestNeighbors(n_neighbors=k).fit(np.asarray(X_train, dtype=np.float32))
-    d, _ = nn.kneighbors(X_query)
+    if index is None:
+        index = build_ood_index(X_train, gate)
+        if index is None:
+            return np.zeros(len(X_query), dtype=bool)
+    d, _ = index.kneighbors(X_query)
     mean_d = d.mean(axis=1)
     per = gate.get("per_label_thr", {})
     g = float(gate.get("global_thr", np.inf))

src/tracer/runtime/router.py

@@ -35,14 +35,17 @@ class Router:
     """
 
     def __init__(self, stages: list, label_space: list, manifest, embedder=None,
-                 ood_gate=None, train_embeddings=None):
+                 ood_gate=None, train_embeddings=None, ood_index=None):
         self._stages = stages
         self._label_space = label_space
         self._idx_to_label = {i: l for i, l in enumerate(label_space)}
         self.manifest = manifest
         self.embedder = embedder
         self._ood_gate = ood_gate
         self._train_embeddings = train_embeddings
+        # Pre-fit k-NN index over the training embeddings (see fit.ood.build_ood_index):
+        # built once here rather than re-fit from scratch on every predict() call.
+        self._ood_index = ood_index
 
     @classmethod
     def load(cls, artifact_dir: Union[str, Path], embedder=None) -> "Router":
@@ -67,26 +70,41 @@ def load(cls, artifact_dir: Union[str, Path], embedder=None) -> "Router":
         label_space = bundle["label_space"]
 
         # Optional OOD safety net: defer inputs far from the training distribution.
+        # train_emb comes from its own dedicated array (ood_train_embeddings.npy),
+        # not the FAISS index's embeddings.npy -- the index build normalizes its
+        # own copy for cosine similarity, a different representation than what
+        # fit_ood_gate calibrated its distance thresholds on. Falls back to the
+        # index's embeddings for artifacts written before this file existed.
         ood_gate = None
         train_emb = None
+        ood_index = None
         ood_path = artifact_dir / "ood.json"
         if ood_path.exists():
             try:
+                import numpy as _np
                 ood_gate = _json.loads(ood_path.read_text())
-                from tracer.embeddings.index import EmbeddingIndex
-                train_emb = EmbeddingIndex.load(artifact_dir / "index").embeddings
+                ood_emb_path = artifact_dir / "ood_train_embeddings.npy"
+                if ood_emb_path.exists():
+                    train_emb = _np.load(ood_emb_path)
+                else:
+                    from tracer.embeddings.index import EmbeddingIndex
+                    train_emb = EmbeddingIndex.load(artifact_dir / "index").embeddings
+                from tracer.fit.ood import build_ood_index
+                ood_index = build_ood_index(train_emb, ood_gate)
             except Exception:
-                ood_gate, train_emb = None, None
+                ood_gate, train_emb, ood_index = None, None, None
         return cls(stages=stages, label_space=label_space, manifest=manifest,
-                   embedder=embedder, ood_gate=ood_gate, train_embeddings=train_emb)
+                   embedder=embedder, ood_gate=ood_gate, train_embeddings=train_emb,
+                   ood_index=ood_index)
 
     def _ood_flags(self, X: np.ndarray, preds) -> np.ndarray:
         """True where each row is out-of-distribution (force-deferred)."""
         if self._ood_gate is None or self._train_embeddings is None:
             return np.zeros(len(X), dtype=bool)
         from tracer.fit.ood import ood_mask
         labels = [self._idx_to_label.get(int(p), "?") for p in preds]
-        return ood_mask(X, self._train_embeddings, labels, self._ood_gate)
+        return ood_mask(X, self._train_embeddings, labels, self._ood_gate,
+                        index=self._ood_index)
 
     def _to_embedding(self, input) -> np.ndarray:
         """Convert input (text or array) to a (dim,) float32 embedding."""

tests/test_embedding_index.py

@@ -0,0 +1,56 @@
+"""Tests for EmbeddingIndex.build (embeddings/index.py).
+
+Regression coverage for a real bug: build() used to normalize its input
+array IN PLACE for the cosine metric, silently rescaling whatever array the
+caller passed in (e.g. api.py reused the same `X` for OOD-gate calibration
+right before building the index). It must not mutate the caller's array, and
+`.embeddings` must match exactly what was indexed.
+"""
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+from tracer.embeddings.index import EmbeddingIndex
+
+
+def _unnormalized(n=20, dim=8, seed=0):
+    rng = np.random.RandomState(seed)
+    return (rng.normal(0, 1, (n, dim)) * 5.0).astype(np.float32)
+
+
+def test_build_does_not_mutate_caller_array():
+    pytest.importorskip("faiss")
+    X = _unnormalized()
+    original = X.copy()
+    EmbeddingIndex.build(X, metric="cosine")
+    np.testing.assert_array_equal(X, original)
+
+
+def test_build_embeddings_attr_matches_what_was_indexed():
+    pytest.importorskip("faiss")
+    X = _unnormalized()
+    idx = EmbeddingIndex.build(X, metric="cosine")
+    # .embeddings should already be L2-normalized (what the FAISS index holds),
+    # not the raw un-normalized input -- so the two never silently diverge.
+    norms = np.linalg.norm(idx.embeddings, axis=1)
+    np.testing.assert_allclose(norms, 1.0, atol=1e-5)
+
+    # The nearest neighbour of a training point (searched via the same
+    # normalization) should be itself.
+    I, _D = idx.search(X[0], k=1)
+    assert I[0] == 0
+
+
+def test_build_no_faiss_falls_back_without_mutating(monkeypatch):
+    monkeypatch.setitem(sys.modules, "faiss", None)  # forces `import faiss` to raise ImportError
+
+    X = _unnormalized()
+    original = X.copy()
+    idx = EmbeddingIndex.build(X, metric="cosine")
+    np.testing.assert_array_equal(X, original)
+    np.testing.assert_array_equal(idx.embeddings, original)

tests/test_gate.py

@@ -51,9 +51,9 @@ def test_nan_probs_do_not_crash_acceptor():
 
 def test_routing_robust_to_extreme_ood_inputs():
     """Out-of-distribution / extreme embeddings must route without crashing and
-    produce finite accept scores. (The OSS library defers via the acceptor
-    threshold; it does not ship a separate distance-based OOD gate, which lives
-    in the hosted gateway.)"""
+    produce finite accept scores, even with no OOD gate calibrated (this test
+    exercises apply_stage/the acceptor threshold directly, not the separate
+    distance-based OOD gate in fit/ood.py -- see tests/test_ood.py for that)."""
     from tracer.fit.pipeline import build_l2d, apply_stage
     split = _make_split(800, 3.0)
     res = build_l2d(split, target_ta=0.9, alpha=0.1)

tests/test_ood.py

@@ -12,7 +12,7 @@
 
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
 
-from tracer.fit.ood import fit_ood_gate, ood_mask
+from tracer.fit.ood import fit_ood_gate, ood_mask, build_ood_index
 
 
 def _train(n=500, dim=16, seed=0):
@@ -52,3 +52,25 @@ def test_ood_mask_no_gate_is_noop():
     X = np.zeros((3, 4), dtype=np.float32)
     out = ood_mask(X, X, ["a", "a", "a"], None)
     assert out.shape == (3,) and not out.any()
+
+
+def test_prebuilt_index_matches_fresh_index():
+    """build_ood_index() + ood_mask(..., index=...) (the path Router now uses,
+    built once at load time) must give identical results to the fresh-fit
+    path, since it's the same computation with the fit step hoisted out."""
+    X, labels = _train()
+    gate = fit_ood_gate(X, labels)
+    rng = np.random.RandomState(3)
+    X_q = rng.normal(0, 1, (25, 16)).astype(np.float32)
+    q = ["a"] * 25
+
+    fresh = ood_mask(X_q, X, q, gate)
+    idx = build_ood_index(X, gate)
+    cached = ood_mask(X_q, X, q, gate, index=idx)
+    assert (fresh == cached).all()
+
+
+def test_build_ood_index_none_cases():
+    assert build_ood_index(None, {"k": 5}) is None
+    assert build_ood_index(np.zeros((0, 4)), {"k": 5}) is None
+    assert build_ood_index(np.zeros((5, 4)), None) is None

tests/test_router.py

@@ -183,3 +183,59 @@ def test_fallback_called_on_deferred(fitted):
             break
     if not deferred_found:
         pytest.skip("no deferred prediction found, policy handled everything")
+
+
+# ── OOD gate: decoupled from the FAISS index (regression) ──────────────────────
+
+def test_ood_train_embeddings_saved_as_dedicated_artifact(fitted):
+    """fit() must save the OOD gate's calibration embeddings to their own file,
+    exactly matching the array it calibrated on -- not reuse the FAISS index's
+    embeddings.npy, which build() may have separately normalized for cosine
+    similarity. See fit/ood.py's module docstring for why the two must stay
+    decoupled."""
+    artifact_dir, X = fitted
+    ood_emb_path = artifact_dir / "ood_train_embeddings.npy"
+    if not ood_emb_path.exists():
+        pytest.skip("no OOD gate calibrated for this fixture")
+    saved = np.load(ood_emb_path)
+    np.testing.assert_array_equal(saved, X)
+
+
+def test_ood_gate_correct_after_fit_and_load_with_faiss():
+    """Regression test for a real bug: with faiss installed, EmbeddingIndex.build
+    L2-normalizes its own copy of the embeddings for the cosine index. Router
+    used to load the OOD gate's train_embeddings from that same (normalized)
+    array, while fit_ood_gate had calibrated its thresholds on the original raw
+    embeddings -- a scale mismatch between calibration and inference. A full
+    fit() -> load_router() round trip must still correctly flag clearly
+    off-distribution inputs and pass through in-distribution ones."""
+    pytest.importorskip("faiss")
+    import tempfile
+    import tracer
+
+    with tempfile.TemporaryDirectory() as tmp:
+        traces_path, X = _make_traces(tmp)  # defaults match the `fitted` fixture above (known-deployable)
+        artifact_dir = Path(tmp) / ".tracer"
+        result = tracer.fit(traces_path, artifact_dir, embeddings=X,
+                            config=tracer.FitConfig(verbose=False))
+        if result.manifest.selected_method is None:
+            pytest.skip("no deployable policy on this synthetic split")
+
+        router = tracer.load_router(artifact_dir)
+        if router._ood_gate is None:
+            pytest.skip("no OOD gate calibrated for this fixture")
+
+        # Check the OOD gate's own flag directly (isolated from the acceptor's
+        # separate handled/deferred decision, which router.predict() conflates
+        # it with): a clearly off-distribution point must be flagged, and the
+        # original training points -- the exact rows the gate calibrated on --
+        # must not be, since a scale mismatch would corrupt both directions.
+        from tracer.fit.pipeline import _predict as _stage_predict
+        rng = np.random.RandomState(7)
+        dim = X.shape[1]
+        far = (rng.randn(1, dim) * 1000).astype(np.float32)
+        far_preds, _ = _stage_predict(router._stages[0]["clf"], far)
+        assert router._ood_flags(far, far_preds)[0]
+
+        train_preds, _ = _stage_predict(router._stages[0]["clf"], X)
+        assert not router._ood_flags(X, train_preds).any()