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

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()