From 32ecb8d2bbc21bfbb7a33fdff542ea93c2569dd4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 13:12:39 +0000 Subject: [PATCH 1/4] feat(profiling): add a compile-time axis to campaign coverage The Profiling Agent could not see the compile-time corpus autolens_profiling is already producing. AGENTS.md listed "JAX compilation-time profiling of likelihood functions" under Future modes, _profiling.py had no occurrence of "compile", and all three modes read only results/runtime/ -- while scripts/misc/jax_compile/ held 93 committed probe records that nothing cross-referenced against the science grid. campaign --axis compile answers how much of the grid has compile data on a tier. Records are placed by their own (dataset_class, model_type, instrument) rather than by path, since results are filed under / and the path drops the class and instrument entirely. Tier mapping is deliberately NOT TIER_CONFIGS. That map keys off sweep config names which fold precision into the name (local_cpu_fp64 / local_cpu_mp), whereas a compile record carries a raw hardware string plus a separate mixed_precision bool; reusing it would mis-bucket every row. "other" is a real answer rather than a fallback -- the corpus holds RTX-2060 rows belonging to neither tier. Off-grid records (knn, delaunay_matern, the datacube_img* multi-band classes) and non-tier hardware get their own buckets: real measurements that are neither grid coverage nor noise. Malformed records are reported with file and index rather than skipped, surfacing the 4 that carry null hardware, class and instrument. The mode reports coverage only and never compares two timings. Compile timings are host-load-sensitive -- jax_compile/README.md records the first measurements being wrong by up to 7x (851s vs 117s for the same compile) because XLA compiles on the host cores -- so rows are comparable only within (hardware, jax_version, mixed_precision, cache state). Comparison waits on the pins in phase 2. ingest and triage reject --axis compile with exit 5 rather than ignoring it, so a compile flag can never silently return a runtime answer. The transform axis is read from probe.py's TRANSFORMS literal via the same ast route load_grid uses for CELLS, so the Brain cannot drift from the instrument. Adds tests/test_profiling_conductor.py -- profiling was the only conductor without a test file, so the "runtime axis unchanged" requirement had nothing to assert against. Hermetic: synthetic workspace fixtures, no real checkout. Runtime-axis output verified byte-identical against the real workspace. Phase 1 of 3; PyAutoMind draft/feature/profiling/. Closes #218 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01L4STU81pQP1GkMzZVvsMsv --- agents/conductors/profiling/AGENTS.md | 26 ++ agents/conductors/profiling/_profiling.py | 224 ++++++++++++++++- tests/test_profiling_conductor.py | 280 ++++++++++++++++++++++ 3 files changed, 528 insertions(+), 2 deletions(-) create mode 100644 tests/test_profiling_conductor.py diff --git a/agents/conductors/profiling/AGENTS.md b/agents/conductors/profiling/AGENTS.md index 4b14ccd..9d5d963 100644 --- a/agents/conductors/profiling/AGENTS.md +++ b/agents/conductors/profiling/AGENTS.md @@ -31,6 +31,32 @@ pyauto-brain profiling triage pyauto-brain profiling --json ``` +### The measurement axes + +`--axis runtime` (the default) is steady-state per-call cost, filed under +`results/runtime/` and bucketed by sweep **config** name (`local_cpu_fp64`, +`local_cpu_mp`, …). `--axis compile` is the one-off cost — trace, XLA compile, +first call — filed under `scripts/misc/jax_compile/results//` and +bucketed by **hardware**, with `mixed_precision` a separate field. The two +vocabularies do not interchange, so the compile axis maps tiers itself rather +than reusing `TIER_CONFIGS`. + +`--axis compile` currently serves `campaign` (coverage); `ingest` and `triage` +reject it with exit 5 until the compile pins land, so a compile flag can never +silently return a runtime answer. + +**Compile timings are host-load-sensitive** — the first measurements in +`jax_compile/README.md` were wrong by up to **7×** (851 s vs 117 s for the same +compile) purely from host load, because XLA compiles on the host cores. Rows are +comparable only within `(hardware, jax_version, mixed_precision, cache state)`. +`campaign --axis compile` therefore reports **coverage only** and never compares +two timings; comparison waits on the pins. + +Records whose cell is not in the sweep grid (`knn`, `delaunay_matern`, the +`datacube_img*` multi-band classes) are reported in an **off-grid** bucket, and +non-tier hardware in an **other-hardware** bucket. Both are real measurements — +neither counts as grid coverage, and neither is silently dropped. + ## Fundamental principles - **The classification is the result** for CPU-unusable cells (the usability diff --git a/agents/conductors/profiling/_profiling.py b/agents/conductors/profiling/_profiling.py index a37e2ec..9021a08 100644 --- a/agents/conductors/profiling/_profiling.py +++ b/agents/conductors/profiling/_profiling.py @@ -51,6 +51,13 @@ } DEFAULT_PER_RUN_TIMEOUT = 3600 +# Only used when jax_compile/probe.py cannot be read; the live list is taken +# from the instrument itself (see load_transforms). +FALLBACK_TRANSFORMS = ("jit", "grad", "vag", "vmap", "vmap_vag", "laxmap_vag", "pyloop_vag") + +# Fields a compile record needs before it can be placed on the grid at all. +COMPILE_KEY_FIELDS = ("hardware", "dataset_class", "model_type", "instrument") + def workspace_root(explicit: str | None = None) -> Path: if explicit: @@ -92,6 +99,64 @@ def load_tables(ws: Path) -> dict[str, Any]: } +def compile_dir(ws: Path) -> Path: + return ws / "scripts" / "misc" / "jax_compile" + + +def load_transforms(ws: Path) -> tuple[str, ...]: + """probe.py's transform axis, read from the workspace rather than copied. + + The list is a module-level literal in `jax_compile/probe.py`, so the same + ast route `load_grid` uses keeps the Brain from drifting out of sync with + the instrument. `FALLBACK_TRANSFORMS` only covers probe.py being absent.""" + got = _module_literal(compile_dir(ws) / "probe.py", "TRANSFORMS") + return tuple(got) if got else FALLBACK_TRANSFORMS + + +def load_compile_corpus(ws: Path) -> "list[tuple[str, int, dict[str, Any]]]": + """(relative-path, index-in-file, record) for every compile probe record. + + Each `results//.json` is an append-only LIST of flat + records. Unreadable or non-list files are skipped the way `ingest` already + skips malformed probe JSON — a corrupt file must not take the mode down.""" + root = compile_dir(ws) / "results" + out: list[tuple[str, int, dict[str, Any]]] = [] + if not root.is_dir(): + return out + for p in sorted(root.glob("*/*.json")): + try: + data = json.loads(p.read_text()) + except (OSError, ValueError): + continue + if not isinstance(data, list): + continue + rel = str(p.relative_to(root)) + for i, rec in enumerate(data): + if isinstance(rec, dict): + out.append((rel, i, rec)) + return out + + +def compile_tier_of(hardware: str | None) -> str: + """Which campaign tier a compile record belongs to. + + Deliberately NOT `TIER_CONFIGS`: that keys off sweep *config* names which + fold precision into the name (`local_cpu_fp64` / `local_cpu_mp`), whereas a + compile record carries a raw `hardware` string plus a SEPARATE + `mixed_precision` bool. Reusing the runtime map would mis-bucket every row. + + `other` is a real answer, not a fallback: the corpus holds RTX-2060 rows + that belong to neither tier, and folding them into one would misreport + coverage on hardware nobody asked about.""" + if not hardware: + return "other" + if hardware == "local_cpu": + return "local" + if "A100" in hardware: + return "a100" + return "other" + + # --------------------------------------------------------------------------- # campaign # --------------------------------------------------------------------------- @@ -157,6 +222,117 @@ def campaign(ws: Path, tier: str) -> dict[str, Any]: } +def campaign_compile(ws: Path, tier: str) -> dict[str, Any]: + """Coverage of the jax_compile corpus against the science grid. + + Answers "how much of the grid has compile data on this tier", which nothing + did before: the compile corpus and the runtime grid live in separate trees + with no cross-reference. Read-only, like every campaign — it never runs + probe.py, it emits the invocations a human would run. + """ + if tier not in TIER_CONFIGS: + return {"agent": "profiling", "mode": "campaign", "error": f"unknown tier {tier!r}"} + + grid = load_grid(ws) + transforms = load_transforms(ws) + on_grid = { + (cls, model, inst) + for cls, model, instruments in grid + for inst in (instruments or (None,)) + } + + # A record is placed by its OWN (dataset_class, model_type, instrument), not + # by its path: results are filed under /, which drops + # the class and the instrument entirely. + covered: set[tuple[tuple[str, str, str | None], str]] = set() + off_grid: dict[str, int] = {} + other_hw: dict[str, int] = {} + malformed: list[dict[str, Any]] = [] + + for rel, idx, rec in load_compile_corpus(ws): + if any(rec.get(f) in (None, "") for f in COMPILE_KEY_FIELDS): + malformed.append( + { + "record": f"{rel}[{idx}]", + "missing": [f for f in COMPILE_KEY_FIELDS if rec.get(f) in (None, "")], + "tag": rec.get("tag"), + } + ) + continue + rec_tier = compile_tier_of(rec.get("hardware")) + if rec_tier != tier: + if rec_tier == "other": + other_hw[str(rec["hardware"])] = other_hw.get(str(rec["hardware"]), 0) + 1 + continue + cell = (rec["dataset_class"], rec["model_type"], rec["instrument"]) + cell_id = "/".join(str(k) for k in cell) + if cell not in on_grid: + # Real measurements, not noise: knn / delaunay_matern are mesh + # variants from the Prodigy census, and the datacube_img* classes + # are the multi-band compile experiment. Reported, never counted as + # grid coverage and never silently dropped. + off_grid[cell_id] = off_grid.get(cell_id, 0) + 1 + continue + covered.add((cell, str(rec.get("transform")))) + + done: list[str] = [] + missing: list[str] = [] + missing_by_cell: dict[tuple[str, str, str | None], list[str]] = {} + for cls, model, instruments in grid: + for inst in instruments or (None,): + cell = (cls, model, inst) + cell_id = f"{cls}/{model}/{inst}" if inst else f"{cls}/{model}" + for tf in transforms: + run_id = f"{cell_id} [{tf}]" + if (cell, tf) in covered: + done.append(run_id) + else: + missing.append(run_id) + missing_by_cell.setdefault(cell, []).append(tf) + + dispatch: list[str] = [] + if tier == "local": + for (cls, model, inst), tfs in sorted(missing_by_cell.items(), key=lambda kv: str(kv[0])): + dispatch.append( + f"python3 scripts/misc/jax_compile/probe.py --dataset-class {cls} " + f"--model-type {model}" + + (f" --instrument {inst}" if inst else "") + + f" --transforms {','.join(tfs)} --cache-dir --tag " + ) + else: + submits = sorted(p.name for p in (compile_dir(ws) / "hpc").glob("submit_*")) or sorted( + p.name for p in (ws / "hpc" / "batch_gpu").glob("submit_*") + ) + dispatch = [f"sbatch hpc/batch_gpu/{s} (on the RAL checkout, post-pull)" for s in submits] + + return { + "agent": "profiling", + "mode": "campaign", + "axis": "compile", + "tier": tier, + "transforms": list(transforms), + "grid_cells": len(on_grid), + "runs_done": len(done), + "runs_missing": len(missing), + "missing": missing, + "off_grid": [{"cell": c, "records": n} for c, n in sorted(off_grid.items())], + "other_hardware": [{"hardware": h, "records": n} for h, n in sorted(other_hw.items())], + "malformed": malformed, + "policy": ( + "Compile timings are host-load-sensitive (the first measurements were " + "wrong by up to 7x from host load alone), so rows are only comparable " + "within (hardware, jax_version, mixed_precision, cache state). This " + "mode reports COVERAGE only — it never compares two timings." + ), + "dispatch_plan": dispatch, + "next_action": ( + "compile grid fully covered on this tier" + if not missing + else f"dispatch the {tier} compile plan ({len(missing)} cell/transform runs outstanding)" + ), + } + + # --------------------------------------------------------------------------- # ingest # --------------------------------------------------------------------------- @@ -288,7 +464,34 @@ def emit_human(d: dict[str, Any]) -> None: if d.get("error"): print(f"ERROR: {d['error']}") return - if d["mode"] == "campaign": + if d["mode"] == "campaign" and d.get("axis") == "compile": + print(f"Tier: {d['tier']}") + print(f"Transforms: {', '.join(d['transforms'])}") + print( + f"Cell/transform runs: {d['runs_done']} done · {d['runs_missing']} missing " + f"({d['grid_cells']} grid cells × {len(d['transforms'])} transforms)" + ) + for r in d["missing"][:10]: + print(f" missing: {r}") + if len(d["missing"]) > 10: + print(f" ... +{len(d['missing']) - 10} more") + if d["off_grid"]: + print("Off-grid records (real measurements, not grid cells):") + for o in d["off_grid"]: + print(f" {o['cell']}: {o['records']} record(s)") + if d["other_hardware"]: + print("Other hardware (neither tier):") + for o in d["other_hardware"]: + print(f" {o['hardware']}: {o['records']} record(s)") + if d["malformed"]: + print(f"Malformed records: {len(d['malformed'])}") + for m in d["malformed"][:10]: + print(f" {m['record']}: missing {', '.join(m['missing'])} (tag={m['tag']!r})") + print(f"Policy: {d['policy']}") + print("Dispatch plan:") + for s in d["dispatch_plan"]: + print(f" - {s}") + elif d["mode"] == "campaign": print(f"Tier: {d['tier']}") print( f"Runs: {d['runs_done']} done · " @@ -329,17 +532,34 @@ def main(argv=None) -> int: ap = argparse.ArgumentParser(prog="profiling") ap.add_argument("mode", nargs="?", default="campaign", choices=["campaign", "ingest", "triage"]) ap.add_argument("--tier", default="local", help="campaign tier: local | a100") + ap.add_argument( + "--axis", + default="runtime", + choices=["runtime", "compile"], + help="what is measured: runtime (steady-state per-call cost) | compile (jax_compile probe)", + ) ap.add_argument("--workspace", default=None, help="override the autolens_profiling path") ap.add_argument("--json", action="store_true", dest="as_json") a = ap.parse_args(argv) + # ingest/triage own the compile axis in later phases of the arc (pins, then + # drift classification). Refusing now is deliberate: a mode that silently + # ignored --axis would report runtime findings under a compile flag. + if a.axis == "compile" and a.mode != "campaign": + print( + f"profiling: --axis compile is not implemented for {a.mode!r} yet " + "(campaign only; ingest/triage land with the compile pins)", + file=sys.stderr, + ) + return 5 + ws = workspace_root(a.workspace) if not ws.is_dir(): print(f"profiling: workspace not found: {ws}", file=sys.stderr) return 4 if a.mode == "campaign": - d = campaign(ws, a.tier) + d = campaign_compile(ws, a.tier) if a.axis == "compile" else campaign(ws, a.tier) elif a.mode == "ingest": d = ingest(ws) else: diff --git a/tests/test_profiling_conductor.py b/tests/test_profiling_conductor.py new file mode 100644 index 0000000..884a19a --- /dev/null +++ b/tests/test_profiling_conductor.py @@ -0,0 +1,280 @@ +"""Contract tests for the profiling conductor's CLI footing. + +Hermetic: every test builds a synthetic `autolens_profiling` fixture in a temp +dir and passes it via `--workspace`, so the assertions never depend on the state +of the real checkout (whose corpus grows every campaign). + +Profiling was the only conductor without a test file when the compile axis was +added; the runtime-axis cases here exist to hold that surface still while the +compile axis grows beside it. +""" + +import json +import os +import subprocess +import sys +from pathlib import Path + +BRAIN_HOME = Path(__file__).resolve().parents[1] +BRAIN = BRAIN_HOME / "bin" / "pyauto-brain" + +# Two cells, deliberately asymmetric: one with instruments spanning two entries +# and one single-instrument cell, so cell-expansion bugs cannot hide. +FIXTURE_CELLS = """ +CELLS: list[tuple[str, str, tuple[str, ...]]] = [ + ("imaging", "mge", ("hst", "jwst")), + ("interferometer", "pixelization", ("sma",)), +] +""" + +FIXTURE_TRANSFORMS = 'TRANSFORMS = ("jit", "vag")\n' + + +def _run(args, workspace, root=None): + env = {**os.environ, "PYAUTO_ROOT": str(root or workspace.parent)} + return subprocess.run( + [sys.executable, str(BRAIN_HOME / "agents" / "conductors" / "profiling" / "_profiling.py"), *args, + "--workspace", str(workspace)], + capture_output=True, text=True, env=env, + ) + + +def _record(**kw): + base = { + "transform": "jit", + "trace_s": 1.0, + "compile_s": 2.0, + "first_s": 0.1, + "steady_s": 0.01, + "dataset_class": "imaging", + "model_type": "mge", + "instrument": "hst", + "hardware": "local_cpu", + "jax_version": "0.10.2", + "cache_dir": "", + "mixed_precision": False, + "tag": "t", + } + base.update(kw) + return base + + +def _workspace(tmp_path, records_by_file=None): + ws = tmp_path / "autolens_profiling" + lr = ws / "scripts" / "misc" / "likelihood_runtime" + lr.mkdir(parents=True) + (lr / "sweep.py").write_text(FIXTURE_CELLS) + + jc = ws / "scripts" / "misc" / "jax_compile" + jc.mkdir(parents=True) + (jc / "probe.py").write_text(FIXTURE_TRANSFORMS) + + for rel, records in (records_by_file or {}).items(): + p = jc / "results" / rel + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(json.dumps(records)) + + (ws / "results" / "runtime").mkdir(parents=True) + vram = ws / "scripts" / "misc" / "vram" + vram.mkdir(parents=True) + (vram / "config.py").write_text("VMAP_BATCH = {}\nVMAP_BATCH_SPARSE = {}\nPROVENANCE = {}\n") + (ws / "hpc" / "batch_gpu").mkdir(parents=True) + return ws + + +# --------------------------------------------------------------------------- +# compile axis +# --------------------------------------------------------------------------- + + +def test_compile_axis_counts_grid_coverage(tmp_path): + """Records land on the grid via their own fields, not their file path.""" + ws = _workspace(tmp_path, { + # Filed under mge.json but carrying an interferometer/pixelization cell: + # placing by path would put this on the wrong cell entirely. + "local_cpu/mge.json": [ + _record(transform="jit"), + _record(transform="vag"), + _record(dataset_class="interferometer", model_type="pixelization", + instrument="sma", transform="jit"), + ], + }) + r = _run(["campaign", "--axis", "compile", "--json"], ws) + assert r.returncode == 0, r.stderr + d = json.loads(r.stdout) + + assert d["axis"] == "compile" + assert d["grid_cells"] == 3 # (imaging,mge,hst/jwst) + (interf,pix,sma) + assert d["transforms"] == ["jit", "vag"] + assert d["runs_done"] == 3 # 3 of 3x2 cell/transform runs + assert d["runs_missing"] == 3 + assert "imaging/mge/jwst [jit]" in d["missing"] + assert "interferometer/pixelization/sma [vag]" in d["missing"] + + +def test_off_grid_records_are_reported_not_dropped(tmp_path): + """knn / delaunay_matern are real measurements, not noise, and not coverage.""" + ws = _workspace(tmp_path, { + "local_cpu/knn.json": [_record(model_type="knn"), _record(model_type="knn", transform="vag")], + }) + d = json.loads(_run(["campaign", "--axis", "compile", "--json"], ws).stdout) + + assert d["runs_done"] == 0, "off-grid records must not count as grid coverage" + assert d["off_grid"] == [{"cell": "imaging/knn/hst", "records": 2}] + + +def test_malformed_records_are_bucketed_with_their_location(tmp_path): + ws = _workspace(tmp_path, { + "local_cpu/export_probe.json": [ + {"transform": "jit", "model_type": "mge", "tag": "census"}, # no hardware/class/instrument + _record(), + ], + }) + d = json.loads(_run(["campaign", "--axis", "compile", "--json"], ws).stdout) + + assert len(d["malformed"]) == 1 + m = d["malformed"][0] + assert m["record"] == "local_cpu/export_probe.json[0]" + assert set(m["missing"]) == {"hardware", "dataset_class", "instrument"} + assert d["runs_done"] == 1, "the well-formed sibling record still counts" + + +def test_tier_split_does_not_reuse_the_runtime_config_names(tmp_path): + """hardware+mixed_precision is a different vocabulary from TIER_CONFIGS. + + An fp64 and an mp record on the same hardware are the SAME compile tier — + folding precision into the tier (as the runtime config names do) would + double-count them. + """ + ws = _workspace(tmp_path, { + "local_cpu/mge.json": [_record(), _record(mixed_precision=True)], + "local_gpu_NVIDIA_A100_80GB_PCIe/mge.json": [ + _record(hardware="local_gpu_NVIDIA_A100_80GB_PCIe"), + ], + "local_gpu_RTX_2060/mge.json": [_record(hardware="local_gpu_RTX_2060")], + }) + + local = json.loads(_run(["campaign", "--axis", "compile", "--json"], ws).stdout) + assert local["runs_done"] == 1, "fp64 + mp on one hardware is one cell/transform run" + assert local["other_hardware"] == [{"hardware": "local_gpu_RTX_2060", "records": 1}] + + a100 = json.loads(_run(["campaign", "--axis", "compile", "--tier", "a100", "--json"], ws).stdout) + assert a100["runs_done"] == 1 + # The RTX row is not the A100 tier either — it must not be absorbed by it. + assert a100["other_hardware"] == [{"hardware": "local_gpu_RTX_2060", "records": 1}] + + +def test_dispatch_plan_names_the_missing_cell_and_transforms(tmp_path): + ws = _workspace(tmp_path, {"local_cpu/mge.json": [_record(transform="jit")]}) + d = json.loads(_run(["campaign", "--axis", "compile", "--json"], ws).stdout) + + hst = [s for s in d["dispatch_plan"] if "--instrument hst" in s] + assert len(hst) == 1 + assert "--dataset-class imaging" in hst[0] and "--model-type mge" in hst[0] + assert "--transforms vag" in hst[0], "only the missing transform is dispatched" + + +def test_transforms_come_from_the_workspace_not_a_copy(tmp_path): + """probe.py owns the transform axis; the Brain must not carry a stale copy.""" + ws = _workspace(tmp_path) + (ws / "scripts" / "misc" / "jax_compile" / "probe.py").write_text( + 'TRANSFORMS = ("jit", "vag", "some_future_transform")\n' + ) + d = json.loads(_run(["campaign", "--axis", "compile", "--json"], ws).stdout) + assert "some_future_transform" in d["transforms"] + + +def test_unreadable_corpus_file_does_not_take_the_mode_down(tmp_path): + ws = _workspace(tmp_path, {"local_cpu/mge.json": [_record()]}) + (ws / "scripts" / "misc" / "jax_compile" / "results" / "local_cpu" / "broken.json").write_text("{{") + r = _run(["campaign", "--axis", "compile", "--json"], ws) + assert r.returncode == 0 + assert json.loads(r.stdout)["runs_done"] == 1 + + +def test_absent_compile_corpus_reports_zero_coverage(tmp_path): + ws = _workspace(tmp_path) + d = json.loads(_run(["campaign", "--axis", "compile", "--json"], ws).stdout) + assert d["runs_done"] == 0 + assert d["runs_missing"] == 6 + assert "outstanding" in d["next_action"] + + +def test_compile_axis_never_compares_timings(tmp_path): + """Coverage only. Comparing rows across the comparability key is phase 2/3.""" + ws = _workspace(tmp_path, { + "local_cpu/mge.json": [_record(compile_s=2.0), _record(compile_s=900.0, tag="loaded")], + }) + d = json.loads(_run(["campaign", "--axis", "compile", "--json"], ws).stdout) + assert "COVERAGE only" in d["policy"] + blob = json.dumps(d) + assert "900" not in blob, "a timing value must not leak into a coverage decision" + + +def test_bad_tier_is_an_error(tmp_path): + ws = _workspace(tmp_path) + d = json.loads(_run(["campaign", "--axis", "compile", "--tier", "nope", "--json"], ws).stdout) + assert "unknown tier" in d["error"] + + +# --------------------------------------------------------------------------- +# axis routing +# --------------------------------------------------------------------------- + + +def test_compile_axis_is_refused_for_ingest_and_triage(tmp_path): + """Better a usage error than runtime findings reported under a compile flag.""" + ws = _workspace(tmp_path) + for mode in ("ingest", "triage"): + r = _run([mode, "--axis", "compile"], ws) + assert r.returncode == 5, f"{mode}: {r.stdout}{r.stderr}" + assert "not implemented" in r.stderr + + +def test_missing_workspace_exits_4(tmp_path): + r = _run(["campaign", "--axis", "compile"], tmp_path / "nope") + assert r.returncode == 4 + + +# --------------------------------------------------------------------------- +# runtime axis — held still while the compile axis grows beside it +# --------------------------------------------------------------------------- + + +def test_runtime_axis_is_the_default_and_unaffected(tmp_path): + ws = _workspace(tmp_path, {"local_cpu/mge.json": [_record()]}) + default = _run(["campaign", "--json"], ws) + explicit = _run(["campaign", "--axis", "runtime", "--json"], ws) + assert default.returncode == 0 + assert default.stdout == explicit.stdout + + d = json.loads(default.stdout) + assert "axis" not in d, "the runtime decision shape must not change" + assert d["grid_cells"] == 3 + # 3 cells x 2 local configs, none present in the empty runtime tree. + assert d["runs_missing"] == 6 + assert d["runs_done"] == 0 + assert d["runs_unusable"] == 0 + # The runtime axis buckets by sweep CONFIG name; a compile transform name + # appearing here would mean the two vocabularies had been crossed. + assert all("local_cpu_" in run_id for run_id in d["missing"]) + assert not any(tf in run_id for run_id in d["missing"] for tf in ("[jit]", "[vag]")) + + +def test_runtime_ingest_and_triage_still_run(tmp_path): + ws = _workspace(tmp_path) + for mode in ("ingest", "triage"): + r = _run([mode, "--json"], ws) + assert r.returncode == 0, r.stderr + assert json.loads(r.stdout)["mode"] == mode + + +def test_cli_dispatcher_exposes_the_axis_flag(tmp_path): + """The flag must reach the agent through bin/pyauto-brain, not just directly.""" + ws = _workspace(tmp_path, {"local_cpu/mge.json": [_record()]}) + r = subprocess.run( + [str(BRAIN), "profiling", "campaign", "--axis", "compile", "--json", "--workspace", str(ws)], + capture_output=True, text=True, env={**os.environ, "PYAUTO_ROOT": str(tmp_path)}, + ) + assert r.returncode == 0, r.stderr + assert json.loads(r.stdout)["axis"] == "compile" From 854706ab03ced7f72942700efd8e81f701f43d03 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 13:13:13 +0000 Subject: [PATCH 2/4] fix(profiling): drop the dead jax_compile/hpc glob from the a100 dispatch The first glob never matched (jax_compile/ has no hpc/ dir) so the fallback always won, but had it matched it would have printed hpc/batch_gpu/ for a file that lives elsewhere. Use the same path the runtime campaign does. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01L4STU81pQP1GkMzZVvsMsv --- agents/conductors/profiling/_profiling.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/agents/conductors/profiling/_profiling.py b/agents/conductors/profiling/_profiling.py index 9021a08..359ea30 100644 --- a/agents/conductors/profiling/_profiling.py +++ b/agents/conductors/profiling/_profiling.py @@ -300,9 +300,7 @@ def campaign_compile(ws: Path, tier: str) -> dict[str, Any]: + f" --transforms {','.join(tfs)} --cache-dir --tag " ) else: - submits = sorted(p.name for p in (compile_dir(ws) / "hpc").glob("submit_*")) or sorted( - p.name for p in (ws / "hpc" / "batch_gpu").glob("submit_*") - ) + submits = sorted(p.name for p in (ws / "hpc" / "batch_gpu").glob("submit_*")) dispatch = [f"sbatch hpc/batch_gpu/{s} (on the RAL checkout, post-pull)" for s in submits] return { From d17ecb38518be2449f7ae6523d7deaa1e5067a11 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 14:18:22 +0000 Subject: [PATCH 3/4] fix(profiling): don't report sibling-instrument records as malformed jax_compile/ hosts export_probe.py and trace_profile.py, which append their own schema into the SAME results// tree probe.py writes to. Their records have no hardware/dataset_class/instrument because they are a different record kind, not because they are corrupt -- so the 4 the mode was reporting as malformed would have sent someone to fix two files that work correctly. Split the two: missing the whole identity triple is a sibling instrument (reported per file, under its own bucket); missing only some key fields is genuine corruption and stays malformed. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01L4STU81pQP1GkMzZVvsMsv --- agents/conductors/profiling/AGENTS.md | 6 +++++ agents/conductors/profiling/_profiling.py | 26 +++++++++++++++----- tests/test_profiling_conductor.py | 30 +++++++++++++++++++---- 3 files changed, 51 insertions(+), 11 deletions(-) diff --git a/agents/conductors/profiling/AGENTS.md b/agents/conductors/profiling/AGENTS.md index 9d5d963..593b714 100644 --- a/agents/conductors/profiling/AGENTS.md +++ b/agents/conductors/profiling/AGENTS.md @@ -57,6 +57,12 @@ Records whose cell is not in the sweep grid (`knn`, `delaunay_matern`, the non-tier hardware in an **other-hardware** bucket. Both are real measurements — neither counts as grid coverage, and neither is silently dropped. +`jax_compile/` also hosts sibling instruments (`export_probe.py`, +`trace_profile.py`) that append their own schema into the same +`results//` tree. Records missing the whole `(hardware, dataset_class, +instrument)` identity triple are reported as **sibling-instrument** records, not +as malformed — only a record missing *some* of its key fields is corruption. + ## Fundamental principles - **The classification is the result** for CPU-unusable cells (the usability diff --git a/agents/conductors/profiling/_profiling.py b/agents/conductors/profiling/_profiling.py index 359ea30..35b880f 100644 --- a/agents/conductors/profiling/_profiling.py +++ b/agents/conductors/profiling/_profiling.py @@ -57,6 +57,9 @@ # Fields a compile record needs before it can be placed on the grid at all. COMPILE_KEY_FIELDS = ("hardware", "dataset_class", "model_type", "instrument") +# Absent *together*, these mark a record written by a sibling instrument sharing +# the results tree rather than a corrupt probe record. +COMPILE_IDENTITY_FIELDS = ("hardware", "dataset_class", "instrument") def workspace_root(explicit: str | None = None) -> Path: @@ -247,16 +250,22 @@ def campaign_compile(ws: Path, tier: str) -> dict[str, Any]: covered: set[tuple[tuple[str, str, str | None], str]] = set() off_grid: dict[str, int] = {} other_hw: dict[str, int] = {} + foreign: dict[str, int] = {} malformed: list[dict[str, Any]] = [] for rel, idx, rec in load_compile_corpus(ws): - if any(rec.get(f) in (None, "") for f in COMPILE_KEY_FIELDS): + absent = [f for f in COMPILE_KEY_FIELDS if rec.get(f) in (None, "")] + if set(absent) >= set(COMPILE_IDENTITY_FIELDS): + # Not corruption: jax_compile/ hosts sibling instruments + # (export_probe.py, trace_profile.py) that append their own schema + # into the SAME results// tree. Missing the whole identity + # triple means "another instrument's record", and calling that + # malformed would send someone to fix a file that is working. + foreign[rel] = foreign.get(rel, 0) + 1 + continue + if absent: malformed.append( - { - "record": f"{rel}[{idx}]", - "missing": [f for f in COMPILE_KEY_FIELDS if rec.get(f) in (None, "")], - "tag": rec.get("tag"), - } + {"record": f"{rel}[{idx}]", "missing": absent, "tag": rec.get("tag")} ) continue rec_tier = compile_tier_of(rec.get("hardware")) @@ -314,6 +323,7 @@ def campaign_compile(ws: Path, tier: str) -> dict[str, Any]: "runs_missing": len(missing), "missing": missing, "off_grid": [{"cell": c, "records": n} for c, n in sorted(off_grid.items())], + "foreign_records": [{"file": f, "records": n} for f, n in sorted(foreign.items())], "other_hardware": [{"hardware": h, "records": n} for h, n in sorted(other_hw.items())], "malformed": malformed, "policy": ( @@ -481,6 +491,10 @@ def emit_human(d: dict[str, Any]) -> None: print("Other hardware (neither tier):") for o in d["other_hardware"]: print(f" {o['hardware']}: {o['records']} record(s)") + if d["foreign_records"]: + print("Sibling-instrument records (not probe.py's schema):") + for f in d["foreign_records"]: + print(f" {f['file']}: {f['records']} record(s)") if d["malformed"]: print(f"Malformed records: {len(d['malformed'])}") for m in d["malformed"][:10]: diff --git a/tests/test_profiling_conductor.py b/tests/test_profiling_conductor.py index 884a19a..f967c40 100644 --- a/tests/test_profiling_conductor.py +++ b/tests/test_profiling_conductor.py @@ -123,19 +123,39 @@ def test_off_grid_records_are_reported_not_dropped(tmp_path): assert d["off_grid"] == [{"cell": "imaging/knn/hst", "records": 2}] -def test_malformed_records_are_bucketed_with_their_location(tmp_path): +def test_sibling_instrument_records_are_not_called_malformed(tmp_path): + """export_probe.py / trace_profile.py share the results tree with probe.py. + + Their records lack the whole identity triple because they are a different + schema, not because they are corrupt — reporting them as malformed would + send someone to fix a file that is working correctly. + """ ws = _workspace(tmp_path, { "local_cpu/export_probe.json": [ - {"transform": "jit", "model_type": "mge", "tag": "census"}, # no hardware/class/instrument - _record(), + {"transform": "jit", "model_type": "mge", "tag": "census"}, + {"transform": "vag", "model_type": "mge", "tag": "census"}, ], + "local_cpu/mge.json": [_record()], + }) + d = json.loads(_run(["campaign", "--axis", "compile", "--json"], ws).stdout) + + assert d["malformed"] == [] + assert d["foreign_records"] == [{"file": "local_cpu/export_probe.json", "records": 2}] + assert d["runs_done"] == 1, "the real probe record still counts" + + +def test_a_genuinely_incomplete_probe_record_is_still_malformed(tmp_path): + """One field missing is corruption; the whole identity triple is a sibling.""" + ws = _workspace(tmp_path, { + "local_cpu/mge.json": [_record(instrument=None), _record()], }) d = json.loads(_run(["campaign", "--axis", "compile", "--json"], ws).stdout) + assert d["foreign_records"] == [] assert len(d["malformed"]) == 1 m = d["malformed"][0] - assert m["record"] == "local_cpu/export_probe.json[0]" - assert set(m["missing"]) == {"hardware", "dataset_class", "instrument"} + assert m["record"] == "local_cpu/mge.json[0]" + assert m["missing"] == ["instrument"] assert d["runs_done"] == 1, "the well-formed sibling record still counts" From 73dac1422fc414d218d71030a167e27871e21193 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 14:41:00 +0000 Subject: [PATCH 4/4] =?UTF-8?q?feat(profiling):=20ingest=20--axis=20compil?= =?UTF-8?q?e=20=E2=80=94=20warm-pin=20drift=20detection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Completes the Brain side of phase 2, and is the surveillance the arc exists for: the persistent cache and --xla_gpu_autotune_level=0 are SETTINGS, so a config drift or an XLA_FLAGS clobber (the PyAutoNerves#127 failure that went undetected for two months) puts the worst case back with nothing failing. Reads the workspace's jax_compile/pins.json and reports warm rows that are unpinned or have drifted. Every comparison happens strictly inside one (hardware, hostname, jax_version, mixed_precision, cache_state); cross-key pairs are never a regression. A jax_version bump recompiles ONCE BY DESIGN, so it surfaces as a new unpinned key rather than as drift. Two corrections found by running it against the real corpus rather than trusting the design: 1. Rows PREDATING their pin are not drift. The first run flagged four, all of them July-16 measurements the July-28 pin had been chosen over -- i.e. it reported the improvement that set the pin as though it were a regression. Drift now requires a row newer than its pin. 2. That exposed the deeper flaw, fixed in the workspace: pins must be sticky. With "most recent wins", re-deriving pins after a cache regression would have quietly baked the regression in and the surveillance would report all-clear forever. Thresholds are deliberately generous and require BOTH gates -- >= 2.0x the pin AND >= 1.0s absolute. The ratio alone screams about sub-second cells where 100ms of jitter is 3x; the floor alone misses a cheap cell degrading by an order of magnitude. Host load alone has produced 7x errors in this corpus, and an alarm that cries wolf gets ignored. Against the real corpus: 25 pins, 0 drifted, 0 unpinned. 11 more tests, including a synthetic warm-reverting-to-cold row proving the alarm fires, each comparability field proving it does not fire across the key, and a guard that the Brain's mirrored key definition matches the workspace's pins.py (mirrored rather than imported, since importing would drag the JAX stack into the Brain). Stacked on feature/compile-axis-campaign-coverage (needs its load_compile_corpus/compile_tier_of). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01L4STU81pQP1GkMzZVvsMsv --- agents/conductors/profiling/AGENTS.md | 22 ++- agents/conductors/profiling/_profiling.py | 165 +++++++++++++++++++- tests/test_profiling_conductor.py | 182 +++++++++++++++++++++- 3 files changed, 356 insertions(+), 13 deletions(-) diff --git a/agents/conductors/profiling/AGENTS.md b/agents/conductors/profiling/AGENTS.md index 593b714..cff64ef 100644 --- a/agents/conductors/profiling/AGENTS.md +++ b/agents/conductors/profiling/AGENTS.md @@ -21,6 +21,7 @@ prompt, PyAutoMind `issued/profiling_agent.md`. |------|----------|-------| | `campaign` | Which grid runs are done / CPU-unusable / missing on this tier, and how do I dispatch the rest? | dispatch plan (local sweep flags incl. the per-run timeout; A100 submit list) | | `ingest` | Which probe JSONs aren't in the vram tables yet, and which results have no pin? | table-update rows, pin list, baseline + dashboard steps | +| `ingest --axis compile` | Which warm compile rows are unpinned, and which have drifted from their pin? | drifted rows (with pinned vs observed), unpinned keys, confirm/classify/re-pin steps | | `triage` | What do the pinned-drift findings mean? | per-finding classification: stale pin → re-pin here; library regression → `bug/` via intake | ``` @@ -41,9 +42,24 @@ bucketed by **hardware**, with `mixed_precision` a separate field. The two vocabularies do not interchange, so the compile axis maps tiers itself rather than reusing `TIER_CONFIGS`. -`--axis compile` currently serves `campaign` (coverage); `ingest` and `triage` -reject it with exit 5 until the compile pins land, so a compile flag can never -silently return a runtime answer. +`--axis compile` serves `campaign` (coverage) and `ingest` (warm-pin drift); +`triage` rejects it with exit 5 until drift classification lands, so a compile +flag can never silently return a runtime answer. + +**Drift is deliberately hard to trigger.** A row counts only if it is *newer* +than its pin, at least `2.0x` the pinned value, **and** at least `1.0 s` above it +in absolute terms. Rows predating the pin are the history the pin was chosen +over — flagging them would report the improvement that set the pin as a +regression. The ratio alone screams about sub-second cells where 100 ms of +jitter is 3x; the absolute floor alone misses a cheap cell degrading by an order +of magnitude. Both gates, generous, because host load alone has produced 7x +errors in this corpus and an alarm that cries wolf gets ignored. + +Pins live in the workspace (`jax_compile/pins.json`) and are **sticky** — the +workspace's `update_pins.py` will not move an existing pin without `--repin`. If +pins auto-followed the newest measurement, re-deriving them after a cache +regression would bake the regression in and the surveillance would report +all-clear forever. **Compile timings are host-load-sensitive** — the first measurements in `jax_compile/README.md` were wrong by up to **7×** (851 s vs 117 s for the same diff --git a/agents/conductors/profiling/_profiling.py b/agents/conductors/profiling/_profiling.py index 35b880f..0aff8f6 100644 --- a/agents/conductors/profiling/_profiling.py +++ b/agents/conductors/profiling/_profiling.py @@ -61,6 +61,20 @@ # the results tree rather than a corrupt probe record. COMPILE_IDENTITY_FIELDS = ("hardware", "dataset_class", "instrument") +# Mirrors autolens_profiling/scripts/misc/jax_compile/pins.py. Duplicated rather +# than imported for the same reason the grid is read via ast: importing the +# workspace would drag the JAX stack into the Brain. Kept honest by a test that +# reads the workspace's own definition. +COMPARABILITY_FIELDS = ("hardware", "hostname", "jax_version", "mixed_precision", "cache_state") +CELL_FIELDS = ("dataset_class", "model_type", "instrument", "transform") +PIN_FIELDS = COMPARABILITY_FIELDS + CELL_FIELDS + +# Drift thresholds. Generous on purpose: host load alone has produced 7x errors +# in this corpus, so a tight bound would flag a busy laptop as a regression and +# teach people to ignore the alarm. +COMPILE_DRIFT_RATIO = 2.0 +COMPILE_DRIFT_FLOOR_S = 1.0 + def workspace_root(explicit: str | None = None) -> Path: if explicit: @@ -140,6 +154,31 @@ def load_compile_corpus(ws: Path) -> "list[tuple[str, int, dict[str, Any]]]": return out +def load_pins(ws: Path) -> list[dict[str, Any]]: + """The workspace's warm-compile pins (`jax_compile/pins.json`).""" + path = compile_dir(ws) / "pins.json" + if not path.is_file(): + return [] + try: + data = json.loads(path.read_text()) + except (OSError, ValueError): + return [] + pins = data.get("pins") if isinstance(data, dict) else data + return [p for p in pins if isinstance(p, dict)] if isinstance(pins, list) else [] + + +def pin_key_str(key: tuple) -> str: + parts = dict(zip(PIN_FIELDS, key)) + cell = "/".join( + str(parts[f]) for f in ("dataset_class", "model_type", "instrument") if parts.get(f) + ) + return ( + f"{cell} [{parts.get('transform')}] " + f"@ {parts.get('hardware')}/{parts.get('hostname')} jax{parts.get('jax_version')}" + f"{' mp' if parts.get('mixed_precision') else ''} {parts.get('cache_state')}" + ) + + def compile_tier_of(hardware: str | None) -> str: """Which campaign tier a compile record belongs to. @@ -341,6 +380,106 @@ def campaign_compile(ws: Path, tier: str) -> dict[str, Any]: } +def ingest_compile(ws: Path) -> dict[str, Any]: + """Which warm compile rows are unpinned, and which have drifted from a pin. + + The surveillance the arc exists for: the persistent cache and + `--xla_gpu_autotune_level=0` are *settings*, so a config drift or an + `XLA_FLAGS` clobber puts the worst case back with nothing failing. + + Every comparison here happens strictly inside one comparability key. Rows + from different hardware, hosts, jax versions, precisions or cache states are + never paired — that is not conservatism, it is the difference between a + signal and noise: compile timings are host-load-sensitive to a measured 7x, + and a `jax_version` bump recompiles once BY DESIGN rather than regressing. + """ + pins = load_pins(ws) + if not pins: + return { + "agent": "profiling", + "mode": "ingest", + "axis": "compile", + "pins": 0, + "unpinned": [], + "drifted": [], + "next_action": ( + "no compile pins — run `python3 scripts/misc/jax_compile/update_pins.py --write` " + "in autolens_profiling first" + ), + } + + by_key = {tuple(p.get(f) for f in PIN_FIELDS): p for p in pins} + unpinned: list[dict[str, Any]] = [] + drifted: list[dict[str, Any]] = [] + seen: set[tuple] = set() + + for rel, idx, rec in load_compile_corpus(ws): + if rec.get("cache_state") != "warm" or "compile_s" not in rec: + continue + key = tuple(rec.get(f) for f in PIN_FIELDS) + if any(k in (None, "") for k in key if k is not False): + continue + pin = by_key.get(key) + if pin is None: + if key not in seen: + seen.add(key) + unpinned.append({"record": f"{rel}[{idx}]", "pin": pin_key_str(key)}) + continue + # Only rows NEWER than the pin can be drift. Every warm row predating + # the pin is the history the pin was chosen over — flagging those + # reports the improvement that set the pin as though it were a + # regression, which is how an alarm earns its way into being ignored. + if str(rec.get("timestamp") or "") <= str(pin.get("source_timestamp") or ""): + continue + expected, got = pin.get("compile_s"), rec.get("compile_s") + if not isinstance(expected, (int, float)) or not isinstance(got, (int, float)): + continue + if expected <= 0: + continue + ratio = got / expected + # Both gates, deliberately. The ratio alone screams about sub-second + # cells where a 100 ms jitter is 3x; the absolute delta alone misses a + # cheap cell degrading by an order of magnitude. Generous because host + # load alone has produced 7x errors in this corpus. + if ratio >= COMPILE_DRIFT_RATIO and abs(got - expected) >= COMPILE_DRIFT_FLOOR_S: + drifted.append( + { + "record": f"{rel}[{idx}]", + "pin": pin_key_str(key), + "pinned_s": expected, + "observed_s": got, + "ratio": round(ratio, 2), + "tag": rec.get("tag"), + } + ) + + return { + "agent": "profiling", + "mode": "ingest", + "axis": "compile", + "pins": len(pins), + "unpinned": unpinned, + "drifted": drifted, + "policy": ( + f"Drift = a warm row NEWER than its pin, >= {COMPILE_DRIFT_RATIO}x the " + f"pinned value AND >= {COMPILE_DRIFT_FLOOR_S}s absolute, compared ONLY " + f"within {'/'.join(COMPARABILITY_FIELDS)}. Cross-key pairs and rows " + "predating the pin are never a regression." + ), + "steps": [ + "re-run the drifted cell warm to confirm it is not host load " + "(check the record's host_state against the pin's)", + "if confirmed, classify it — `pyauto-brain profiling triage --axis compile`", + "pin the unpinned rows: `python3 scripts/misc/jax_compile/update_pins.py --write`", + ], + "next_action": ( + "compile pins current — no warm drift" + if not drifted and not unpinned + else f"{len(drifted)} drifted, {len(unpinned)} unpinned warm key(s)" + ), + } + + # --------------------------------------------------------------------------- # ingest # --------------------------------------------------------------------------- @@ -517,6 +656,21 @@ def emit_human(d: dict[str, Any]) -> None: print("Dispatch plan:") for s in d["dispatch_plan"]: print(f" - {s}") + elif d["mode"] == "ingest" and d.get("axis") == "compile": + print(f"Compile pins: {d['pins']}") + print(f"Drifted: {len(d['drifted'])}") + for x in d["drifted"][:10]: + print( + f" {x['pin']}: pinned {x['pinned_s']}s -> observed " + f"{x['observed_s']}s ({x['ratio']}x, tag={x['tag']!r})" + ) + print(f"Unpinned warm keys: {len(d['unpinned'])}") + for x in d["unpinned"][:10]: + print(f" {x['pin']}") + if d.get("policy"): + print(f"Policy: {d['policy']}") + for s in d.get("steps", []): + print(f" - {s}") elif d["mode"] == "ingest": print(f"Provenance: {d['provenance']}") print(f"Probe updates: {len(d['probe_updates'])}") @@ -557,10 +711,13 @@ def main(argv=None) -> int: # ingest/triage own the compile axis in later phases of the arc (pins, then # drift classification). Refusing now is deliberate: a mode that silently # ignored --axis would report runtime findings under a compile flag. - if a.axis == "compile" and a.mode != "campaign": + # triage owns the compile axis in phase 3 (drift CLASSIFICATION). Refusing + # is deliberate: a mode that silently ignored --axis would report runtime + # findings under a compile flag. + if a.axis == "compile" and a.mode == "triage": print( - f"profiling: --axis compile is not implemented for {a.mode!r} yet " - "(campaign only; ingest/triage land with the compile pins)", + "profiling: --axis compile is not implemented for 'triage' yet " + "(campaign + ingest only; classification lands next)", file=sys.stderr, ) return 5 @@ -573,7 +730,7 @@ def main(argv=None) -> int: if a.mode == "campaign": d = campaign_compile(ws, a.tier) if a.axis == "compile" else campaign(ws, a.tier) elif a.mode == "ingest": - d = ingest(ws) + d = ingest_compile(ws) if a.axis == "compile" else ingest(ws) else: d = triage(ws) diff --git a/tests/test_profiling_conductor.py b/tests/test_profiling_conductor.py index f967c40..6b19b37 100644 --- a/tests/test_profiling_conductor.py +++ b/tests/test_profiling_conductor.py @@ -242,13 +242,16 @@ def test_bad_tier_is_an_error(tmp_path): # --------------------------------------------------------------------------- -def test_compile_axis_is_refused_for_ingest_and_triage(tmp_path): - """Better a usage error than runtime findings reported under a compile flag.""" +def test_compile_axis_is_refused_for_triage(tmp_path): + """Better a usage error than runtime findings reported under a compile flag. + + `ingest` gained the axis with the pins; `triage` classifies drift and lands + with phase 3. + """ ws = _workspace(tmp_path) - for mode in ("ingest", "triage"): - r = _run([mode, "--axis", "compile"], ws) - assert r.returncode == 5, f"{mode}: {r.stdout}{r.stderr}" - assert "not implemented" in r.stderr + r = _run(["triage", "--axis", "compile"], ws) + assert r.returncode == 5, f"{r.stdout}{r.stderr}" + assert "not implemented" in r.stderr def test_missing_workspace_exits_4(tmp_path): @@ -298,3 +301,170 @@ def test_cli_dispatcher_exposes_the_axis_flag(tmp_path): ) assert r.returncode == 0, r.stderr assert json.loads(r.stdout)["axis"] == "compile" + + +# --------------------------------------------------------------------------- +# ingest --axis compile (warm pins) +# --------------------------------------------------------------------------- + + +def _pinned(ws, pins): + (ws / "scripts" / "misc" / "jax_compile" / "pins.json").write_text( + json.dumps({"schema": 1, "pins": pins}) + ) + + +def _pin(**kw): + base = { + "hardware": "local_cpu", + "hostname": "laptop", + "jax_version": "0.10.2", + "mixed_precision": False, + "cache_state": "warm", + "dataset_class": "imaging", + "model_type": "mge", + "instrument": "hst", + "transform": "vag", + "compile_s": 2.3, + "source_tag": "census-warm", + "source_timestamp": "2026-07-01T00:00:00", + } + base.update(kw) + return base + + +def _warm(**kw): + base = { + "cache_state": "warm", + "hostname": "laptop", + "timestamp": "2026-08-01T00:00:00", + "transform": "vag", # matches _pin's default, so the keys line up + } + base.update(kw) + return _record(**base) + + +def test_a_warm_row_reverting_toward_cold_is_drift(tmp_path): + """The alarm the whole arc exists for: the cache stopped being hit.""" + ws = _workspace(tmp_path, {"local_cpu/mge.json": [_warm(compile_s=117.0)]}) + _pinned(ws, [_pin(compile_s=2.3)]) + d = json.loads(_run(["ingest", "--axis", "compile", "--json"], ws).stdout) + + assert len(d["drifted"]) == 1 + x = d["drifted"][0] + assert (x["pinned_s"], x["observed_s"]) == (2.3, 117.0) + assert x["ratio"] > 50 + + +def test_rows_predating_the_pin_are_not_drift(tmp_path): + """The pin was CHOSEN over this history; flagging it reports the + improvement that set the pin as a regression.""" + ws = _workspace(tmp_path, { + "local_cpu/mge.json": [_warm(compile_s=117.0, timestamp="2026-06-01T00:00:00")], + }) + _pinned(ws, [_pin(compile_s=2.3, source_timestamp="2026-07-01T00:00:00")]) + d = json.loads(_run(["ingest", "--axis", "compile", "--json"], ws).stdout) + assert d["drifted"] == [] + + +def test_drift_never_pairs_across_the_comparability_key(tmp_path): + """A slow row on ANOTHER host/version/precision is not this pin's business.""" + ws = _workspace(tmp_path, { + "local_cpu/mge.json": [ + _warm(compile_s=117.0, hostname="euclid-ral-compute-22"), + _warm(compile_s=117.0, jax_version="0.11.0"), + _warm(compile_s=117.0, mixed_precision=True), + ], + }) + _pinned(ws, [_pin(compile_s=2.3)]) + d = json.loads(_run(["ingest", "--axis", "compile", "--json"], ws).stdout) + + assert d["drifted"] == [], "cross-key rows must never be reported as drift" + assert len(d["unpinned"]) == 3, "they are unpinned keys of their own, not silence" + + +def test_a_jax_version_bump_is_a_new_key_not_a_regression(tmp_path): + """Cache keys include the jax version, so a bump recompiles once BY DESIGN.""" + ws = _workspace(tmp_path, { + "local_cpu/mge.json": [_warm(compile_s=117.0, jax_version="0.11.0")], + }) + _pinned(ws, [_pin(compile_s=2.3, jax_version="0.10.2")]) + d = json.loads(_run(["ingest", "--axis", "compile", "--json"], ws).stdout) + assert d["drifted"] == [] + assert len(d["unpinned"]) == 1 + + +def test_small_absolute_moves_on_cheap_cells_are_not_drift(tmp_path): + """0.05s -> 0.30s is 6x and completely uninteresting; the absolute floor + exists so sub-second jitter does not train people to ignore the alarm.""" + ws = _workspace(tmp_path, {"local_cpu/mge.json": [_warm(compile_s=0.30)]}) + _pinned(ws, [_pin(compile_s=0.05)]) + d = json.loads(_run(["ingest", "--axis", "compile", "--json"], ws).stdout) + assert d["drifted"] == [] + + +def test_large_absolute_move_below_the_ratio_is_not_drift(tmp_path): + ws = _workspace(tmp_path, {"local_cpu/mge.json": [_warm(compile_s=130.0)]}) + _pinned(ws, [_pin(compile_s=100.0)]) + d = json.loads(_run(["ingest", "--axis", "compile", "--json"], ws).stdout) + assert d["drifted"] == [] + + +def test_cold_rows_are_never_compared_against_a_warm_pin(tmp_path): + ws = _workspace(tmp_path, { + "local_cpu/mge.json": [_record(cache_state="cold", compile_s=117.0, hostname="laptop")], + }) + _pinned(ws, [_pin(compile_s=2.3)]) + d = json.loads(_run(["ingest", "--axis", "compile", "--json"], ws).stdout) + assert d["drifted"] == [] and d["unpinned"] == [] + + +def test_unpinned_warm_keys_are_reported_once_each(tmp_path): + ws = _workspace(tmp_path, { + "local_cpu/mge.json": [_warm(compile_s=2.0), _warm(compile_s=2.1)], + }) + _pinned(ws, []) + d = json.loads(_run(["ingest", "--axis", "compile", "--json"], ws).stdout) + assert d["next_action"].startswith("no compile pins") + + +def test_absent_pins_file_says_so_rather_than_reporting_all_clear(tmp_path): + ws = _workspace(tmp_path, {"local_cpu/mge.json": [_warm()]}) + d = json.loads(_run(["ingest", "--axis", "compile", "--json"], ws).stdout) + assert d["pins"] == 0 + assert "update_pins.py" in d["next_action"] + + +def test_triage_still_refuses_the_compile_axis(tmp_path): + ws = _workspace(tmp_path) + r = _run(["triage", "--axis", "compile"], ws) + assert r.returncode == 5 + assert "not implemented" in r.stderr + + +def test_brain_comparability_key_matches_the_workspace_definition(tmp_path): + """The Brain mirrors pins.py rather than importing it (importing the + workspace would drag the JAX stack in), so pin the two together.""" + import ast as _ast + + ws = _workspace(tmp_path) + pins_py = ws / "scripts" / "misc" / "jax_compile" / "pins.py" + pins_py.write_text( + 'COMPARABILITY_FIELDS = ("hardware", "hostname", "jax_version", ' + '"mixed_precision", "cache_state")\n' + 'CELL_FIELDS = ("dataset_class", "model_type", "instrument", "transform")\n' + ) + tree = _ast.parse(pins_py.read_text()) + found = { + t.id: _ast.literal_eval(n.value) + for n in tree.body + if isinstance(n, _ast.Assign) + for t in n.targets + if isinstance(t, _ast.Name) + } + + sys.path.insert(0, str(BRAIN_HOME / "agents" / "conductors" / "profiling")) + import _profiling # noqa: PLC0415 + + assert _profiling.COMPARABILITY_FIELDS == found["COMPARABILITY_FIELDS"] + assert _profiling.CELL_FIELDS == found["CELL_FIELDS"]