Profile SLaM stage-resume overhead + judge checkpointing vs speed-up

[Jammy2211]

Overview

Resuming a partially-complete SLaM pipeline can take significant wall time even though every completed stage skips its search: inter-stage work (result loading, max-likelihood fit recomputation, adapt-image construction, visualization re-output, JAX JIT compile) runs again on every resume, slowing the science loop. This task adds a pipeline_resume/ profiling tier to autolens_profiling that measures where resume time goes across a full 5-stage SLaM pipeline, then delivers a judgment: targeted speed-ups vs a checkpoint system (weighing the adapt-image data-flow complexity a checkpoint design would need to handle).

Plan

• Add a new pipeline_resume/ tier to autolens_profiling that runs a full 5-stage SLaM pipeline to completion, then re-runs it to profile the resume path.
• Instrument both runs with per-stage and inter-stage timers: samples load, result rebuild, max-likelihood fit recompute, adapt-image construction, visualization/output, JIT compile.
• Write versioned JSON + PNG results under results/ per repo convention, and register the tier in the README dashboard.
• Run cold + resume on CPU (optionally laptop GPU) and capture the breakdown.
• Deliver the judgment on this issue: whether resume is slow, where the time goes, and whether the fix is targeted speed-ups (e.g. persist adapt images to files/ at stage completion and load on resume) or a full checkpoint system.
• Any resulting library change (PyAutoFit/PyAutoLens caching or checkpointing) is filed as a follow-up prompt, not implemented here.

Detailed implementation plan

Work Classification

Workspace

Worktree root

~/Code/PyAutoLabs-wt/slam-resume-profiling/ (created by /start_workspace)

Affected Repositories

• autolens_profiling (primary)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_profiling	main	clean

Note: stale unregistered branch feature/jax-optimizer-tuning exists from the merged optimizer-benchmark work — warning only, no worktree claim conflicts.

Suggested branch: feature/slam-resume-profiling

Implementation Steps

1. pipeline_resume/slam_resume.py — a full 5-stage SLaM chain mirroring autolens_workspace/scripts/guides/modeling/slam_start_here.py (SOURCE LP → SOURCE PIX (adaptive pixelization via al.AdaptImages from the prior result) → LIGHT LP → MASS TOTAL → SUBHALO), so a bottleneck in any real stage boundary is captured rather than trimmed away. Small auto-simulated dataset, fast search settings so the cold run completes in reasonable time. Reuses _profile_cli.py conventions (device-info capture, output-path resolver, auto-simulate hook) without forcing the script into the JIT-sweep mold.
2. Instrumentation — wall-clock timers around each search.fit() and each inter-stage block (result → galaxy_name_image_dict → AdaptImages, model composition). On the resume run, decompose inside the completed-fit path (PyAutoFit/autofit/non_linear/search/abstract_search.py → paths.is_complete → result_via_completed_fit): samples CSV load, result rebuild, max-LH fit recompute, visualization re-output — via targeted timer wrappers in the script (no library edits).
3. Output — versioned summary JSON + stacked-bar PNG (stage × component, cold vs resume) under results/, shapes per results/README.md; add the tier to scripts/build_readme.py and the README dashboard.
4. Measurement — run cold + resume on CPU (optionally laptop GPU; cap JAX GPU prealloc). Hypotheses to test: per-stage JAX JIT recompile and adapt-image fit recomputation dominate resume; samples parsing and visualization are secondary. Adjacent to the filed jax_compile_time_profiling intake — this task quantifies compile's share of resume specifically.
5. Judgment (deliverable) — written assessment on this issue: is resume overhead material; component breakdown; recommendation of targeted speed-ups (e.g. persist adapt images / galaxy_name_image_dict into files/ at stage completion, load on resume; skip visualization re-output on completed fits; lazy samples load) vs a full checkpoint system — explicitly weighing the adapt-image pass-through problem: checkpointed resume must load adapt images from the right stage's output, which may make a general checkpoint system unnecessarily complex.
6. Follow-up seam — any library change lands as a new PyAutoMind prompt via intake; out of scope here.

Key Files

• autolens_profiling/pipeline_resume/slam_resume.py — new profiling script (this task's main artifact)
• autolens_profiling/_profile_cli.py — shared helper conventions to reuse
• autolens_profiling/results/README.md — result JSON/PNG shape conventions
• autolens_profiling/scripts/build_readme.py — README dashboard generator to extend
• autolens_workspace/scripts/guides/modeling/slam_start_here.py — canonical 5-stage SLaM chain to mirror (read-only reference)
• PyAutoFit/autofit/non_linear/search/abstract_search.py — resume path (is_complete → result_via_completed_fit) being profiled (read-only reference)

Testing

• Run the script twice (cold, then resume) and verify the JSON schema, presence and sanity of per-stage/per-component timings.
• ruff check . — the repo's PR lint gate (lint.yml, Python 3.12).

Original Prompt

Click to expand starting prompt

It is common for a user to set off a lens model, get so far, and then resume that model. For SLaM pipelines,
this can take time if it has to do tasks between each stage, which makes doing the science slower.

First, can you extend autolens_profiling to have some testing and profiling code for this, its part of the
core run time and thus falls under profiling rather than hygeine (e.g. in the context of agent work).

Then, can you assess if it is slow, whether the best way forward is to try and speed it up or if it would
be feasily to build a checkpoing system, where resuming goes straight to the latest result to
resume? Normally this would be fine, the issue is that stuff like adapt images, which get passed through
the pipeline, would need to be carefully thought about so the resume has their data loaded from the right place.
This may make a full chckpointing system unnecessary complex. You judge.

[Jammy2211]

Progress update (2026-07-16, end of day)

Harness built and pushed — branch feature/slam-resume-profiling, commits 9129d49 + 0af2d6e:

• pipeline_resume/slam_resume.py: full 5-stage SLaM chain (mirrors slam_start_here.py), cold/partial/resume mode detection via .completed markers, per-stage spans (search_fit / adapt_images / positions / model_compose) plus runtime wrappers decomposing the completed-fit path (samples summary load, full samples load) — no library edits.
• Versioned JSON+PNG summaries under results/pipeline_resume/; build_readme.py gained the pipeline-resume auto-table; tier README + root/results README index rows.

Traps hit and fixed

• Nautilus default n_batch=100 OOMs a contended 16GB machine on the source_lp MGE model — the script now mirrors slam_start_here.py (50/20; 4 under --fast). n_batch is not an identifier field, so tuning it preserves completed stages.
• Ruff's import sort moved jax_wrapper below autofit/autolens (silently breaks JAX env setup) — pinned with # isort: off.

Findings so far (pre-measurement, verified in source)

• The per-stage "checkpoint" largely already exists: save_attributes writes each stage's consumed adapt images into files/, and agg_util.adapt_images_from already reconstructs AdaptImages from them. The resume cost comes from the SLaM script recomputing adapt images from the upstream result (max-LH fit rebuild → JIT compile + inversion) instead of loading the saved copy.
• Resume also pays a per-stage zip→unzip restore cycle (post-fit output zips the search dir).
• Side project filed and picked up by a parallel session: PYAUTO_TEST_MODE_SAMPLES size-realistic bypass outputs (PyAutoFit#1378), which will make cold runs of this tier near-instant.

In flight overnight — the --fast (n_live=50, n_batch=4) chain is running detached on CPU: stage 1 complete, stage 2 sampling (~4.8k calls), stages 3–5 to follow.

Next session — pure-resume pass → per-component decomposition → full-fidelity cold+resume runs → the speed-up-vs-checkpoint judgment here.

[Jammy2211]

Strategy change (2026-07-17): sampling-based cold runs are paused (two OOM kills under machine contention + a WSL reboot ate the overnight chain mid-stage-2). Cold-run generation will instead use the test-mode feature this task spawned: PYAUTO_TEST_MODE=2 PYAUTO_TEST_MODE_SAMPLES=~10000 (PyAutoFit#1378 design / #1379 implementation, now in progress) — instant 5-stage completion with production-size samples.csv, so the resume measurements stay honest.

Harness is already test-mode-ready (96acde6): mode/reset detection follows the output/test_mode/ namespacing, test-mode runs write to a _testmode-suffixed artifact with the level recorded per run, and the resume invocation must keep the env vars set. Resuming here once #1379 ships.

[Jammy2211]

UNPAUSED — PYAUTO_TEST_MODE_SAMPLES is live on main (PyAutoLabs/PyAutoNerves#126 + PyAutoLabs/PyAutoFit#1381 merged 2026-07-17; design PyAutoLabs/PyAutoFit#1378). The test-mode epic's phase 3 is folded into this task (its draft prompt is retired) — this task owns the pipeline_resume subtree and its harness is already test-mode-ready.

The resume recipe in the active.md note stands; phase 3 adds these deliverables on top:

1. Size-parity validation: samples.csv row count + byte size of the bypass-completed stages vs the production target measured on #1378 (source_lp[1] = 10,187 rows × 21 cols = 9.07 MB; pick N accordingly, ~10k). Record the comparison with the results.
2. README recipe for the pipeline_resume tier: cold PYAUTO_TEST_MODE=2 PYAUTO_TEST_MODE_SAMPLES=<N> ... --reset, resume with the same env vars (output/test_mode/ namespacing trap), plus the known deltas — latents auto-skipped, no search-internal checkpoint, adapt-image FITS values are prior-median (size right, values wrong; timing-honest not science-honest).
3. Benchmarks on an idle machine; then the decomposition → judgment deliverable already scoped here (load-not-recompute vs checkpoint).

Phase 4 (library-side docs) stays a separate small task, issued after this ships.

[Jammy2211]

Measurement + judgment (2026-07-17)

Is resuming slow? Yes — and it is almost entirely recomputation, not loading.

Instant test-mode chain (PYAUTO_TEST_MODE=2 PYAUTO_TEST_MODE_SAMPLES=10000, production-fidelity models, CPU, v2026.7.9.1): cold 205s → resume of the fully-complete 5-stage chain = 148s of pure overhead. Size parity vs production: source_lp[1] samples.csv 10,000×21 = 8.81 MB (target 10,187×21 = 9.07 MB).

Resume component	Time	Share
positions_likelihood_from (source_pix_1 + mass_total)	70.0s	47%
Adapt-image reconstruction (source_pix_1 + source_pix_2)	54.7s	37%
Completed-fit path, all 5 stages (unzip, samples summary, result rebuild; stage 1 carries the check_likelihood_function recompute + first JAX compile at 9.3s)	12.2s	8%
Imports + model composition	~10.5s	7%
Full samples CSV load (10k rows / 8.8 MB)	0.5s	<1%

Both dominant components rebuild the upstream result's max-likelihood fit / tracer — i.e. they pay a fresh JAX compile + inversion/point-solve per resume. On this laptop-scale dataset that is ~2.5 min per resume; on production data (Delaunay meshes, GPU compile in the minutes, cluster scale) it scales up and is paid on every resume of a SLaM campaign. Two adapt legs were even free here only because the in-process cached_property was warm — a fresh resume process pays all four.

Judgment: targeted speed-ups. A checkpoint system is not warranted.

The prompt's worry — "adapt images get passed through the pipeline, resume would need their data loaded from the right place" — turns out to be already solved by existing persistence: every stage writes the adapt images it consumed into its own files/ (AnalysisDataset.save_attributes), and autogalaxy/aggregator/agg_util.adapt_images_from already reconstructs AdaptImages from that artifact. The "right place" is the consuming stage's own output — provenance-free. A general checkpoint system would re-solve exactly this per-object data-flow problem at much higher complexity, for components that are already ~free (samples load 0.5s, zip restore ~0.7s/stage).

Recommended fixes (filed as follow-up prompts, not implemented here):

1. Adapt images disk-first (~55s here, more in production): when composing stage N's analysis on a resume, load files/adapt_images from stage N's own output if present instead of recomputing from the upstream result. Loader machinery exists.
2. Persist + reload the positions likelihood (~70s here): positions_likelihood_from re-runs the point solver from the upstream tracer on every resume; the object is small and serializable — save it into the consuming stage's files/ at composition time, load when present.

Together these attack ~85% of measured resume overhead. The remaining tail (imports, model compose, the stage-1 consistency check) is either fixed-cost or a deliberate integrity feature.

Extras from this measurement session

• New config/general.yaml (workspace parity): the 5-stage chain previously crashed at source_pix[2] on the pixelization-without-positions guard — the real-sampling runs had never reached stage 3 to expose it; the instant path found it in 2 minutes.
• PyAutoFit gap found + filed (draft/bug/autofit/): _fit_bypass_test_mode returns without paths.completed(), so bypassed stages were never resumable; the harness carries a one-line runtime stopgap.
• README now documents the instant recipe, the same-env-vars trap, size-parity numbers and known deltas (phase-3 deliverables).

Benchmarks ran under moderate background load (load avg ~4); the decomposition shape is robust but absolute numbers are a floor. Next: /ship_workspace for the PR, plus the follow-up intake for the two fixes.

[Jammy2211]

Shipped — 2026-07-17

PRs

• pipeline_resume tier: SLaM 5-stage resume-overhead profiler + instant test-mode recipe (#70) #75

Summary

pipeline_resume/ tier shipped: full 5-stage SLaM resume-overhead profiler with per-component decomposition, instant test-mode cold path (phase 3 deliverables: recipe, size parity, reference benchmarks), and the workspace-parity config/general.yaml. Measurement + judgment are in the comment above: resume ≈148s pure overhead, 85% recomputation; recommendation is the two targeted load-not-recompute fixes (follow-up prompt filed), not a checkpoint system.

Session Notes

• Shipped through Heart RED on human ack (5 pre-existing reasons, quoted in-session; none from this branch).
• PyAutoFit gap found: test-mode bypass never writes .completed → bug prompt draft/bug/autofit/test_mode_bypass_never_writes_the_completed.md; harness carries a one-line runtime stopgap to remove when it ships.
• Follow-up feature: draft/feature/autogalaxy/slam_resume_fast_path_load_persisted_adapt.md (adapt-images disk-first + persist/reload positions likelihood; validated before/after via this tier's instant recipe).
• Real-sampling partial outputs (stage 1 + stage-2 checkpoint) remain under output/pipeline_resume/hst_fast/ locally if a sampled cross-check is ever wanted.