feat(profiling): add a compile-time axis to campaign coverage

[Jammy2211]

Overview

The Profiling Agent cannot see the compile-time corpus that autolens_profiling
is already producing. agents/conductors/profiling/AGENTS.md still lists "JAX
compilation-time profiling of likelihood functions" under Future modes, and
_profiling.py (353 lines) contains zero occurrences of "compile" — all three
modes read only results/runtime/. Meanwhile scripts/misc/jax_compile/ holds
93 committed probe records that nothing cross-references against the science grid.

This is phase 1 of the compile-axis arc: a read-only coverage answer. It is the
cheapest of the three legs and produces the number that justifies the other two.

Plan

• Add --axis compile to the existing campaign mode; runtime stays the default
  and is untouched.
• Load the compile corpus from scripts/misc/jax_compile/results/*/*.json with
  stdlib json, honouring the never-import-the-workspace rule.
• Resolve each record to a grid cell from its in-record (dataset_class, model_type, instrument) — not its path, which drops class and instrument.
• Report coverage over cell × transform × hardware tier, with off-grid and
  malformed records in their own buckets rather than silently counted or dropped.
• Emit a probe.py dispatch plan for what is missing, mirroring how the runtime
  campaign emits sweep.py lines.
• Add tests/test_profiling_conductor.py — it does not exist yet; profiling is
  currently the only conductor without one.

Detailed implementation plan

Affected Repositories

• PyAutoBrain (primary)

Work Classification

Library

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoBrain	main	clean

Suggested branch: feature/compile-axis-campaign-coverage
Worktree root: ~/Code/PyAutoLabs-wt/compile-axis-campaign-coverage/

Implementation Steps

1. agents/conductors/profiling/_profiling.py — add load_compile_corpus(ws):
   glob scripts/misc/jax_compile/results/*/*.json, each an append-only list of
   flat records. Tolerate unreadable/malformed files the way ingest already
   does (except (OSError, ValueError): continue).

2. Add TRANSFORMS, the seven axes from probe.py's docstring: jit, grad,
   vag, vmap, vmap_vag, laxmap_vag, pyloop_vag.

3. Add campaign_compile(ws, tier). Resolve records to grid cells from
   load_grid(ws) (existing, reads sweep.py::CELLS via ast). Bucket into
   done / missing / off_grid / malformed:

   • off-grid: knn and delaunay_matern are real Prodigy-census
     measurements that are not CELLS entries — they must not read as noise.
   • malformed: 4 records currently carry hardware/dataset_class/
     instrument all null; report file + index so they can be fixed
     deliberately.

4. Map hardware to tier. Runtime uses TIER_CONFIGS config names
   (local_cpu_fp64 …); compile records carry a hardware string
   (local_cpu, local_gpu_NVIDIA_A100_80GB_PCIe, …) plus a separate
   mixed_precision bool. These are different vocabularies — needs an
   explicit mapping, not a reuse of TIER_CONFIGS.

5. Dispatch plan: python jax_compile/probe.py --model-type <m> --transforms <list> for local; the hpc/batch_gpu/submit_* form for --tier a100.

6. main() — add --axis {runtime,compile} defaulting to runtime; dispatch
   campaign accordingly. Reject --axis compile for ingest/triage with
   exit 5 (usage) so phases 2–3 land it deliberately rather than half-working.

7. emit_human() — a compile branch printing the four buckets and the plan.

8. agents/conductors/profiling/AGENTS.md — document the axis under campaign.
   Leave the Future modes entry until phase 3 closes the arc.

9. tests/test_profiling_conductor.py — new. Follow the hermetic pattern in
   test_hygiene_conductor.py: PYAUTO_ROOT at a temp dir, drive the real CLI
   via subprocess, build a synthetic workspace fixture (a sweep.py carrying
   CELLS + a jax_compile/results/ tree) so no real checkout is needed.
   Cover: bucket counts, off-grid separation, malformed handling, --json
   shape, exit 5 on ingest --axis compile, and a runtime-axis regression guard.

Key Files

• agents/conductors/profiling/_profiling.py — all logic
• agents/conductors/profiling/AGENTS.md — mode documentation
• tests/test_profiling_conductor.py — new
• autolens_profiling/scripts/misc/jax_compile/probe.py — read-only reference
  for the record schema and transform names

Out of scope

• Any judgement about whether a compile time is good (phase 3).
• Any pin, baseline or dashboard (phase 2).
• Executing probe.py — the agent reasons and emits a plan, never runs the
  workspace (ProfilingDecision contract).

Arc context

Phase 1 of three, re-scoped 2026-08-10 from the 2026-07-14 prompt
draft/feature/profiling/profiling_agent_jax_compile_time_scope.md. The original
"measure compile time so we can speed it up" framing is spent: the speed-up
shipped (persistent compilation cache 117.0s → 2.3s; --xla_gpu_autotune_level=0
17× on the pathological cold probe), so the surviving task is regression
surveillance — both wins are settings, and nothing watches for them reverting.

• Phase 2 — compile_warm_baseline_dashboard.md (autolens_profiling + PyAutoBrain)
• Phase 3 — compile_axis_triage_drift.md (PyAutoBrain)

Original Prompt

Click to expand starting prompt

Profiling Agent phase 1 — campaign --axis compile: what compile coverage do we actually have?

Type: feature
Target: PyAutoBrain
Repos:

• PyAutoBrain
  Difficulty: small
  Autonomy: supervised
  Priority: high
  Status: formalised

Phase 1 of profiling_agent_jax_compile_time_scope.md (read that first — it
carries the re-scope, the evidence and the comparability constraint).

Why this leg first

It is the cheapest of the three, it is pure-read, and it produces the number that
justifies the other two. Right now nobody can say how much of the science grid the
compile corpus covers, because nothing cross-references the two trees.

Measured by hand 2026-08-10 — this is what the mode should compute automatically:

• 93 records, all under scripts/misc/jax_compile/results/<hardware>/<model_type>.json.
• Distinct (dataset_class, model_type, instrument) touched: imaging/mge/hst,
  imaging/pixelization/hst, imaging/delaunay_matern/hst, imaging/knn/hst,
  plus two synthetic multi-band classes (datacube_img, datacube_img_hetero).
• The runtime grid (scripts/misc/likelihood_runtime/sweep.py::CELLS) is 24
  (class, model, instrument) cells across imaging / interferometer /
  datacube × hst,jwst,ao / sma,alma,alma_high,jvla.
• So: no interferometer row, no datacube row, no jwst, no ao — one
  instrument out of seven, and delaunay_matern / knn are mesh variants that are
  not grid cells at all.
• Hardware: local_cpu (82 rows), A100 (9, all pixelization/mge × jit/vag),
  RTX 2060 (2). Plus 4 malformed rows with hardware/dataset_class/
  instrument all null.

Scope

Add --axis compile to the existing campaign mode in
agents/conductors/profiling/_profiling.py. The runtime axis stays the default and
is untouched.

1. Read the compile corpus. Load scripts/misc/jax_compile/results/*/*.json
   (append-only lists of flat records). Stdlib json only — the same
   never-import-the-workspace rule the runtime path already follows via ast.
2. Resolve each record to a grid cell using its in-record
   (dataset_class, model_type, instrument), not its file path — the results tree
   is filed by <hardware>/<model_type> and drops class/instrument from the path.
3. Report coverage over cell × transform × hardware-tier. The transform axis is
   probe.py's seven: jit, grad, vag, vmap, vmap_vag, laxmap_vag,
   pyloop_vag. Cells not in CELLS are reported in their own off-grid bucket
   rather than silently counted or silently dropped — knn and delaunay_matern
   are real measurements from the Prodigy census and must not read as noise.
4. Report malformed records (missing hardware / dataset_class /
   instrument) as their own bucket with file + index, so they can be fixed or
   deleted deliberately.
5. Emit a dispatch plan of concrete probe.py invocations for what is missing,
   mirroring how the runtime campaign emits sweep.py lines:
   python jax_compile/probe.py --model-type <m> --transforms <list> for local,
   and the hpc/batch_gpu/submit_* form for --tier a100.

Explicitly not in scope

• Any judgement about whether a compile time is good — that is phase 3.
• Any pin, baseline or dashboard — that is phase 2.
• Running probe.py. The agent reasons and emits a plan; it never executes the
  workspace (ProfilingDecision contract).

Acceptance

• pyauto-brain profiling campaign --axis compile and
  ... --axis compile --tier a100 both run against a real autolens_profiling
  checkout and report done / missing / off-grid / malformed counts.
• The reported on-grid coverage matches the hand count above (3–4 cells of 24,
  hst only) on the corpus as it stands.
• --json emits the same structure, consistent with the existing modes.
• The runtime axis output is byte-identical to before the change (regression test).
• No file in autolens_profiling is written or executed by the agent.