maint: PreOptimizationTimes campaign — local-CPU leg (polish phase 3)

[Jammy2211]

Overview

Phase 3 of the polish series: the PreOptimizationTimes baseline campaign — the last full likelihood_runtime profiling before the optimization push, so these are the run times everything after is compared to. Phases 1–2 locked the design, validated all 9 cells on current source, and staged the A100 probes. RAL is down, so this run executes the laptop-CPU leg of the matrix and stages the rest: HPC-CPU and A100 rows (and the probe ingest) join when RAL returns; laptop-GPU rows are the user's follow-up. Autonomy: --auto, effective level supervised.

Plan

• Sweep driver instrument axis — sweep.py's CELLS grid is (class, model) with per-script default instruments; the campaign needs (class, model, instrument): per-cell campaign instrument lists (imaging × ao/jwst/hst; interferometer and datacube × their instrument sets), --instrument passthrough, and per-instrument output dirs (results/runtime/<class>/<model>/<instrument>/ — aggregate.py already supports the 3-level layout).
• Baseline tooling — scripts/build_baseline.py: snapshot the campaign's comparison.jsons into results/baselines/PreOptimizationTimes/ plus the rendered one-page .md (convention from results/notes/design_lock_in.md); the dashboard's baseline column picks it up automatically.
• Run the local-CPU matrix (local_cpu_fp64 + local_cpu_mp), cheapest→heaviest in the background, imaging sparse rows included; vmap rows run at the non-GPU clamp (production vmap numbers are the A100 rows, later).
• Pin the unpinned — first clean campaign runs print their likelihood/evidence values (record-and-flag design); pin them per-instrument so the Heart drift leg has teeth.
• Aggregate + publish — aggregate.py per cell, dashboard refresh, results note.
• Deferred (RAL): HPC-CPU rows, A100 probe ingest → VMAP_BATCH/VMAP_BATCH_SPARSE refresh, A100 fp64+mp rows, baseline finalisation for those columns.

Detail

Affected Repositories

• autolens_profiling (primary, only)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_profiling	main (post PR #55)	untracked legacy artifacts only

Suggested branch: feature/profiling-preopt-campaign · Classification: Workspace · Worktree: ~/Code/PyAutoLabs-wt/profiling-preopt-campaign/

Key Files

• likelihood_runtime/sweep.py (instrument axis), scripts/build_baseline.py (new), likelihood_runtime/*/*.py (pins), results/baselines/PreOptimizationTimes/, results/notes/

Original Prompt

PyAutoMind maintenance/autolens_profiling/polish_phase_3_runtimes.md (phase 3 of polish.md; laptop-CPU leg scoped in-session 2026-07-08 — RAL down).

[Jammy2211]

Milestone 1 — campaign tooling shipped to branch; local-CPU matrix running

On feature/profiling-preopt-campaign (commit 1): sweep.py gained the instrument axis — CELLS is now (class, model, instruments) with --instrument passthrough and per-instrument output dirs (results/runtime/<class>/<model>/<instrument>/, the 3-level layout aggregate.py already reads); --only/--skip accept class/model or class/model/instrument. scripts/build_baseline.py snapshots the campaign comparisons into results/baselines/PreOptimizationTimes/ + a rendered one-page .md; the dashboard baseline column activates automatically. 25/25 tests, ruff clean, dry-run verified.

Now running (background, cheapest→heaviest): the dense pass — 24 cells (imaging×{hst,jwst,ao}×3 models, interferometer×{sma,alma,alma_high,jvla}×3 models, datacube×{sma,alma,alma_high}) × local_cpu_fp64+local_cpu_mp = 48 runs — then the imaging sparse pass (18 runs), then per-cell aggregation. Datasets auto-simulate on first need; vmap rows run at the non-GPU clamp (production vmap numbers are the A100 rows, post-RAL). Expect this to grind for many hours; unpinned instruments print their likelihood/evidence values for pinning at ingest.

Deferred to RAL's return: HPC-CPU rows, probe-only submits → VMAP_BATCH/VMAP_BATCH_SPARSE ingest, A100 fp64+mp rows, baseline finalisation. Laptop-GPU rows are the user's follow-up.

[Jammy2211]

Interruption resilience (laptop may shut down overnight): every completed run's JSON persists immediately in the worktree, so a shutdown loses only the single in-flight run. sweep.py gained --skip-existing (commit 2) — resume skips completed (cell, config) pairs, verified live against the 9 already-finished runs.

Cold-resume instructions (any session, any morning): the worktree is ~/Code/PyAutoLabs-wt/profiling-preopt-campaign/; run the two passes with --skip-gpu --skip-existing (dense, then --sparse --only imaging/...), then aggregate.py — or just re-run scratchpad/preopt_campaign_resume.sh from the session scratchpad if it still exists. Then: pin unpinned instruments from the printed values, build_baseline.py --name PreOptimizationTimes, build_readme.py, results note, ship gate.

[Jammy2211]

Overnight interruption + morning resume (scoped). The laptop went down as anticipated; 18/48 dense runs persisted (imaging mge×3 + pixelization/delaunay×{hst,jwst}, interferometer mge/pixelization×sma). The failure logs from the first pass produced real findings:

• interferometer/mge at alma+: inherently blocked everywhere, exactly as vram/config.py documents — the dense MGE mapping matrix demanded a 62.72 GB gather at alma. Removed from the campaign grid (commit 3); these were never valid campaign cells.
• Laptop-infeasible (deferred to the HPC-CPU/A100 legs): imaging pixelization/delaunay @ ao (385k-pixel mask → silent OOM-kill on 15 GB); interferometer pixelization @ alma+ (NUFFT precision operator, "may exceed hours" + OOM risk); anything @ jvla (the dataset simulation itself was SIGKILLed). Note imaging/mge @ ao completed fine (3.11 s/call) — it's the mesh models that blow up.
• Resume relaunched with --skip-existing + the infeasible-cell skip list: ~12 dense runs remain (interferometer delaunay sma/alma/alma_high, datacube ×3) plus the 12-run hst/jwst imaging sparse pass.

The infeasibility map is itself campaign data — it goes in the results note and pre-scopes the HPC-CPU dispatch when RAL returns.

[Jammy2211]

Campaign-validity finding from the phase-4 breakdown work (#59): ~/venv/PyAuto/bin/activate:74 globally exports XLA_FLAGS="--xla_disable_hlo_passes=constant_folding" (leftover from the pure_callback const-fold debugging). Every local run that activates the PyAuto venv inherits it — which includes this campaign's overnight local-CPU matrix (2026-07-08/09).

Measured effect on a breakdown cell: imaging/mge step-sum 0.946 s with the flag vs 0.292 s without (3.2×). A with/without control pair on likelihood_runtime/imaging/mge.py is running now (results will be posted here) to quantify the effect on this campaign's numbers specifically.

Implications for the PreOptimizationTimes runtime baseline if the control confirms:

1. Local-CPU rows are inflated by a hardware-independent constant-ish factor — fine as a relative pre-optimization reference among themselves, but not production-representative in absolute terms.
2. Local-vs-A100 ratios are skewed — the HPC submits source the HPC checkout's activate.sh, which does not set the flag, so A100 rows are honest while local rows are taxed.
3. The vmap-vs-single-JIT comparisons within local rows are internally consistent (both taxed).

Result JSONs now record xla_flags/omp_num_threads/cpu_count provenance going forward (phase-4 branch, 7edf06a). Whether to remove the venv export (repo convention is vmap-for-honesty, not the pass-disable) is batched with the human on #59 — if it goes, the local-CPU matrix likely wants a rerun before the baseline is tagged; if it stays, it needs documenting and consistent application on the HPC side too.

[Jammy2211]

Follow-up to the validity flag — mostly good news for phase 3. The XLA flag turns out to be deliberate library policy (autoconf/jax_wrapper.py sets it at import and recommends the terminal export), so the campaign's local-CPU rows were measured under the canonical env, not a poisoned one — no rerun needed on that account.

Quantified effects (mge control pair, equal load): single-JIT 1.54× slower with the flag; vmap per-call insensitive (0.110 s both ways). So within the local matrix, single-JIT columns carry the flag's compile-time-vs-steady-state trade, vmap columns don't — internally consistent either way.

One residual to make explicit: the A100 submits don't export XLA_FLAGS — they inherit autoconf's in-process set, whose effectiveness depends on import-vs-backend-init order. Worth adding the explicit export to hpc/batch_gpu/submit_* templates (the phase-4 breakdown submits will carry it) so local-vs-A100 columns are guaranteed uniform.