From d40970937d77d642499adcf74e359dd95e5b4b55 Mon Sep 17 00:00:00 2001 From: aithal007 Date: Sat, 29 Aug 2026 18:07:11 +0530 Subject: [PATCH] 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 --- src/tracer/api.py | 10 ++++++ src/tracer/embeddings/index.py | 10 ++++-- src/tracer/fit/ood.py | 50 ++++++++++++++++++++++++------ src/tracer/runtime/router.py | 30 ++++++++++++++---- tests/test_embedding_index.py | 56 ++++++++++++++++++++++++++++++++++ tests/test_gate.py | 6 ++-- tests/test_ood.py | 24 ++++++++++++++- tests/test_router.py | 56 ++++++++++++++++++++++++++++++++++ 8 files changed, 221 insertions(+), 21 deletions(-) create mode 100644 tests/test_embedding_index.py diff --git a/src/tracer/api.py b/src/tracer/api.py index 0c6cf15..e8d7a88 100644 --- a/src/tracer/api.py +++ b/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)") diff --git a/src/tracer/embeddings/index.py b/src/tracer/embeddings/index.py index 9e75f70..eb6dc7b 100644 --- a/src/tracer/embeddings/index.py +++ b/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: diff --git a/src/tracer/fit/ood.py b/src/tracer/fit/ood.py index 35aafbd..6ff9009 100644 --- a/src/tracer/fit/ood.py +++ b/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)) diff --git a/src/tracer/runtime/router.py b/src/tracer/runtime/router.py index e2dd959..b0f3a81 100644 --- a/src/tracer/runtime/router.py +++ b/src/tracer/runtime/router.py @@ -35,7 +35,7 @@ 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)} @@ -43,6 +43,9 @@ def __init__(self, stages: list, label_space: list, manifest, embedder=None, 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,18 +70,32 @@ 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).""" @@ -86,7 +103,8 @@ def _ood_flags(self, X: np.ndarray, preds) -> np.ndarray: 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.""" diff --git a/tests/test_embedding_index.py b/tests/test_embedding_index.py new file mode 100644 index 0000000..7efa2a6 --- /dev/null +++ b/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) diff --git a/tests/test_gate.py b/tests/test_gate.py index 6a41d73..d357f64 100644 --- a/tests/test_gate.py +++ b/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) diff --git a/tests/test_ood.py b/tests/test_ood.py index 2271c72..9a48648 100644 --- a/tests/test_ood.py +++ b/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 diff --git a/tests/test_router.py b/tests/test_router.py index 9ea53ce..8ec04b0 100644 --- a/tests/test_router.py +++ b/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()