Skip to content

feat(jax_compile): make warm compile machine-identifiable via cache_state - #104

Merged
Jammy2211 merged 5 commits into
mainfrom
feature/compile-warm-baseline-dashboard
Aug 10, 2026
Merged

feat(jax_compile): make warm compile machine-identifiable via cache_state#104
Jammy2211 merged 5 commits into
mainfrom
feature/compile-warm-baseline-dashboard

Conversation

@Jammy2211

Copy link
Copy Markdown
Contributor

Part of #103 (phase 2 of the compile-axis arc). This is the schema leg — the
prerequisite that unblocks pins and the dashboard; those remain open on the issue.

Why

Tracking warm compile per cell is the point of the arc: the persistent cache
turned 117.0 s into 2.3 s (CPU MGE vag) and 5517.8 s into 937.1 s (A100
end-to-end), and nothing watches for that reverting.

But warmness was recoverable only from the free-text tag, which carries ~40
ad-hoc spellings across this corpus — census-warm, census-warm2,
prodigy-census-warm-retry, cache-warm, mb_homo_cold — so the dashboard that
has to track warm compile could not tell which rows were warm.

What changed

probe.py derives cache_state from what the compile did: count cache entries
either side of lowered.compile().

value meaning
cold wrote a new entry — a MISS
warm wrote nothing into a populated cache — a HIT
none no --cache-dir
unknown configured, empty, and nothing written — not silently warm

The trap this avoids, and tests against: cache_dir is non-empty on cold rows
too, because the cold run is the one that populates the cache. Any substring check
on the tag or a truthiness check on cache_dir gets this backwards.

Per transform, not per run. Each transform compiles its own module, so one
invocation can legitimately miss on one and hit on another — which the old per-run
tag could not express at all.

host_state (cpu_count + 1m load average) is added because XLA compiles on the
host cores. The README's measurement-discipline section records the first
numbers being wrong by up to (851 s vs 117 s for the same compile) purely
from host load, and the corpus already mixes a 32-core RAL allocation with laptop
rows. Without it, "is this compile slower?" is unanswerable after the fact.

record_kind="compile_probe" makes the schema self-describing — export_probe.py
and trace_profile.py append their own shape into the same results tree.

Backfill

backfill_cache_state.py migrates the 89 existing compile records: 33 cold, 34
warm, 19 none, 3 unknown
.

  • Exact where cache_dir was empty → none.
  • Inferred only from an end-anchored cold/warm tag.
  • Everything else left unknown rather than guessed.

The anchoring is the point: mb_homo_cold_laxmap_gpu contains "cold" mid-tag and
is left unknown rather than mislabelled — precisely the false match an unanchored
parse would have made, and precisely the error cache_state exists to remove. The
4 sibling-instrument records are left untouched.

Verification

  • 9 new tests in scripts/misc/test/test_jax_compile_cache_state.py, JAX-free.
  • ruff check . and ruff format --check . clean across the repo.
  • build_readme.py --check passes (dashboard idempotence gate).

CI caveat, worth knowing before merge: scripts/misc/test/ is not run by any
workflow — lint.yml has no pytest step — so these tests pass locally but are not
a gate here. Relatedly, test_vram_config.py::test_aggregate_matches_cell_prefixed_stems
is currently failing on main unnoticed. Both are flagged on #103 rather than
folded into this PR.


Generated by Claude Code

claude added 5 commits August 10, 2026 14:22
…tate

Tracking WARM compile per cell is the point of the compile-axis arc: the
persistent cache turned 117.0s into 2.3s (CPU MGE vag) and 5517.8s into 937.1s
(A100 end to end), and nothing watches for that reverting. But warmness was
recoverable only from the free-text tag, which carries ~40 ad-hoc spellings
across this corpus -- census-warm, census-warm2, prodigy-census-warm-retry,
cache-warm, mb_homo_cold and so on -- so the dashboard that has to track warm
compile could not tell which rows were warm.

probe.py now derives cache_state from what the compile DID: count cache entries
either side of lowered.compile(), so a compile that wrote a new entry was a
MISS (cold) and one that wrote nothing into a populated cache was a HIT (warm).
No cache dir is none; configured-but-empty-and-nothing-written is unknown
rather than silently warm.

The obvious shortcut is a trap and is tested against: cache_dir is non-empty on
COLD rows too, because the cold run is the one that populates the cache.

Derived per transform rather than per run -- each transform compiles its own
module, so one invocation can legitimately miss on one and hit on another,
which a per-run tag could not express at all.

Also adds host_state (cpu_count + 1m load average). XLA compiles on the host
cores, and README.md records the first measurements being wrong by up to 7x
(851s vs 117s for the same compile) purely from host load, with the corpus
already mixing a 32-core RAL allocation against laptop rows. Without it, "is
this compile slower?" is unanswerable after the fact.

record_kind="compile_probe" makes the schema self-describing: export_probe.py
and trace_profile.py append their own shape into the same results tree.

backfill_cache_state.py migrates the 89 existing compile records: exact where
cache_dir was empty (none), inferred only from an END-ANCHORED cold/warm tag,
and left unknown otherwise -- 33 cold, 34 warm, 19 none, 3 unknown. The
anchoring matters: mb_homo_cold_laxmap_gpu contains "cold" mid-tag and is left
unknown rather than mislabelled, which is the false match an unanchored parse
would have made. The 4 sibling-instrument records are left untouched.

Phase 2 (workspace leg) of the compile-axis arc; PyAutoMind
draft/feature/profiling/compile_warm_baseline_dashboard.md.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L4STU81pQP1GkMzZVvsMsv
Completes the workspace side of phase 2: the schema leg made warm rows
identifiable, this makes them a regression surface.

pins.json holds the expected warm compile per cell/transform, derived by
update_pins.py from the corpus. 25 pins from 34 warm records. Only warm rows
are pinnable -- a cold row is the cost the cache exists to remove, so pinning
one would enshrine exactly what this arc watches for, and none/unknown rows
cannot be placed on either side of that line. Where several warm rows share a
pin identity the most recent wins: a pin states what warm costs NOW, and
averaging would blend a pre- and post-regression world into a number that
describes neither. The chosen row's tag and timestamp travel with the pin.

hostname is part of the comparability key, and that is a measured decision
rather than completeness. `hardware` is only ever local_cpu / local_gpu_<dev>,
so the single `local_cpu` label currently spans a laptop (66 records) AND a
32-core RAL node (12). The first rendering of this table proved the harm: it
put imaging/mge/hst [laxmap_vag] pinned on the RAL node directly beside [vag]
pinned on the laptop, under one heading, as though they were comparable --
which is the very 7x host-load hazard README.md warns about.

The dashboard renders through the existing sentinel-block mechanism
(build_readme.py auto-table:jax-compile-warm, jax_compile/README.md added to
TARGET_READMES), grouped by comparability key and never merged across it: one
ranked table would invite the cross-key comparison the pins exist to prevent.
build_readme.py --check stays idempotent, so the CI dashboard gate covers it.

update_pins.py --check exits 1 when pins.json is stale, so a corpus edit that
lands without re-deriving is catchable.

9 more tests: only-warm-is-pinnable, each comparability field splitting the
key (hostname included), most-recent-wins, provenance, determinism, and a
guard that the committed pins match a fresh derivation.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L4STU81pQP1GkMzZVvsMsv
"Most recent warm row wins" was wrong in the way that defeats the whole
purpose. With it, running update_pins.py after a cache regression would move
the pin ONTO the regressed value, and every subsequent check would report
all-clear forever.

Pins are now committed expectations, matching how the runtime results'
pinned_expected behaves: a new key is pinned from its most recent warm row,
but an existing pin never moves without an explicit --repin. That flag is how
a confirmed-benign change is accepted, and it is named and documented as the
deliberate act it is.

Found by running the Brain's ingest --axis compile against the real corpus:
it flagged four rows as drifted that were all OLDER than their pin -- the
history the pin had been chosen over -- which made the auto-following pin
visible.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L4STU81pQP1GkMzZVvsMsv
The sticky behaviour was fixed but only asserted in prose. Adds the three
cases that matter: a regression landing after the pin does not move it,
--repin is the deliberate way to accept a change, and stickiness still lets a
NEW key be pinned.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L4STU81pQP1GkMzZVvsMsv
These tests ran in NO workflow. lint.yml covered ruff, build_readme --check,
lychee and an import smoke, so nothing checked the vram tables, the
compile-probe cache_state derivation or the warm-compile pins on a PR -- an
odd gap for an arc whose whole subject is regression surveillance.

The job already installs the full dependency chain, which is what these tests
need: aggregate.py imports matplotlib at module scope, so a bare checkout
cannot run them (that missing dep, not any defect, is what made the suite look
red outside CI). Runs in ~1s.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L4STU81pQP1GkMzZVvsMsv
@Jammy2211
Jammy2211 merged commit 355d555 into main Aug 10, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants