perf: cut the three slowest smoke-gate scripts below a 6-minute py3.12 leg

[Jammy2211]

Overview

The per-PR smoke gate's py3.12 leg spends 553.0s executing scripts, and three
entries — imaging/subhalo_recovery.py (120.7s), misc/database/scrape/general.py
(63.4s) and point_source/jax_likelihood/point.py (47.7s) — are 42% of that.
This task makes each one materially faster or demotes it to the weekly/release
channels, per script, without weakening what the gate discriminates. Cost
attribution below was measured in a CI-equivalent environment (py3.12, source
chain install, smoke-profile env); the sibling prompt for running the gate less
often is out of scope here.

Plan

• imaging/subhalo_recovery.py: split the script — keep the one-shot engine as
  the PR-gate entry, move the warm-started iterative-LM confirmation leg to a
  new script covered by the weekly/release channels (measured: the iterative leg
  is ~80% of runtime and reproduces the one-shot metrics to 3 decimal places;
  the iterative engine class stays PR-gated via the interferometer sibling).
• misc/database/scrape/general.py: demote to the weekly/release channels with
  the required smoke_tests.txt comment (its unique coverage — scraping a real
  Nautilus run's outputs — is release-level, and the aggregator trio keeps the
  scrape API covered on the gate); fallback if demotion is declined: shrink the
  real search.
• point_source/jax_likelihood/point.py: measured compile-dominated (~44s of
  55s local is two JAX JIT compiles; likelihood evals are 0.03s), and problem
  size is a weak lever (10x looser solver precision saves only ~15%) — instead
  drop the FitPositionsImagePairAllSolved block, whose coverage the weekly
  channel already runs verbatim in image_plane.py; or demote the whole
  script if the ~6:00 acceptance is strict (see arithmetic below).
• Re-measure the gate on the PR and record per-entry timings; verify 23/23 pass
  and the py3.12 leg lands under ~6 minutes.

Detailed implementation plan

Affected Repositories

• autolens_workspace_test (primary, only)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace_test	main	clean

Suggested branch: feature/smoke-gate-slowest-scripts

Measured baseline (this assessment, 2026-08-23)

CI-equivalent environment: py3.12 venv, chain installed from source checkouts
per .github/scripts/smoke_install.sh, smoke-profile env per
config/build/profile_smoke.yaml + each script's ENV: declaration. Local
total for subhalo_recovery.py cold-cache is 148.0s vs 120.7s CI (factor
~0.82); warm-cache local is 134.7s (the ~13s delta is numba compile, which CI
always pays — fresh runners).

imaging/subhalo_recovery.py phase split (local warm / cold):

phase	warm	cold
import autolens	1.4s	1.6s
simulate + arc mask (5858 px)	2.1s	6.7s
one-shot joint solve (fit.log_evidence + best_fit_dpsi)	21.1s	37.5s
iterative LM solve_joint_optimization (n_iter=3, warm-start)	92.5s	83.9s
iterative log_evidence() (Laplace) — printed, never asserted	17.6s	18.2s
metrics + gauge asserts	0.0s	0.0s

Both engines recover essentially identical metrics: one-shot corr 0.8228 /
peak dist 0.06", iterative corr 0.8233 / 0.06". The warm-started iterative leg
is a stationarity confirmation of the one-shot optimum, not independent
discrimination.

Levers measured (all still assertion-green variants):

variant	effect	verdict
grid 120→100	mask 5858→5870 px (S/N-driven), no time change	dead end
dpsi mesh factor 2→3	no time change (cost is data x source dense algebra)	dead end
arc-mask threshold 3.0→4.0	mask −3.5%, no time change	dead end
n_iter 3→2	iterative solve 92.5→67.8s, corr 0.8231, evidence −0.001 dex	works
drop unasserted Laplace print	−17.6s	free
one-shot only	134.7→~25s warm, ~46s cold → est. ~35–40s CI	biggest lever

point_source/jax_likelihood/point.py phase split (local, 55.4s total vs 47.7s CI):

phase	time
imports	1.8s
dataset + PointSolver.for_grid construction	0.5s
first _vmap — FitPositionsImagePairAll JIT compile	11.3s
cached _vmap + NumPy fit + jit(fit_from) round-trip	0.5s
first _vmap — FitPositionsImagePairAllSolved JIT compile	33.0s
remainder (NumPy solved cross-check)	8.2s

Compile-dominated: ~44s of 55s is XLA compilation of the two pinned code
paths; each actual likelihood evaluation is 0.03s. Problem size was measured
as a WEAK lever: loosening pixel_scale_precision 0.001→0.01 (10x, ≈3 fewer
traced refinement levels) only cuts the compiles to 8.9s + 28.4s (total
55.4→47.3s, −15%) — the XLA cost is mostly a per-program floor, not
refinement depth. The strong lever is that the 33s Solved-block compile
duplicates coverage image_plane.py already provides on the weekly channel.

misc/database/scrape/general.py phase split (local, ~78s total vs 63.4s CI,
from run log timestamps):

phase	time
import + dataset bootstrap/load + mask (2828 px)	~6s
real Nautilus search (n_like_max=300, JAX): sampling	~54s
— of which JAX compiles (vmap 3.3s + per-sample latent 14.3s)	~18s in the above
database scrape + ALL aggregator assertions	~3s

The search is ~72s of ~78s; the scrape + assertions the script uniquely
gate-tests are ~3s. The misc/aggregator/ trio on the gate covers the same
scrape/aggregator API under PYAUTO_TEST_MODE=1 fixtures; general.py's unique
coverage is scraping a REAL search's output directory.

Implementation Steps

1. imaging/subhalo_recovery.py → split (recommended).

   • Trim imaging/subhalo_recovery.py to simulate + arc mask + one-shot
     FitDpsiSrcImaging + dkappa corr/dist assertions (thresholds unchanged:
     corr > 0.5, dist < 0.5"). Keep the ENV: full_datasets declaration and the
     docstring's parity-anchor framing; note in the docstring where the
     iterative leg went.
   • Create imaging/subhalo_recovery_iterative.py carrying the iterative-LM
     leg verbatim (warm start recomputed from a fresh one-shot fit, or
     restructured to re-derive x0 locally), with its existing corr/dist/gauge
     assertions. Do NOT add it to smoke_tests.txt; the weekly
     workspace-smoke.yml and release-integrate run every script, so it is
     covered there automatically — but note it in smoke_tests.txt's comment
     for the record.
   • Either way, delete the unasserted iter_fit.log_evidence() print from
     whichever script keeps the iterative leg (−15s+; nothing discriminates on it).
   • Discrimination demonstration (constraint 1): break the guarded thing —
     e.g. zero the recovered dkappa or simulate without the subhalo — and show
     the one-shot corr assertion fail; record the demonstration in the PR.
   • Engine-class coverage note for the PR: the PR gate still runs BOTH pc
     engines via interferometer/subhalo_recovery_interferometer.py (one-shot
     sparse+dense parity + IterFitDpsiSrcInterferometer n_iter=5 cold start,
     ~28s); imaging's iterative leg additionally runs weekly + at release.
   • Fallback if the human wants both engines on the PR gate: keep one script,
     drop the Laplace print and set n_iter=2 (measured corr 0.8231 vs 0.8233,
     gauge asserts unchanged) → est. ~85s CI. State which option was chosen.

2. misc/database/scrape/general.py → demote (recommended).

   • Remove from smoke_tests.txt with the required comment naming the
     channels that still run it: Heart's weekly workspace-smoke.yml (Mondays
     03:00 UTC, library main, same profile) and release-integrate
     (PYAUTO_TEST_MODE=0, full resolution).
   • Rationale to record: its unique gate coverage is scraping a REAL Nautilus
     run's output directory into the SQLite database (the misc/aggregator/
     trio on the gate covers the same scrape/aggregator API under
     PYAUTO_TEST_MODE=1 fixtures); a lens-modelling PR that breaks real-run
     database serialization without tripping the aggregator trio, the JAX
     likelihood pins, or the library unit tests is the least likely regression
     of the three, and it remains covered twice weekly.
   • Fallback if demotion is declined: keep it but shrink the real search
     (n_like_max 300→100 and/or a smaller in-script circular mask) — its
     assertions are relational (> 0, name/info checks), not pinned literals,
     so a smaller fit preserves them; verify the scrape still ingests > 0
     samples and every agg.values(...) generator still yields.

3. point_source/jax_likelihood/point.py → drop the gate-redundant Solved
   block (or demote the script).

   • Primary: remove the __Model: Solved Source (Parameter-Free)__ block
     (the second _vmap compile, 33s measured, plus its ~8s NumPy
     cross-check) from point.py. Its coverage — FitPositionsImagePairAllSolved
     with a pinned literal — already exists in image_plane.py, which the
     weekly workspace-smoke.yml and release-integrate run; update
     point.py's docstring (it currently says the solved variant is here
     BECAUSE this script is the gate entry — record the demotion decision and
     name the covering channel). The base FitPositionsImagePairAll pin
     (-83.38049778), the NumPy cross-check and the jit(fit_from) round-trip
     stay on the gate untouched — no literal re-pins needed.
     Est. 47.7 → ~18s CI.
   • Do NOT bother shrinking pixel_scale_precision / solver grid: measured a
     10x loosening saves only ~15% (XLA floor), and it would force re-pinning
     both literals for nothing (cf. test: repin vmap literals halved by the PositionsLH penalty fix (PyAutoLens#700) #257 repin pain).
   • Alternative if the ~6:00 acceptance is read strictly: demote the whole
     script to the weekly channel (comment in smoke_tests.txt naming
     image_plane.py + weekly/release as the covering surface) — but this
     removes the last point-source entry from the PR gate; flag that trade
     explicitly to the human.

4. Re-measure + close out.

   • Run the full smoke gate locally (python .github/scripts/run_smoke.py)
     and on the PR; grep the [PASS] <name> — <n>s lines; confirm 23/23 (or
     22/23+new script count if split) and py3.12 script time < ~360s.
   • Projected arithmetic from measurements (CI seconds): 553.0 − 84
     (subhalo split, 120.7→37) − 63.4 (general demote) − 30 (point Solved
     block, 47.7→18) ≈ 376s ≈ 6m16s; with point.py demoted entirely
     instead: ≈ 358s ≈ 5m58s. The recommended set lands ~16s above a
     strict 360s; whether "~6 minutes" tolerates that, or point.py should be
     fully demoted, is the one open acceptance call — flag it on the PR.

Key Files

• scripts/imaging/subhalo_recovery.py — split: one-shot stays; iterative leg
  moves to scripts/imaging/subhalo_recovery_iterative.py (new)
• smoke_tests.txt — general.py demotion comment; no new-script addition
• scripts/misc/database/scrape/general.py — only if fallback-shrink chosen
• scripts/point_source/jax_likelihood/point.py — remove the Solved block
  (coverage stays in image_plane.py on the weekly channel); no re-pins
• config/build/profile_smoke.yaml — unchanged (no cap changes; constraint 3)

Work Classification

Workspace

Worktree root

~/Code/PyAutoLabs-wt/smoke-gate-slowest-scripts/ (created later by /start_workspace)

Original Prompt

Click to expand starting prompt

Speed up the three slowest autolens_workspace_test smoke-gate scripts

Type: test
Target: workspaces
Repos:

• autolens_workspace_test
  Difficulty: medium
  Autonomy: supervised
  Priority: normal
  Status: formalised

The per-PR smoke gate in autolens_workspace_test costs ~11m20s wall-clock, of
which ~9m13s is script execution. Three entries are 42% of that. This prompt is
the make them cheaper half; cutting how often the gate runs at all is the
sibling prompt draft/test/pyautoheart/smoke_relevance_gate.md. Do not merge
the two — one edits scripts, the other edits a workflow, and they land in
different repos.

Measured

CI run 32605025472 (2026-08-22, main), py3.12 leg (the critical path; py3.13
is ~6% faster). 23 entries, 23/23 pass, 553.0s total — which reconciles to
the step wall-clock exactly, so the runner adds no measurable overhead and the
scripts are the cost.

py3.12	share	script
120.7s	21.8%	imaging/subhalo_recovery.py
63.4s	11.5%	misc/database/scrape/general.py
47.7s	8.6%	point_source/jax_likelihood/point.py
33.9s	6.1%	imaging/jax_likelihood/rectangular.py
31.5s	5.7%	misc/jax_assertions/delaunay_nn.py
30.2s	5.5%	imaging/jax_likelihood/mge.py

The remaining 17 entries are 4.5–29.2s each and are not in scope. Reproduce
with gh api on the job log and grep the runner's [PASS] <name> — <n>s
lines; the runner prints one per entry.

Note the tail is short: after these three the curve flattens, so this prompt
can win ~4 minutes and no more. Do not chase entries below ~30s.

Task

Per script, either make it materially faster or demote it — both are
acceptable outcomes, and the choice is per script, not global.

1. imaging/subhalo_recovery.py (120.7s). Already the subject of one
   speed-up pass: complete/2026/08/potential-correction-validation.md leg 1
   recorded it at 232s/224s against the 300s cap, and it now runs at 120.7s, so
   half the work is done and the why is documented there — read it before
   re-deriving. It asserts end-to-end dkappa recovery of a simulated 1e10
   Msun subhalo for both the one-shot and iterative engines. Ask whether the
   PR gate needs both engines or whether one belongs on the weekly channel.
2. misc/database/scrape/general.py (63.4s). Un-parked on 2026-07-21
   (chore(no_run): un-park database/scrape/general, autolens_workspace_test#192).
   A database-scrape regression is the least likely of the three to be broken
   by a typical lens-modelling PR, so it is the strongest demotion candidate —
   check what it uniquely covers before deciding.
3. point_source/jax_likelihood/point.py (47.7s). Check whether the cost
   is PointSolver iterations or JAX compile time; if compile-dominated, the
   lever is problem size, not sample count.

The cap does not apply to these scripts, and cannot be made to

Read this before proposing "just turn PYAUTO_SMALL_DATASETS on".

config/build/profile_smoke.yaml has set PYAUTO_SMALL_DATASETS: "1" in its
defaults since the file was created (2026-04-08, then spelled
PYAUTO_WORKSPACE_SMALL_DATASETS; renamed 2026-04-30). It has never been
missing. What has grown is the set of entries exempted from it:

date	entries exempted from the cap
2026-04-30	10 profile overrides
2026-05-28	23
2026-07-23	27 — then migrated into in-file ENV: declarations (PyAutoHands#187), leaving 1 override
today	90 ENV: … full_datasets declarations repo-wide

Of the 23 live smoke entries, 16 declare ENV: [jax] full_datasets and a
17th (misc/database/scrape/general.py) is exempted by the one surviving
profile override. That is 516.5s of the 553.0s — 93% of the gate — running
uncapped. Only six entries actually run capped (3 aggregator + 2 latent at
9-pixel masks, multi_galaxy/model_fit.py at 80), totalling 36.5s.

Measured mask sizes from the same CI run, against autolens_workspace's
80–208 pixels:

entry	mask
imaging/subhalo_recovery.py	5858 px
misc/database/scrape/general.py	2828 px
multi_galaxy/jax_likelihood/lp.py	2828 px
imaging/jax_likelihood/potential_correction.py	1466 px
imaging/jax_likelihood/{rectangular,mge}.py	952 px
imaging/jax_likelihood/{lp,smbh}.py	716 px

The exemptions are deliberate and load-bearing: the jax_likelihood scripts
assert hardcoded full-resolution likelihood literals at rtol 1e-4, so a capped
dataset changes the likelihood and the assertion fails. That failure has been
hit and recorded — see complete/2026/08/sph-transform-name-check.md
("ran with capped datasets against full-resolution reference assertions and
failed") and complete/2026/08/mge-sigma-min-workspace-sweep.md.

So the real question this prompt has to answer is not "why is the cap off"
but "should the per-PR gate be running full-resolution parity assertions at
all?" For each of the three scripts, the options are: re-derive the literals
at a capped size (buys the cap but re-pins every literal, and those pins have
already had to be regenerated once — autolens_workspace_test#257, "repin vmap
literals halved by the PositionsLH penalty fix", which blocked the nightly
release at Stage 3 on 08-18/08-19), shrink the problem some other way, or
demote to the weekly channel that is meant for full-fidelity work. Choose per
script and say which.

Constraints

• A faster script that no longer tests anything is a regression, not a win.
  This repo has three recorded instances of exactly that failure mode: the
  vacuous JAX assertions (complete/2026/07/vacuous-jax-assertions.md), the
  NUFFT parity legs that compared nufftax against itself and reported
  max |Δ| = 0.0000e+00, and latent/latent_nan_robustness passing vacuously
  under the smoke profile (see planned.md). For every reduction, state what
  the assertion still discriminates against and show it failing when the thing
  it guards is broken.
• Coverage given up is not coverage lost. Heart's weekly workspace-smoke.yml
  (Mondays 03:00 UTC) already runs every script in this repo against library
  main under the same profile_smoke.yaml, and release-integrate re-runs
  the matrix at PYAUTO_TEST_MODE=0, full resolution. The curated PR list is a
  strict subset of both. Demotion = removing the entry from smoke_tests.txt
  with a comment saying which channel still covers it.
• Do not raise BUILD_SCRIPT_TIMEOUT for these. The 300s smoke cap is a
  runaway detector; three entries sitting under it is the problem, not the cap.
• Keep the entries that stay inside the cap with real margin — 120.7s against
  300s already flaked into timeout historically under sweep-load contention
  (planned.md records 252s uncontended vs a 300s timeout under load).

Acceptance

• Smoke step wall-clock for the py3.12 leg drops below ~6 minutes.
• 23/23 still pass, and every touched script has a stated, demonstrated
  discriminating assertion.
• Any demoted entry is commented in smoke_tests.txt naming the channel that
  still runs it.

[Jammy2211]

Shipped — #269 merged (a6965a5).

Outcome: py3.12 smoke leg 553.0s → 301.3s (5m01s), 22/22 pass; py3.13 351.1s, 22/22. Acceptance (<~6 min) met with margin, without demoting point.py.

Per script: imaging/subhalo_recovery.py 120.7s → 13.2s (split; iterative leg → subhalo_recovery_iterative.py on the weekly/release channels), point.py 47.7s → 22.7s (Solved block dropped, pinned in image_plane.py weekly), general.py demoted with the channel-naming comment in smoke_tests.txt. Every touched script's surviving assertion was demonstrated to fail when the guarded thing is broken (details on the PR).

Completion recorded in PyAutoMind (complete/2026/08/smoke-gate-slowest-scripts.md). The sibling gate-frequency prompt (smoke_relevance_gate, PyAutoHeart) remains a separate task.

---

Generated by Claude Code