From fd5c727d6a3ee19bb3f539e91a479f9a61c2002e Mon Sep 17 00:00:00 2001 From: drewOrc <36374426+drewOrc@users.noreply.github.com> Date: Wed, 23 Sep 2026 11:01:01 +0800 Subject: [PATCH 1/4] Add logits archive, training cost records and make ac2 - archive.py: per-run .npz with float32 validation/test logits, gold intent ids and required metadata (model and revision, seed, k, OOS training rows, dataset revision, label-space hash, git commit, time). load_logits refuses missing keys, wrong types, shape mismatches, NaN, out-of-range labels and a different label space. Archives are gitignored; their SHA-256 goes into results/logits-manifest.json. - efficiency.py: parameter counts, wall time, steps and peak memory (MPS sampled driver/tensor memory, CUDA allocator peak, CPU ru_maxrss with the macOS/Linux unit difference handled), recorded in train_summary.json and copied into the results JSON. - ac2.py and `make ac2`: seeds 42/43/44, resumable, FORCE=1 reruns, deletes weights of seeds other than 42 after archiving, writes results/ac2.json and exits 1 on FAIL. - Results JSON moves to results/runs/ so report.py does not read ac2.json or the manifest as run records. - make smoke now writes and re-reads a logits archive. --- .gitignore | 8 +- Makefile | 12 +- README.md | 13 +- src/tinyrouter/ac2.py | 185 ++++++++++++++++++++++++ src/tinyrouter/archive.py | 265 +++++++++++++++++++++++++++++++++++ src/tinyrouter/efficiency.py | 118 ++++++++++++++++ src/tinyrouter/evaluate.py | 93 ++++++++++-- src/tinyrouter/labels.py | 14 ++ src/tinyrouter/report.py | 7 +- src/tinyrouter/smoke.py | 18 ++- src/tinyrouter/train.py | 48 +++++-- tests/archive_fakes.py | 41 ++++++ tests/test_ac2.py | 174 +++++++++++++++++++++++ tests/test_archive.py | 170 ++++++++++++++++++++++ tests/test_efficiency.py | 39 ++++++ tests/test_train_smoke.py | 42 ++++++ 16 files changed, 1209 insertions(+), 38 deletions(-) create mode 100644 src/tinyrouter/ac2.py create mode 100644 src/tinyrouter/archive.py create mode 100644 src/tinyrouter/efficiency.py create mode 100644 tests/archive_fakes.py create mode 100644 tests/test_ac2.py create mode 100644 tests/test_archive.py create mode 100644 tests/test_efficiency.py diff --git a/.gitignore b/.gitignore index 8ea391a..3df792d 100644 --- a/.gitignore +++ b/.gitignore @@ -26,8 +26,14 @@ checkpoints/ *.pt *.onnx -# Run outputs (results/*.json are small and will be committed once real runs exist) +# Run outputs. Committed: results/runs/*.json, results/ac2.json and +# results/logits-manifest.json (SHA-256 of every logits archive). Not +# committed: the archives themselves (about 5 MB each, about 40 runs); +# they are attached to a GitHub Release and checked with `make verify-logits`. results/summary.md +results/logits/ +*.npz +*.tmp # iCloud conflict copies ("name 2") * [0-9]* diff --git a/Makefile b/Makefile index 36f32d2..887af1f 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: setup lint format test test-network smoke train evaluate report clean-checkpoints +.PHONY: setup lint format test test-network smoke train evaluate ac2 verify-logits report clean-checkpoints CONFIG ?= configs/bert-base.yaml SEED ?= 42 @@ -54,6 +54,16 @@ train: evaluate: uv run $(UV_ENV) python -m tinyrouter.evaluate --config $(CONFIG) --seed $(SEED) +# AC2: bert-base-uncased, full data, seeds 42/43/44 -> results/ac2.json; exit 1 on FAIL. +# Resumes: seeds already scored and archived are skipped. FORCE=1 reruns all. +# Weights of seeds 43 and 44 are deleted after their logits are archived. +ac2: + uv run $(UV_ENV) python -m tinyrouter.ac2 --config configs/bert-base.yaml $(if $(filter 1,$(FORCE)),--force,) + +# Every archive listed in results/logits-manifest.json is present and matches its SHA-256. +verify-logits: + uv run python -m tinyrouter.archive + report: uv run python -m tinyrouter.report diff --git a/README.md b/README.md index ff399ac..218c160 100644 --- a/README.md +++ b/README.md @@ -27,8 +27,10 @@ Other targets: | `make lint` | ruff check, ruff format check, em dash check | | `make test-network` | tests that download from the Hugging Face Hub (dataset checksums) | | `make train CONFIG=configs/bert-base.yaml SEED=42` | fine-tune; final weights in `checkpoints//final` | -| `make evaluate CONFIG=... SEED=...` | score validation and test, write `results/.json` | -| `make report` | build `results/summary.md` from `results/*.json` | +| `make evaluate CONFIG=... SEED=...` | score validation and test, write `results/runs/.json` and the logits archive | +| `make ac2` | bert-base-uncased, full data, seeds 42/43/44; writes `results/ac2.json`, exits 1 on FAIL; resumes, `FORCE=1` reruns | +| `make verify-logits` | check every archive in `results/logits/` against `results/logits-manifest.json` | +| `make report` | build `results/summary.md` from `results/runs/*.json` | | `make clean-checkpoints` | delete all trained weights | Put `HF_TOKEN` and `ANTHROPIC_API_KEY` in a `.env` (copy `.env.example`) if you need them; the Makefile passes it to `uv run`. Nothing in the tests or the smoke run calls a paid API. @@ -38,6 +40,8 @@ Put `HF_TOKEN` and `ANTHROPIC_API_KEY` in a `.env` (copy `.env.example`) if you - **Labels.** The model is trained on all 151 CLINC150 intents (150 plus `oos`). Each intent maps to one of 8 routing targets (7 agents plus `oos`) through `src/tinyrouter/resources/intent_to_agent.json`. - **Data.** `clinc/clinc_oos`, `plus` config, from the Hugging Face Hub at a pinned commit. Each parquet file is checked against its SHA-256 and row count (15,250 / 3,100 / 5,500), and the label names inside it must match the committed `intent_names.json`. - **Calibration.** Temperature scaling, fitted on validation logits only. `fit_temperature` raises `LeakageError` if given anything else, and a test checks that. +- **Logits archive.** Every evaluation stores per-example validation and test logits (float32) with gold labels and metadata (model and revision, seed, k, OOS training rows, dataset revision, label-space hash, git commit, time) in `results/logits/.npz`. Later analysis reads only these files. They are not committed (about 5 MB per bert-base run); their SHA-256 goes into the committed `results/logits-manifest.json`, and the files are attached to a GitHub Release. +- **Training cost.** Each results JSON records parameter counts, training wall time, steps, device and peak memory. How memory is measured depends on the device (sampled Metal driver memory on MPS, the allocator peak on CUDA, process peak RSS on CPU) and is written next to the number; see `src/tinyrouter/efficiency.py`. - **Metrics.** 8-way accuracy, 150-way in-scope accuracy and OOS recall (as defined in Larson et al. 2019), ECE, and NLL. They are computed before and after calibration on both splits. ## Layout @@ -54,6 +58,9 @@ src/tinyrouter/ calibrate.py temperature scaling (validation only) train.py HF Trainer wrapper, device auto-select (cuda > mps > cpu) evaluate.py logits, temperature fit, metrics, results JSON + archive.py per-example logits archive (.npz) and SHA-256 manifest + efficiency.py parameter counts, wall time, peak memory per device + ac2.py AC2 run over three seeds and PASS/FAIL verdict llm.py Claude Haiku zero-shot router (baseline and fallback; not run yet) report.py results/*.json -> results/summary.md smoke.py end-to-end wiring check @@ -64,7 +71,7 @@ tests/ pytest; `network` marker for Hub downloads This repository sits in a folder synced by iCloud Drive. iCloud would otherwise try to upload the virtualenv (about 1 GB of torch) and every checkpoint (about 440 MB per bert-base run), and it can evict local files to save space in the middle of a run. iCloud skips any path ending in `.nosync`, so `make setup` creates `.venv.nosync/` and `checkpoints.nosync/` and puts symlinks at `.venv` and `checkpoints`. uv and the training code use the usual names and never notice. If you clone this somewhere outside iCloud, the symlinks do no harm. -Disk is also tight, so training keeps at most one checkpoint (`save_total_limit=1`) and deletes it once the final weights are saved. +Disk is also tight, so training keeps at most one checkpoint (`save_total_limit=1`) and deletes it once the final weights are saved. `make ac2` also deletes the weights of seeds 43 and 44 after their logits are archived; only seed 42's are kept. ## Reproducibility notes diff --git a/src/tinyrouter/ac2.py b/src/tinyrouter/ac2.py new file mode 100644 index 0000000..dd0d1b1 --- /dev/null +++ b/src/tinyrouter/ac2.py @@ -0,0 +1,185 @@ +"""AC2 pipeline check: bert-base-uncased on the full OOS+ train split, seeds 42/43/44. + +Pass means every seed's 150-way in-scope test accuracy is at least 95.7% +(docs/PLAN.md AC2; derived from Larson et al.'s 96.7%, not a claim of +exact reproduction). Writes ``results/ac2.json`` and exits 1 on FAIL. + +Resumable: a seed whose results JSON exists and whose logits archive +matches the manifest is not rerun, unless ``--force``. Trained weights +of every seed except 42 are deleted once that seed's logits are archived +(disk is tight; seed 42 is kept as the RQ7 ONNX fallback). +""" + +from __future__ import annotations + +import argparse +import json +import shutil +from collections.abc import Callable +from pathlib import Path + +from tinyrouter.archive import ArchiveError, check_against_manifest, load_logits +from tinyrouter.config import RunConfig, load_config + +SEEDS = (42, 43, 44) +KEEP_WEIGHTS_SEEDS = frozenset({42}) +THRESHOLD = 0.957 +EXPECTED_MODEL = "google-bert/bert-base-uncased" +EXPECTED_TRAIN_ROWS = 15_250 +EXPECTED_OOS_TRAIN_ROWS = 250 +EXPECTED_TEST_ROWS = 5_500 + +TrainFn = Callable[[RunConfig], Path] +EvaluateFn = Callable[[RunConfig, Path], dict[str, object]] +Log = Callable[[str], None] + + +class SetupError(ValueError): + """A results record is not the AC2 setup (wrong model, subsampled data, ...).""" + + +def run_dir(config: RunConfig) -> Path: + return Path(config.checkpoint_root) / config.run_name + + +def is_done(config: RunConfig) -> bool: + """Results JSON present and its logits archive loads and matches the manifest.""" + from tinyrouter.evaluate import RunPaths + + paths = RunPaths.of(config) + if not (paths.results_json.exists() and paths.logits.exists()): + return False + try: + check_against_manifest(paths.manifest, paths.logits) + load_logits(paths.logits) + except ArchiveError: + return False + return True + + +def has_final_weights(config: RunConfig) -> bool: + final = run_dir(config) / "final" + return (final / "model.safetensors").exists() and (final / "train_summary.json").exists() + + +def run_seed( + config: RunConfig, train_fn: TrainFn, evaluate_fn: EvaluateFn, force: bool, log: Log +) -> None: + if is_done(config) and not force: + log(f"seed {config.seed}: results and logits already archived, skipping") + else: + if has_final_weights(config) and not force: + log(f"seed {config.seed}: reusing trained weights in {run_dir(config)}/final") + final_dir = run_dir(config) / "final" + else: + if run_dir(config).exists(): + # Stale or partial output from an earlier run must not mix with the new one. + shutil.rmtree(run_dir(config)) + log(f"seed {config.seed}: removed earlier output in {run_dir(config)}") + log(f"seed {config.seed}: training {config.run_name}") + final_dir = train_fn(config) + evaluate_fn(config, final_dir) + log(f"seed {config.seed}: evaluated, logits archived") + if config.seed not in KEEP_WEIGHTS_SEEDS and run_dir(config).exists(): + shutil.rmtree(run_dir(config)) + log(f"seed {config.seed}: deleted weights {run_dir(config)} (only seed 42 is kept)") + + +def check_setup(record: dict[str, object]) -> None: + config, training = record["config"], record["training"] + assert isinstance(config, dict) and isinstance(training, dict) + test_rows = record["metrics"]["test"]["raw"]["n"] # type: ignore[index] + problems = [] + if config["model_name"] != EXPECTED_MODEL: + problems.append(f"model {config['model_name']} != {EXPECTED_MODEL}") + if config["per_intent"] is not None or config["eval_per_intent"] is not None: + problems.append("data was subsampled") + if training["train_rows"] != EXPECTED_TRAIN_ROWS: + problems.append(f"train_rows {training['train_rows']} != {EXPECTED_TRAIN_ROWS}") + if training["oos_train_rows"] != EXPECTED_OOS_TRAIN_ROWS: + problems.append(f"oos_train_rows {training['oos_train_rows']} != 250") + if test_rows != EXPECTED_TEST_ROWS: + problems.append(f"test rows {test_rows} != {EXPECTED_TEST_ROWS}") + if problems: + raise SetupError(f"{record['run_name']} is not the AC2 setup: {'; '.join(problems)}") + + +def judge(records: dict[int, dict[str, object]]) -> dict: + """PASS iff every seed in SEEDS has test in-scope accuracy >= THRESHOLD.""" + missing = [s for s in SEEDS if s not in records] + if missing: + raise SetupError(f"no results for seeds {missing}") + per_seed = {} + for seed in SEEDS: + record = records[seed] + check_setup(record) + test = record["metrics"]["test"]["raw"] # type: ignore[index] + accuracy = float(test["in_scope_accuracy_150"]) + per_seed[str(seed)] = { + "run_name": record["run_name"], + "in_scope_accuracy_150": accuracy, + "oos_recall_151": float(test["oos_recall_151"]), + "accuracy_8": float(test["accuracy_8"]), + "oos_recall_8": float(test["oos_recall_8"]), + "passed": accuracy >= THRESHOLD, + } + passed = all(entry["passed"] for entry in per_seed.values()) + return { + "criterion": "AC2: test in_scope_accuracy_150 >= threshold for every seed", + "split": "test", + "threshold": THRESHOLD, + "seeds": per_seed, + "verdict": "PASS" if passed else "FAIL", + } + + +def run_ac2( + base: RunConfig, + train_fn: TrainFn, + evaluate_fn: EvaluateFn, + force: bool = False, + log: Log = print, +) -> dict: + from tinyrouter.evaluate import RunPaths + + records: dict[int, dict[str, object]] = {} + for seed in SEEDS: + config = base.with_seed(seed) + run_seed(config, train_fn, evaluate_fn, force, log) + path = RunPaths.of(config).results_json + records[seed] = json.loads(path.read_text(encoding="utf-8")) + result = judge(records) + result["weights_kept"] = {str(s): run_dir(base.with_seed(s)).exists() for s in SEEDS} + out = Path(base.results_root) / "ac2.json" + out.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") + log(f"wrote {out}: {result['verdict']}") + return result + + +def default_train(config: RunConfig) -> Path: + from tinyrouter.train import prepare_train_split, train + + return train(config, prepare_train_split(config), run_dir(config)) + + +def default_evaluate(config: RunConfig, model_dir: Path) -> dict[str, object]: + from tinyrouter.evaluate import evaluate + + return evaluate(config, model_dir) + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", default="configs/bert-base.yaml") + parser.add_argument("--force", action="store_true", help="retrain and re-score every seed") + args = parser.parse_args(argv) + result = run_ac2(load_config(args.config), default_train, default_evaluate, args.force) + for seed, entry in result["seeds"].items(): + mark = "pass" if entry["passed"] else "FAIL" + print(f" seed {seed}: in-scope {entry['in_scope_accuracy_150']:.4f} {mark}") + if result["verdict"] != "PASS": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/src/tinyrouter/archive.py b/src/tinyrouter/archive.py new file mode 100644 index 0000000..8e83097 --- /dev/null +++ b/src/tinyrouter/archive.py @@ -0,0 +1,265 @@ +"""Per-example logits archive: the only input the analysis steps read (docs/PLAN.md AC3). + +One ``.npz`` per run holds validation and test logits, gold intent ids, +and a JSON metadata block. Logits are stored as float32, the dtype the +model produced them in, so nothing downstream (temperature fit, ECE, +threshold sweeps, risk-coverage ranking) is computed on rounded numbers. + +Archives live in ``results/logits/`` and are not committed. Their SHA-256 +is recorded in ``results/logits-manifest.json``, which is committed, so a +copy downloaded from a GitHub Release can be checked byte for byte against +what the run produced. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path + +import numpy as np + +from tinyrouter.calibrate import SplitLogits +from tinyrouter.data import SplitName, sha256_of +from tinyrouter.labels import load_label_space + +FORMAT_VERSION = 1 +SPLITS: tuple[SplitName, ...] = ("validation", "test") +LOGITS_DTYPE = np.float32 +MANIFEST_NAME = "logits-manifest.json" + +# Every key must be present. Values may be None only where listed in NULLABLE. +REQUIRED_METADATA: dict[str, type | tuple[type, ...]] = { + "format_version": int, + "run_name": str, + "model_name": str, + "model_revision": str, + "seed": int, + "per_intent": int, + "train_rows": int, + "oos_train_rows": int, + "eval_per_intent": int, + "dataset_revision": str, + "label_space_sha256": str, + "git_commit": str, + "git_dirty": bool, + "created_at": str, +} +NULLABLE = {"per_intent", "eval_per_intent"} + + +class ArchiveError(ValueError): + """A logits archive is incomplete, inconsistent, or does not match the manifest.""" + + +@dataclass(frozen=True) +class LogitsArchive: + metadata: dict[str, object] + splits: dict[str, SplitLogits] + + @property + def validation(self) -> SplitLogits: + return self.splits["validation"] + + @property + def test(self) -> SplitLogits: + return self.splits["test"] + + +def git_state(repo_dir: Path | None = None) -> tuple[str, bool]: + """Current commit and whether the working tree has uncommitted changes. + + Returns ``("unknown", True)`` outside a git checkout, so the archive + still says it cannot be traced to a commit instead of omitting the key. + """ + cwd = repo_dir or Path(__file__).resolve().parent + try: + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], cwd=cwd, capture_output=True, text=True, check=True + ).stdout.strip() + status = subprocess.run( + ["git", "status", "--porcelain", "--untracked-files=no"], + cwd=cwd, + capture_output=True, + text=True, + check=True, + ).stdout + except (OSError, subprocess.CalledProcessError): + return "unknown", True + return commit, bool(status.strip()) + + +def utc_now() -> str: + return datetime.now(UTC).isoformat(timespec="seconds") + + +def validate_metadata(metadata: dict[str, object]) -> None: + missing = sorted(set(REQUIRED_METADATA) - set(metadata)) + if missing: + raise ArchiveError(f"metadata is missing {missing}") + for key, expected in REQUIRED_METADATA.items(): + value = metadata[key] + if value is None and key in NULLABLE: + continue + # bool is a subclass of int; do not let True pass as a seed. + wrong_bool = isinstance(value, bool) and expected is not bool + if wrong_bool or not isinstance(value, expected): + raise ArchiveError(f"metadata[{key!r}] = {value!r} is not {expected}") + if isinstance(value, str) and not value.strip(): + raise ArchiveError(f"metadata[{key!r}] is empty") + if metadata["format_version"] != FORMAT_VERSION: + raise ArchiveError( + f"format_version {metadata['format_version']} is not supported (expected " + f"{FORMAT_VERSION})" + ) + expected_hash = load_label_space().sha256 + if metadata["label_space_sha256"] != expected_hash: + raise ArchiveError( + "label space fingerprint differs from resources/*.json; the logit columns in this " + "archive may not mean what the current code assumes" + ) + + +def validate_split(split: SplitLogits) -> None: + num_intents = load_label_space().num_intents + logits, labels = split.logits, split.labels + if logits.shape[1] != num_intents: + raise ArchiveError( + f"{split.split}: {logits.shape[1]} logit columns, expected {num_intents}" + ) + if logits.dtype != LOGITS_DTYPE: + raise ArchiveError(f"{split.split}: logits are {logits.dtype}, expected {LOGITS_DTYPE}") + if not np.issubdtype(labels.dtype, np.integer): + raise ArchiveError(f"{split.split}: labels are {labels.dtype}, expected integers") + if labels.size and (labels.min() < 0 or labels.max() >= num_intents): + raise ArchiveError(f"{split.split}: label ids outside [0, {num_intents})") + if not np.isfinite(logits).all(): + raise ArchiveError(f"{split.split}: logits contain NaN or inf") + + +def save_logits(path: Path, splits: dict[str, SplitLogits], metadata: dict[str, object]) -> Path: + """Validate and write an archive atomically (a crash never leaves a half-written file).""" + if set(splits) != set(SPLITS): + raise ArchiveError(f"need exactly the splits {SPLITS}, got {sorted(splits)}") + arrays: dict[str, np.ndarray] = {} + for name in SPLITS: + split = splits[name] + if split.split != name: + raise ArchiveError(f"splits[{name!r}] holds the '{split.split}' split") + stored = SplitLogits(name, split.logits.astype(LOGITS_DTYPE), split.labels.astype(np.int64)) + validate_split(stored) + arrays[f"{name}_logits"] = stored.logits + arrays[f"{name}_labels"] = stored.labels + validate_metadata(metadata) + arrays["metadata"] = np.array(json.dumps(metadata, sort_keys=True)) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + with tmp.open("wb") as fh: + np.savez_compressed(fh, **arrays) + os.replace(tmp, path) + return path + + +def load_logits(path: Path) -> LogitsArchive: + """Read an archive back, refusing it if metadata or shapes are incomplete or inconsistent.""" + expected_keys = {"metadata"} | {f"{s}_{part}" for s in SPLITS for part in ("logits", "labels")} + with np.load(path, allow_pickle=False) as npz: + missing = sorted(expected_keys - set(npz.files)) + if missing: + raise ArchiveError(f"{path}: arrays missing {missing}") + metadata = json.loads(str(npz["metadata"])) + arrays = {key: npz[key] for key in expected_keys - {"metadata"}} + if not isinstance(metadata, dict): + raise ArchiveError(f"{path}: metadata is not a JSON object") + validate_metadata(metadata) + splits: dict[str, SplitLogits] = {} + for name in SPLITS: + logits, labels = arrays[f"{name}_logits"], arrays[f"{name}_labels"] + if logits.ndim != 2 or labels.ndim != 1 or logits.shape[0] != labels.shape[0]: + raise ArchiveError( + f"{path}: {name} logits {logits.shape} and labels {labels.shape} do not line up" + ) + split = SplitLogits(name, logits, labels) + validate_split(split) + splits[name] = split + return LogitsArchive(metadata=metadata, splits=splits) + + +def read_manifest(manifest_path: Path) -> dict[str, dict[str, object]]: + if not manifest_path.exists(): + return {} + raw = json.loads(manifest_path.read_text(encoding="utf-8")) + return dict(raw["files"]) + + +def record_in_manifest(manifest_path: Path, archive_path: Path) -> dict[str, object]: + """Add or replace the entry for ``archive_path``; keys are file names relative to its dir.""" + archive = load_logits(archive_path) + entry = { + "sha256": sha256_of(archive_path), + "bytes": archive_path.stat().st_size, + "run_name": archive.metadata["run_name"], + "git_commit": archive.metadata["git_commit"], + "created_at": archive.metadata["created_at"], + "rows": {name: int(split.labels.shape[0]) for name, split in archive.splits.items()}, + } + files = read_manifest(manifest_path) + files[archive_path.name] = entry + body = {"format_version": FORMAT_VERSION, "files": dict(sorted(files.items()))} + manifest_path.parent.mkdir(parents=True, exist_ok=True) + tmp = manifest_path.with_name(manifest_path.name + ".tmp") + tmp.write_text(json.dumps(body, indent=2) + "\n", encoding="utf-8") + os.replace(tmp, manifest_path) + return entry + + +def check_against_manifest(manifest_path: Path, archive_path: Path) -> None: + """Raise unless the manifest lists ``archive_path`` with this exact SHA-256.""" + entry = read_manifest(manifest_path).get(archive_path.name) + if entry is None: + raise ArchiveError(f"{archive_path.name} is not listed in {manifest_path}") + actual = sha256_of(archive_path) + if actual != entry["sha256"]: + raise ArchiveError( + f"{archive_path.name}: sha256 {actual} != manifest {entry['sha256']}; the file " + "changed after it was recorded, or it is a different run's archive" + ) + + +def verify_all(manifest_path: Path, logits_dir: Path) -> list[str]: + """Problems found checking every manifest entry against the files in ``logits_dir``.""" + problems = [] + for name in read_manifest(manifest_path): + path = logits_dir / name + if not path.exists(): + problems.append(f"{name}: listed in the manifest but not in {logits_dir}") + continue + try: + check_against_manifest(manifest_path, path) + load_logits(path) + except ArchiveError as exc: + problems.append(str(exc)) + return problems + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description="Verify logits archives against the manifest.") + parser.add_argument("--results-root", default="results") + args = parser.parse_args(argv) + root = Path(args.results_root) + manifest = root / MANIFEST_NAME + problems = verify_all(manifest, root / "logits") + for problem in problems: + print(f"FAIL {problem}") + count = len(read_manifest(manifest)) + if problems: + raise SystemExit(1) + print(f"OK {count} archive(s) match {manifest}") + + +if __name__ == "__main__": + main() diff --git a/src/tinyrouter/efficiency.py b/src/tinyrouter/efficiency.py new file mode 100644 index 0000000..b93648d --- /dev/null +++ b/src/tinyrouter/efficiency.py @@ -0,0 +1,118 @@ +"""Training cost numbers for the AC5 efficiency table: parameters, wall time, peak memory. + +Peak memory is measured differently per device, and the numbers are not +comparable across devices. Each result names what was measured: + +- cuda: ``torch.cuda.max_memory_allocated`` after ``reset_peak_memory_stats``. + The allocator tracks the true peak of tensor memory, including + activations inside a step. +- mps: torch 2.14 has no peak counter for MPS, so the callback samples + ``torch.mps.driver_allocated_memory()`` (everything Metal holds for the + process, cached allocator pools included) and + ``torch.mps.current_allocated_memory()`` (live tensors only) after each + backward pass and each optimizer step, and keeps the maxima. The driver + number is the headline: freed activation blocks stay in the pool, so it + approximates the high-water mark from above. The tensor number is + sampled between passes and misses activations that are freed before the + sample, so it is a lower bound. +- cpu: ``resource.getrusage(RUSAGE_SELF).ru_maxrss``, the process's peak + resident set size since start, which also counts everything loaded + before training (Python, torch, the dataset). ru_maxrss is bytes on + macOS and kilobytes on Linux; ``peak_rss_bytes`` converts. +""" + +from __future__ import annotations + +import sys +from collections.abc import Callable +from dataclasses import dataclass, field + +Probe = Callable[[], dict[str, int]] + + +def count_parameters(model: object) -> dict[str, int]: + params = list(model.parameters()) # type: ignore[attr-defined] + return { + "total": int(sum(p.numel() for p in params)), + "trainable": int(sum(p.numel() for p in params if p.requires_grad)), + } + + +def peak_rss_bytes(ru_maxrss: int, platform: str = sys.platform) -> int: + return ru_maxrss if platform == "darwin" else ru_maxrss * 1024 + + +def mps_probe() -> dict[str, int]: + import torch + + return { + "driver_allocated_bytes": int(torch.mps.driver_allocated_memory()), + "tensor_allocated_bytes": int(torch.mps.current_allocated_memory()), + } + + +def cpu_probe() -> dict[str, int]: + import resource + + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return {"process_peak_rss_bytes": peak_rss_bytes(rss)} + + +def cuda_probe() -> dict[str, int]: + import torch + + return {"max_memory_allocated_bytes": int(torch.cuda.max_memory_allocated())} + + +MEASURES = { + "mps": "sampled after each backward and optimizer step: torch.mps.driver_allocated_memory " + "(Metal driver, includes allocator cache; headline) and current_allocated_memory " + "(live tensors; lower bound)", + "cuda": "torch.cuda.max_memory_allocated after reset_peak_memory_stats (tensor peak)", + "cpu": "resource.getrusage ru_maxrss: process peak RSS since start, not training only", +} +PROBES: dict[str, Probe] = {"mps": mps_probe, "cuda": cuda_probe, "cpu": cpu_probe} + + +@dataclass +class MemorySampler: + """Keeps the per-key maximum of whatever ``probe`` returns each time ``sample`` runs.""" + + probe: Probe + peaks: dict[str, int] = field(default_factory=dict) + samples: int = 0 + + def sample(self) -> None: + for key, value in self.probe().items(): + self.peaks[key] = max(self.peaks.get(key, 0), int(value)) + self.samples += 1 + + +def start_measurement(device: str) -> MemorySampler: + if device == "cuda": + import torch + + torch.cuda.reset_peak_memory_stats() + return MemorySampler(PROBES[device]) + + +def memory_report(device: str, sampler: MemorySampler) -> dict[str, object]: + measured = {"device": device, "measures": MEASURES[device], "samples": sampler.samples} + return {**measured, **sampler.peaks} + + +def sampler_callback(sampler: MemorySampler) -> object: + """HF Trainer callback that samples after backward and after each optimizer step.""" + from transformers import TrainerCallback + + class SampleMemory(TrainerCallback): + def on_train_begin(self, args, state, control, **kwargs): # noqa: ANN001 + sampler.sample() + + def on_pre_optimizer_step(self, args, state, control, **kwargs): # noqa: ANN001 + sampler.sample() + + def on_step_end(self, args, state, control, **kwargs): # noqa: ANN001 + sampler.sample() + + return SampleMemory() diff --git a/src/tinyrouter/evaluate.py b/src/tinyrouter/evaluate.py index f8a443d..d57d5ba 100644 --- a/src/tinyrouter/evaluate.py +++ b/src/tinyrouter/evaluate.py @@ -1,9 +1,15 @@ """Score a trained model on validation and test, raw and temperature-scaled. The temperature is fitted on validation logits only (see calibrate.py) and -then applied unchanged to test. Results go to ``results/.json`` -with the config, versions, and metrics, so tables are rebuilt from JSON, -never typed by hand. +then applied unchanged to test. Each run writes three things under +``results_root``: + +- ``runs/.json``: config, versions, training cost, metrics + (committed; tables are rebuilt from these, never typed by hand); +- ``logits/.npz``: per-example validation and test logits + (not committed, see archive.py); +- an entry in ``logits-manifest.json`` with that archive's SHA-256 + (committed). """ from __future__ import annotations @@ -11,17 +17,41 @@ import argparse import json import platform -from dataclasses import asdict +from dataclasses import asdict, dataclass from pathlib import Path import numpy as np +from tinyrouter.archive import ( + MANIFEST_NAME, + git_state, + load_logits, + record_in_manifest, + save_logits, + utc_now, +) from tinyrouter.calibrate import SplitLogits, apply_temperature, fit_temperature from tinyrouter.config import RunConfig, load_config from tinyrouter.data import DATASET_REVISION, Split, SplitName, load_split, subsample_per_intent from tinyrouter.labels import load_label_space from tinyrouter.metrics import routing_metrics, softmax -from tinyrouter.train import pick_device +from tinyrouter.train import SUMMARY_NAME, pick_device + + +@dataclass(frozen=True) +class RunPaths: + results_json: Path + logits: Path + manifest: Path + + @classmethod + def of(cls, config: RunConfig) -> RunPaths: + root = Path(config.results_root) + return cls( + results_json=root / "runs" / f"{config.run_name}.json", + logits=root / "logits" / f"{config.run_name}.npz", + manifest=root / MANIFEST_NAME, + ) def predict_logits(model_dir: Path, split: Split, config: RunConfig) -> SplitLogits: @@ -78,17 +108,56 @@ def eval_split(name: SplitName, config: RunConfig) -> Split: return subsample_per_intent(load_split(name), config.eval_per_intent, config.seed) -def evaluate(config: RunConfig, model_dir: Path, results_path: Path) -> dict[str, object]: +def read_training_summary(model_dir: Path) -> dict[str, object]: + path = model_dir / SUMMARY_NAME + if not path.exists(): + raise FileNotFoundError( + f"{path} not found; evaluate needs the summary train() writes next to the weights " + "(training set size goes into the logits metadata)" + ) + return json.loads(path.read_text(encoding="utf-8")) + + +def archive_metadata(config: RunConfig, training: dict[str, object]) -> dict[str, object]: + commit, dirty = git_state() + return { + "format_version": 1, + "run_name": config.run_name, + "model_name": config.model_name, + "model_revision": config.model_revision, + "seed": config.seed, + "per_intent": config.per_intent, + "train_rows": training["train_rows"], + "oos_train_rows": training["oos_train_rows"], + "eval_per_intent": config.eval_per_intent, + "dataset_revision": DATASET_REVISION, + "label_space_sha256": load_label_space().sha256, + "git_commit": commit, + "git_dirty": dirty, + "created_at": utc_now(), + } + + +def evaluate(config: RunConfig, model_dir: Path) -> dict[str, object]: + """Score the model, archive its logits, and write the results JSON; return the record.""" + paths = RunPaths.of(config) + training = read_training_summary(model_dir) val = predict_logits(model_dir, eval_split("validation", config), config) test = predict_logits(model_dir, eval_split("test", config), config) + save_logits(paths.logits, {"validation": val, "test": test}, archive_metadata(config, training)) + entry = record_in_manifest(paths.manifest, paths.logits) + # Metrics come from the archive as written, so they match what later analysis reads. + archived = load_logits(paths.logits) record = { "run_name": config.run_name, "config": asdict(config), "environment": environment(), - "metrics": score(val, test), + "training": training, + "logits": {"file": paths.logits.name, "sha256": entry["sha256"], "bytes": entry["bytes"]}, + "metrics": score(archived.validation, archived.test), } - results_path.parent.mkdir(parents=True, exist_ok=True) - results_path.write_text(json.dumps(record, indent=2) + "\n", encoding="utf-8") + paths.results_json.parent.mkdir(parents=True, exist_ok=True) + paths.results_json.write_text(json.dumps(record, indent=2) + "\n", encoding="utf-8") return record @@ -102,9 +171,9 @@ def main(argv: list[str] | None = None) -> None: if args.seed is not None: config = config.with_seed(args.seed) model_dir = Path(args.model_dir or Path(config.checkpoint_root) / config.run_name / "final") - results_path = Path(config.results_root) / f"{config.run_name}.json" - evaluate(config, model_dir, results_path) - print(f"wrote {results_path}") + evaluate(config, model_dir) + paths = RunPaths.of(config) + print(f"wrote {paths.results_json} and {paths.logits}") if __name__ == "__main__": diff --git a/src/tinyrouter/labels.py b/src/tinyrouter/labels.py index 4fd32ec..1a5a66f 100644 --- a/src/tinyrouter/labels.py +++ b/src/tinyrouter/labels.py @@ -9,6 +9,7 @@ from __future__ import annotations +import hashlib import json from dataclasses import dataclass from functools import cache @@ -58,6 +59,19 @@ def agents_of(self, intent_ids: np.ndarray) -> np.ndarray: """Map integer intent ids (any shape) to integer agent ids.""" return self.intent_to_agent_id[np.asarray(intent_ids, dtype=np.int64)] + @property + def sha256(self) -> str: + """Fingerprint of intent order plus the intent-to-agent map. + + Two archives with the same fingerprint agree on what every logit + column means and which agent it rolls up to. + """ + canonical = json.dumps( + [[name, self.intent_to_agent[name]] for name in self.intent_names], + separators=(",", ":"), + ) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + def aggregate_probs(self, intent_probs: np.ndarray) -> np.ndarray: """Sum 151-way probabilities into 8-way agent probabilities. diff --git a/src/tinyrouter/report.py b/src/tinyrouter/report.py index 5c4b0a0..cd7aa64 100644 --- a/src/tinyrouter/report.py +++ b/src/tinyrouter/report.py @@ -1,4 +1,4 @@ -"""Build results/summary.md from results/*.json. Minimal for now. +"""Build results/summary.md from results/runs/*.json. Minimal for now. Tables are generated, never typed by hand (docs/PLAN.md AC7). With no results yet this says so and exits 0; the README section it will feed is @@ -31,9 +31,10 @@ def main(argv: list[str] | None = None) -> None: parser.add_argument("--results-dir", default="results") args = parser.parse_args(argv) results_dir = Path(args.results_dir) - paths = sorted(results_dir.glob("*.json")) if results_dir.is_dir() else [] + runs_dir = results_dir / "runs" + paths = sorted(runs_dir.glob("*.json")) if runs_dir.is_dir() else [] if not paths: - print(f"no results in {results_dir}/ yet; run `make train evaluate` first") + print(f"no results in {runs_dir}/ yet; run `make train evaluate` first") return records = [json.loads(p.read_text(encoding="utf-8")) for p in paths] out = results_dir / "summary.md" diff --git a/src/tinyrouter/smoke.py b/src/tinyrouter/smoke.py index 701319d..ac4b17d 100644 --- a/src/tinyrouter/smoke.py +++ b/src/tinyrouter/smoke.py @@ -1,7 +1,8 @@ """End-to-end smoke run: real data, tiny random model, one training step, full scoring path. Checks that every piece connects (download and checksum, label mapping, -training, logits, temperature fit, metrics, JSON output). The numbers it +training, logits, logits archive and manifest, temperature fit, metrics, +JSON output). The numbers it prints are meaningless; the model is random and sees about 150 rows. Everything is written to a temporary directory and removed afterwards. """ @@ -15,8 +16,9 @@ from dataclasses import replace from pathlib import Path +from tinyrouter.archive import check_against_manifest, load_logits from tinyrouter.config import load_config -from tinyrouter.evaluate import evaluate +from tinyrouter.evaluate import RunPaths, evaluate from tinyrouter.train import prepare_train_split, train @@ -33,7 +35,12 @@ def main(argv: list[str] | None = None) -> None: ) train_split = prepare_train_split(config) final_dir = train(config, train_split, Path(config.checkpoint_root) / config.run_name) - record = evaluate(config, final_dir, Path(config.results_root) / "smoke.json") + record = evaluate(config, final_dir) + paths = RunPaths.of(config) + check_against_manifest(paths.manifest, paths.logits) + archive = load_logits(paths.logits) + logits_bytes = paths.logits.stat().st_size + logits_cells = sum(s.logits.size for s in archive.splits.values()) metrics = record["metrics"] assert isinstance(metrics, dict) temperature = float(metrics["temperature"]) @@ -52,6 +59,11 @@ def main(argv: list[str] | None = None) -> None: f"smoke OK: train_rows={len(train_split)} val_rows={val_n} test_rows={test_n} " f"T={temperature:.3f} elapsed={time.monotonic() - started:.1f}s" ) + print( + f"smoke logits archive: {logits_bytes} bytes for {logits_cells} float32 logits " + f"({logits_bytes / logits_cells:.2f} bytes each), sha256 in manifest, load_logits OK; " + f"peak memory {record['training']['peak_memory']}" + ) if __name__ == "__main__": diff --git a/src/tinyrouter/train.py b/src/tinyrouter/train.py index 35b094a..ddad6f4 100644 --- a/src/tinyrouter/train.py +++ b/src/tinyrouter/train.py @@ -10,12 +10,21 @@ import argparse import json import shutil +import time from pathlib import Path from tinyrouter.config import RunConfig, load_config from tinyrouter.data import Split, load_split, subsample_per_intent +from tinyrouter.efficiency import ( + count_parameters, + memory_report, + sampler_callback, + start_measurement, +) from tinyrouter.labels import LabelSpace, load_label_space +SUMMARY_NAME = "train_summary.json" + def pick_device(requested: str = "auto") -> str: """Resolve 'auto' to cuda, then mps, then cpu; validate an explicit choice.""" @@ -88,7 +97,12 @@ def training_arguments(config: RunConfig, output_dir: Path, device: str) -> obje def train(config: RunConfig, train_split: Split, output_dir: Path) -> Path: - """Train on ``train_split`` and return the directory holding the final weights.""" + """Train on ``train_split`` and return the directory holding the final weights. + + ``final/train_summary.json`` records the training set size, parameter + counts, wall time, steps and peak memory; evaluation copies it into the + results JSON, because the weights (and this file) may be deleted later. + """ from transformers import DataCollatorWithPadding, Trainer, set_seed if train_split.name != "train": @@ -97,32 +111,36 @@ def train(config: RunConfig, train_split: Split, output_dir: Path) -> Path: labels = load_label_space() device = pick_device(config.device) model, tokenizer = load_model_and_tokenizer(config, labels) + sampler = start_measurement(device) trainer = Trainer( model=model, # type: ignore[arg-type] args=training_arguments(config, output_dir, device), # type: ignore[arg-type] train_dataset=to_hf_dataset(train_split, tokenizer, config.max_length), # type: ignore[arg-type] data_collator=DataCollatorWithPadding(tokenizer), # type: ignore[arg-type] processing_class=tokenizer, # type: ignore[arg-type] + callbacks=[sampler_callback(sampler)], # type: ignore[list-item] ) + started = time.perf_counter() result = trainer.train() + wall_seconds = time.perf_counter() - started + sampler.sample() final_dir = output_dir / "final" trainer.save_model(str(final_dir)) for leftover in output_dir.glob("checkpoint-*"): shutil.rmtree(leftover) - (final_dir / "train_summary.json").write_text( - json.dumps( - { - "run_name": config.run_name, - "device": device, - "train_rows": len(train_split), - "global_step": result.global_step, - "train_loss": result.training_loss, - }, - indent=2, - ) - + "\n", - encoding="utf-8", - ) + summary = { + "run_name": config.run_name, + "device": device, + "train_rows": len(train_split), + "oos_train_rows": int((train_split.intents == labels.oos_intent_id).sum()), + "per_intent": config.per_intent, + "global_step": result.global_step, + "train_loss": result.training_loss, + "parameters": count_parameters(model), + "train_wall_seconds": round(wall_seconds, 3), + "peak_memory": memory_report(device, sampler), + } + (final_dir / SUMMARY_NAME).write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") return final_dir diff --git a/tests/archive_fakes.py b/tests/archive_fakes.py new file mode 100644 index 0000000..fbe48c6 --- /dev/null +++ b/tests/archive_fakes.py @@ -0,0 +1,41 @@ +"""Synthetic logits archives shared by the archive and AC2 tests.""" + +import numpy as np + +from tinyrouter.calibrate import SplitLogits +from tinyrouter.labels import load_label_space + +NUM_INTENTS = 151 + + +def fake_splits(seed: int = 0, n_val: int = 7, n_test: int = 11) -> dict[str, SplitLogits]: + rng = np.random.default_rng(seed) + return { + name: SplitLogits( + name, + rng.normal(scale=4.0, size=(n, NUM_INTENTS)).astype(np.float32), + rng.integers(0, NUM_INTENTS, size=n), + ) + for name, n in (("validation", n_val), ("test", n_test)) + } + + +def fake_metadata(**overrides: object) -> dict[str, object]: + meta: dict[str, object] = { + "format_version": 1, + "run_name": "bert-base-uncased-full-seed42", + "model_name": "google-bert/bert-base-uncased", + "model_revision": "86b5e0934494bd15c9632b12f734a8a67f723594", + "seed": 42, + "per_intent": None, + "train_rows": 15_250, + "oos_train_rows": 250, + "eval_per_intent": None, + "dataset_revision": "155b9c710419136e17307b80d0a13e68cd46b4ec", + "label_space_sha256": load_label_space().sha256, + "git_commit": "0" * 40, + "git_dirty": False, + "created_at": "2026-09-23T00:00:00+00:00", + } + meta.update(overrides) + return meta diff --git a/tests/test_ac2.py b/tests/test_ac2.py new file mode 100644 index 0000000..6be062a --- /dev/null +++ b/tests/test_ac2.py @@ -0,0 +1,174 @@ +import json +from pathlib import Path + +import numpy as np +import pytest + +from archive_fakes import fake_metadata, fake_splits +from tinyrouter import ac2 +from tinyrouter.ac2 import SEEDS, THRESHOLD, SetupError, judge, run_ac2 +from tinyrouter.archive import record_in_manifest, save_logits +from tinyrouter.config import RunConfig +from tinyrouter.evaluate import RunPaths + + +def record(seed: int, accuracy: float, **config_overrides: object) -> dict[str, object]: + config = {"model_name": ac2.EXPECTED_MODEL, "per_intent": None, "eval_per_intent": None} + config.update(config_overrides) + test = { + "in_scope_accuracy_150": accuracy, + "oos_recall_151": 0.5, + "accuracy_8": 0.9, + "oos_recall_8": 0.5, + "n": ac2.EXPECTED_TEST_ROWS, + } + return { + "run_name": f"bert-base-uncased-full-seed{seed}", + "config": config, + "training": {"train_rows": 15_250, "oos_train_rows": 250}, + "metrics": {"test": {"raw": test}}, + } + + +def records(*accuracies: float) -> dict[int, dict[str, object]]: + return {seed: record(seed, acc) for seed, acc in zip(SEEDS, accuracies, strict=True)} + + +def test_accuracy_exactly_at_the_threshold_passes(): + # 957 of 1000 correct, computed the way metrics.accuracy computes it. + exact = float(np.mean(np.arange(1000) < 957)) + assert exact == THRESHOLD + assert judge(records(exact, exact, exact))["verdict"] == "PASS" + + +def test_accuracy_just_below_the_threshold_fails(): + below = 4306 / 4500 # one query short of 95.7% on the 4,500 in-scope test rows + result = judge(records(below, below, below)) + assert result["verdict"] == "FAIL" + assert not any(entry["passed"] for entry in result["seeds"].values()) + + +def test_one_seed_below_fails_even_when_the_mean_is_above(): + result = judge(records(0.99, 0.95, 0.99)) + assert result["verdict"] == "FAIL" + assert [result["seeds"][str(s)]["passed"] for s in SEEDS] == [True, False, True] + + +def test_all_seeds_above_pass_and_report_every_number(): + result = judge(records(0.97, 0.968, 0.971)) + assert result["verdict"] == "PASS" + assert result["threshold"] == THRESHOLD and result["split"] == "test" + assert set(result["seeds"]["42"]) >= {"in_scope_accuracy_150", "oos_recall_151", "accuracy_8"} + + +def test_missing_seed_is_an_error_not_a_pass(): + partial = records(0.99, 0.99, 0.99) + del partial[44] + with pytest.raises(SetupError, match="44"): + judge(partial) + + +@pytest.mark.parametrize( + "overrides", + [{"model_name": "answerdotai/ModernBERT-base"}, {"per_intent": 10}, {"eval_per_intent": 1}], +) +def test_a_record_from_a_different_setup_is_refused(overrides): + bad = records(0.99, 0.99, 0.99) + bad[43] = record(43, 0.99, **overrides) + with pytest.raises(SetupError, match="seed43"): + judge(bad) + + +class FakePipeline: + """Stands in for train() and evaluate(): writes weights, an archive and a results JSON.""" + + def __init__(self, accuracy: float = 0.97): + self.accuracy = accuracy + self.trained: list[int] = [] + self.evaluated: list[int] = [] + + def train(self, config: RunConfig) -> Path: + self.trained.append(config.seed) + final = Path(config.checkpoint_root) / config.run_name / "final" + final.mkdir(parents=True) + (final / "model.safetensors").write_bytes(b"weights") + (final / "train_summary.json").write_text("{}") + return final + + def evaluate(self, config: RunConfig, model_dir: Path) -> dict[str, object]: + self.evaluated.append(config.seed) + paths = RunPaths.of(config) + splits = fake_splits(seed=config.seed) + save_logits(paths.logits, splits, fake_metadata(seed=config.seed)) + record_in_manifest(paths.manifest, paths.logits) + rec = record(config.seed, self.accuracy) + paths.results_json.parent.mkdir(parents=True, exist_ok=True) + paths.results_json.write_text(json.dumps(rec)) + return rec + + +@pytest.fixture +def base(tmp_path) -> RunConfig: + return RunConfig( + model_name=ac2.EXPECTED_MODEL, + model_revision="r", + checkpoint_root=str(tmp_path / "ckpt"), + results_root=str(tmp_path / "results"), + ) + + +def test_run_trains_all_seeds_writes_ac2_json_and_keeps_only_seed_42_weights(base, tmp_path): + fake, logs = FakePipeline(), [] + result = run_ac2(base, fake.train, fake.evaluate, log=logs.append) + assert fake.trained == fake.evaluated == list(SEEDS) + assert result["weights_kept"] == {"42": True, "43": False, "44": False} + assert (tmp_path / "ckpt" / "bert-base-uncased-full-seed42" / "final").exists() + assert not (tmp_path / "ckpt" / "bert-base-uncased-full-seed43").exists() + assert sum("deleted weights" in line for line in logs) == 2 + on_disk = json.loads((tmp_path / "results" / "ac2.json").read_text()) + assert on_disk["verdict"] == "PASS" + + +def test_rerun_skips_seeds_already_archived(base): + fake = FakePipeline() + run_ac2(base, fake.train, fake.evaluate, log=lambda _: None) + again = FakePipeline() + run_ac2(base, again.train, again.evaluate, log=lambda _: None) + assert again.trained == again.evaluated == [] + + +def test_force_retrains_every_seed(base): + fake = FakePipeline() + run_ac2(base, fake.train, fake.evaluate, log=lambda _: None) + forced = FakePipeline() + run_ac2(base, forced.train, forced.evaluate, force=True, log=lambda _: None) + assert forced.trained == list(SEEDS) + + +def test_a_seed_whose_archive_no_longer_matches_the_manifest_is_rerun(base): + fake = FakePipeline() + run_ac2(base, fake.train, fake.evaluate, log=lambda _: None) + tampered = RunPaths.of(base.with_seed(43)).logits + tampered.write_bytes(tampered.read_bytes() + b"x") + again = FakePipeline() + run_ac2(base, again.train, again.evaluate, log=lambda _: None) + assert again.trained == [43] + + +def test_weights_left_by_a_crashed_evaluation_are_reused_not_retrained(base): + crashed = FakePipeline() + crashed.train(base.with_seed(42)) + fake = FakePipeline() + run_ac2(base, fake.train, fake.evaluate, log=lambda _: None) + assert fake.trained == [43, 44] + assert fake.evaluated == list(SEEDS) + + +def test_fail_verdict_makes_the_command_exit_nonzero(base, monkeypatch): + fake = FakePipeline(accuracy=0.9) + monkeypatch.setattr(ac2, "default_train", fake.train) + monkeypatch.setattr(ac2, "default_evaluate", fake.evaluate) + monkeypatch.setattr(ac2, "load_config", lambda _: base) + with pytest.raises(SystemExit) as exc: + ac2.main(["--config", "unused.yaml"]) + assert exc.value.code == 1 diff --git a/tests/test_archive.py b/tests/test_archive.py new file mode 100644 index 0000000..e39bea4 --- /dev/null +++ b/tests/test_archive.py @@ -0,0 +1,170 @@ +import json + +import numpy as np +import pytest + +from archive_fakes import NUM_INTENTS, fake_metadata, fake_splits +from tinyrouter.archive import ( + REQUIRED_METADATA, + ArchiveError, + check_against_manifest, + git_state, + load_logits, + read_manifest, + record_in_manifest, + save_logits, + verify_all, +) +from tinyrouter.calibrate import SplitLogits +from tinyrouter.labels import build_label_space, load_label_space + + +def test_round_trip_preserves_logits_labels_and_metadata_exactly(tmp_path): + splits, meta = fake_splits(), fake_metadata() + path = save_logits(tmp_path / "run.npz", splits, meta) + archive = load_logits(path) + assert archive.metadata == meta + for name in ("validation", "test"): + got = archive.splits[name] + assert got.split == name + assert got.logits.dtype == np.float32 + np.testing.assert_array_equal(got.logits, splits[name].logits) + np.testing.assert_array_equal(got.labels, splits[name].labels) + assert archive.validation.logits.shape == (7, NUM_INTENTS) + assert archive.test.logits.shape == (11, NUM_INTENTS) + + +def test_float64_logits_are_stored_as_float32(tmp_path): + splits = fake_splits() + wide = splits["test"].logits.astype(np.float64) + splits["test"] = SplitLogits("test", wide, splits["test"].labels) + archive = load_logits(save_logits(tmp_path / "r.npz", splits, fake_metadata())) + assert archive.test.logits.dtype == np.float32 + + +@pytest.mark.parametrize("key", sorted(REQUIRED_METADATA)) +def test_save_refuses_metadata_missing_any_required_key(tmp_path, key): + meta = fake_metadata() + del meta[key] + with pytest.raises(ArchiveError, match=key): + save_logits(tmp_path / "r.npz", fake_splits(), meta) + assert not (tmp_path / "r.npz").exists() + + +def rewrite_metadata(path, mutate) -> None: + with np.load(path, allow_pickle=False) as npz: + arrays = {k: npz[k] for k in npz.files} + meta = json.loads(str(arrays["metadata"])) + mutate(meta) + arrays["metadata"] = np.array(json.dumps(meta)) + with path.open("wb") as fh: + np.savez_compressed(fh, **arrays) + + +def test_load_refuses_an_archive_whose_metadata_lost_a_key(tmp_path): + path = save_logits(tmp_path / "r.npz", fake_splits(), fake_metadata()) + rewrite_metadata(path, lambda m: m.pop("model_revision")) + with pytest.raises(ArchiveError, match="model_revision"): + load_logits(path) + + +@pytest.mark.parametrize( + ("key", "value"), + [("seed", True), ("seed", "42"), ("git_commit", ""), ("oos_train_rows", None)], +) +def test_metadata_with_wrong_type_or_empty_value_is_refused(tmp_path, key, value): + with pytest.raises(ArchiveError, match=key): + save_logits(tmp_path / "r.npz", fake_splits(), fake_metadata(**{key: value})) + + +def test_metadata_for_a_different_label_space_is_refused(tmp_path): + other = build_label_space(["a", "oos"], {"a": "finance_agent", "oos": "oos"}).sha256 + with pytest.raises(ArchiveError, match="label space"): + save_logits(tmp_path / "r.npz", fake_splits(), fake_metadata(label_space_sha256=other)) + + +def test_label_space_fingerprint_changes_when_an_intent_moves_agent(): + labels = load_label_space() + mapping = dict(labels.intent_to_agent) + name = next(n for n, a in mapping.items() if a == "finance_agent") + mapping[name] = "travel_agent" + moved = build_label_space(list(labels.intent_names), mapping) + assert moved.sha256 != labels.sha256 + assert len(labels.sha256) == 64 + + +def test_load_refuses_logits_and_labels_of_different_length(tmp_path): + path = save_logits(tmp_path / "r.npz", fake_splits(), fake_metadata()) + with np.load(path, allow_pickle=False) as npz: + arrays = {k: npz[k] for k in npz.files} + arrays["test_labels"] = arrays["test_labels"][:-1] + with path.open("wb") as fh: + np.savez_compressed(fh, **arrays) + with pytest.raises(ArchiveError, match="line up"): + load_logits(path) + + +def test_save_refuses_wrong_column_count_nan_and_out_of_range_labels(tmp_path): + base = fake_splits() + narrow = dict(base, test=SplitLogits("test", base["test"].logits[:, :150], base["test"].labels)) + with pytest.raises(ArchiveError, match="columns"): + save_logits(tmp_path / "a.npz", narrow, fake_metadata()) + nan_logits = base["test"].logits.copy() + nan_logits[0, 0] = np.nan + with pytest.raises(ArchiveError, match="NaN"): + save_logits( + tmp_path / "b.npz", + dict(base, test=SplitLogits("test", nan_logits, base["test"].labels)), + fake_metadata(), + ) + bad_labels = base["test"].labels.copy() + bad_labels[0] = NUM_INTENTS + with pytest.raises(ArchiveError, match="label ids"): + save_logits( + tmp_path / "c.npz", + dict(base, test=SplitLogits("test", base["test"].logits, bad_labels)), + fake_metadata(), + ) + + +def test_save_refuses_mislabelled_or_missing_splits(tmp_path): + base = fake_splits() + swapped = {"validation": base["test"], "test": base["validation"]} + with pytest.raises(ArchiveError, match="holds the 'test' split"): + save_logits(tmp_path / "r.npz", swapped, fake_metadata()) + with pytest.raises(ArchiveError, match="exactly the splits"): + save_logits(tmp_path / "r.npz", {"test": base["test"]}, fake_metadata()) + + +def test_manifest_records_sha256_and_detects_a_changed_file(tmp_path): + manifest = tmp_path / "logits-manifest.json" + path = save_logits(tmp_path / "logits" / "run.npz", fake_splits(), fake_metadata()) + entry = record_in_manifest(manifest, path) + assert read_manifest(manifest)["run.npz"] == entry + assert entry["rows"] == {"validation": 7, "test": 11} + check_against_manifest(manifest, path) + assert verify_all(manifest, tmp_path / "logits") == [] + + save_logits(path, fake_splits(seed=1), fake_metadata()) + with pytest.raises(ArchiveError, match="sha256"): + check_against_manifest(manifest, path) + assert len(verify_all(manifest, tmp_path / "logits")) == 1 + + +def test_manifest_check_fails_for_an_unlisted_or_missing_file(tmp_path): + manifest = tmp_path / "logits-manifest.json" + a = save_logits(tmp_path / "logits" / "a.npz", fake_splits(), fake_metadata()) + b = save_logits(tmp_path / "logits" / "b.npz", fake_splits(seed=2), fake_metadata()) + record_in_manifest(manifest, a) + with pytest.raises(ArchiveError, match="not listed"): + check_against_manifest(manifest, b) + record_in_manifest(manifest, b) + assert sorted(read_manifest(manifest)) == ["a.npz", "b.npz"] + a.unlink() + assert verify_all(manifest, tmp_path / "logits") == [ + f"a.npz: listed in the manifest but not in {tmp_path / 'logits'}" + ] + + +def test_git_state_outside_a_repository_says_unknown(tmp_path): + assert git_state(tmp_path) == ("unknown", True) diff --git a/tests/test_efficiency.py b/tests/test_efficiency.py new file mode 100644 index 0000000..4e72e25 --- /dev/null +++ b/tests/test_efficiency.py @@ -0,0 +1,39 @@ +import pytest + +from tinyrouter.efficiency import ( + MEASURES, + MemorySampler, + count_parameters, + memory_report, + peak_rss_bytes, +) + + +def test_sampler_keeps_the_maximum_of_each_key_across_samples(): + readings = iter( + [ + {"driver_allocated_bytes": 100, "tensor_allocated_bytes": 50}, + {"driver_allocated_bytes": 300, "tensor_allocated_bytes": 20}, + {"driver_allocated_bytes": 200, "tensor_allocated_bytes": 80}, + ] + ) + sampler = MemorySampler(lambda: next(readings)) + for _ in range(3): + sampler.sample() + assert sampler.peaks == {"driver_allocated_bytes": 300, "tensor_allocated_bytes": 80} + report = memory_report("mps", sampler) + assert report["samples"] == 3 + assert report["measures"] == MEASURES["mps"] + + +@pytest.mark.parametrize(("platform", "expected"), [("darwin", 2048), ("linux", 2048 * 1024)]) +def test_ru_maxrss_unit_differs_between_macos_and_linux(platform, expected): + assert peak_rss_bytes(2048, platform) == expected + + +def test_count_parameters_separates_frozen_ones(): + import torch + + model = torch.nn.Sequential(torch.nn.Linear(3, 4), torch.nn.Linear(4, 2)) + model[0].weight.requires_grad_(False) + assert count_parameters(model) == {"total": 12 + 4 + 8 + 2, "trainable": 4 + 8 + 2} diff --git a/tests/test_train_smoke.py b/tests/test_train_smoke.py index 66932cc..5931f35 100644 --- a/tests/test_train_smoke.py +++ b/tests/test_train_smoke.py @@ -110,3 +110,45 @@ def test_smoke_config_file_uses_cpu_and_one_step(): config = load_config(Path(__file__).parent.parent / "configs" / "smoke.yaml") assert (config.device, config.max_steps, config.warmup_ratio) == ("cpu", 1, 0.0) assert replace(config, seed=7).run_name.endswith("seed7") + + +@pytest.mark.filterwarnings("ignore::tinyrouter.calibrate.TemperatureBoundWarning") +def test_train_then_evaluate_archives_logits_and_records_training_cost( + tiny_model_dir, tmp_path, monkeypatch +): + import json + + from tinyrouter import evaluate as evaluate_module + from tinyrouter.archive import check_against_manifest, load_logits + from tinyrouter.evaluate import RunPaths, evaluate + + config = replace(smoke_config(tiny_model_dir, tmp_path), results_root=str(tmp_path / "res")) + final_dir = train(config, synthetic_split("train"), tmp_path / "ckpt" / config.run_name) + + summary = json.loads((final_dir / "train_summary.json").read_text()) + assert summary["parameters"]["total"] == summary["parameters"]["trainable"] > 0 + assert summary["global_step"] == 2 and summary["train_wall_seconds"] > 0 + assert summary["oos_train_rows"] == 1 and summary["train_rows"] == 32 + memory = summary["peak_memory"] + assert memory["device"] == "cpu" and memory["process_peak_rss_bytes"] > 10_000_000 + assert memory["samples"] >= 3 + + monkeypatch.setattr(evaluate_module, "eval_split", lambda name, _: synthetic_split(name)) + record = evaluate(config, final_dir) + paths = RunPaths.of(config) + check_against_manifest(paths.manifest, paths.logits) + archive = load_logits(paths.logits) + assert archive.metadata["seed"] == config.seed + assert archive.metadata["oos_train_rows"] == 1 + assert archive.test.logits.shape == (32, 151) + assert record["training"] == summary + assert record["logits"]["file"] == f"{config.run_name}.npz" + on_disk = json.loads(paths.results_json.read_text()) + assert on_disk["metrics"]["test"]["raw"]["n"] == 32 + + +def test_evaluate_refuses_a_model_dir_without_a_training_summary(tmp_path): + from tinyrouter.evaluate import read_training_summary + + with pytest.raises(FileNotFoundError, match="train_summary.json"): + read_training_summary(tmp_path) From 010fa400be89746a22a0eebab2b66b5c5c9d1bee Mon Sep 17 00:00:00 2001 From: drewOrc <36374426+drewOrc@users.noreply.github.com> Date: Wed, 23 Sep 2026 11:10:58 +0800 Subject: [PATCH 2/4] Tie AC2 resume to the config and check run identity Review follow-up for the AC2 runner and the logits archive. - Resume only what the current config produced: results JSON and train_summary.json both record the config, and a seed counts as done only when the config matches (output paths aside) and the logits SHA-256 in the results JSON, the manifest and the file all agree. Weights trained with another config are cleared, not reused. - FORCE clears every seed's weights, results JSON, archive and manifest entry before training anything; evaluate removes the old results JSON before writing a new archive. - judge checks each record's run_name, config seed and training seed against its slot, and refuses two seeds sharing one archive. - ac2.json and the printout list validation numbers per seed; README and docstring say tuning uses validation only. - evaluate refuses a model whose id2label differs from intent_names. - git_dirty ignores results/; train_summary records the git state. - environment() records device, device name and deterministic mode. - Tests for seed propagation, logit row alignment, git_dirty, the crash and FORCE scenarios, and weights kept after a failed evaluation. --- README.md | 6 +- src/tinyrouter/ac2.py | 134 ++++++++++++++++++--------- src/tinyrouter/archive.py | 21 ++++- src/tinyrouter/config.py | 15 ++- src/tinyrouter/efficiency.py | 10 +- src/tinyrouter/evaluate.py | 50 +++++++++- src/tinyrouter/train.py | 7 ++ tests/test_ac2.py | 174 +++++++++++++++++++++++++++++------ tests/test_archive.py | 25 +++++ tests/test_train_smoke.py | 97 +++++++++++++++++++ 10 files changed, 456 insertions(+), 83 deletions(-) diff --git a/README.md b/README.md index 218c160..8ab7d93 100644 --- a/README.md +++ b/README.md @@ -28,11 +28,15 @@ Other targets: | `make test-network` | tests that download from the Hugging Face Hub (dataset checksums) | | `make train CONFIG=configs/bert-base.yaml SEED=42` | fine-tune; final weights in `checkpoints//final` | | `make evaluate CONFIG=... SEED=...` | score validation and test, write `results/runs/.json` and the logits archive | -| `make ac2` | bert-base-uncased, full data, seeds 42/43/44; writes `results/ac2.json`, exits 1 on FAIL; resumes, `FORCE=1` reruns | +| `make ac2` | bert-base-uncased, full data, seeds 42/43/44; writes `results/ac2.json`, exits 1 on FAIL; resumes only runs made with the current config, `FORCE=1` clears and reruns all three | | `make verify-logits` | check every archive in `results/logits/` against `results/logits-manifest.json` | | `make report` | build `results/summary.md` from `results/runs/*.json` | | `make clean-checkpoints` | delete all trained weights | +Run every target from the repository root: `checkpoint_root` and `results_root` in the configs are relative paths. + +**Tuning after an AC2 FAIL uses validation only.** `results/ac2.json` and the `make ac2` printout list each seed's validation in-scope accuracy, 8-way accuracy and OOS recall for that purpose; the test numbers are the verdict and are not looked at while choosing hyperparameters. + Put `HF_TOKEN` and `ANTHROPIC_API_KEY` in a `.env` (copy `.env.example`) if you need them; the Makefile passes it to `uv run`. Nothing in the tests or the smoke run calls a paid API. ## How it works diff --git a/src/tinyrouter/ac2.py b/src/tinyrouter/ac2.py index dd0d1b1..ee2b7d3 100644 --- a/src/tinyrouter/ac2.py +++ b/src/tinyrouter/ac2.py @@ -4,10 +4,19 @@ (docs/PLAN.md AC2; derived from Larson et al.'s 96.7%, not a claim of exact reproduction). Writes ``results/ac2.json`` and exits 1 on FAIL. -Resumable: a seed whose results JSON exists and whose logits archive -matches the manifest is not rerun, unless ``--force``. Trained weights -of every seed except 42 are deleted once that seed's logits are archived -(disk is tight; seed 42 is kept as the RQ7 ONNX fallback). +Tuning after a FAIL uses the validation numbers only; ``ac2.json`` and the +printout list them for that purpose. The test numbers are the verdict and +are not a tuning signal. + +Resumable. A seed is skipped only when its results JSON was produced by +the current config (output paths aside) and its logits archive has the +same SHA-256 in the results JSON, in the manifest and on disk. Trained +weights are reused only when their ``train_summary.json`` records the +current config. Anything else is cleared and redone. ``--force`` clears +every seed's outputs before training anything, so a crash halfway never +leaves old and new runs mixed. Weights of every seed except 42 are deleted +once that seed's logits are archived (disk is tight; seed 42 is kept as +the RQ7 ONNX fallback). """ from __future__ import annotations @@ -18,8 +27,10 @@ from collections.abc import Callable from pathlib import Path -from tinyrouter.archive import ArchiveError, check_against_manifest, load_logits +from tinyrouter.archive import ArchiveError, load_logits, read_manifest, remove_from_manifest from tinyrouter.config import RunConfig, load_config +from tinyrouter.data import sha256_of +from tinyrouter.evaluate import RunPaths SEEDS = (42, 43, 44) KEEP_WEIGHTS_SEEDS = frozenset({42}) @@ -28,6 +39,7 @@ EXPECTED_TRAIN_ROWS = 15_250 EXPECTED_OOS_TRAIN_ROWS = 250 EXPECTED_TEST_ROWS = 5_500 +REPORTED = ("in_scope_accuracy_150", "oos_recall_151", "accuracy_8", "oos_recall_8") TrainFn = Callable[[RunConfig], Path] EvaluateFn = Callable[[RunConfig, Path], dict[str, object]] @@ -35,7 +47,11 @@ class SetupError(ValueError): - """A results record is not the AC2 setup (wrong model, subsampled data, ...).""" + """A results record is not the AC2 setup (wrong model, seed, data, or a duplicate).""" + + +def expected_run_name(seed: int) -> str: + return f"bert-base-uncased-full-seed{seed}" def run_dir(config: RunConfig) -> Path: @@ -43,39 +59,57 @@ def run_dir(config: RunConfig) -> Path: def is_done(config: RunConfig) -> bool: - """Results JSON present and its logits archive loads and matches the manifest.""" - from tinyrouter.evaluate import RunPaths - + """This config's results JSON exists and its archive is intact and the one it was scored on.""" paths = RunPaths.of(config) if not (paths.results_json.exists() and paths.logits.exists()): return False + record = json.loads(paths.results_json.read_text(encoding="utf-8")) + if not config.matches(record.get("config")): + return False + recorded = record.get("logits", {}).get("sha256") + listed = read_manifest(paths.manifest).get(paths.logits.name, {}).get("sha256") + if not recorded or not recorded == listed == sha256_of(paths.logits): + return False try: - check_against_manifest(paths.manifest, paths.logits) - load_logits(paths.logits) + archive = load_logits(paths.logits) except ArchiveError: return False - return True + return archive.metadata["seed"] == config.seed -def has_final_weights(config: RunConfig) -> bool: +def reusable_weights(config: RunConfig) -> bool: + """Final weights exist and were trained with exactly this config.""" final = run_dir(config) / "final" - return (final / "model.safetensors").exists() and (final / "train_summary.json").exists() + summary = final / "train_summary.json" + if not ((final / "model.safetensors").exists() and summary.exists()): + return False + return config.matches(json.loads(summary.read_text(encoding="utf-8")).get("config")) + +def clear_outputs(config: RunConfig, log: Log) -> None: + """Remove this seed's weights, results JSON, logits archive and manifest entry.""" + paths = RunPaths.of(config) + removed = [str(p) for p in (paths.results_json, paths.logits) if p.exists()] + for name in removed: + Path(name).unlink() + if remove_from_manifest(paths.manifest, paths.logits.name): + removed.append(f"manifest entry {paths.logits.name}") + if run_dir(config).exists(): + shutil.rmtree(run_dir(config)) + removed.append(str(run_dir(config))) + if removed: + log(f"seed {config.seed}: cleared {', '.join(removed)}") -def run_seed( - config: RunConfig, train_fn: TrainFn, evaluate_fn: EvaluateFn, force: bool, log: Log -) -> None: - if is_done(config) and not force: - log(f"seed {config.seed}: results and logits already archived, skipping") + +def run_seed(config: RunConfig, train_fn: TrainFn, evaluate_fn: EvaluateFn, log: Log) -> None: + if is_done(config): + log(f"seed {config.seed}: results and logits for this config already archived, skipping") else: - if has_final_weights(config) and not force: - log(f"seed {config.seed}: reusing trained weights in {run_dir(config)}/final") + if reusable_weights(config): + log(f"seed {config.seed}: reusing weights trained with this config") final_dir = run_dir(config) / "final" else: - if run_dir(config).exists(): - # Stale or partial output from an earlier run must not mix with the new one. - shutil.rmtree(run_dir(config)) - log(f"seed {config.seed}: removed earlier output in {run_dir(config)}") + clear_outputs(config, log) log(f"seed {config.seed}: training {config.run_name}") final_dir = train_fn(config) evaluate_fn(config, final_dir) @@ -85,11 +119,15 @@ def run_seed( log(f"seed {config.seed}: deleted weights {run_dir(config)} (only seed 42 is kept)") -def check_setup(record: dict[str, object]) -> None: +def check_setup(record: dict[str, object], seed: int) -> None: config, training = record["config"], record["training"] assert isinstance(config, dict) and isinstance(training, dict) test_rows = record["metrics"]["test"]["raw"]["n"] # type: ignore[index] problems = [] + if record["run_name"] != expected_run_name(seed): + problems.append(f"run_name {record['run_name']} != {expected_run_name(seed)}") + if config["seed"] != seed or training.get("seed") != seed: + problems.append(f"config seed {config['seed']}, training seed {training.get('seed')}") if config["model_name"] != EXPECTED_MODEL: problems.append(f"model {config['model_name']} != {EXPECTED_MODEL}") if config["per_intent"] is not None or config["eval_per_intent"] is not None: @@ -101,7 +139,12 @@ def check_setup(record: dict[str, object]) -> None: if test_rows != EXPECTED_TEST_ROWS: problems.append(f"test rows {test_rows} != {EXPECTED_TEST_ROWS}") if problems: - raise SetupError(f"{record['run_name']} is not the AC2 setup: {'; '.join(problems)}") + raise SetupError(f"seed {seed} is not the AC2 setup: {'; '.join(problems)}") + + +def split_numbers(record: dict[str, object], split: str) -> dict[str, float]: + raw = record["metrics"][split]["raw"] # type: ignore[index] + return {key: float(raw[key]) for key in REPORTED} def judge(records: dict[int, dict[str, object]]) -> dict: @@ -109,25 +152,26 @@ def judge(records: dict[int, dict[str, object]]) -> dict: missing = [s for s in SEEDS if s not in records] if missing: raise SetupError(f"no results for seeds {missing}") + for seed in SEEDS: + check_setup(records[seed], seed) + shas = [records[s]["logits"]["sha256"] for s in SEEDS] # type: ignore[index] + if len(set(shas)) != len(SEEDS): + raise SetupError("two seeds point at the same logits archive; they are not separate runs") per_seed = {} for seed in SEEDS: - record = records[seed] - check_setup(record) - test = record["metrics"]["test"]["raw"] # type: ignore[index] - accuracy = float(test["in_scope_accuracy_150"]) + test = split_numbers(records[seed], "test") per_seed[str(seed)] = { - "run_name": record["run_name"], - "in_scope_accuracy_150": accuracy, - "oos_recall_151": float(test["oos_recall_151"]), - "accuracy_8": float(test["accuracy_8"]), - "oos_recall_8": float(test["oos_recall_8"]), - "passed": accuracy >= THRESHOLD, + "run_name": records[seed]["run_name"], + "validation": split_numbers(records[seed], "validation"), + "test": test, + "passed": test["in_scope_accuracy_150"] >= THRESHOLD, } passed = all(entry["passed"] for entry in per_seed.values()) return { "criterion": "AC2: test in_scope_accuracy_150 >= threshold for every seed", "split": "test", "threshold": THRESHOLD, + "tuning_note": "tune on the validation numbers only; test is the verdict", "seeds": per_seed, "verdict": "PASS" if passed else "FAIL", } @@ -140,12 +184,13 @@ def run_ac2( force: bool = False, log: Log = print, ) -> dict: - from tinyrouter.evaluate import RunPaths - + if force: + for seed in SEEDS: + clear_outputs(base.with_seed(seed), log) records: dict[int, dict[str, object]] = {} for seed in SEEDS: config = base.with_seed(seed) - run_seed(config, train_fn, evaluate_fn, force, log) + run_seed(config, train_fn, evaluate_fn, log) path = RunPaths.of(config).results_json records[seed] = json.loads(path.read_text(encoding="utf-8")) result = judge(records) @@ -171,12 +216,17 @@ def default_evaluate(config: RunConfig, model_dir: Path) -> dict[str, object]: def main(argv: list[str] | None = None) -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", default="configs/bert-base.yaml") - parser.add_argument("--force", action="store_true", help="retrain and re-score every seed") + parser.add_argument("--force", action="store_true", help="clear and rerun every seed") args = parser.parse_args(argv) result = run_ac2(load_config(args.config), default_train, default_evaluate, args.force) + print("seed val in-scope val OOS recall test in-scope verdict (test only)") for seed, entry in result["seeds"].items(): + val, test = entry["validation"], entry["test"] mark = "pass" if entry["passed"] else "FAIL" - print(f" seed {seed}: in-scope {entry['in_scope_accuracy_150']:.4f} {mark}") + print( + f"{seed:>4} {val['in_scope_accuracy_150']:.4f} {val['oos_recall_151']:.4f}" + f" {test['in_scope_accuracy_150']:.4f} {mark}" + ) if result["verdict"] != "PASS": raise SystemExit(1) diff --git a/src/tinyrouter/archive.py b/src/tinyrouter/archive.py index 8e83097..7749572 100644 --- a/src/tinyrouter/archive.py +++ b/src/tinyrouter/archive.py @@ -71,8 +71,10 @@ def test(self) -> SplitLogits: def git_state(repo_dir: Path | None = None) -> tuple[str, bool]: - """Current commit and whether the working tree has uncommitted changes. + """Current commit and whether tracked files outside ``results/`` have uncommitted changes. + ``results/`` is excluded because runs rewrite it themselves (the + manifest, results JSON); counting that would mark every later run dirty. Returns ``("unknown", True)`` outside a git checkout, so the archive still says it cannot be traced to a commit instead of omitting the key. """ @@ -82,7 +84,7 @@ def git_state(repo_dir: Path | None = None) -> tuple[str, bool]: ["git", "rev-parse", "HEAD"], cwd=cwd, capture_output=True, text=True, check=True ).stdout.strip() status = subprocess.run( - ["git", "status", "--porcelain", "--untracked-files=no"], + ["git", "status", "--porcelain", "--untracked-files=no", "--", ":(top,exclude)results"], cwd=cwd, capture_output=True, text=True, @@ -209,12 +211,25 @@ def record_in_manifest(manifest_path: Path, archive_path: Path) -> dict[str, obj } files = read_manifest(manifest_path) files[archive_path.name] = entry + write_manifest(manifest_path, files) + return entry + + +def write_manifest(manifest_path: Path, files: dict[str, dict[str, object]]) -> None: body = {"format_version": FORMAT_VERSION, "files": dict(sorted(files.items()))} manifest_path.parent.mkdir(parents=True, exist_ok=True) tmp = manifest_path.with_name(manifest_path.name + ".tmp") tmp.write_text(json.dumps(body, indent=2) + "\n", encoding="utf-8") os.replace(tmp, manifest_path) - return entry + + +def remove_from_manifest(manifest_path: Path, archive_name: str) -> bool: + """Drop one entry; returns whether it was there.""" + files = read_manifest(manifest_path) + if files.pop(archive_name, None) is None: + return False + write_manifest(manifest_path, files) + return True def check_against_manifest(manifest_path: Path, archive_path: Path) -> None: diff --git a/src/tinyrouter/config.py b/src/tinyrouter/config.py index ba1d89e..9620b7e 100644 --- a/src/tinyrouter/config.py +++ b/src/tinyrouter/config.py @@ -2,7 +2,7 @@ from __future__ import annotations -from dataclasses import dataclass, fields, replace +from dataclasses import asdict, dataclass, fields, replace from pathlib import Path from typing import Literal @@ -10,6 +10,9 @@ Device = Literal["auto", "cpu", "mps", "cuda"] +# Where output goes, not what it is: two configs differing only here produce the same run. +LOCATION_FIELDS = frozenset({"checkpoint_root", "results_root"}) + @dataclass(frozen=True) class RunConfig: @@ -46,6 +49,16 @@ def run_name(self) -> str: short = self.model_name.rstrip("/").split("/")[-1] return f"{short}-{size}-seed{self.seed}" + def identity(self) -> dict[str, object]: + """Every field that can change the trained model or its scores.""" + return {k: v for k, v in asdict(self).items() if k not in LOCATION_FIELDS} + + def matches(self, recorded: object) -> bool: + """Whether a config dict saved with earlier output describes this same run.""" + if not isinstance(recorded, dict): + return False + return {k: v for k, v in recorded.items() if k not in LOCATION_FIELDS} == self.identity() + def with_seed(self, seed: int) -> RunConfig: return replace(self, seed=seed) diff --git a/src/tinyrouter/efficiency.py b/src/tinyrouter/efficiency.py index b93648d..6572540 100644 --- a/src/tinyrouter/efficiency.py +++ b/src/tinyrouter/efficiency.py @@ -11,8 +11,10 @@ process, cached allocator pools included) and ``torch.mps.current_allocated_memory()`` (live tensors only) after each backward pass and each optimizer step, and keeps the maxima. The driver - number is the headline: freed activation blocks stay in the pool, so it - approximates the high-water mark from above. The tensor number is + number is the headline. Freed activation blocks usually stay in the + allocator pool, so it is usually at or above the true tensor peak, but + that is not guaranteed: the pool can release memory between samples, and + a short-lived peak can fall between two samples. The tensor number is sampled between passes and misses activations that are freed before the sample, so it is a lower bound. - cpu: ``resource.getrusage(RUSAGE_SELF).ru_maxrss``, the process's peak @@ -66,8 +68,8 @@ def cuda_probe() -> dict[str, int]: MEASURES = { "mps": "sampled after each backward and optimizer step: torch.mps.driver_allocated_memory " - "(Metal driver, includes allocator cache; headline) and current_allocated_memory " - "(live tensors; lower bound)", + "(Metal driver, includes allocator cache; headline, usually at or above the tensor " + "peak but not guaranteed) and current_allocated_memory (live tensors; lower bound)", "cuda": "torch.cuda.max_memory_allocated after reset_peak_memory_stats (tensor peak)", "cpu": "resource.getrusage ru_maxrss: process peak RSS since start, not training only", } diff --git a/src/tinyrouter/evaluate.py b/src/tinyrouter/evaluate.py index d57d5ba..377e603 100644 --- a/src/tinyrouter/evaluate.py +++ b/src/tinyrouter/evaluate.py @@ -17,6 +17,7 @@ import argparse import json import platform +import subprocess from dataclasses import asdict, dataclass from pathlib import Path @@ -61,6 +62,7 @@ def predict_logits(model_dir: Path, split: Split, config: RunConfig) -> SplitLog device = pick_device(config.device) tokenizer = AutoTokenizer.from_pretrained(model_dir) model = AutoModelForSequenceClassification.from_pretrained(model_dir).to(device).eval() + check_model_labels(model.config.id2label, model_dir) chunks: list[np.ndarray] = [] with torch.inference_mode(): for start in range(0, len(split), config.eval_batch_size): @@ -76,6 +78,26 @@ def predict_logits(model_dir: Path, split: Split, config: RunConfig) -> SplitLog return SplitLogits(split.name, np.concatenate(chunks), split.intents) +class LabelMismatchError(ValueError): + """The model's output columns are not the intents the current label space lists.""" + + +def check_model_labels(id2label: dict, model_dir: Path) -> None: + """Refuse a model whose id2label differs from resources/intent_names.json, in order.""" + names = load_label_space().intent_names + model_names = tuple(str(label) for _, label in sorted((int(k), v) for k, v in id2label.items())) + if model_names != names: + diff = next( + (i for i, (a, b) in enumerate(zip(model_names, names, strict=False)) if a != b), + min(len(model_names), len(names)), + ) + raise LabelMismatchError( + f"{model_dir}: model id2label ({len(model_names)} labels) differs from " + f"intent_names.json ({len(names)}) first at column {diff}; logit column i " + "would not mean intent i" + ) + + def score(val: SplitLogits, test: SplitLogits) -> dict[str, object]: """Metrics for both splits before and after temperature scaling fitted on ``val``.""" labels = load_label_space() @@ -91,15 +113,37 @@ def score(val: SplitLogits, test: SplitLogits) -> dict[str, object]: return out -def environment() -> dict[str, str]: +def device_name(device: str) -> str: + import torch + + if device == "cuda": + return torch.cuda.get_device_name() + if platform.system() == "Darwin": + try: + return subprocess.run( + ["sysctl", "-n", "machdep.cpu.brand_string"], + capture_output=True, + text=True, + check=True, + ).stdout.strip() + except (OSError, subprocess.CalledProcessError): + pass + return platform.processor() or platform.machine() + + +def environment(config: RunConfig) -> dict[str, object]: import torch import transformers + device = pick_device(config.device) return { "python": platform.python_version(), "torch": torch.__version__, "transformers": transformers.__version__, "platform": platform.platform(), + "device": device, + "device_name": device_name(device), + "deterministic_algorithms": torch.are_deterministic_algorithms_enabled(), "dataset_revision": DATASET_REVISION, } @@ -141,6 +185,8 @@ def archive_metadata(config: RunConfig, training: dict[str, object]) -> dict[str def evaluate(config: RunConfig, model_dir: Path) -> dict[str, object]: """Score the model, archive its logits, and write the results JSON; return the record.""" paths = RunPaths.of(config) + # A crash below must not leave the previous run's JSON next to this run's archive. + paths.results_json.unlink(missing_ok=True) training = read_training_summary(model_dir) val = predict_logits(model_dir, eval_split("validation", config), config) test = predict_logits(model_dir, eval_split("test", config), config) @@ -151,7 +197,7 @@ def evaluate(config: RunConfig, model_dir: Path) -> dict[str, object]: record = { "run_name": config.run_name, "config": asdict(config), - "environment": environment(), + "environment": environment(config), "training": training, "logits": {"file": paths.logits.name, "sha256": entry["sha256"], "bytes": entry["bytes"]}, "metrics": score(archived.validation, archived.test), diff --git a/src/tinyrouter/train.py b/src/tinyrouter/train.py index ddad6f4..bd71875 100644 --- a/src/tinyrouter/train.py +++ b/src/tinyrouter/train.py @@ -11,8 +11,10 @@ import json import shutil import time +from dataclasses import asdict from pathlib import Path +from tinyrouter.archive import git_state from tinyrouter.config import RunConfig, load_config from tinyrouter.data import Split, load_split, subsample_per_intent from tinyrouter.efficiency import ( @@ -128,8 +130,13 @@ def train(config: RunConfig, train_split: Split, output_dir: Path) -> Path: trainer.save_model(str(final_dir)) for leftover in output_dir.glob("checkpoint-*"): shutil.rmtree(leftover) + commit, dirty = git_state() summary = { "run_name": config.run_name, + "seed": config.seed, + "config": asdict(config), + "git_commit": commit, + "git_dirty": dirty, "device": device, "train_rows": len(train_split), "oos_train_rows": int((train_split.intents == labels.oos_intent_id).sum()), diff --git a/tests/test_ac2.py b/tests/test_ac2.py index 6be062a..55854d4 100644 --- a/tests/test_ac2.py +++ b/tests/test_ac2.py @@ -1,4 +1,5 @@ import json +from dataclasses import asdict, replace from pathlib import Path import numpy as np @@ -6,27 +7,36 @@ from archive_fakes import fake_metadata, fake_splits from tinyrouter import ac2 -from tinyrouter.ac2 import SEEDS, THRESHOLD, SetupError, judge, run_ac2 +from tinyrouter.ac2 import SEEDS, THRESHOLD, SetupError, is_done, judge, run_ac2 from tinyrouter.archive import record_in_manifest, save_logits from tinyrouter.config import RunConfig from tinyrouter.evaluate import RunPaths -def record(seed: int, accuracy: float, **config_overrides: object) -> dict[str, object]: - config = {"model_name": ac2.EXPECTED_MODEL, "per_intent": None, "eval_per_intent": None} - config.update(config_overrides) - test = { +def metrics(accuracy: float) -> dict[str, object]: + return { "in_scope_accuracy_150": accuracy, "oos_recall_151": 0.5, "accuracy_8": 0.9, "oos_recall_8": 0.5, "n": ac2.EXPECTED_TEST_ROWS, } + + +def record(seed: int, accuracy: float, **config_overrides: object) -> dict[str, object]: + config = { + "model_name": ac2.EXPECTED_MODEL, + "seed": seed, + "per_intent": None, + "eval_per_intent": None, + } + config.update(config_overrides) return { - "run_name": f"bert-base-uncased-full-seed{seed}", + "run_name": ac2.expected_run_name(seed), "config": config, - "training": {"train_rows": 15_250, "oos_train_rows": 250}, - "metrics": {"test": {"raw": test}}, + "training": {"seed": seed, "train_rows": 15_250, "oos_train_rows": 250}, + "logits": {"sha256": f"{seed:064d}"}, + "metrics": {"validation": {"raw": metrics(0.96)}, "test": {"raw": metrics(accuracy)}}, } @@ -54,11 +64,14 @@ def test_one_seed_below_fails_even_when_the_mean_is_above(): assert [result["seeds"][str(s)]["passed"] for s in SEEDS] == [True, False, True] -def test_all_seeds_above_pass_and_report_every_number(): +def test_verdict_lists_validation_numbers_for_tuning_next_to_test(): result = judge(records(0.97, 0.968, 0.971)) assert result["verdict"] == "PASS" assert result["threshold"] == THRESHOLD and result["split"] == "test" - assert set(result["seeds"]["42"]) >= {"in_scope_accuracy_150", "oos_recall_151", "accuracy_8"} + seed42 = result["seeds"]["42"] + assert seed42["validation"]["in_scope_accuracy_150"] == 0.96 + assert set(seed42["validation"]) == set(seed42["test"]) == set(ac2.REPORTED) + assert "validation" in result["tuning_note"] def test_missing_seed_is_an_error_not_a_pass(): @@ -75,38 +88,79 @@ def test_missing_seed_is_an_error_not_a_pass(): def test_a_record_from_a_different_setup_is_refused(overrides): bad = records(0.99, 0.99, 0.99) bad[43] = record(43, 0.99, **overrides) - with pytest.raises(SetupError, match="seed43"): + with pytest.raises(SetupError, match="seed 43"): + judge(bad) + + +def test_one_result_reused_for_all_three_seeds_is_refused(): + same = record(42, 0.99) + with pytest.raises(SetupError, match="seed 43"): + judge({42: same, 43: same, 44: same}) + + +def test_records_whose_seed_fields_disagree_with_their_slot_are_refused(): + bad = records(0.99, 0.99, 0.99) + bad[44]["training"]["seed"] = 42 # type: ignore[index] + with pytest.raises(SetupError, match="training seed 42"): + judge(bad) + + +def test_two_seeds_sharing_one_logits_archive_are_refused(): + bad = records(0.99, 0.99, 0.99) + bad[44]["logits"] = bad[43]["logits"] + with pytest.raises(SetupError, match="same logits archive"): + judge(bad) + + +@pytest.mark.parametrize(("key", "value"), [("train_rows", 15_249), ("oos_train_rows", 100)]) +def test_training_set_size_other_than_full_oos_plus_is_refused(key, value): + bad = records(0.99, 0.99, 0.99) + bad[42]["training"][key] = value # type: ignore[index] + with pytest.raises(SetupError, match=key): judge(bad) class FakePipeline: - """Stands in for train() and evaluate(): writes weights, an archive and a results JSON.""" + """Stands in for train() and evaluate() with the same on-disk contract.""" - def __init__(self, accuracy: float = 0.97): + def __init__(self, accuracy: float = 0.97, fail_train: int | None = None): self.accuracy = accuracy + self.fail_train = fail_train + self.fail_evaluate: int | None = None self.trained: list[int] = [] self.evaluated: list[int] = [] def train(self, config: RunConfig) -> Path: + if config.seed == self.fail_train: + raise RuntimeError("simulated crash during training") self.trained.append(config.seed) final = Path(config.checkpoint_root) / config.run_name / "final" final.mkdir(parents=True) (final / "model.safetensors").write_bytes(b"weights") - (final / "train_summary.json").write_text("{}") + summary = {"seed": config.seed, "config": asdict(config)} + (final / "train_summary.json").write_text(json.dumps(summary)) return final def evaluate(self, config: RunConfig, model_dir: Path) -> dict[str, object]: + if config.seed == self.fail_evaluate: + raise RuntimeError("simulated crash during evaluation") self.evaluated.append(config.seed) paths = RunPaths.of(config) - splits = fake_splits(seed=config.seed) - save_logits(paths.logits, splits, fake_metadata(seed=config.seed)) - record_in_manifest(paths.manifest, paths.logits) + paths.results_json.unlink(missing_ok=True) + save_logits(paths.logits, fake_splits(seed=config.seed), fake_metadata(seed=config.seed)) + entry = record_in_manifest(paths.manifest, paths.logits) rec = record(config.seed, self.accuracy) + rec["config"] = asdict(config) + rec["logits"] = {"sha256": entry["sha256"]} paths.results_json.parent.mkdir(parents=True, exist_ok=True) paths.results_json.write_text(json.dumps(rec)) return rec +def quiet(_: str) -> None: + pass + + @pytest.fixture def base(tmp_path) -> RunConfig: return RunConfig( @@ -130,41 +184,100 @@ def test_run_trains_all_seeds_writes_ac2_json_and_keeps_only_seed_42_weights(bas def test_rerun_skips_seeds_already_archived(base): - fake = FakePipeline() - run_ac2(base, fake.train, fake.evaluate, log=lambda _: None) + run_ac2(base, FakePipeline().train, FakePipeline().evaluate, log=quiet) again = FakePipeline() - run_ac2(base, again.train, again.evaluate, log=lambda _: None) + run_ac2(base, again.train, again.evaluate, log=quiet) assert again.trained == again.evaluated == [] -def test_force_retrains_every_seed(base): +def test_rerun_after_changing_a_hyperparameter_retrains_every_seed(base): + """S1: tuning after a FAIL must not be judged on the old config's results.""" fake = FakePipeline() - run_ac2(base, fake.train, fake.evaluate, log=lambda _: None) + run_ac2(base, fake.train, fake.evaluate, log=quiet) + tuned = replace(base, learning_rate=3e-5) + again = FakePipeline() + run_ac2(tuned, again.train, again.evaluate, log=quiet) + assert again.trained == list(SEEDS) + stored = json.loads(RunPaths.of(tuned.with_seed(42)).results_json.read_text()) + assert stored["config"]["learning_rate"] == 3e-5 + + +def test_weights_trained_with_another_config_are_not_reused(base): + FakePipeline().train(base.with_seed(42)) + fake = FakePipeline() + run_ac2(replace(base, num_train_epochs=3.0), fake.train, fake.evaluate, log=quiet) + assert fake.trained == list(SEEDS) + + +def test_changing_only_output_locations_does_not_count_as_a_new_config(base): + assert base.matches(asdict(replace(base, results_root="elsewhere", checkpoint_root="x"))) + assert not base.matches(asdict(replace(base, model_revision="other"))) + + +def test_force_retrains_every_seed(base): + run_ac2(base, FakePipeline().train, FakePipeline().evaluate, log=quiet) forced = FakePipeline() - run_ac2(base, forced.train, forced.evaluate, force=True, log=lambda _: None) + run_ac2(base, forced.train, forced.evaluate, force=True, log=quiet) assert forced.trained == list(SEEDS) +def test_force_run_that_crashes_leaves_no_old_results_to_mix_in(base): + """S2: after FORCE crashes at seed 43, a plain rerun must redo 43 and 44, not reuse old 44.""" + run_ac2(base, FakePipeline().train, FakePipeline().evaluate, log=quiet) + crashing = FakePipeline(fail_train=43) + with pytest.raises(RuntimeError, match="simulated"): + run_ac2(base, crashing.train, crashing.evaluate, force=True, log=quiet) + assert not RunPaths.of(base.with_seed(44)).results_json.exists() + assert ( + "bert-base-uncased-full-seed44.npz" + not in json.loads(RunPaths.of(base).manifest.read_text())["files"] + ) + resumed = FakePipeline() + run_ac2(base, resumed.train, resumed.evaluate, log=quiet) + assert resumed.trained == [43, 44] + + +def test_old_results_json_next_to_a_newer_archive_is_not_done(base): + """S3: evaluation wrote archive and manifest, then crashed before the results JSON.""" + run_ac2(base, FakePipeline().train, FakePipeline().evaluate, log=quiet) + config = base.with_seed(43) + paths = RunPaths.of(config) + assert is_done(config) + save_logits(paths.logits, fake_splits(seed=99), fake_metadata(seed=43)) + record_in_manifest(paths.manifest, paths.logits) + assert not is_done(config) + again = FakePipeline() + run_ac2(base, again.train, again.evaluate, log=quiet) + assert again.trained == [43] + + def test_a_seed_whose_archive_no_longer_matches_the_manifest_is_rerun(base): - fake = FakePipeline() - run_ac2(base, fake.train, fake.evaluate, log=lambda _: None) + run_ac2(base, FakePipeline().train, FakePipeline().evaluate, log=quiet) tampered = RunPaths.of(base.with_seed(43)).logits tampered.write_bytes(tampered.read_bytes() + b"x") again = FakePipeline() - run_ac2(base, again.train, again.evaluate, log=lambda _: None) + run_ac2(base, again.train, again.evaluate, log=quiet) assert again.trained == [43] def test_weights_left_by_a_crashed_evaluation_are_reused_not_retrained(base): - crashed = FakePipeline() - crashed.train(base.with_seed(42)) + FakePipeline().train(base.with_seed(42)) fake = FakePipeline() - run_ac2(base, fake.train, fake.evaluate, log=lambda _: None) + run_ac2(base, fake.train, fake.evaluate, log=quiet) assert fake.trained == [43, 44] assert fake.evaluated == list(SEEDS) -def test_fail_verdict_makes_the_command_exit_nonzero(base, monkeypatch): +def test_weights_survive_when_evaluation_of_that_seed_raises(base): + fake = FakePipeline() + fake.fail_evaluate = 43 + with pytest.raises(RuntimeError, match="evaluation"): + run_ac2(base, fake.train, fake.evaluate, log=quiet) + final = Path(base.checkpoint_root) / "bert-base-uncased-full-seed43" / "final" + assert (final / "model.safetensors").exists() + + +def test_fail_verdict_makes_the_command_exit_nonzero(base, monkeypatch, capsys): fake = FakePipeline(accuracy=0.9) monkeypatch.setattr(ac2, "default_train", fake.train) monkeypatch.setattr(ac2, "default_evaluate", fake.evaluate) @@ -172,3 +285,4 @@ def test_fail_verdict_makes_the_command_exit_nonzero(base, monkeypatch): with pytest.raises(SystemExit) as exc: ac2.main(["--config", "unused.yaml"]) assert exc.value.code == 1 + assert "val in-scope" in capsys.readouterr().out diff --git a/tests/test_archive.py b/tests/test_archive.py index e39bea4..b31b66b 100644 --- a/tests/test_archive.py +++ b/tests/test_archive.py @@ -168,3 +168,28 @@ def test_manifest_check_fails_for_an_unlisted_or_missing_file(tmp_path): def test_git_state_outside_a_repository_says_unknown(tmp_path): assert git_state(tmp_path) == ("unknown", True) + + +def test_git_state_ignores_changes_under_results_but_not_elsewhere(tmp_path): + import subprocess + + def git(*args): + subprocess.run( + ["git", "-c", "user.email=t@example.com", "-c", "user.name=t", *args], + cwd=tmp_path, + check=True, + capture_output=True, + ) + + (tmp_path / "results").mkdir() + (tmp_path / "results" / "logits-manifest.json").write_text("{}") + (tmp_path / "code.py").write_text("x = 1\n") + git("init", "-q") + git("add", ".") + git("commit", "-q", "-m", "init") + commit, dirty = git_state(tmp_path) + assert len(commit) == 40 and dirty is False + (tmp_path / "results" / "logits-manifest.json").write_text('{"files": {}}') + assert git_state(tmp_path)[1] is False + (tmp_path / "code.py").write_text("x = 2\n") + assert git_state(tmp_path)[1] is True diff --git a/tests/test_train_smoke.py b/tests/test_train_smoke.py index 5931f35..a50b18c 100644 --- a/tests/test_train_smoke.py +++ b/tests/test_train_smoke.py @@ -126,6 +126,8 @@ def test_train_then_evaluate_archives_logits_and_records_training_cost( final_dir = train(config, synthetic_split("train"), tmp_path / "ckpt" / config.run_name) summary = json.loads((final_dir / "train_summary.json").read_text()) + assert summary["seed"] == config.seed and config.matches(summary["config"]) + assert len(summary["git_commit"]) in (7, 40) or summary["git_commit"] == "unknown" assert summary["parameters"]["total"] == summary["parameters"]["trainable"] > 0 assert summary["global_step"] == 2 and summary["train_wall_seconds"] > 0 assert summary["oos_train_rows"] == 1 and summary["train_rows"] == 32 @@ -142,6 +144,9 @@ def test_train_then_evaluate_archives_logits_and_records_training_cost( assert archive.metadata["oos_train_rows"] == 1 assert archive.test.logits.shape == (32, 151) assert record["training"] == summary + env = record["environment"] + assert env["device"] == "cpu" and env["device_name"] + assert env["deterministic_algorithms"] in (True, False) assert record["logits"]["file"] == f"{config.run_name}.npz" on_disk = json.loads(paths.results_json.read_text()) assert on_disk["metrics"]["test"]["raw"]["n"] == 32 @@ -152,3 +157,95 @@ def test_evaluate_refuses_a_model_dir_without_a_training_summary(tmp_path): with pytest.raises(FileNotFoundError, match="train_summary.json"): read_training_summary(tmp_path) + + +def test_training_arguments_carry_the_config_seed_for_both_init_and_data_order(tmp_path): + config = replace(smoke_config(Path("unused"), tmp_path), seed=7) + args = __import__("tinyrouter.train", fromlist=["x"]).training_arguments( + config, tmp_path, "cpu" + ) + assert args.seed == args.data_seed == 7 + + +def test_global_seed_is_set_from_the_config_before_the_model_is_built(tmp_path, monkeypatch): + import torch + + import tinyrouter.train as train_module + + seen: list[int] = [] + + def capture(config, labels): + seen.append(torch.initial_seed()) + raise RuntimeError("stop after model construction point") + + monkeypatch.setattr(train_module, "load_model_and_tokenizer", capture) + config = replace(smoke_config(Path("unused"), tmp_path), seed=1234) + with pytest.raises(RuntimeError, match="stop"): + train(config, synthetic_split("train"), tmp_path / "x") + assert seen == [1234] + + +class StubTokenizer: + """Encodes 'q' as the single id N, so the model can echo it back.""" + + def __call__(self, texts, **_): + import torch + from transformers import BatchEncoding + + ids = torch.tensor([[int(t[1:])] for t in texts]) + return BatchEncoding({"input_ids": ids}) + + +class StubModel: + """Puts all logit mass on the column named by the input id: row i predicts intent N_i.""" + + def __init__(self): + from types import SimpleNamespace + + names = load_label_space().intent_names + self.config = SimpleNamespace(id2label=dict(enumerate(names))) + + def to(self, _): + return self + + def eval(self): + return self + + def __call__(self, input_ids): + from types import SimpleNamespace + + import torch + + logits = torch.full((input_ids.shape[0], 151), -5.0) + logits[torch.arange(input_ids.shape[0]), input_ids[:, 0]] = 5.0 + return SimpleNamespace(logits=logits) + + +def test_predicted_logit_rows_stay_aligned_with_their_labels_across_batches(tmp_path, monkeypatch): + import transformers + + monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", lambda *_: StubTokenizer()) + monkeypatch.setattr( + transformers.AutoModelForSequenceClassification, "from_pretrained", lambda *_: StubModel() + ) + intents = np.array([150, 3, 42, 7, 0, 99, 42, 11, 5, 120, 64], dtype=np.int64) + split = Split("test", tuple(f"q{i}" for i in intents), intents) + config = replace(smoke_config(Path("unused"), tmp_path), eval_batch_size=4) + out = predict_logits(tmp_path, split, config) + np.testing.assert_array_equal(out.logits.argmax(axis=1), out.labels) + np.testing.assert_array_equal(out.labels, intents) + + +def test_model_whose_id2label_differs_from_intent_names_is_refused(tiny_model_dir, tmp_path): + from tinyrouter.evaluate import LabelMismatchError, check_model_labels + + names = load_label_space().intent_names + check_model_labels(dict(enumerate(names)), tmp_path) + swapped = dict(enumerate(names)) + swapped[0], swapped[1] = swapped[1], swapped[0] + with pytest.raises(LabelMismatchError, match="column 0"): + check_model_labels(swapped, tmp_path) + # The fixture model was built without id2label (LABEL_0, LABEL_1, ...). + config = smoke_config(tiny_model_dir, tmp_path) + with pytest.raises(LabelMismatchError): + predict_logits(tiny_model_dir, synthetic_split("test"), config) From 4b40fe43dd73b2a75bfbf737d5c6579225965f13 Mon Sep 17 00:00:00 2001 From: drewOrc <36374426+drewOrc@users.noreply.github.com> Date: Wed, 23 Sep 2026 11:12:27 +0800 Subject: [PATCH 3/4] Test the run_name identity check on its own --- tests/test_ac2.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_ac2.py b/tests/test_ac2.py index 55854d4..da7d3ad 100644 --- a/tests/test_ac2.py +++ b/tests/test_ac2.py @@ -98,6 +98,13 @@ def test_one_result_reused_for_all_three_seeds_is_refused(): judge({42: same, 43: same, 44: same}) +def test_record_with_another_runs_name_is_refused_even_if_its_seeds_match(): + bad = records(0.99, 0.99, 0.99) + bad[43]["run_name"] = "bert-base-uncased-k10-seed43" + with pytest.raises(SetupError, match="run_name"): + judge(bad) + + def test_records_whose_seed_fields_disagree_with_their_slot_are_refused(): bad = records(0.99, 0.99, 0.99) bad[44]["training"]["seed"] = 42 # type: ignore[index] From e535792a449077a1bb9aee07c0cef02b4cf147ad Mon Sep 17 00:00:00 2001 From: drewOrc <36374426+drewOrc@users.noreply.github.com> Date: Wed, 23 Sep 2026 11:19:18 +0800 Subject: [PATCH 4/4] Clear stale AC2 verdict and record what produced it - run_ac2 deletes results/ac2.json before doing anything, so a run with a changed config (or FORCE) that fails partway leaves no old PASS. - ac2.json records the config (seed excluded) and each seed's logits SHA-256. - Tests for: no verdict after a failed changed-config run, manifest sha disagreeing with the results JSON, an archive carrying another seed, evaluate removing the old results JSON before predicting, and the training summary seed with a seed other than 42. - Document that FORCE deletes seed 42's kept weights and that only FORCE brings them back after clean-checkpoints; note the conservative config identity in RunConfig.identity(). --- Makefile | 4 ++++ README.md | 2 +- src/tinyrouter/ac2.py | 12 ++++++++-- src/tinyrouter/config.py | 8 ++++++- tests/test_ac2.py | 50 +++++++++++++++++++++++++++++++++++++++ tests/test_train_smoke.py | 31 ++++++++++++++++++++++-- 6 files changed, 101 insertions(+), 6 deletions(-) diff --git a/Makefile b/Makefile index 887af1f..f250813 100644 --- a/Makefile +++ b/Makefile @@ -57,6 +57,10 @@ evaluate: # AC2: bert-base-uncased, full data, seeds 42/43/44 -> results/ac2.json; exit 1 on FAIL. # Resumes: seeds already scored and archived are skipped. FORCE=1 reruns all. # Weights of seeds 43 and 44 are deleted after their logits are archived. +# FORCE=1 also deletes seed 42's kept weights before retraining. After +# `make clean-checkpoints`, seed 42's weights are not rebuilt by a plain +# `make ac2` (its results already exist); getting them back takes FORCE=1, +# which retrains all three seeds. ac2: uv run $(UV_ENV) python -m tinyrouter.ac2 --config configs/bert-base.yaml $(if $(filter 1,$(FORCE)),--force,) diff --git a/README.md b/README.md index 8ab7d93..e54f46c 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ Other targets: | `make test-network` | tests that download from the Hugging Face Hub (dataset checksums) | | `make train CONFIG=configs/bert-base.yaml SEED=42` | fine-tune; final weights in `checkpoints//final` | | `make evaluate CONFIG=... SEED=...` | score validation and test, write `results/runs/.json` and the logits archive | -| `make ac2` | bert-base-uncased, full data, seeds 42/43/44; writes `results/ac2.json`, exits 1 on FAIL; resumes only runs made with the current config, `FORCE=1` clears and reruns all three | +| `make ac2` | bert-base-uncased, full data, seeds 42/43/44; writes `results/ac2.json`, exits 1 on FAIL; resumes only runs made with the current config, `FORCE=1` clears and reruns all three, including deleting seed 42's kept weights; after `make clean-checkpoints` only `FORCE=1` brings seed 42's weights back | | `make verify-logits` | check every archive in `results/logits/` against `results/logits-manifest.json` | | `make report` | build `results/summary.md` from `results/runs/*.json` | | `make clean-checkpoints` | delete all trained weights | diff --git a/src/tinyrouter/ac2.py b/src/tinyrouter/ac2.py index ee2b7d3..dbfbc23 100644 --- a/src/tinyrouter/ac2.py +++ b/src/tinyrouter/ac2.py @@ -2,7 +2,9 @@ Pass means every seed's 150-way in-scope test accuracy is at least 95.7% (docs/PLAN.md AC2; derived from Larson et al.'s 96.7%, not a claim of -exact reproduction). Writes ``results/ac2.json`` and exits 1 on FAIL. +exact reproduction). Writes ``results/ac2.json`` (with the config and +each seed's logits SHA-256) and exits 1 on FAIL. Any earlier ``ac2.json`` +is deleted first, so a run that fails partway leaves no verdict behind. Tuning after a FAIL uses the validation numbers only; ``ac2.json`` and the printout list them for that purpose. The test numbers are the verdict and @@ -162,6 +164,7 @@ def judge(records: dict[int, dict[str, object]]) -> dict: test = split_numbers(records[seed], "test") per_seed[str(seed)] = { "run_name": records[seed]["run_name"], + "logits_sha256": records[seed]["logits"]["sha256"], # type: ignore[index] "validation": split_numbers(records[seed], "validation"), "test": test, "passed": test["in_scope_accuracy_150"] >= THRESHOLD, @@ -184,6 +187,11 @@ def run_ac2( force: bool = False, log: Log = print, ) -> dict: + out = Path(base.results_root) / "ac2.json" + if out.exists(): + # A verdict from an earlier config must not survive a run that fails partway. + out.unlink() + log(f"removed previous {out}") if force: for seed in SEEDS: clear_outputs(base.with_seed(seed), log) @@ -194,8 +202,8 @@ def run_ac2( path = RunPaths.of(config).results_json records[seed] = json.loads(path.read_text(encoding="utf-8")) result = judge(records) + result["config"] = {k: v for k, v in base.identity().items() if k != "seed"} result["weights_kept"] = {str(s): run_dir(base.with_seed(s)).exists() for s in SEEDS} - out = Path(base.results_root) / "ac2.json" out.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") log(f"wrote {out}: {result['verdict']}") return result diff --git a/src/tinyrouter/config.py b/src/tinyrouter/config.py index 9620b7e..dd56bc1 100644 --- a/src/tinyrouter/config.py +++ b/src/tinyrouter/config.py @@ -50,7 +50,13 @@ def run_name(self) -> str: return f"{short}-{size}-seed{self.seed}" def identity(self) -> dict[str, object]: - """Every field that can change the trained model or its scores.""" + """Every field that can change the trained model or its scores. + + Deliberately conservative: fields that only affect evaluation + (``eval_batch_size``, ``eval_per_intent``) also count, so changing + one of them retrains instead of just re-scoring. This could later be + split into a training identity and an evaluation identity. + """ return {k: v for k, v in asdict(self).items() if k not in LOCATION_FIELDS} def matches(self, recorded: object) -> bool: diff --git a/tests/test_ac2.py b/tests/test_ac2.py index da7d3ad..a3453eb 100644 --- a/tests/test_ac2.py +++ b/tests/test_ac2.py @@ -293,3 +293,53 @@ def test_fail_verdict_makes_the_command_exit_nonzero(base, monkeypatch, capsys): ac2.main(["--config", "unused.yaml"]) assert exc.value.code == 1 assert "val in-scope" in capsys.readouterr().out + + +def test_ac2_json_records_the_config_and_each_seeds_logits_sha(base): + result = run_ac2(base, FakePipeline().train, FakePipeline().evaluate, log=quiet) + on_disk = json.loads((Path(base.results_root) / "ac2.json").read_text()) + assert on_disk == json.loads(json.dumps(result)) + assert "seed" not in on_disk["config"] + assert on_disk["config"]["learning_rate"] == base.learning_rate + for seed in SEEDS: + stored = json.loads(RunPaths.of(base.with_seed(seed)).results_json.read_text()) + assert on_disk["seeds"][str(seed)]["logits_sha256"] == stored["logits"]["sha256"] + + +def test_changed_config_that_fails_partway_leaves_no_ac2_verdict(base): + run_ac2(base, FakePipeline().train, FakePipeline().evaluate, log=quiet) + verdict = Path(base.results_root) / "ac2.json" + assert json.loads(verdict.read_text())["verdict"] == "PASS" + crashing = FakePipeline(fail_train=43) + with pytest.raises(RuntimeError, match="simulated"): + run_ac2(replace(base, learning_rate=3e-5), crashing.train, crashing.evaluate, log=quiet) + assert not verdict.exists() + + +def rewrite_results_sha(config: RunConfig, sha: str) -> None: + path = RunPaths.of(config).results_json + rec = json.loads(path.read_text()) + rec["logits"]["sha256"] = sha + path.write_text(json.dumps(rec)) + + +def test_not_done_when_the_manifest_disagrees_with_json_and_file(base): + run_ac2(base, FakePipeline().train, FakePipeline().evaluate, log=quiet) + config = base.with_seed(43) + paths = RunPaths.of(config) + manifest = json.loads(paths.manifest.read_text()) + manifest["files"][paths.logits.name]["sha256"] = "f" * 64 + paths.manifest.write_text(json.dumps(manifest)) + assert not is_done(config) + + +def test_not_done_when_the_archive_belongs_to_another_seed(base): + from tinyrouter.data import sha256_of + + run_ac2(base, FakePipeline().train, FakePipeline().evaluate, log=quiet) + config = base.with_seed(43) + paths = RunPaths.of(config) + save_logits(paths.logits, fake_splits(seed=43), fake_metadata(seed=42)) + record_in_manifest(paths.manifest, paths.logits) + rewrite_results_sha(config, sha256_of(paths.logits)) + assert not is_done(config) diff --git a/tests/test_train_smoke.py b/tests/test_train_smoke.py index a50b18c..b32a823 100644 --- a/tests/test_train_smoke.py +++ b/tests/test_train_smoke.py @@ -122,11 +122,14 @@ def test_train_then_evaluate_archives_logits_and_records_training_cost( from tinyrouter.archive import check_against_manifest, load_logits from tinyrouter.evaluate import RunPaths, evaluate - config = replace(smoke_config(tiny_model_dir, tmp_path), results_root=str(tmp_path / "res")) + # A seed other than 42, so a summary that hardcodes 42 is caught. + config = replace( + smoke_config(tiny_model_dir, tmp_path), results_root=str(tmp_path / "res"), seed=43 + ) final_dir = train(config, synthetic_split("train"), tmp_path / "ckpt" / config.run_name) summary = json.loads((final_dir / "train_summary.json").read_text()) - assert summary["seed"] == config.seed and config.matches(summary["config"]) + assert summary["seed"] == 43 and config.matches(summary["config"]) assert len(summary["git_commit"]) in (7, 40) or summary["git_commit"] == "unknown" assert summary["parameters"]["total"] == summary["parameters"]["trainable"] > 0 assert summary["global_step"] == 2 and summary["train_wall_seconds"] > 0 @@ -249,3 +252,27 @@ def test_model_whose_id2label_differs_from_intent_names_is_refused(tiny_model_di config = smoke_config(tiny_model_dir, tmp_path) with pytest.raises(LabelMismatchError): predict_logits(tiny_model_dir, synthetic_split("test"), config) + + +def test_evaluate_removes_the_old_results_json_before_it_can_fail(tmp_path, monkeypatch): + import json + + from tinyrouter import evaluate as evaluate_module + from tinyrouter.evaluate import RunPaths, evaluate + + config = replace(smoke_config(Path("unused"), tmp_path), results_root=str(tmp_path / "res")) + old = RunPaths.of(config).results_json + old.parent.mkdir(parents=True) + old.write_text(json.dumps({"from": "an earlier run"})) + model_dir = tmp_path / "model" + model_dir.mkdir() + (model_dir / "train_summary.json").write_text('{"train_rows": 1, "oos_train_rows": 0}') + + def boom(*_): + raise RuntimeError("prediction failed") + + monkeypatch.setattr(evaluate_module, "predict_logits", boom) + monkeypatch.setattr(evaluate_module, "eval_split", lambda name, _: synthetic_split(name)) + with pytest.raises(RuntimeError, match="prediction failed"): + evaluate(config, model_dir) + assert not old.exists()