From eb7591935bef52df1102db553e06584b4b09b8c8 Mon Sep 17 00:00:00 2001 From: Adam Date: Fri, 25 Sep 2026 18:49:47 +0200 Subject: [PATCH] Fix routing correctness and harden updates, tracing, and reports --- .github/workflows/ci.yml | 4 + AGENTS.md | 6 +- CHANGELOG.md | 26 ++++++ README.md | 2 +- benchmarks/README.md | 38 ++++++++ benchmarks/routing_reliability.py | 70 ++++++++++++++ docs/api.md | 9 +- docs/artifacts.md | 2 +- docs/cold_start.md | 2 +- docs/javascript.md | 7 +- docs/watch.md | 6 ++ js/src/index.ts | 3 + js/test/watch.test.ts | 23 +++++ llms.txt | 2 +- src/tracer/analysis/html_report.py | 29 +++--- src/tracer/analysis/sankey.py | 5 +- src/tracer/api.py | 110 ++++++++++++++++------ src/tracer/embeddings/index.py | 4 +- src/tracer/fit/ood.py | 7 +- src/tracer/fit/pipeline.py | 55 ++++++----- src/tracer/fit/surrogate.py | 25 ++--- src/tracer/runtime/router.py | 11 ++- src/tracer/runtime/serve.py | 46 +++++++--- src/tracer/scanner.py | 8 +- src/tracer/traces/loader.py | 4 +- src/tracer/watch.py | 97 ++++++++++++++----- tests/test_report_safety.py | 82 +++++++++++++++++ tests/test_routing_regressions.py | 140 ++++++++++++++++++++++++++++ tests/test_serve.py | 76 ++++++++++++++- tests/test_update_safety.py | 143 +++++++++++++++++++++++++++++ tests/test_watch_safety.py | 126 +++++++++++++++++++++++++ 31 files changed, 1025 insertions(+), 143 deletions(-) create mode 100644 benchmarks/README.md create mode 100644 benchmarks/routing_reliability.py create mode 100644 tests/test_report_safety.py create mode 100644 tests/test_routing_regressions.py create mode 100644 tests/test_update_safety.py create mode 100644 tests/test_watch_safety.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f35400e..08fe9b7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,6 +17,8 @@ jobs: include: - python-version: "3.12" dependencies: minimum + - python-version: "3.12" + dependencies: faiss steps: - uses: actions/checkout@v4 @@ -31,6 +33,8 @@ jobs: python -m pip install --upgrade pip if [ "${{ matrix.dependencies }}" = "minimum" ]; then python -m pip install -e ".[dev]" "numpy==1.26.4" "scikit-learn==1.4.2" "joblib==1.3.2" + elif [ "${{ matrix.dependencies }}" = "faiss" ]; then + python -m pip install -e ".[dev,faiss]" plotly else python -m pip install -e ".[dev]" fi diff --git a/AGENTS.md b/AGENTS.md index 61a40c6..ba93d4f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -32,7 +32,7 @@ out = router.predict(embedding_vector) # numpy array out = router.predict("raw text", embedder=embedder) # or text directly # 3. UPDATE - continual learning (coverage grows over time) -tracer.update("new_traces.jsonl", embeddings=X_new) +tracer.update("new_traces.jsonl", new_embeddings=X_new) # 4. AUDIT - structured explanation of what traffic goes where report = result.qualitative_report # slices, boundary pairs, examples @@ -148,7 +148,7 @@ out = router.predict(text, fallback=lambda: call_gpt4(text)) ### Pattern 3: Continual learning loop ```python # Collect new traces from production -tracer.update("new_traces.jsonl", embeddings=X_new) +tracer.update("new_traces.jsonl", new_embeddings=X_new) # Reload router with improved policy router = tracer.load_router(".tracer", embedder=my_embedder) ``` @@ -190,7 +190,7 @@ tracer.serve(".tracer", port=8000) | `tracer.load_router(dir, embedder=e)` | Load for inference | No | | `router.predict(text_or_emb)` | Route one input | No | | `router.predict_batch(inputs)` | Route a batch | No | -| `tracer.update(new_traces, embeddings=X)` | Refit with new data | No | +| `tracer.update(new_traces, new_embeddings=X)` | Refit with new data | No | | `tracer.embed(texts)` | Compute embeddings | Needs `pip install tracer-llm[embeddings]` | | `Embedder.from_endpoint(url)` | External embedding API | Needs URL + auth from human | | `tracer.generate_html_report(dir)` | Visual audit | No | diff --git a/CHANGELOG.md b/CHANGELOG.md index d57cfc2..8479c34 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,33 @@ versioning. ## Unreleased +### Fixed +- Preserve the previous artifact generation when an update fails validation, + fitting, or publication. Honor explicit update configuration without mutating + the caller, retain saved configuration by default, and increment the fit count. +- Preserve classifier label IDs in residual stages and select the largest + eligible acceptance set during small-data calibration. +- Keep original embedding scales when building FAISS indexes and reuse the OOD + nearest-neighbor index across predictions. Existing FAISS artifacts trained + with unnormalized inputs need refitting to recover the original scale. +- Apply configured seeds to subsampling, splits, calibration, and model training. +- Record async outputs, errors, cancellation, timing, and parent spans correctly; + isolate tracing storage/export failures from application calls. +- Escape trace and artifact text in reports and reject unsafe watcher filenames + in Python and JavaScript. +- Preserve deferred confidence scores in reports; validate prediction shapes and + finite values; report actual file lines for null teacher labels. +- Handle concurrent HTTP requests and client disconnects, reject malformed or + oversized request bodies, and close the server cleanly on interruption. +- Correct update examples to use `new_embeddings=`. + ### Changed (breaking) +- The prediction server binds to `127.0.0.1` by default and no longer sends a + wildcard CORS header. Explicitly configure a remote bind address when needed; + browser deployments should configure CORS at their authenticated proxy. +- Watcher filenames must start with a letter or digit and use only ASCII + letters, digits, dots, underscores, or hyphens, up to 128 characters. +- Invalid prediction inputs return HTTP 400 rather than 500. - Require Python 3.12 or newer; drop support for Python 3.9–3.11. CI tests Python 3.12, 3.13, and 3.14, plus the minimum core dependencies on 3.12. - Allow NumPy 1.26.4 through 2.x so Python 3.14 can install compatible wheels. diff --git a/README.md b/README.md index 900a8dc..70b1a7b 100644 --- a/README.md +++ b/README.md @@ -142,7 +142,7 @@ Day 5: 10,000 traces → 92% coverage → 9,200 calls/day saved ``` ```python -tracer.update("new_traces.jsonl", embeddings=X_new) # refit with new production traces +tracer.update("new_traces.jsonl", new_embeddings=X_new) # refit with new production traces ``` The parity gate re-calibrates on each update, so coverage only increases when the surrogate actually earns it. diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000..ee1108a --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,38 @@ +# Routing reliability ablation + +`routing_reliability.py` exercises three deterministic failure cases without +external data, model downloads, or API calls. The inputs and seeds are in the +script. These are targeted regression measurements, not general workload claims. + +Run the same script against each checkout with identical dependencies: + +```bash +pip install -e '.[faiss]' +OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 LOKY_MAX_CPU_COUNT=1 \ + python benchmarks/routing_reliability.py +``` + +Measured with Python 3.12.13, NumPy 2.5.3, scikit-learn 1.9.1, and FAISS 1.15.1. +The baseline is commit `0c918712c09a677feae80eede7069442c4e259a9`. + +| Case | Baseline | Fixed | +|---|---:|---:| +| Small calibration: accepted fraction | 30.56% | 91.67% | +| Small calibration: agreement on accepted rows | 100% | 90.91% | +| Small calibration: Clopper–Pearson lower bound | 81.11% | 80.86% | +| Model trained on a label subset: label agreement | 0% | 100% | +| FAISS fit/reload: held-out coverage | 0% | 99% | +| FAISS fit/reload: agreement on handled held-out rows | Not applicable (none handled) | 100% | +| FAISS fit: caller input unchanged | No | Yes | + +The small-calibration case has 36 rows, a target of 80%, and alpha 0.1. +Both selections satisfy the same lower-bound test; the correction chooses the +larger eligible set. The label-subset case compares six predictions against +known labels `{1, 3}`. The FAISS case uses 240 training rows and **600 separate, +newly generated held-out rows**, target agreement 90%, and the same linear model +candidate in both versions. OOD detection remains Euclidean and retains its +existing threshold rule. + +The regression suite also covers staged-update rollback, retained recovery +backups, explicit configuration, reproducible seeds, async tracing and +cancellation, report escaping, path validation, and HTTP error/concurrency paths. diff --git a/benchmarks/routing_reliability.py b/benchmarks/routing_reliability.py new file mode 100644 index 0000000..c613f9a --- /dev/null +++ b/benchmarks/routing_reliability.py @@ -0,0 +1,70 @@ +"""Small deterministic routing ablation; no model downloads or API calls. + +Run with `python benchmarks/routing_reliability.py`. Install the optional faiss +extra to include the fit/reload case. Output is JSON suitable for comparing two +checkouts with the same Python and dependencies. +""" +import json +import tempfile +from pathlib import Path + +import numpy as np +from sklearn.linear_model import LogisticRegression + +from tracer.api import fit, load_router +from tracer.config import FitConfig +from tracer.fit.pipeline import _calibrate_threshold, _predict + + +def run(): + scores = np.linspace(.50, .99, 36) + teacher = np.zeros(36, dtype=int) + predictions = teacher.copy() + predictions[:6] = 1 + chosen = _calibrate_threshold(scores, predictions, teacher, .8) + X = np.array([[-3.], [-2.], [-1.], [1.], [2.], [3.]]) + y = np.array([1, 1, 1, 3, 3, 3]) + model = LogisticRegression(random_state=42).fit(X, y) + results = { + 'small_calibration': {key: chosen[key] for key in + ('coverage', 'teacher_agreement', 'teacher_agreement_lower')}, + 'subset_label_agreement': float((_predict(model, X)[0] == y).mean()), + } + try: + import faiss + except ImportError: + results['fit_reload'] = 'skipped: install tracer-llm[faiss]' + return results + + rng = np.random.RandomState(24) + centers = rng.normal(size=(3, 8)) * 10 + y_train = np.repeat(np.arange(3), 80) + y_test = np.repeat(np.arange(3), 200) + train = (centers[y_train] + rng.normal(size=(240, 8)) * .2).astype(np.float32) + heldout = (centers[y_test] + rng.normal(size=(600, 8)) * .2).astype(np.float32) + original = train.copy() + skip = ('logreg_c10', 'sgd_log', 'mlp_1h', 'mlp_2h', 'dt', 'rf', 'et', 'gbt', 'xgb') + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + traces = root / 'traces.jsonl' + traces.write_text(''.join(json.dumps({'input': str(i), 'teacher': str(label)}) + '\n' + for i, label in enumerate(y_train))) + fitted = fit(traces, root / 'model', train, + FitConfig(frontier_targets=(.9,), skip_candidates=skip, verbose=False)) + router = load_router(root / 'model') + output = router.predict_batch(heldout) + accepted = output['handled'] + labels = np.asarray(output['labels']) + results['fit_reload'] = { + 'faiss_version': faiss.__version__, + 'training_input_unchanged': bool(np.array_equal(train, original)), + 'calibration_coverage': fitted.manifest.coverage_cal, + 'heldout_coverage': float(accepted.mean()), + 'heldout_teacher_agreement': float((labels[accepted] == y_test[accepted].astype(str)).mean()) if accepted.any() else None, + 'heldout_rows': len(y_test), + } + return results + + +if __name__ == '__main__': + print(json.dumps(run(), indent=2)) diff --git a/docs/api.md b/docs/api.md index b146580..82e6da9 100644 --- a/docs/api.md +++ b/docs/api.md @@ -118,10 +118,17 @@ tracer.update( | `new_trace_path` | `str \| Path` | required | Path to NEW traces JSONL file | | `artifact_dir` | `str \| Path` | `".tracer"` | Existing artifact directory to update | | `new_embeddings` | `np.ndarray \| None` | `None` | Embeddings for the NEW traces only `(n_new, dim)` | -| `config` | `FitConfig \| None` | `None` | If None, re-uses `target_teacher_agreement` from the existing manifest | +| `config` | `FitConfig \| None` | `None` | If None, reuses the saved fit configuration. An explicit configuration takes precedence and is not mutated. | **Returns:** `FitResult` (same as `fit()`) +The replacement model is fitted in a staging directory. Validation or fitting +failures leave existing artifacts intact; a failed publication restores the +previous directory. If restoration itself fails, the exception identifies the +retained backup. Keep one writer per artifact directory and reload readers only +after `update()` returns: directory publication is not a concurrent-reader or +power-loss transaction. Allow disk space for the staged generation and backup. + **Example:** ```python diff --git a/docs/artifacts.md b/docs/artifacts.md index c50974d..3f4bf30 100644 --- a/docs/artifacts.md +++ b/docs/artifacts.md @@ -48,7 +48,7 @@ The top-level summary. Human-readable and machine-readable. | `selected_method` | Which pipeline was deployed: `"global"`, `"l2d"`, `"rsb"`, or `null` (parity gate blocked) | | `coverage_cal` | Fraction of calibration-set traffic handled by surrogate. Proxy for production coverage. | | `teacher_agreement_cal` | On calibration-set handled traffic, fraction where surrogate agreed with teacher. Should be ≥ `target_teacher_agreement`. | -| `n_retrains` | How many times this artifact has been updated via `tracer.update()` | +| `n_retrains` | Number of fits in this artifact's history: 1 after initial fitting, incremented by each successful `tracer.update()` | **Null method:** If `selected_method` is `null`, coverage is 0%. All traffic routes to the teacher. This happens when no pipeline could reach the target teacher agreement on the calibration set. It's safe -- not an error. diff --git a/docs/cold_start.md b/docs/cold_start.md index 26d606b..69c1762 100644 --- a/docs/cold_start.md +++ b/docs/cold_start.md @@ -99,7 +99,7 @@ As your LLM handles more production traffic, append new traces and re-fit: ```python # After collecting more traces -result = tracer.update("new_traces.jsonl", ".tracer", embeddings=X_new) +result = tracer.update("new_traces.jsonl", ".tracer", new_embeddings=X_new) ``` Each update round adds deferred inputs (the ones the surrogate wasn't confident about) back into the training pool, automatically expanding coverage over time. This is TRACER's core feedback loop. diff --git a/docs/javascript.md b/docs/javascript.md index 8c32ed6..97e4b42 100644 --- a/docs/javascript.md +++ b/docs/javascript.md @@ -94,6 +94,11 @@ tracer.serve(".tracer", host="127.0.0.1", port=8000) Run this as a sidecar next to your Node app, same server, same docker-compose, same machine. This example binds to `127.0.0.1:8000`. For a container sidecar, bind to `0.0.0.0` inside the container and limit network access to your application. +The server defaults to loopback, accepts JSON bodies up to 8 MiB, and uses a +10-second socket timeout. It handles requests concurrently. Invalid input returns +HTTP 400. It provides no authentication or CORS headers; for browser or remote +access, configure authentication, TLS, and CORS at your reverse proxy. + --- ## Predicting from your JS app @@ -168,7 +173,7 @@ Every deferred input that reaches your LLM is a new trace. Accumulate them and r # Run this script on a schedule; supply matching embeddings or a sibling .npy. import tracer -tracer.update("new_traces.jsonl", embeddings=X_new) +tracer.update("new_traces.jsonl", new_embeddings=X_new) ``` Then restart the `serve_policy.py` process to pick up the updated policy. Coverage typically grows from ~84% at day 1 to 90%+ within a week of production traffic. diff --git a/docs/watch.md b/docs/watch.md index 3234b42..e955afd 100644 --- a/docs/watch.md +++ b/docs/watch.md @@ -7,6 +7,12 @@ OpenTelemetry GenAI schema), and the same object that feeds ## Quick start +Decorators support both regular functions and `async def`; async spans finish +after the awaited call returns or raises. Recording failures do not replace the +application's result or exception. Set `TRACER_WATCH_DEBUG=1` to diagnose them. +Local watcher names must start with an ASCII letter or digit and contain only +letters, digits, dots, underscores, or hyphens (maximum 128 characters). + ```python import tracer diff --git a/js/src/index.ts b/js/src/index.ts index 4127a70..d9b07fd 100644 --- a/js/src/index.ts +++ b/js/src/index.ts @@ -375,6 +375,9 @@ function getEnv(name: string): string | undefined { export class LocalFileSink implements Sink { readonly path: string; constructor(name: string, dir: string = ".tracer/watch") { + if (typeof name !== "string" || name.length > 128 || !/^[A-Za-z0-9]/.test(name) || /[^A-Za-z0-9._-]/.test(name)) { + throw new Error("Watcher names must start with a letter or digit and contain only letters, digits, '.', '_' or '-' (max 128 characters)"); + } this.path = path.join(dir, `${name}.jsonl`); try { fs.mkdirSync(dir, { recursive: true }); diff --git a/js/test/watch.test.ts b/js/test/watch.test.ts index 9458d7a..b3edac1 100644 --- a/js/test/watch.test.ts +++ b/js/test/watch.test.ts @@ -372,3 +372,26 @@ describe("watch factory", () => { expect(w).toBeInstanceOf(Watcher); }); }); + +// Names are filenames on both POSIX and Windows. +describe("LocalFileSink path boundaries", () => { + it.each(["../escape", "/absolute", "..\\escape", ".", "..", "", "trailing\n", "line\nbreak", "nul\0byte", "x".repeat(129)])("rejects unsafe name %j", (name) => { + const dir = tmpDir(); + try { + expect(() => new LocalFileSink(name, path.join(dir, "watch"))).toThrow(/Watcher names/); + expect(fs.existsSync(path.join(dir, "escape.jsonl"))).toBe(false); + } finally { + fs.rmSync(dir, { recursive: true, force: true }); + } + }); + it("accepts a portable service name", () => { + const dir = tmpDir(); + try { + const sink = new LocalFileSink("service-1.production_trace", dir); + sink.emit(new GenAISpan({ inputText: "ok" })); + expect(fs.existsSync(path.join(dir, "service-1.production_trace.jsonl"))).toBe(true); + } finally { + fs.rmSync(dir, { recursive: true, force: true }); + } + }); +}); diff --git a/llms.txt b/llms.txt index 3515c84..69f988d 100644 --- a/llms.txt +++ b/llms.txt @@ -11,7 +11,7 @@ pip install tracer-llm tracer.fit("traces.jsonl", embeddings=X) → FitResult tracer.load_router(".tracer", embedder=e) → Router router.predict("text" or embedding) → {label, decision, accept_score} -tracer.update("new.jsonl", embeddings=X_new) → FitResult +tracer.update("new.jsonl", new_embeddings=X_new) → FitResult ## Trace format (JSONL, one per line) {"input": "user text", "teacher": "label_from_llm"} diff --git a/src/tracer/analysis/html_report.py b/src/tracer/analysis/html_report.py index 76e9ba8..522dc10 100644 --- a/src/tracer/analysis/html_report.py +++ b/src/tracer/analysis/html_report.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from html import escape from pathlib import Path from typing import Union @@ -278,7 +279,7 @@ def generate_html_report(

Routing Policy

{n_traces:,} traces  ·  {n_labels} labels -  ·  {emb_dim}-dim embeddings +  ·  {escape(str(emb_dim))}-dim embeddings  ·  target TA = {target_ta:.0%}
@@ -296,7 +297,7 @@ def generate_html_report(
Method
- {method.upper()} + {escape(str(method).upper())}
selected pipeline
@@ -389,7 +390,7 @@ def generate_html_report( """ for s in sorted(label_slices, key=lambda x: -x["handled_rate"]): - label = s["slice_name"].replace("label:", "") + label = escape(str(s["slice_name"]).replace("label:", "")) hr_v = s["handled_rate"] ta_s = s.get("teacher_agreement_handled") ta_str = f"{ta_s:.1%}" if ta_s is not None else "-" @@ -397,7 +398,7 @@ def generate_html_report( f'{label}' f'{hr_v:.1%}' f'{_bar_html(hr_v)}' - f'{s["count"]}' + f'{escape(str(s["count"]))}' f'{ta_str}\n' ) html += " \n \n\n" @@ -408,10 +409,10 @@ def generate_html_report( html += 'BucketCoverageCount\n' for s in length_slices: html += ( - f'{s["slice_name"]}' + f'{escape(str(s["slice_name"]))}' f'{s["handled_rate"]:.1%}' f'{_bar_html(s["handled_rate"])}' - f'{s["count"]}\n' + f'{escape(str(s["count"]))}\n' ) html += '\n\n' @@ -423,14 +424,14 @@ def generate_html_report( hs = _score_str(p.get("handled_score")) ds = _score_str(p.get("deferred_score")) html += f"""
-
{p["teacher_label"]}
+
{escape(str(p["teacher_label"]))}
SURROGATE - {p["handled_preview"]}{hs} + {escape(str(p["handled_preview"]))}{hs}
→ LLM - {p["deferred_preview"]}{ds} + {escape(str(p["deferred_preview"]))}{ds}
\n""" html += '\n' @@ -443,18 +444,18 @@ def generate_html_report( for ex in handled_ex[:6]: score_str = f'score {ex["accept_score"]:.2f}' if ex.get("accept_score") else "" html += (f'
' - f'
{ex["input_preview"]}
' + f'
{escape(str(ex["input_preview"]))}
' f'
' - f'{ex["teacher_label"]}' + f'{escape(str(ex["teacher_label"]))}' f'{score_str}
\n') html += '\n' html += '
\n

Deferred to teacher

\n' for ex in deferred_ex[:6]: html += (f'
' - f'
{ex["input_preview"]}
' + f'
{escape(str(ex["input_preview"]))}
' f'
' - f'{ex["teacher_label"]}' + f'{escape(str(ex["teacher_label"]))}' f'
\n') html += '
\n' @@ -470,7 +471,7 @@ def generate_html_report( prev_bar = int(d["previous_handled_rate"] * 60) cur_bar = int(d["current_handled_rate"] * 60) html += ( - f'{d["label"]}' + f'{escape(str(d["label"]))}' f'{d["previous_handled_rate"]:.1%}' f'
' f'{d["current_handled_rate"]:.1%}' diff --git a/src/tracer/analysis/sankey.py b/src/tracer/analysis/sankey.py index efd0518..1c40394 100644 --- a/src/tracer/analysis/sankey.py +++ b/src/tracer/analysis/sankey.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +from html import escape from pathlib import Path from typing import Optional, Union @@ -84,7 +85,7 @@ def generate_sankey( fig.update_layout( title=dict( - text=title, + text=escape(str(title)), font=dict(size=15, color=_TEXT, family="Inter, system-ui, sans-serif"), x=0.01, y=0.97, ), @@ -131,7 +132,7 @@ def _build_sankey_figure(manifest: dict, qr: dict, top_k: int, dark: bool): n_other_labels = len(manifest.get("label_space", [])) - n_top has_other = (other_handled + other_deferred) > 0.5 - node_labels = [s["slice_name"].replace("label:", "").replace("_", " ") for s in top] + node_labels = [escape(s["slice_name"].replace("label:", "").replace("_", " ")) for s in top] if has_other: node_labels.append(f"other ({n_other_labels} labels)") surrogate_idx = len(node_labels) diff --git a/src/tracer/api.py b/src/tracer/api.py index 0c6cf15..c7c5f78 100644 --- a/src/tracer/api.py +++ b/src/tracer/api.py @@ -9,8 +9,11 @@ from __future__ import annotations import json +import shutil import sys +import tempfile import time +from dataclasses import asdict, replace from pathlib import Path from typing import Callable, Optional, Union @@ -73,10 +76,7 @@ def fit( raise FileNotFoundError( f"No embeddings found. Pass embeddings= or place a .npy file at {emb_path}") - if len(embeddings) != len(dataset): - raise ValueError(f"Trace/embedding mismatch: {len(dataset)} traces vs {len(embeddings)} embeddings") - - X = embeddings.astype(np.float32, copy=False) + X = _embedding_matrix(embeddings, len(dataset)) y_teacher = np.array([label_to_idx[r.teacher_label] for r in dataset.records], dtype=int) y_true = None has_gt = all(r.ground_truth is not None for r in dataset.records) @@ -102,7 +102,7 @@ def log_fn(msg: str) -> None: # noqa: E301, local helper frontier, split = fit_frontier(X, y_teacher, targets, max_fit_labels=config.max_fit_labels, min_coverage=config.min_deploy_coverage, - log=log_fn, skip=config.skip_candidates) + log=log_fn, skip=config.skip_candidates, seed=config.seed) # Select best pipeline at target TA selected = None @@ -137,11 +137,13 @@ def log_fn(msg: str) -> None: # noqa: E301, local helper decisions = ["handled" if h else "deferred" for h in handled] scores = np.zeros(len(X)) - for i, stage in enumerate(best["stages"]): - mask = stage_id == i - if mask.any(): - _, _, s = apply_stage(stage, X[mask]) - scores[mask] = s + remaining = np.ones(len(X), dtype=bool) + for stage in best["stages"]: + if remaining.any(): + _, accepted, stage_scores = apply_stage(stage, X[remaining]) + rows = np.flatnonzero(remaining) + scores[rows] = stage_scores + remaining[rows[accepted]] = False qual_report = build_qualitative_report( texts=texts, teacher_labels=teacher_labels_str, @@ -179,13 +181,7 @@ def log_fn(msg: str) -> None: # noqa: E301, local helper # Save config config_path = artifact_dir / "config.json" - config_path.write_text(json.dumps({ - "target_teacher_agreement": config.target_teacher_agreement, - "frontier_targets": list(config.frontier_targets), - "min_deploy_coverage": config.min_deploy_coverage, - "max_fit_labels": config.max_fit_labels, - "seed": config.seed, - }, indent=2), encoding="utf-8") + config_path.write_text(json.dumps(asdict(config), indent=2), encoding="utf-8") # Save frontier summary frontier_path = artifact_dir / "frontier.json" @@ -216,6 +212,19 @@ def log_fn(msg: str) -> None: # noqa: E301, local helper qualitative_report=qual_report, notes=notes) +def _embedding_matrix(values, n_rows, dim=None): + X = np.asarray(values, dtype=np.float32) + if X.ndim != 2 or X.shape[1] == 0: + raise ValueError("Embeddings must be a nonempty-width 2-D matrix (n_traces, dim)") + if len(X) != n_rows: + raise ValueError(f"Trace/embedding mismatch: {n_rows} traces vs {len(X)} embeddings") + if dim is not None and X.shape[1] != dim: + raise ValueError(f"Embedding dimension mismatch: expected {dim}, got {X.shape[1]}") + if not np.isfinite(X).all(): + raise ValueError("Embeddings must contain only finite values") + return X + + def update( new_trace_path: Union[str, Path], artifact_dir: Union[str, Path] = ".tracer", @@ -225,7 +234,10 @@ def update( """Refit a TRACER policy with additional traces (continual learning). Loads the existing traces from the artifact dir, appends the new ones, - and re-fits. The .tracer directory is updated in place. + and re-fits in a staging directory. Failed validation or fitting leaves the + existing artifacts intact. Use one writer per artifact directory; reload + readers after update() returns. An explicit config takes precedence over + the saved configuration and is never mutated. """ artifact_dir = Path(artifact_dir) manifest = load_manifest(artifact_dir / "manifest.json") @@ -242,10 +254,6 @@ def update( else: raise FileNotFoundError(f"No embeddings for new traces at {emb_path}") - # Load existing embeddings - existing_index = EmbeddingIndex.load(artifact_dir / "index") - X_combined = np.vstack([existing_index.embeddings, new_embeddings.astype(np.float32)]) - # Load existing traces so we can re-save the combined set. fit() always # writes all_traces.jsonl, so it should be present. If it is missing (e.g. # it was deleted), the original records cannot be recovered from the index @@ -263,17 +271,61 @@ def update( from tracer.traces.loader import load_traces as _lt existing_ds = _lt(existing_traces_path) combined_records = existing_ds.records + new_ds.records + existing_index = EmbeddingIndex.load(artifact_dir / "index") + X_existing = _embedding_matrix(existing_index.embeddings, len(existing_ds), manifest.embedding_dim) + X_new = _embedding_matrix(new_embeddings, len(new_ds), X_existing.shape[1]) + X_combined = np.vstack([X_existing, X_new]) from tracer.types import TraceDataset from tracer.traces.loader import save_traces combined_ds = TraceDataset(records=combined_records) - save_traces(combined_ds, existing_traces_path) - - config = config or FitConfig() - if manifest.target_teacher_agreement: - config.target_teacher_agreement = manifest.target_teacher_agreement - - return fit(existing_traces_path, artifact_dir, embeddings=X_combined, config=config) + if config is None: + from tracer.config import EmbeddingConfig + config_path = artifact_dir / "config.json" + saved = json.loads(config_path.read_text()) if config_path.exists() else {} + if "embedding" in saved: + saved["embedding"] = EmbeddingConfig(**saved["embedding"]) + saved.setdefault("target_teacher_agreement", manifest.target_teacher_agreement) + config = FitConfig(**saved) + else: + config = replace(config) + + # Keep the temporary generation on the same filesystem as the destination. + with tempfile.TemporaryDirectory(prefix=f".{artifact_dir.name}-update-", + dir=artifact_dir.absolute().parent) as temp: + root = Path(temp) + staged = root / "next" + shutil.copytree(artifact_dir, staged) + # A refit must not inherit reports or gates from the previous model. + for filename in ("pipeline.joblib", "qualitative_report.json", "ood.json", + "index.faiss", "report.html", "sankey.html"): + (staged / filename).unlink(missing_ok=True) + combined_path = root / "combined.jsonl" + save_traces(combined_ds, combined_path) + result = fit(combined_path, staged, embeddings=X_combined, config=config) + result.artifact_dir = str(artifact_dir) + result.manifest.n_retrains = manifest.n_retrains + 1 + for field_name in ("pipeline_path", "index_path", "config_path", "qualitative_report_path"): + path = getattr(result.manifest, field_name) + if path is not None: + setattr(result.manifest, field_name, str(artifact_dir / Path(path).relative_to(staged))) + write_manifest(staged / "manifest.json", result.manifest) + # Outside the staging directory so a failed rollback cannot discard it. + backup = root.with_name(root.name + "-previous") + artifact_dir.rename(backup) + try: + staged.rename(artifact_dir) + except BaseException: + try: + backup.rename(artifact_dir) + except OSError as exc: + raise RuntimeError(f"Previous artifacts are preserved at {backup}; restore them before retrying") from exc + raise + try: + shutil.rmtree(backup) + except OSError: + result.notes.append(f"Update completed; previous artifact backup retained at {backup}") + return result def report(artifact_dir: Union[str, Path] = ".tracer") -> ArtifactManifest: diff --git a/src/tracer/embeddings/index.py b/src/tracer/embeddings/index.py index ea5e95e..96d79d7 100644 --- a/src/tracer/embeddings/index.py +++ b/src/tracer/embeddings/index.py @@ -22,7 +22,9 @@ def build(cls, embeddings: np.ndarray, metric: str = "cosine") -> "EmbeddingInde except ImportError: return cls(embeddings, index=None) - X = embeddings.astype(np.float32, copy=False) + # FAISS normalizes in place. Keep the caller and persisted training data + # at their original scale for refitting and the Euclidean OOD gate. + X = np.array(embeddings, dtype=np.float32, order="C", copy=True) if metric == "cosine": faiss.normalize_L2(X) idx = faiss.IndexFlatIP(X.shape[1]) diff --git a/src/tracer/fit/ood.py b/src/tracer/fit/ood.py index 35aafbd..10547cf 100644 --- a/src/tracer/fit/ood.py +++ b/src/tracer/fit/ood.py @@ -52,7 +52,7 @@ 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: +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 @@ -64,8 +64,9 @@ def ood_mask(X_query: np.ndarray, X_train: np.ndarray, query_labels, gate) -> np k = int(min(gate.get("k", 10), len(X_train))) if k < 1: 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 = NearestNeighbors(n_neighbors=k).fit(np.asarray(X_train, dtype=np.float32)) + d, _ = index.kneighbors(X_query, n_neighbors=k) 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/fit/pipeline.py b/src/tracer/fit/pipeline.py index 4b0f056..fe63549 100644 --- a/src/tracer/fit/pipeline.py +++ b/src/tracer/fit/pipeline.py @@ -119,7 +119,7 @@ def _accept_features(probs: np.ndarray) -> np.ndarray: return np.column_stack([top1, top2, margin, entropy]) -def _fit_acceptor(probs, y_pred, y_teacher): +def _fit_acceptor(probs, y_pred, y_teacher, seed=42): correct = (y_pred == y_teacher).astype(int) feats = _accept_features(probs) ok = np.isfinite(feats).all(axis=1) @@ -128,7 +128,7 @@ def _fit_acceptor(probs, y_pred, y_teacher): return None clf = Pipeline([ ("scale", StandardScaler()), - ("clf", LogisticRegression(max_iter=1000, solver="lbfgs", random_state=42)), + ("clf", LogisticRegression(max_iter=1000, solver="lbfgs", random_state=seed)), ]) clf.fit(feats, correct) return clf @@ -140,7 +140,8 @@ def _accept_scores(acceptor, probs): return acceptor.predict_proba(_accept_features(probs))[:, 1] -def _calibrate_threshold(scores, preds, y_teacher, target_ta, alpha=0.1, min_accept=10): +def _calibrate_threshold(scores, preds, y_teacher, target_ta, alpha=0.1, min_accept=10, + seed=42): """Choose the accept threshold on a held-out lower bound, then check it generalises. Earlier iterations sat at two extremes. The original gate selected the @@ -172,12 +173,11 @@ def _calibrate_threshold(scores, preds, y_teacher, target_ta, alpha=0.1, min_acc n = len(scores) holdout = False if n >= 40: - s_idx, v_idx = _holdout_indices(y_teacher, 0.7) # 70% select / 30% verify + s_idx, v_idx = _holdout_indices(y_teacher, 0.7, seed=seed) if len(s_idx) >= 20 and len(v_idx) >= 12: holdout = True if not holdout: # Too little data to hold out: single-set CP lower-bound gate over all rows. - best = None for t in np.unique(np.sort(scores)): acc = scores >= t n_acc = int(acc.sum()) @@ -185,13 +185,11 @@ def _calibrate_threshold(scores, preds, y_teacher, target_ta, alpha=0.1, min_acc continue k_acc = int((preds[acc] == y_teacher[acc]).sum()) if _cp_lower(k_acc, n_acc, alpha) >= target_ta: - best = (float(t), k_acc, n_acc) # keep lowest-threshold (max coverage) - if best is None: - return None - t, k_acc, n_acc = best - return {"threshold": t, "teacher_agreement": float(k_acc / n_acc), - "teacher_agreement_lower": float(_cp_lower(k_acc, n_acc, alpha)), - "coverage": float((scores >= t).mean()), "holdout": False} + # Thresholds are ascending: the first eligible set has maximum coverage. + return {"threshold": float(t), "teacher_agreement": float(k_acc / n_acc), + "teacher_agreement_lower": float(_cp_lower(k_acc, n_acc, alpha)), + "coverage": float(acc.mean()), "holdout": False} + return None s_sel, p_sel, y_sel = scores[s_idx], preds[s_idx], y_teacher[s_idx] s_ver, p_ver, y_ver = scores[v_idx], preds[v_idx], y_teacher[v_idx] @@ -233,12 +231,13 @@ def _calibrate_threshold(scores, preds, y_teacher, target_ta, alpha=0.1, min_acc def _predict(clf, X): probs = clf.predict_proba(X) - preds = probs.argmax(axis=1).astype(int) + # A residual stage may only have seen a subset of the global label ids. + preds = np.take(clf.classes_, probs.argmax(axis=1)).astype(int) return preds, probs def build_global(split, target_ta, alpha: float = 0.1, log: Optional[LogFn] = None, - skip: Iterable[str] = ()): + skip: Iterable[str] = (), seed=42): """Global pipeline: one surrogate, accept all if the agreement lower bound clears the target. Gating on the Clopper-Pearson lower bound (not the raw point estimate) stops a small or lucky calibration set from certifying an @@ -251,7 +250,7 @@ def build_global(split, target_ta, alpha: float = 0.1, log: Optional[LogFn] = No log(f" build_global: surrogate sweep on {len(split['X_train'])} train / {len(split['X_val'])} val") clf, name, val_m = search_best_surrogate( split["X_train"], split["y_train"], split["X_val"], split["y_val"], - on_candidate=_candidate_log(log), skip=skip) + on_candidate=_candidate_log(log), skip=skip, seed=seed) if clf is None: log(" build_global: all surrogate candidates failed (no model trained)") return {"method": "global", "stages": [], "summary": { @@ -293,23 +292,23 @@ def _cb(name: str, val_f1: float, elapsed: Optional[float] = None) -> None: def _build_accepting_stage(X_tr, y_tr, X_val, y_val, X_cal, y_cal, target_ta, stage_name, alpha: float = 0.1, - log: Optional[LogFn] = None, skip: Iterable[str] = ()): + log: Optional[LogFn] = None, skip: Iterable[str] = (), seed=42): log = log or _noop_log if len(np.unique(y_tr)) < 2 or len(X_val) == 0 or len(X_cal) == 0: return None t0 = time.perf_counter() log(f" {stage_name}: surrogate sweep on {len(X_tr)} train / {len(X_val)} val") clf, name, val_m = search_best_surrogate( - X_tr, y_tr, X_val, y_val, on_candidate=_candidate_log(log), skip=skip) + X_tr, y_tr, X_val, y_val, on_candidate=_candidate_log(log), skip=skip, seed=seed) if clf is None: log(f" {stage_name}: all surrogate candidates failed (no model trained)") return None log(f" {stage_name}: surrogate done in {time.perf_counter()-t0:.1f}s best={name} f1={val_m['teacher_f1']:.3f}") preds_val, probs_val = _predict(clf, X_val) preds_cal, probs_cal = _predict(clf, X_cal) - acceptor = _fit_acceptor(probs_val, preds_val, y_val) + acceptor = _fit_acceptor(probs_val, preds_val, y_val, seed=seed) scores_cal = _accept_scores(acceptor, probs_cal) - ti = _calibrate_threshold(scores_cal, preds_cal, y_cal, target_ta, alpha) + ti = _calibrate_threshold(scores_cal, preds_cal, y_cal, target_ta, alpha, seed=seed) if ti is None: return None return {"stage_name": stage_name, "model_name": name, "clf": clf, @@ -324,14 +323,14 @@ def _build_accepting_stage(X_tr, y_tr, X_val, y_val, X_cal, y_cal, target_ta, st def build_l2d(split, target_ta, alpha: float = 0.1, log: Optional[LogFn] = None, - skip: Iterable[str] = ()): + skip: Iterable[str] = (), seed=42): """L2D: surrogate + acceptor-gated deferral.""" log = log or _noop_log log(f"build_l2d: target_TA={target_ta:.2f}") s1 = _build_accepting_stage( split["X_train"], split["y_train"], split["X_val"], split["y_val"], split["X_cal"], split["y_cal"], target_ta, "stage_1", - alpha=alpha, log=log, skip=skip) + alpha=alpha, log=log, skip=skip, seed=seed) if s1 is None: return {"method": "l2d", "stages": [], "summary": {"status": "no_stage", "coverage_cal_total": 0.0}} stages = [s1] @@ -345,14 +344,14 @@ def build_l2d(split, target_ta, alpha: float = 0.1, log: Optional[LogFn] = None, def build_rsb(split, target_ta, alpha: float = 0.1, log: Optional[LogFn] = None, - skip: Iterable[str] = ()): + skip: Iterable[str] = (), seed=42): """RSB: residual two-stage cascade.""" log = log or _noop_log log(f"build_rsb: target_TA={target_ta:.2f}") s1 = _build_accepting_stage( split["X_train"], split["y_train"], split["X_val"], split["y_val"], split["X_cal"], split["y_cal"], target_ta, "stage_1", - alpha=alpha, log=log, skip=skip) + alpha=alpha, log=log, skip=skip, seed=seed) if s1 is None: return {"method": "rsb", "stages": [], "summary": {"status": "no_stage", "coverage_cal_total": 0.0}} stages = [s1] @@ -369,7 +368,7 @@ def build_rsb(split, target_ta, alpha: float = 0.1, log: Optional[LogFn] = None, split["X_val"][rej_val], split["y_val"][rej_val], split["X_cal"][rej_cal], split["y_cal"][rej_cal], target_ta, "stage_2", - alpha=alpha, log=log, skip=skip) + alpha=alpha, log=log, skip=skip, seed=seed) if s2 is not None: stages.append(s2) @@ -449,7 +448,7 @@ def _pipeline_cal_summary(method, stages, X_cal, y_cal): def fit_frontier(X, y_teacher, targets, max_fit_labels=8000, min_coverage=0.05, alpha: float = 0.1, - log: Optional[LogFn] = None, skip: Iterable[str] = ()): + log: Optional[LogFn] = None, skip: Iterable[str] = (), seed=42): """Build global/l2d/rsb for each target TA, return best per target. `alpha` is the confidence level for the deployment lower bound: a candidate @@ -460,8 +459,8 @@ def fit_frontier(X, y_teacher, targets, max_fit_labels=8000, min_coverage=0.05, log(f"fit_frontier: X={X.shape} targets={sorted(set(float(t) for t in targets))} " f"max_fit_labels={max_fit_labels} alpha={alpha}" + (f" skip={tuple(skip)}" if tuple(skip) else "")) - X_fit, y_fit = _subsample(X, y_teacher, max_fit_labels) - split = _split_buffer(X_fit, y_fit) + X_fit, y_fit = _subsample(X, y_teacher, max_fit_labels, seed=seed) + split = _split_buffer(X_fit, y_fit, seed=seed) log(f"fit_frontier: split -> {len(split['X_train'])} train / " f"{len(split['X_val'])} val / {len(split['X_cal'])} cal") builders = {"global": build_global, "l2d": build_l2d, "rsb": build_rsb} @@ -470,7 +469,7 @@ def fit_frontier(X, y_teacher, targets, max_fit_labels=8000, min_coverage=0.05, for target in sorted(set(float(t) for t in targets)): candidates = [] for method_name, builder in builders.items(): - pipeline = builder(split, target, alpha=alpha, log=log, skip=skip) + pipeline = builder(split, target, alpha=alpha, log=log, skip=skip, seed=seed) pipeline["summary"]["method"] = method_name if pipeline["summary"].get("coverage_cal_total", 0.0) < min_coverage: if pipeline["summary"].get("status") == "ok": diff --git a/src/tracer/fit/surrogate.py b/src/tracer/fit/surrogate.py index 1182084..b84f9fa 100644 --- a/src/tracer/fit/surrogate.py +++ b/src/tracer/fit/surrogate.py @@ -31,7 +31,7 @@ SEED = 42 -def _candidates(n_samples: int, skip: Iterable[str] = ()) -> dict: +def _candidates(n_samples: int, skip: Iterable[str] = (), seed: int = SEED) -> dict: """Return model factory dict. Factories are callables → fitted-ready estimator. Pass ``skip=("name", ...)`` to exclude candidates from the sweep, useful @@ -45,17 +45,17 @@ def _candidates(n_samples: int, skip: Iterable[str] = ()) -> dict: "logreg_c1": lambda: Pipeline([ ("scale", StandardScaler()), ("clf", LogisticRegression(C=1.0, max_iter=1000, solver="lbfgs", - random_state=SEED)), + random_state=seed)), ]), "logreg_c10": lambda: Pipeline([ ("scale", StandardScaler()), ("clf", LogisticRegression(C=10.0, max_iter=1000, solver="lbfgs", - random_state=SEED)), + random_state=seed)), ]), "sgd_log": lambda: Pipeline([ ("scale", StandardScaler()), ("clf", SGDClassifier(loss="log_loss", max_iter=200, tol=1e-3, - random_state=SEED)), + random_state=seed)), ]), } @@ -65,33 +65,33 @@ def _candidates(n_samples: int, skip: Iterable[str] = ()) -> dict: ("scale", StandardScaler()), ("clf", MLPClassifier(hidden_layer_sizes=(256,), alpha=1e-4, max_iter=140, early_stopping=True, - random_state=SEED)), + random_state=seed)), ]), "mlp_2h": lambda: Pipeline([ ("scale", StandardScaler()), ("clf", MLPClassifier(hidden_layer_sizes=(256, 96), alpha=1e-4, max_iter=140, early_stopping=True, - random_state=SEED)), + random_state=seed)), ]), } # ── tree-based (no scaling needed) ──────────────────────────────────────── trees = { "dt": lambda: DecisionTreeClassifier( - max_depth=16, min_samples_leaf=2, random_state=SEED), + max_depth=16, min_samples_leaf=2, random_state=seed), "rf": lambda: RandomForestClassifier( n_estimators=200, max_features="sqrt", - min_samples_leaf=2, n_jobs=-1, random_state=SEED), + min_samples_leaf=2, n_jobs=-1, random_state=seed), "et": lambda: ExtraTreesClassifier( n_estimators=200, max_features="sqrt", - min_samples_leaf=2, n_jobs=-1, random_state=SEED), + min_samples_leaf=2, n_jobs=-1, random_state=seed), } # GradientBoosting is slow on large datasets -- skip above 4k if n_samples <= 4_000: trees["gbt"] = lambda: GradientBoostingClassifier( n_estimators=150, max_depth=4, learning_rate=0.1, - subsample=0.8, random_state=SEED) + subsample=0.8, random_state=seed) # ── xgboost (optional) ─────────────────────────────────────────────────── try: @@ -102,7 +102,7 @@ def _candidates(n_samples: int, skip: Iterable[str] = ()) -> dict: n_estimators=200, max_depth=6, learning_rate=0.1, subsample=0.8, colsample_bytree=0.8, use_label_encoder=False, eval_metric="mlogloss", - n_jobs=-1, random_state=SEED, verbosity=0, + n_jobs=-1, random_state=seed, verbosity=0, ) except Exception: pass @@ -138,6 +138,7 @@ def search_best_surrogate( y_val: np.ndarray, on_candidate: Optional[Callable[..., None]] = None, skip: Iterable[str] = (), + seed: int = SEED, ) -> Tuple[Any, str, dict]: """Train all candidates and return (best_clf, model_name, metrics). @@ -156,7 +157,7 @@ def search_best_surrogate( dominate wall-time, e.g. ``skip=("gbt",)``. """ n = len(X_train) - candidates = _candidates(n, skip=skip) + candidates = _candidates(n, skip=skip, seed=seed) best_clf = None best_name = None diff --git a/src/tracer/runtime/router.py b/src/tracer/runtime/router.py index e2dd959..932bfc7 100644 --- a/src/tracer/runtime/router.py +++ b/src/tracer/runtime/router.py @@ -43,6 +43,10 @@ def __init__(self, stages: list, label_space: list, manifest, embedder=None, self.embedder = embedder self._ood_gate = ood_gate self._train_embeddings = train_embeddings + self._ood_index = None + if ood_gate is not None and train_embeddings is not None and len(train_embeddings): + from sklearn.neighbors import NearestNeighbors + self._ood_index = NearestNeighbors().fit(np.asarray(train_embeddings, dtype=np.float32)) @classmethod def load(cls, artifact_dir: Union[str, Path], embedder=None) -> "Router": @@ -86,7 +90,7 @@ 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.""" @@ -131,6 +135,8 @@ def predict( dict with keys: label, decision ("handled"/"deferred"), accept_score, stage """ embedding = self._to_embedding(input) + if embedding.ndim != 1 or not np.isfinite(embedding).all(): + raise ValueError("A single embedding must be a finite 1-D vector") expected_dim = self.manifest.embedding_dim if expected_dim is not None and embedding.shape[-1] != expected_dim: raise ValueError( @@ -175,6 +181,8 @@ def predict_batch( """ from tracer.fit.pipeline import route_pipeline X = self._to_embeddings(inputs) + if X.ndim != 2 or not np.isfinite(X).all(): + raise ValueError("Batch embeddings must be a finite 2-D matrix") expected_dim = self.manifest.embedding_dim if expected_dim is not None and X.shape[-1] != expected_dim: raise ValueError( @@ -186,6 +194,7 @@ def predict_batch( if self._ood_gate is not None and np.asarray(handled).any(): ood = self._ood_flags(X, preds) handled = np.asarray(handled) & ~ood + stage_id = np.where(handled, stage_id, -1) labels = [] decisions = [] for i in range(len(X)): diff --git a/src/tracer/runtime/serve.py b/src/tracer/runtime/serve.py index d284f88..69ca9c0 100644 --- a/src/tracer/runtime/serve.py +++ b/src/tracer/runtime/serve.py @@ -14,7 +14,7 @@ from __future__ import annotations import json -from http.server import HTTPServer, BaseHTTPRequestHandler +from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler from pathlib import Path from typing import Union @@ -27,6 +27,10 @@ class _Handler(BaseHTTPRequestHandler): + def setup(self): + super().setup() + self.connection.settimeout(10) + def do_GET(self): if self.path == "/health": self._json_response(200, { @@ -65,6 +69,8 @@ def _handle_predict(self, body: dict): x = np.asarray(emb, dtype=np.float32) out = _router.predict(x) self._json_response(200, out) + except (ValueError, TypeError) as e: + self._json_response(400, {"error": str(e)}) except Exception as e: self._json_response(500, {"error": str(e)}) @@ -83,34 +89,43 @@ def _handle_predict_batch(self, body: dict): "handled": out["handled"].tolist(), } self._json_response(200, result) + except (ValueError, TypeError) as e: + self._json_response(400, {"error": str(e)}) except Exception as e: self._json_response(500, {"error": str(e)}) def _read_body(self) -> dict: length = int(self.headers.get("Content-Length", 0)) - if length == 0: - raise ValueError("empty request body") + if length <= 0 or length > 8 * 1024 * 1024: + raise ValueError("request body must be between 1 byte and 8 MiB") raw = self.rfile.read(length) - return json.loads(raw) + if len(raw) != length: + raise ValueError("incomplete request body") + body = json.loads(raw) + if not isinstance(body, dict): + raise ValueError("request body must be a JSON object") + return body def _json_response(self, code: int, data: dict): body = json.dumps(data, default=str).encode("utf-8") - self.send_response(code) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.send_header("Access-Control-Allow-Origin", "*") - self.end_headers() - self.wfile.write(body) + try: + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + except (BrokenPipeError, ConnectionResetError, TimeoutError): + self.close_connection = True def log_message(self, fmt, *args): # Quiet logging - only errors - if args and str(args[1]).startswith("5"): + if len(args) > 1 and str(args[1]).startswith("5"): super().log_message(fmt, *args) def serve( artifact_dir: Union[str, Path] = ".tracer", - host: str = "0.0.0.0", + host: str = "127.0.0.1", port: int = 8000, ): """Start a prediction server for a fitted TRACER policy. @@ -118,7 +133,7 @@ def serve( Parameters ---------- artifact_dir : path to .tracer/ directory - host : bind address (default: 0.0.0.0) + host : bind address (default: 127.0.0.1; set explicitly for remote access) port : listen port (default: 8000) """ global _router, _manifest @@ -130,7 +145,7 @@ def serve( _manifest = load_manifest(artifact_dir / "manifest.json") _router = Router.load(artifact_dir) - server = HTTPServer((host, port), _Handler) + server = ThreadingHTTPServer((host, port), _Handler) method = _manifest.selected_method or "none" cov = f"{_manifest.coverage_cal:.1%}" if _manifest.coverage_cal else "n/a" print(f"\n TRACER serve") @@ -147,4 +162,5 @@ def serve( server.serve_forever() except KeyboardInterrupt: print("\n Shutting down.") - server.shutdown() + finally: + server.server_close() diff --git a/src/tracer/scanner.py b/src/tracer/scanner.py index 740d798..0dd2c21 100644 --- a/src/tracer/scanner.py +++ b/src/tracer/scanner.py @@ -569,7 +569,8 @@ def rule(char="-", width=56): var ids=Object.keys(C).map(Number).sort(function(a,b){return (C[b].share||0)-(C[a].share||0);}); for(var j=0;j'+esc(lab)+''); + var swatch=/^#[0-9a-fA-F]{6}$/.test(cl.lc)?cl.lc:'#94a3b8'; + items.push(''+esc(lab)+''); if(items.length>=8){ break; } } if(ids.length>items.length){ items.push('+ more'); } legend.innerHTML=items.join(''); @@ -685,7 +686,10 @@ def scan_html(r: ScanResult, source_name: str = "traces") -> str: has_viz = bool(r.projection and r.projection.get("points")) viz_block = _VIZ_HTML if has_viz else "" - script = _VIZ_SCRIPT.replace("__VIZ_DATA__", _json.dumps(r.projection)) if has_viz else "" + # JSON quoting alone does not prevent a label from closing a script element. + viz_json = (_json.dumps(r.projection).replace("<", "\\u003c") + .replace(">", "\\u003e").replace("&", "\\u0026")) if has_viz else "" + script = _VIZ_SCRIPT.replace("__VIZ_DATA__", viz_json) if has_viz else "" save_html = f"

{money}

" if money else "" return f""" diff --git a/src/tracer/traces/loader.py b/src/tracer/traces/loader.py index 9ae45a4..24057b2 100644 --- a/src/tracer/traces/loader.py +++ b/src/tracer/traces/loader.py @@ -35,6 +35,8 @@ def load_traces(path: Union[str, Path]) -> TraceDataset: raise ValueError( f"Malformed JSON in {path} at line {lineno}: {exc}" ) from exc + if not isinstance(row, dict): + raise ValueError(f"Trace at line {lineno} must be a JSON object") # Accept common key aliases so middleware dumps load without # reshaping (mirrors the scan loader). Canonical keys are # 'input' and 'teacher'. @@ -49,7 +51,7 @@ def load_traces(path: Union[str, Path]) -> TraceDataset: teacher_val = row[label_key] if teacher_val is None or (isinstance(teacher_val, float) and math.isnan(teacher_val)): raise ValueError( - f"Trace at line {len(records) + 1} has a null/NaN 'teacher' field. " + f"Trace at line {lineno} has a null/NaN 'teacher' field. " f"This usually happens when json.dumps() serializes float('nan') as " f"unquoted NaN. Filter out rows with missing labels before fitting." ) diff --git a/src/tracer/watch.py b/src/tracer/watch.py index 18c81d4..ea39117 100644 --- a/src/tracer/watch.py +++ b/src/tracer/watch.py @@ -22,8 +22,11 @@ from __future__ import annotations import contextvars +import inspect import json import os +import re +import sys import threading import time import uuid @@ -196,19 +199,35 @@ def emit(self, span: GenAISpan) -> None: ... def close(self) -> None: ... +def _debug_failure(exc: Exception) -> None: + if os.environ.get("TRACER_WATCH_DEBUG"): + try: + print(f"[tracer.watch] recording failed: {exc}", file=sys.stderr) + except Exception: + pass + + class LocalFileSink: """Append spans as JSONL to ``/.jsonl``. No network, no key.""" def __init__(self, name: str, dir: str = ".tracer/watch") -> None: + if not isinstance(name, str) or len(name) > 128 or re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]*", name) is None: + raise ValueError("Watcher names must start with a letter or digit and contain only letters, digits, '.', '_' or '-' (max 128 characters)") self.path = os.path.join(dir, f"{name}.jsonl") - os.makedirs(dir, exist_ok=True) + try: + os.makedirs(dir, exist_ok=True) + except OSError as exc: + _debug_failure(exc) self._lock = threading.Lock() def emit(self, span: GenAISpan) -> None: - line = json.dumps(span.to_dict(), ensure_ascii=False) - with self._lock: - with open(self.path, "a", encoding="utf-8") as f: - f.write(line + "\n") + try: + line = json.dumps(span.to_dict(), ensure_ascii=False) + with self._lock: + with open(self.path, "a", encoding="utf-8") as f: + f.write(line + "\n") + except Exception as exc: + _debug_failure(exc) def close(self) -> None: # nothing buffered return None @@ -259,11 +278,17 @@ def __init__(self, sinks: Sequence[Sink]) -> None: def emit(self, span: GenAISpan) -> None: for s in self.sinks: - s.emit(span) + try: + s.emit(span) + except Exception as exc: + _debug_failure(exc) def close(self) -> None: for s in self.sinks: - s.close() + try: + s.close() + except Exception as exc: + _debug_failure(exc) # Tracks the currently-open span so nested watched calls form a trace tree @@ -445,26 +470,50 @@ def _sink_from_env(name: str) -> Sink: def __call__(self, fn: Callable) -> Callable: import functools + if inspect.iscoroutinefunction(fn): + @functools.wraps(fn) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + span = self._begin(self._input(args, kwargs)) + try: + result = await fn(*args, **kwargs) + self._capture(span, result) + return result + except BaseException as exc: + span.status, span.error = "error", str(exc) + raise + finally: + self._finish(span) + return async_wrapper + @functools.wraps(fn) def wrapper(*args: Any, **kwargs: Any) -> Any: - span = self._begin(self.extract_input(args, kwargs)) + span = self._begin(self._input(args, kwargs)) try: result = fn(*args, **kwargs) - except Exception as e: - span.status = "error" - span.error = str(e) - self._finish(span) + self._capture(span, result) + return result + except BaseException as exc: + span.status, span.error = "error", str(exc) raise - # Auto-capture model / tokens / finish_reason / tool calls if the - # function returned a provider response object; otherwise fall back - # to the plain output extractor. + finally: + self._finish(span) + + return wrapper + + def _input(self, args, kwargs) -> str: + try: + return self.extract_input(args, kwargs) + except Exception as exc: + _debug_failure(exc) + return "" + + def _capture(self, span, result) -> None: + try: extract_response(span, result) if not span.output_text: span.output_text = self.extract_output(result) - self._finish(span) - return result - - return wrapper + except Exception as exc: + _debug_failure(exc) # -- context manager -------------------------------------------------- # def span( @@ -520,10 +569,16 @@ def _finish(self, span: GenAISpan) -> None: _ACTIVE.reset(tok) except (ValueError, LookupError): pass - self.sink.emit(span) + try: + self.sink.emit(span) + except Exception as exc: + _debug_failure(exc) def close(self) -> None: - self.sink.close() + try: + self.sink.close() + except Exception as exc: + _debug_failure(exc) class _SpanCtx: diff --git a/tests/test_report_safety.py b/tests/test_report_safety.py new file mode 100644 index 0000000..b0c4eca --- /dev/null +++ b/tests/test_report_safety.py @@ -0,0 +1,82 @@ +"""Untrusted trace text must remain text in reports, including script data.""" +import json +from dataclasses import asdict +from html.parser import HTMLParser +from pathlib import Path + +import pytest + +from tracer.analysis.html_report import generate_html_report +from tracer.scanner import ScanResult, scan_html +from tracer.types import ArtifactManifest + +PAYLOAD = '' + + +class Elements(HTMLParser): + def __init__(self): + super().__init__() + self.tags = [] + self.scripts = [] + self.in_script = False + def handle_starttag(self, tag, attrs): + self.tags.append((tag, attrs)) + if tag == 'script': + self.in_script = True + self.scripts.append('') + def handle_endtag(self, tag): + if tag == 'script': self.in_script = False + def handle_data(self, text): + if self.in_script: self.scripts[-1] += text + + +def assert_inert(html): + parsed = Elements(); parsed.feed(html) + assert 'window.TRACER_REPORT_TEST=1' not in parsed.scripts + assert not any(tag == 'img' for tag, attrs in parsed.tags) + assert not any(key.lower().startswith('on') for tag, attrs in parsed.tags for key, value in attrs) + + +def test_scan_projection_cannot_end_script_element(): + projection = {'points': [[0, 0, 0, 0]], 'clusters': {'0': {'label': PAYLOAD, 'ex': [PAYLOAD]}}} + result = ScanResult(1, 1, 1, .9, 0, 0, projection=projection) + html = scan_html(result) + assert_inert(html) + assert '\\u003c/script' in html + + +@pytest.fixture +def report_files(tmp_path): + manifest = asdict(ArtifactManifest(version='0.1.0', n_traces=20, label_space=[PAYLOAD], + selected_method=PAYLOAD, embedding_dim=PAYLOAD, + coverage_cal=.5, teacher_agreement_cal=1)) + report = { + 'coverage': .5, + 'slices': [{'slice_name': 'label:' + PAYLOAD, 'count': 20, 'handled_rate': .5, + 'deferred_rate': .5, 'teacher_agreement_handled': 1}], + 'boundary_pairs': [{'teacher_label': PAYLOAD, 'handled_preview': PAYLOAD, + 'deferred_preview': PAYLOAD, 'handled_score': .9, 'deferred_score': .7}], + 'handled_examples': [{'input_preview': PAYLOAD, 'teacher_label': PAYLOAD, 'accept_score': .9}], + 'deferred_examples': [{'input_preview': PAYLOAD, 'teacher_label': PAYLOAD}], + 'temporal_deltas': [{'label': PAYLOAD, 'previous_handled_rate': .4, 'current_handled_rate': .5, 'delta': .1}], + } + (tmp_path / 'manifest.json').write_text(json.dumps(manifest)) + (tmp_path / 'qualitative_report.json').write_text(json.dumps(report)) + return tmp_path, manifest, report + + +def test_fit_report_escapes_trace_and_manifest_fields(report_files): + root, _, _ = report_files + html = Path(generate_html_report(root)).read_text() + assert_inert(html) + assert '<' in html + + +def test_sankey_labels_and_title_are_plain_text(report_files): + pytest.importorskip('plotly') + from tracer.analysis.sankey import generate_sankey, _build_sankey_figure + root, manifest, report = report_files + fig = _build_sankey_figure(manifest, report, 15, False) + assert all('<' not in label for label in fig.data[0].node.label) + html = Path(generate_sankey(root, title=PAYLOAD)).read_text() + assert_inert(html) diff --git a/tests/test_routing_regressions.py b/tests/test_routing_regressions.py new file mode 100644 index 0000000..ab5297d --- /dev/null +++ b/tests/test_routing_regressions.py @@ -0,0 +1,140 @@ +"""Focused checks for label identity, calibration, seeds and persisted embeddings.""" +import json +from types import SimpleNamespace + +import numpy as np +import pytest +from sklearn.linear_model import LogisticRegression +from sklearn.neighbors import NearestNeighbors + +from tracer.api import fit, load_router +from tracer.config import FitConfig +from tracer.embeddings.index import EmbeddingIndex +from tracer.fit.ood import fit_ood_gate, ood_mask +from tracer.fit.pipeline import _predict, apply_stage, _calibrate_threshold, _cp_lower, fit_frontier +from tracer.runtime.router import Router + +SKIP = ('logreg_c10', 'sgd_log', 'mlp_1h', 'mlp_2h', 'dt', 'rf', 'et', 'gbt', 'xgb') + + +def test_residual_stage_preserves_noncontiguous_labels(): + X = np.array([[-3.], [-2.], [-1.], [1.], [2.], [3.]]) + y = np.array([1, 1, 1, 3, 3, 3]) + clf = LogisticRegression().fit(X, y) + expected = clf.predict(X) + np.testing.assert_array_equal(_predict(clf, X)[0], expected) + stage = {'clf': clf, 'accept_all': True} + router = Router([stage], ['zero', 'one', 'two', 'three'], SimpleNamespace(embedding_dim=1)) + assert router.predict(X[-1])['label'] == 'three' + assert router.predict_batch(X)['labels'] == ['one'] * 3 + ['three'] * 3 + + +def test_small_calibration_chooses_largest_eligible_set(): + scores = np.linspace(.50, .99, 36) + teacher = np.zeros(36, dtype=int) + preds = teacher.copy(); preds[:6] = 1 + chosen = _calibrate_threshold(scores, preds, teacher, .8) + eligible = [] + for threshold in np.unique(scores): + mask = scores >= threshold + n = int(mask.sum()); correct = int((preds[mask] == teacher[mask]).sum()) + if n >= 10 and _cp_lower(correct, n, .1) >= .8: + eligible.append(mask.mean()) + assert chosen['coverage'] == max(eligible) + assert chosen['teacher_agreement_lower'] >= .8 + + +def test_seed_changes_split_but_same_seed_repeats(): + X = np.random.RandomState(4).normal(size=(240, 3)) + y = np.arange(len(X)) % 3 + opts = dict(targets=(.9,), max_fit_labels=120, skip=SKIP) + _, a = fit_frontier(X, y, seed=17, **opts) + _, b = fit_frontier(X, y, seed=17, **opts) + _, c = fit_frontier(X, y, seed=93, **opts) + np.testing.assert_array_equal(a['X_train'], b['X_train']) + assert not np.array_equal(a['X_train'], c['X_train']) + + +@pytest.mark.parametrize('value', [.5, [[0, 1]], [0, np.nan], [np.inf, 0]]) +def test_single_prediction_rejects_invalid_shapes_and_values(value): + router = Router([], ['a'], SimpleNamespace(embedding_dim=2)) + with pytest.raises(ValueError, match='finite 1-D'): + router.predict(value) + + +@pytest.mark.parametrize('value', [.5, [0, 1], [[[0, 1]]], [[0, np.nan]]]) +def test_batch_prediction_rejects_invalid_shapes_and_values(value): + router = Router([], ['a'], SimpleNamespace(embedding_dim=2)) + with pytest.raises(ValueError, match='finite 2-D'): + router.predict_batch(value) + + +def test_router_reuses_ood_index(monkeypatch): + X = np.random.RandomState(3).normal(size=(80, 3)).astype(np.float32) + gate = fit_ood_gate(X, ['a'] * len(X)) + queries = np.vstack([X[:2], np.full((1, 3), 100)]) + expected = ood_mask(queries, X, ['a'] * len(queries), gate) + router = Router([], ['a'], SimpleNamespace(embedding_dim=3), ood_gate=gate, train_embeddings=X) + def unexpected_fit(*args, **kwargs): + raise AssertionError('OOD index was rebuilt during prediction') + monkeypatch.setattr(NearestNeighbors, 'fit', unexpected_fit) + for _ in range(2): + np.testing.assert_array_equal(router._ood_flags(queries, np.zeros(3)), expected) + + +def test_faiss_fit_reload_preserves_embedding_scale(tmp_path): + pytest.importorskip('faiss') + rng = np.random.RandomState(24) + y = np.repeat(np.arange(3), 80) + X = (rng.normal(size=(3, 8))[y] * 10 + rng.normal(size=(240, 8)) * .2).astype(np.float32) + original = X.copy() + traces = tmp_path / 'traces.jsonl' + traces.write_text(''.join(json.dumps({'input': str(i), 'teacher': str(label)}) + '\n' for i, label in enumerate(y))) + result = fit(traces, tmp_path / 'model', X, FitConfig(frontier_targets=(.9,), skip_candidates=SKIP, verbose=False)) + np.testing.assert_array_equal(X, original) + np.testing.assert_array_equal(EmbeddingIndex.load(tmp_path / 'model' / 'index').embeddings, original) + router = load_router(tmp_path / 'model') + assert result.manifest.coverage_cal == 1 + assert not router._ood_flags(original, y).any() + assert router.predict_batch(original)['handled'].all() + + +def test_fit_report_keeps_deferred_confidence(tmp_path, monkeypatch): + # The lazy-import test reloads the package; patch the function used below. + from tracer.api import fit + + X = np.array([[-3.], [-1.], [-.1], [.1], [1.], [3.]], dtype=np.float32) + y = np.array([0, 0, 0, 1, 1, 1]) + clf = LogisticRegression().fit(X, y) + stage = {'clf': clf, 'accept_all': False, 'acceptor': None, 'threshold': .9} + best = {'stages': [stage], 'summary': {'method': 'l2d', 'coverage_cal_total': .3, + 'teacher_agreement_cal_total': 1}} + monkeypatch.setattr('tracer.api.fit_frontier', lambda *args, **kwargs: + ([{'target': .9, 'best': best, 'candidates': [best]}], {})) + path = tmp_path / 'traces.jsonl' + path.write_text(''.join(json.dumps({'input':str(i), 'teacher':str(label)})+'\n' for i,label in enumerate(y))) + result = fit(path, tmp_path / 'out', X, FitConfig(frontier_targets=(.9,), verbose=False)) + assert result.qualitative_report.deferred_examples + for example in result.qualitative_report.deferred_examples: + expected = clf.predict_proba(X[int(example.input_preview):int(example.input_preview)+1]).max() + assert example.accept_score == pytest.approx(expected) + + +def test_fit_config_seed_reaches_model(tmp_path): + import joblib + y = np.arange(160) % 2 + X = np.column_stack([y * 10., y * 10.]).astype(np.float32) + path = tmp_path / 'traces.jsonl' + path.write_text(''.join(json.dumps({'input':str(i), 'teacher':str(label)})+'\n' for i,label in enumerate(y))) + result = fit(path, tmp_path / 'out', X, + FitConfig(seed=291, frontier_targets=(.9,), skip_candidates=SKIP, verbose=False)) + model = joblib.load(result.manifest.pipeline_path)['pipeline']['stages'][0]['clf'] + assert model.named_steps['clf'].random_state == 291 + + +def test_null_teacher_error_uses_file_line(tmp_path): + from tracer.traces.loader import load_traces + path = tmp_path / 'traces.jsonl' + path.write_text('{"input":"a","teacher":"yes"}\n\n\n{"input":"b","teacher":null}\n') + with pytest.raises(ValueError, match='line 4'): + load_traces(path) diff --git a/tests/test_serve.py b/tests/test_serve.py index 99ca0c5..5a7317c 100644 --- a/tests/test_serve.py +++ b/tests/test_serve.py @@ -8,7 +8,7 @@ import json import threading -from http.server import HTTPServer +from http.server import ThreadingHTTPServer from pathlib import Path from urllib.error import HTTPError from urllib.request import Request, urlopen @@ -65,7 +65,7 @@ def server(fitted_artifact): # can bind an ephemeral port and shut down cleanly. serve_mod._manifest = load_manifest(artifact_dir / "manifest.json") serve_mod._router = Router.load(artifact_dir) - httpd = HTTPServer(("127.0.0.1", 0), serve_mod._Handler) + httpd = ThreadingHTTPServer(("127.0.0.1", 0), serve_mod._Handler) port = httpd.server_address[1] thread = threading.Thread(target=httpd.serve_forever, daemon=True) thread.start() @@ -119,10 +119,9 @@ def test_predict_missing_field_is_400(server): assert "embedding" in body["error"] -def test_predict_wrong_dim_is_500(server, fitted_artifact): - # Router raises a dimension-mismatch ValueError, surfaced as a 500. +def test_predict_wrong_dim_is_400(server, fitted_artifact): status, body = _post(server, "/predict", {"embedding": [0.1, 0.2, 0.3]}) - assert status == 500 + assert status == 400 assert "dimension" in body["error"].lower() @@ -160,3 +159,70 @@ def test_unknown_get_path_is_404(server): def test_unknown_post_path_is_404(server): status, body = _post(server, "/nope", {"x": 1}) assert status == 404 + + +@pytest.mark.parametrize('body', [[], None, 2, 'string']) +def test_non_object_request_is_400(server, body): + status, result = _post(server, '/predict', body) + assert status == 400 + assert 'JSON object' in result['error'] + assert _get(server, '/health')[0] == 200 + + +def test_slow_prediction_does_not_block_health(server, monkeypatch): + entered = threading.Event() + release = threading.Event() + results = [] + def slow_predict(_): + entered.set() + assert release.wait(3) + return {'label': 'a', 'decision': 'handled', 'accept_score': 1, 'stage': 0} + monkeypatch.setattr(serve_mod._router, 'predict', slow_predict) + thread = threading.Thread(target=lambda: results.append(_post(server, '/predict', {'embedding': [0]}))) + thread.start() + try: + assert entered.wait(2) + with urlopen(server + '/health', timeout=1) as response: + assert response.status == 200 + finally: + release.set() + thread.join(timeout=3) + assert results[0][0] == 200 + + +@pytest.mark.parametrize('stage', ['headers', 'body']) +@pytest.mark.parametrize('error', [BrokenPipeError, ConnectionResetError]) +def test_disconnect_during_response_is_ignored(stage, error): + from types import SimpleNamespace + handler = object.__new__(serve_mod._Handler) + handler.send_response = lambda *_: None + handler.send_header = lambda *_: None + def disconnected(*_): raise error('peer closed') + handler.end_headers = disconnected if stage == 'headers' else lambda: None + handler.wfile = SimpleNamespace(write=disconnected) + handler._json_response(200, {'ok': True}) + assert handler.close_connection + + +def test_invalid_content_length_is_rejected_before_reading(): + import io + handler = object.__new__(serve_mod._Handler) + handler.rfile = io.BytesIO(b'{}') + for value in ('-1', '0', str(8 * 1024 * 1024 + 1)): + handler.headers = {'Content-Length': value} + with pytest.raises(ValueError, match='request body'): + handler._read_body() + assert handler.rfile.tell() == 0 + + +def test_serve_uses_threaded_loopback_and_closes_on_interrupt(fitted_artifact, monkeypatch): + from unittest.mock import Mock + artifact_dir, _ = fitted_artifact + instance = Mock() + instance.serve_forever.side_effect = KeyboardInterrupt + factory = Mock(return_value=instance) + monkeypatch.setattr(serve_mod, 'ThreadingHTTPServer', factory) + serve_mod.serve(artifact_dir) + assert factory.call_args.args[0] == ('127.0.0.1', 8000) + instance.server_close.assert_called_once() + instance.shutdown.assert_not_called() diff --git a/tests/test_update_safety.py b/tests/test_update_safety.py new file mode 100644 index 0000000..c369c1e --- /dev/null +++ b/tests/test_update_safety.py @@ -0,0 +1,143 @@ +"""Refits must preserve the previous generation until a replacement is ready.""" +import json +from pathlib import Path + +import numpy as np +import pytest + +from tracer.api import fit, update, load_router +from tracer.config import FitConfig +from tracer.traces.loader import load_traces + +SKIP = ('logreg_c10', 'sgd_log', 'mlp_1h', 'mlp_2h', 'dt', 'rf', 'et', 'gbt', 'xgb') + + +def write_traces(path, n): + y = np.arange(n) % 2 + path.write_text(''.join(json.dumps({'input': f'row {i}', 'teacher': str(label)}) + '\n' + for i, label in enumerate(y))) + rng = np.random.RandomState(14) + return (np.column_stack([y * 8, y * 8]) + rng.normal(0, .1, (n, 2))).astype(np.float32) + + +@pytest.fixture +def artifacts(tmp_path): + traces = tmp_path / 'initial.jsonl' + X = write_traces(traces, 160) + output = tmp_path / 'artifacts' + fit(traces, output, X, FitConfig(frontier_targets=(.9,), skip_candidates=SKIP, + seed=19, verbose=False)) + (output / 'user-notes.txt').write_text('keep me') + return output, X + + +def snapshot(directory): + return {str(p.relative_to(directory)): p.read_bytes() for p in directory.rglob('*') if p.is_file()} + + +@pytest.mark.parametrize('new_X', [np.zeros((1, 2)), np.zeros((2, 3)), + np.zeros(2), np.full((2, 2), np.nan)]) +def test_invalid_update_leaves_all_artifacts_unchanged(artifacts, tmp_path, new_X): + output, _ = artifacts + traces = tmp_path / 'new.jsonl' + write_traces(traces, 2) + before = snapshot(output) + with pytest.raises(ValueError): + update(traces, output, new_embeddings=new_X) + assert snapshot(output) == before + # A corrected retry must still work and append exactly once. + result = update(traces, output, new_embeddings=write_traces(traces, 2)) + assert result.manifest.n_traces == 162 + assert len(load_traces(output / 'all_traces.jsonl')) == 162 + + +def test_refit_failure_leaves_original_generation(artifacts, tmp_path, monkeypatch): + # The lazy-import test reloads the package; patch the function used below. + from tracer.api import update + + output, _ = artifacts + traces = tmp_path / 'new.jsonl' + X = write_traces(traces, 2) + before = snapshot(output) + + def fail_after_writing(trace_path, artifact_dir, **kwargs): + (artifact_dir / 'pipeline.joblib').write_bytes(b'partial model') + raise OSError('disk write failed') + + monkeypatch.setattr('tracer.api.fit', fail_after_writing) + with pytest.raises(OSError, match='disk write failed'): + update(traces, output, new_embeddings=X) + assert snapshot(output) == before + + +def test_publish_failure_restores_original_generation(artifacts, tmp_path, monkeypatch): + output, _ = artifacts + traces = tmp_path / 'new.jsonl' + X = write_traces(traces, 2) + before = snapshot(output) + rename = Path.rename + + def fail_publish(path, target): + if path.name == 'next': + raise OSError('publish failed') + return rename(path, target) + + monkeypatch.setattr(Path, 'rename', fail_publish) + with pytest.raises(OSError, match='publish failed'): + update(traces, output, new_embeddings=X) + assert snapshot(output) == before + + +def test_explicit_config_is_honored_and_not_mutated(artifacts, tmp_path): + output, _ = artifacts + traces = tmp_path / 'new.jsonl' + X = write_traces(traces, 2) + config = FitConfig(target_teacher_agreement=.99, frontier_targets=(.99,), + skip_candidates=SKIP, seed=27, verbose=False) + result = update(traces, output, new_embeddings=X, config=config) + assert config.target_teacher_agreement == .99 + assert result.manifest.target_teacher_agreement == .99 + assert result.manifest.n_retrains == 2 + assert json.loads((output / 'config.json').read_text())['seed'] == 27 + # A strict refit with no deployable policy must not keep the old model. + if result.manifest.selected_method is None: + assert not (output / 'pipeline.joblib').exists() + assert not (output / 'ood.json').exists() + + +def test_default_update_preserves_config_and_returns_live_paths(artifacts, tmp_path): + output, X_old = artifacts + traces = tmp_path / 'new.jsonl' + X = write_traces(traces, 10) + result = update(traces, output, new_embeddings=X) + saved = json.loads((output / 'config.json').read_text()) + assert saved['seed'] == 19 + assert saved['skip_candidates'] == list(SKIP) + assert result.manifest.n_retrains == 2 + assert (output / 'user-notes.txt').read_text() == 'keep me' + for name in ('pipeline_path', 'index_path', 'config_path', 'qualitative_report_path'): + value = getattr(result.manifest, name) + if value is not None: + assert Path(value).parent == output + if name != 'index_path': + assert Path(value).exists() + np.testing.assert_allclose(np.load(output / 'index.embeddings.npy'), np.vstack([X_old, X])) + assert load_router(output).predict(X_old[0])['decision'] == 'handled' + + +def test_failed_rollback_keeps_recoverable_backup(artifacts, tmp_path, monkeypatch): + output, _ = artifacts + traces = tmp_path / 'new.jsonl' + X = write_traces(traces, 2) + before = snapshot(output) + rename = Path.rename + def fail_publish_and_rollback(path, target): + if path.name == 'next' or path.name.endswith('-previous'): + raise OSError('destination unavailable') + return rename(path, target) + monkeypatch.setattr(Path, 'rename', fail_publish_and_rollback) + with pytest.raises(RuntimeError, match='preserved at'): + update(traces, output, new_embeddings=X) + backups = list(tmp_path.glob('.artifacts-update-*-previous')) + assert len(backups) == 1 + assert snapshot(backups[0]) == before diff --git a/tests/test_watch_safety.py b/tests/test_watch_safety.py new file mode 100644 index 0000000..d5f01b6 --- /dev/null +++ b/tests/test_watch_safety.py @@ -0,0 +1,126 @@ +"""Tracing must observe async work and remain harmless when storage fails.""" +import asyncio +import inspect + +import pytest + +from tracer.watch import GenAISpan, LocalFileSink, MultiSink, Watcher + + +class Capture: + def __init__(self): self.spans = [] + def emit(self, span): self.spans.append(span.to_dict()) + def close(self): pass + + +class Broken: + def emit(self, span): raise OSError('recording unavailable') + def close(self): raise OSError('close unavailable') + + +def test_async_results_errors_and_concurrent_parentage(): + captured = Capture() + watcher = Watcher('async', sink=captured) + + @watcher + async def child(text): + await asyncio.sleep(.01) + return text.upper() + + @watcher + async def parent(text): + return await child(text) + + @watcher + async def failed(text): + await asyncio.sleep(0) + raise ValueError('application error') + + async def run(): + assert await asyncio.gather(parent('a'), parent('b')) == ['A', 'B'] + with pytest.raises(ValueError, match='application error'): + await failed('failure') + + assert inspect.iscoroutinefunction(parent) + asyncio.run(run()) + roots = [s for s in captured.spans if s['parent_span_id'] is None and s['status'] == 'ok'] + children = [s for s in captured.spans if s['parent_span_id'] is not None] + assert len(roots) == len(children) == 2 + assert roots[0]['trace_id'] != roots[1]['trace_id'] + for child_span in children: + parent_span = next(s for s in roots if s['span_id'] == child_span['parent_span_id']) + assert child_span['trace_id'] == parent_span['trace_id'] + assert child_span['output_text'] == child_span['input_text'].upper() + assert child_span['latency_ms'] >= 5 + assert captured.spans[-1]['status'] == 'error' + assert captured.spans[-1]['error'] == 'application error' + + +def test_async_cancellation_is_recorded_and_reraised(): + captured = Capture() + watcher = Watcher('cancel', sink=captured) + @watcher + async def cancelled(): + raise asyncio.CancelledError() + async def run(): + with pytest.raises(asyncio.CancelledError): + await cancelled() + asyncio.run(run()) + assert captured.spans[0]['status'] == 'error' + + +@pytest.mark.parametrize('sink_type', ['local', 'custom', 'fanout']) +def test_sink_failure_preserves_results_and_original_errors(tmp_path, sink_type): + captured = Capture() + if sink_type == 'local': + sink = LocalFileSink('blocked', str(tmp_path)) + (tmp_path / 'blocked.jsonl').mkdir() + elif sink_type == 'fanout': + sink = MultiSink([Broken(), captured]) + else: + sink = Broken() + watcher = Watcher('safe', sink=sink) + assert watcher(lambda: 'success')() == 'success' + original = ValueError('original failure') + def fail(): raise original + with pytest.raises(ValueError) as raised: + watcher(fail)() + assert raised.value is original + watcher.close() + if sink_type == 'fanout': + assert [s['status'] for s in captured.spans] == ['ok', 'error'] + + +def test_extractors_cannot_break_the_application(): + def broken(*args): raise TypeError('bad extractor') + sink = Capture() + watcher = Watcher('extract', sink=sink, extract_input=broken, extract_output=broken) + assert watcher(lambda value: value)(object()) is not None + assert sink.spans[0]['status'] == 'ok' + + +def test_debug_output_failure_cannot_break_the_application(monkeypatch): + import io + import sys + + closed_stderr = io.StringIO() + closed_stderr.close() + monkeypatch.setenv('TRACER_WATCH_DEBUG', '1') + monkeypatch.setattr(sys, 'stderr', closed_stderr) + watcher = Watcher('debug', sink=Broken()) + assert watcher(lambda: 'success')() == 'success' + watcher.close() + + +@pytest.mark.parametrize('name', ['../escape', '/absolute', r'..\escape', '.', '..', '', + 'line\nbreak', 'trailing\n', 'nul\0byte', 'x' * 129]) +def test_watcher_filename_rejects_unsafe_names(tmp_path, name): + with pytest.raises(ValueError, match='Watcher names'): + LocalFileSink(name, str(tmp_path / 'watch')) + assert not (tmp_path / 'escape.jsonl').exists() + + +def test_safe_watcher_filename_stays_in_directory(tmp_path): + sink = LocalFileSink('service-1.production_trace', str(tmp_path)) + sink.emit(GenAISpan(input_text='ok')) + assert (tmp_path / 'service-1.production_trace.jsonl').is_file()