Skip to content

Commit 74d7b1b

Browse files
authored
profiling: ingest --axis compile — warm-pin drift detection (#220)
The surveillance the compile-axis 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 autolens_profiling#104: 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). Phase 2 of 3, Brain side. Pairs with autolens_profiling#104. Prompt: PyAutoMind active/compile_warm_baseline_dashboard.md.
1 parent cd79005 commit 74d7b1b

3 files changed

Lines changed: 356 additions & 13 deletions

File tree

agents/conductors/profiling/AGENTS.md

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ prompt, PyAutoMind `issued/profiling_agent.md`.
2121
|------|----------|-------|
2222
| `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) |
2323
| `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 |
24+
| `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 |
2425
| `triage` | What do the pinned-drift findings mean? | per-finding classification: stale pin → re-pin here; library regression → `bug/` via intake |
2526

2627
```
@@ -41,9 +42,24 @@ bucketed by **hardware**, with `mixed_precision` a separate field. The two
4142
vocabularies do not interchange, so the compile axis maps tiers itself rather
4243
than reusing `TIER_CONFIGS`.
4344

44-
`--axis compile` currently serves `campaign` (coverage); `ingest` and `triage`
45-
reject it with exit 5 until the compile pins land, so a compile flag can never
46-
silently return a runtime answer.
45+
`--axis compile` serves `campaign` (coverage) and `ingest` (warm-pin drift);
46+
`triage` rejects it with exit 5 until drift classification lands, so a compile
47+
flag can never silently return a runtime answer.
48+
49+
**Drift is deliberately hard to trigger.** A row counts only if it is *newer*
50+
than its pin, at least `2.0x` the pinned value, **and** at least `1.0 s` above it
51+
in absolute terms. Rows predating the pin are the history the pin was chosen
52+
over — flagging them would report the improvement that set the pin as a
53+
regression. The ratio alone screams about sub-second cells where 100 ms of
54+
jitter is 3x; the absolute floor alone misses a cheap cell degrading by an order
55+
of magnitude. Both gates, generous, because host load alone has produced 7x
56+
errors in this corpus and an alarm that cries wolf gets ignored.
57+
58+
Pins live in the workspace (`jax_compile/pins.json`) and are **sticky** — the
59+
workspace's `update_pins.py` will not move an existing pin without `--repin`. If
60+
pins auto-followed the newest measurement, re-deriving them after a cache
61+
regression would bake the regression in and the surveillance would report
62+
all-clear forever.
4763

4864
**Compile timings are host-load-sensitive** — the first measurements in
4965
`jax_compile/README.md` were wrong by up to **** (851 s vs 117 s for the same

agents/conductors/profiling/_profiling.py

Lines changed: 161 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,20 @@
6161
# the results tree rather than a corrupt probe record.
6262
COMPILE_IDENTITY_FIELDS = ("hardware", "dataset_class", "instrument")
6363

64+
# Mirrors autolens_profiling/scripts/misc/jax_compile/pins.py. Duplicated rather
65+
# than imported for the same reason the grid is read via ast: importing the
66+
# workspace would drag the JAX stack into the Brain. Kept honest by a test that
67+
# reads the workspace's own definition.
68+
COMPARABILITY_FIELDS = ("hardware", "hostname", "jax_version", "mixed_precision", "cache_state")
69+
CELL_FIELDS = ("dataset_class", "model_type", "instrument", "transform")
70+
PIN_FIELDS = COMPARABILITY_FIELDS + CELL_FIELDS
71+
72+
# Drift thresholds. Generous on purpose: host load alone has produced 7x errors
73+
# in this corpus, so a tight bound would flag a busy laptop as a regression and
74+
# teach people to ignore the alarm.
75+
COMPILE_DRIFT_RATIO = 2.0
76+
COMPILE_DRIFT_FLOOR_S = 1.0
77+
6478

6579
def workspace_root(explicit: str | None = None) -> Path:
6680
if explicit:
@@ -140,6 +154,31 @@ def load_compile_corpus(ws: Path) -> "list[tuple[str, int, dict[str, Any]]]":
140154
return out
141155

142156

157+
def load_pins(ws: Path) -> list[dict[str, Any]]:
158+
"""The workspace's warm-compile pins (`jax_compile/pins.json`)."""
159+
path = compile_dir(ws) / "pins.json"
160+
if not path.is_file():
161+
return []
162+
try:
163+
data = json.loads(path.read_text())
164+
except (OSError, ValueError):
165+
return []
166+
pins = data.get("pins") if isinstance(data, dict) else data
167+
return [p for p in pins if isinstance(p, dict)] if isinstance(pins, list) else []
168+
169+
170+
def pin_key_str(key: tuple) -> str:
171+
parts = dict(zip(PIN_FIELDS, key))
172+
cell = "/".join(
173+
str(parts[f]) for f in ("dataset_class", "model_type", "instrument") if parts.get(f)
174+
)
175+
return (
176+
f"{cell} [{parts.get('transform')}] "
177+
f"@ {parts.get('hardware')}/{parts.get('hostname')} jax{parts.get('jax_version')}"
178+
f"{' mp' if parts.get('mixed_precision') else ''} {parts.get('cache_state')}"
179+
)
180+
181+
143182
def compile_tier_of(hardware: str | None) -> str:
144183
"""Which campaign tier a compile record belongs to.
145184
@@ -341,6 +380,106 @@ def campaign_compile(ws: Path, tier: str) -> dict[str, Any]:
341380
}
342381

343382

383+
def ingest_compile(ws: Path) -> dict[str, Any]:
384+
"""Which warm compile rows are unpinned, and which have drifted from a pin.
385+
386+
The surveillance the arc exists for: the persistent cache and
387+
`--xla_gpu_autotune_level=0` are *settings*, so a config drift or an
388+
`XLA_FLAGS` clobber puts the worst case back with nothing failing.
389+
390+
Every comparison here happens strictly inside one comparability key. Rows
391+
from different hardware, hosts, jax versions, precisions or cache states are
392+
never paired — that is not conservatism, it is the difference between a
393+
signal and noise: compile timings are host-load-sensitive to a measured 7x,
394+
and a `jax_version` bump recompiles once BY DESIGN rather than regressing.
395+
"""
396+
pins = load_pins(ws)
397+
if not pins:
398+
return {
399+
"agent": "profiling",
400+
"mode": "ingest",
401+
"axis": "compile",
402+
"pins": 0,
403+
"unpinned": [],
404+
"drifted": [],
405+
"next_action": (
406+
"no compile pins — run `python3 scripts/misc/jax_compile/update_pins.py --write` "
407+
"in autolens_profiling first"
408+
),
409+
}
410+
411+
by_key = {tuple(p.get(f) for f in PIN_FIELDS): p for p in pins}
412+
unpinned: list[dict[str, Any]] = []
413+
drifted: list[dict[str, Any]] = []
414+
seen: set[tuple] = set()
415+
416+
for rel, idx, rec in load_compile_corpus(ws):
417+
if rec.get("cache_state") != "warm" or "compile_s" not in rec:
418+
continue
419+
key = tuple(rec.get(f) for f in PIN_FIELDS)
420+
if any(k in (None, "") for k in key if k is not False):
421+
continue
422+
pin = by_key.get(key)
423+
if pin is None:
424+
if key not in seen:
425+
seen.add(key)
426+
unpinned.append({"record": f"{rel}[{idx}]", "pin": pin_key_str(key)})
427+
continue
428+
# Only rows NEWER than the pin can be drift. Every warm row predating
429+
# the pin is the history the pin was chosen over — flagging those
430+
# reports the improvement that set the pin as though it were a
431+
# regression, which is how an alarm earns its way into being ignored.
432+
if str(rec.get("timestamp") or "") <= str(pin.get("source_timestamp") or ""):
433+
continue
434+
expected, got = pin.get("compile_s"), rec.get("compile_s")
435+
if not isinstance(expected, (int, float)) or not isinstance(got, (int, float)):
436+
continue
437+
if expected <= 0:
438+
continue
439+
ratio = got / expected
440+
# Both gates, deliberately. The ratio alone screams about sub-second
441+
# cells where a 100 ms jitter is 3x; the absolute delta alone misses a
442+
# cheap cell degrading by an order of magnitude. Generous because host
443+
# load alone has produced 7x errors in this corpus.
444+
if ratio >= COMPILE_DRIFT_RATIO and abs(got - expected) >= COMPILE_DRIFT_FLOOR_S:
445+
drifted.append(
446+
{
447+
"record": f"{rel}[{idx}]",
448+
"pin": pin_key_str(key),
449+
"pinned_s": expected,
450+
"observed_s": got,
451+
"ratio": round(ratio, 2),
452+
"tag": rec.get("tag"),
453+
}
454+
)
455+
456+
return {
457+
"agent": "profiling",
458+
"mode": "ingest",
459+
"axis": "compile",
460+
"pins": len(pins),
461+
"unpinned": unpinned,
462+
"drifted": drifted,
463+
"policy": (
464+
f"Drift = a warm row NEWER than its pin, >= {COMPILE_DRIFT_RATIO}x the "
465+
f"pinned value AND >= {COMPILE_DRIFT_FLOOR_S}s absolute, compared ONLY "
466+
f"within {'/'.join(COMPARABILITY_FIELDS)}. Cross-key pairs and rows "
467+
"predating the pin are never a regression."
468+
),
469+
"steps": [
470+
"re-run the drifted cell warm to confirm it is not host load "
471+
"(check the record's host_state against the pin's)",
472+
"if confirmed, classify it — `pyauto-brain profiling triage --axis compile`",
473+
"pin the unpinned rows: `python3 scripts/misc/jax_compile/update_pins.py --write`",
474+
],
475+
"next_action": (
476+
"compile pins current — no warm drift"
477+
if not drifted and not unpinned
478+
else f"{len(drifted)} drifted, {len(unpinned)} unpinned warm key(s)"
479+
),
480+
}
481+
482+
344483
# ---------------------------------------------------------------------------
345484
# ingest
346485
# ---------------------------------------------------------------------------
@@ -517,6 +656,21 @@ def emit_human(d: dict[str, Any]) -> None:
517656
print("Dispatch plan:")
518657
for s in d["dispatch_plan"]:
519658
print(f" - {s}")
659+
elif d["mode"] == "ingest" and d.get("axis") == "compile":
660+
print(f"Compile pins: {d['pins']}")
661+
print(f"Drifted: {len(d['drifted'])}")
662+
for x in d["drifted"][:10]:
663+
print(
664+
f" {x['pin']}: pinned {x['pinned_s']}s -> observed "
665+
f"{x['observed_s']}s ({x['ratio']}x, tag={x['tag']!r})"
666+
)
667+
print(f"Unpinned warm keys: {len(d['unpinned'])}")
668+
for x in d["unpinned"][:10]:
669+
print(f" {x['pin']}")
670+
if d.get("policy"):
671+
print(f"Policy: {d['policy']}")
672+
for s in d.get("steps", []):
673+
print(f" - {s}")
520674
elif d["mode"] == "ingest":
521675
print(f"Provenance: {d['provenance']}")
522676
print(f"Probe updates: {len(d['probe_updates'])}")
@@ -557,10 +711,13 @@ def main(argv=None) -> int:
557711
# ingest/triage own the compile axis in later phases of the arc (pins, then
558712
# drift classification). Refusing now is deliberate: a mode that silently
559713
# ignored --axis would report runtime findings under a compile flag.
560-
if a.axis == "compile" and a.mode != "campaign":
714+
# triage owns the compile axis in phase 3 (drift CLASSIFICATION). Refusing
715+
# is deliberate: a mode that silently ignored --axis would report runtime
716+
# findings under a compile flag.
717+
if a.axis == "compile" and a.mode == "triage":
561718
print(
562-
f"profiling: --axis compile is not implemented for {a.mode!r} yet "
563-
"(campaign only; ingest/triage land with the compile pins)",
719+
"profiling: --axis compile is not implemented for 'triage' yet "
720+
"(campaign + ingest only; classification lands next)",
564721
file=sys.stderr,
565722
)
566723
return 5
@@ -573,7 +730,7 @@ def main(argv=None) -> int:
573730
if a.mode == "campaign":
574731
d = campaign_compile(ws, a.tier) if a.axis == "compile" else campaign(ws, a.tier)
575732
elif a.mode == "ingest":
576-
d = ingest(ws)
733+
d = ingest_compile(ws) if a.axis == "compile" else ingest(ws)
577734
else:
578735
d = triage(ws)
579736

0 commit comments

Comments
 (0)