maint: vram-first validation sweep + batch-size provenance (polish phase 2)

[Jammy2211]

Overview

Phase 2 of the polish series (parent: PyAutoMind maintenance/autolens_profiling/polish.md; phase 1: #52 / PR #53). Validate the VRAM/vmap story before the PreOptimizationTimes campaign burns A100 hours: the parent prompt's instruction is "run vram first so memory issues don't sidetrack the campaign". Autonomy: --auto at effective level supervised (plan approved interactively in-session; checkpoint-and-continue per PyAutoBrain/AUTONOMY.md).

Plan

• Coverage + staleness audit — vram/config.py fully covers the phase-3 matrix (its Nones are documented intentional blocks), but every row was probed 2026-05-24 against older source; add per-row provenance (probe date + PyAutoLens version) and flag the rows the campaign relies on.
• Local CPU validation sweep — run every phase-3 cell script once on CPU with its cheapest instrument preset to flush functional bugs while it's cheap; exercise --vmap-probe end-to-end locally and confirm vmap_probe.json lands in the new results/runtime/<class>/<model>/ layout.
• Sparse-path batch sizes — extend the batch lookup with an inversion-path axis (dense default, back-compatible vmap_batch_for), since phase 3 vmaps imaging sparse-vs-mapping and the table has no sparse dimension.
• Close the probe→runtime gap — A100 submits re-probe in Phase A but Phase B reads the static table; runtime cells will prefer a fresh same-cell vmap_probe.json over the table (logged), so a stale table can't OOM a job.
• Docs + dispatch prep — fix the stale z_projects probe path in vram/README.md; stage probe-only A100 submits (multi-point (1,4,16) for mge/pixelization, single-point (1,) for delaunay, sparse variants included).
• A100 leg (human/HPC-blocked) — submits + ingest instructions staged; dispatch happens on RAL (HPC checkout must pull the phase-1 merge first); probe JSONs then ingested into the table with provenance.

Detailed implementation plan

Affected Repositories

• autolens_profiling (primary, only)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_profiling	main	clean

Suggested branch: feature/profiling-vram-validation
Work Classification: Workspace
Worktree root: ~/Code/PyAutoLabs-wt/profiling-vram-validation/

Implementation Steps

1. vram/config.py: key (dataset, model, instrument) → (dataset, model, instrument, path) with back-compatible vmap_batch_for(..., path="dense"); add PROBED provenance mapping; extend test/test_vram_probe.py.
2. likelihood_runtime/<class>/<model>.py: batch-size resolution order — fresh vmap_probe.json in the cell output dir → vram/config table → skip-with-reason (log which source won).
3. hpc/batch_gpu/: probe-only submits for the campaign's vmap cells incl. sparse variants.
4. Local CPU sweep of all phase-3 cells (smallest instrument, CPU-pinned, scratch results); triage any failures — fix in-repo if small, else file bug/ prompts without sidetracking (parent-prompt rule).
5. vram/README.md stale-path fix + in-repo probe→config workflow doc.

Key Files

• vram/config.py, vram/probe.py, test/test_vram_probe.py
• likelihood_runtime/{imaging,interferometer,datacube}/*.py
• hpc/batch_gpu/, vram/README.md

Original Prompt

Click to expand starting prompt

polish phase 2 — vram-first validation sweep

Type: maintenance
Target: autolens_profiling
Difficulty: medium
Autonomy: supervised
Priority: normal
Status: formalised

Phase 2 of 4 of polish.md (see parent for full intent). Depends on phase 1
(design locked).

Profiling runs have historically flagged memory/VRAM issues that side-tracked
the campaign into source fixes. To avoid that, run the vram package first,
across the imaging, interferometer and datacube configs that phase 3 will
profile, so memory issues and other bugs surface now rather than mid-campaign.

Deliverables:

• vram results for the phase-3 config matrix, saved per the phase-1 results
  conventions (.json + .md).
• The computed vmap batch sizes per config — phase 3's likelihood_runtime
  sweeps are vmap-only and consume these batch sizes.
• Any memory/VRAM bugs found are triaged: fix here only if small and in-repo;
  otherwise file as separate bug/ prompts and do not side-track the campaign.

Out of scope: runtime/breakdown sweeps (phases 3–4); searches; point_source;
laptop GPU.

[Jammy2211]

Milestone 1 — code + dispatch legs done; local validation sweep running

Implemented on feature/profiling-vram-validation (commit 1):

• Probe-over-table resolution (vram.resolve_vmap_batch): every vmap-capable runtime cell now prefers a fresh, cell-matching vmap_probe_<model>[_sparse].json in its output dir over the static table, logging which source won — closes the gap where A100 submits re-probe in Phase A but Phase B read the stale table (and could OOM).
• Sparse-path table (VMAP_BATCH_SPARSE): own table for w-tilde rows rather than re-keying VMAP_BATCH; unprobed sparse cells fall back to the dense row (conservative — sparse per-replica footprints have measured smaller everywhere so far).
• Provenance (vram.config.PROVENANCE): the dense table's numbers date to 2026-05-24 against older source — recorded explicitly so campaigns check before trusting.
• Probe-only A100 submits: submit_probe_fast_a100 (mge/pixelization, ~minutes) + submit_probe_delaunay_a100 (delaunay, hours) re-probe every campaign vmap cell, dense + sparse imaging variants; ingest workflow documented in vram/README.md (which also loses its stale z_projects path). Probe JSON basenames normalised across all cells.
• 11 new unit tests (19/19 total); ruff + format clean.

Now running: the local CPU validation sweep — each campaign cell once on its cheapest instrument (imaging×euclid dense+sparse, interferometer×sma, datacube×sma) plus an end-to-end check that a fresh probe JSON actually wins the batch-size resolution. Results and any flushed bugs follow in the next update.

For you (HPC leg): after this PR lands, pull it on the RAL checkout and sbatch hpc/batch_gpu/submit_probe_fast_a100 then submit_probe_delaunay_a100 — probe JSONs land under results/runtime/, and ingest instructions are in vram/README.md "Probe-only submits + ingest".

[Jammy2211]

Milestone 2 — sweep flushed three defects; pinned checks redesigned (user-directed)

The local CPU validation sweep did exactly what phase 2 exists for — three real defects in the first cells, all fixed on the branch:

1. A100 batch sizes applied on every backend — runtime cells vmapped at table batch (64) on laptop CPU and got OOM-killed; the probe's 65 GB budget default is A100-specific too. Fix: resolve_vmap_batch takes the running backend and clamps to 3 on non-GPU; probes now record their backend and resolution ignores probes from a different backend.
2. Instrument-naive pins — imaging/{mge,pixelization,delaunay} asserted an HST-only pinned value for every --instrument, failing any non-HST run by construction. Converted to the per-instrument dict pattern interferometer/mge already had.
3. Design defect (user call): pins were hard asserts. They guard benchmark comparability, not library correctness (that's autolens_workspace_test's remit) — and a crash-at-the-final-assert could kill a multi-hour A100 job. All 9 cells now record-and-flag: drift prints a loud warning and lands in the result JSON as pinned_expected + pinned_drift; eager ≡ JIT ≡ vmap cross-consistency asserts stay hard. Boundary rule recorded in results/notes/design_lock_in.md. PyAutoHeart pairing filed (feature/pyautoheart/profiling_drift_check.md): a vitals leg scanning result JSONs for non-empty pinned_drift.

Sweep so far (cheapest instrument per cell, CPU): probe pipeline end-to-end OK; probe-JSON-over-table resolution verified live; post-fix cells all green — imaging_delaunay dense (381s) + sparse (213s), interferometer_mge (23s), interferometer_pixelization (70s); drift fields verified present in the real artifacts. Two cells remain, then the three pre-fix casualties rerun under fixed code. 25/25 unit tests, ruff clean.

[Jammy2211]

Checkpoint — phase 2 ready to ship, awaiting your go (supervised autonomous run: parking per the autonomy contract)

Where things stand. Phase 2 is complete and verified on feature/profiling-vram-validation (local worktree, 3 commits, not yet pushed):

1. vram batch resolution — probe override, sparse table, provenance — resolve_vmap_batch (probe JSON > table), VMAP_BATCH_SPARSE, PROVENANCE, probe-only A100 submits, probe-basename normalisation.
2. pinned likelihood checks record-and-flag, never crash — the user-directed redesign: drift warns + lands in result JSONs (pinned_expected/pinned_drift, the surface the filed PyAutoHeart vitals leg reads); eager ≡ JIT ≡ vmap asserts stay hard; per-instrument pins; non-GPU batch clamp; backend-tagged probes.
3. validation sweep results note + dashboard refresh — all 9 campaign cells green on CPU (results/notes/vram_validation_2026_07_08.md is the durable record), probe pipeline + probe-override verified live, three defects flushed and fixed.

Ship-gate legs. Tests: 25/25 (downstream suites n/a — standalone leaf). Smoke: six-workspace curated smoke n/a (no downstream script surface); repo smoke-import gate passes, and the full validation sweep is a much stronger live check. Review: CLEAN — every changed path was exercised end-to-end by the sweep. Heart: YELLOW — 10 reasons, all organism-wide (today's workspace validation 5 failures, 58 stale parked scripts, 63 script timing regressions, version pin, install/release validation, 4 old PRs). None touch this repo.

The question: ship now — push the branch and open the pending-release PR — acknowledging the YELLOW reason set as unrelated?

• (a) Ship (recommended — same shape as phase 1's YELLOW: all reasons predate/live outside this repo; this repo's own PR gate is lint.yml, green locally) → push, PR with ## Scripts Changed + validation checklist, heart-ack recorded, task → awaiting-merge, calibration row appended. Merge stays with you.
• (b) Hold until the organism is GREEN → nothing pushed, task stays parked.
• (c) Request changes → say what; I iterate on the branch first.

After merge, the HPC leg is yours: pull on the RAL checkout, sbatch hpc/batch_gpu/submit_probe_fast_a100 then submit_probe_delaunay_a100, and the ingest workflow in vram/README.md turns the probe JSONs into refreshed table rows — that ingest + any table updates can be a quick follow-up session. Phase 3 (the PreOptimizationTimes campaign) is then unblocked.

[Jammy2211]

Shipped to PR — you answered (a): PR #55 is open with pending-release, Heart-YELLOW ack recorded, task → awaiting-merge. Merge stays with you. RAL note: the A100 probe leg is deferred until RAL is back (submits + ingest workflow are on the branch, ready to dispatch); nothing else in this phase depends on it.

[Jammy2211]

CI green on PR #55 — six layers deep, all now fixed. The repo's lint gate had never executed past lychee: each failing step masked the next. In order: (1) dashboard divergence — results/simulators/ was still gitignored while the committed README carried its row (un-ignored per the phase-1 design); (2) dead link to the retired likelihood/README.md in simulators/README.md; (3) Slack invite 403s to non-browser clients (lychee exclusion); (4) both workflows still invoked retired likelihood/ paths — profile.yml's release-time sweep would have failed on first dispatch; (5) the smoke step's hand-picked pip list lacked the import chain (pyyaml, …) — replaced with the lib-tests.yml sibling-checkout install pattern; (6) searches/nautilus/simple.py predates the searches redesign → searches/nautilus/imaging/mge.py.

Net effect beyond this PR: merging fixes main's currently-red CI, and the smoke + profile workflows now actually exercise the current package layout. Very much phase 2's job description — flushing latent breakage before it costs campaign time.

[Jammy2211]

Shipped — 2026-07-08

PRs

• maint: vram-first validation — probe override, record-and-flag pins (polish phase 2) #55 (merged)

Summary

Phase 2 (vram-first validation) is merged. All 9 campaign cells verified green on CPU against current source; probe-over-table batch resolution with non-GPU clamp and backend-tagged probes; pinned checks converted to record-and-flag with the drift surface PyAutoHeart consumes; probe-only A100 submits + ingest workflow staged; six layers of latent CI breakage fixed (main's lint gate is green again).

Session Notes

• RAL is down — the A100 probe leg waits: on return, pull on the RAL checkout, sbatch hpc/batch_gpu/submit_probe_fast_a100 then submit_probe_delaunay_a100, ingest per vram/README.md.
• Phase 3 (maintenance/autolens_profiling/polish_phase_3_runtimes.md) is unblocked for its local-CPU rows now; A100 rows follow the probe ingest.
• The Heart pairing (PyAutoHeart#37) ships separately.