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"]