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