feat: SLaM resume fast-path — persisted adapt images + positions

[Jammy2211]

Overview

Implements the judgment of autolens_profiling#70: resuming a fully-complete 5-stage SLaM chain costs ~148s of pure overhead (laptop-scale; larger in production), of which ~85% is inter-stage recomputation — positions_likelihood_from re-running the point solver (70s) and adapt images rebuilt from the upstream result via a max-likelihood fit + fresh JIT compile (55s). Samples loading is <1%. A checkpoint system was rejected: the needed artifacts are already (or trivially can be) persisted per stage. This task adds targeted load-not-recompute paths.

Plan

• PyAutoGalaxy: adapt-images disk-first — on resume, load the persisted adapt-image artifact instead of rebuilding the upstream result's max-likelihood fit.
• PyAutoLens: persist the solved multiple-image positions (or the positions-likelihood object) on stage completion and load them on resume instead of re-running the point solver.
• First-arrival semantics everywhere: artifact absent → compute exactly as today; present → load. Staleness is structurally guarded (upstream model changes produce new search identifiers and fresh output dirs).
• Unit tests per repo (numpy-only), exercising both arrival paths.
• Validate before/after with the merged autolens_profiling pipeline_resume instant recipe (~3 min per decomposition); target is the positions + adapt buckets ≈ 0 on re-arrival.
• Library PRs ordered PyAutoGalaxy → PyAutoLens (dependency direction); workspace guide expected unchanged.

Detailed implementation plan

Work Classification

Library (PyAutoGalaxy primary, PyAutoLens second) — autolens_profiling used read-only for validation (claimed by jax-compile-time-research; no edits there).

Worktree root

~/Code/PyAutoLabs-wt/slam-resume-fastpath/ (created by /start_library)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoGalaxy	main	clean
./PyAutoLens	main	1 modified (paper_jax/paper.md — JOSS draft, unrelated subtree; warning only)

Suggested branch: feature/slam-resume-fastpath

The design decision to make first (step 1)

The expensive calls are eager at SLaM-script level, before the consuming search's paths/identifier exist — so a naive "check the consuming stage's files/ at composition time" cannot work. Two candidate hooks, decide after inspecting what each stage actually persists:

• (A) Result-property disk-first (preferred if artifacts exist upstream): ResultDataset.model_image_galaxy_dict / subtracted_signal_to_noise_map_galaxy_dict are already lazy cached_propertys (autogalaxy/analysis/result.py). When the result was loaded from a completed fit, try loading the upstream stage's persisted galaxy-image / fit artifacts from its files/ before rebuilding max_log_likelihood_fit. Zero script/API change. Requires the upstream stage to persist per-galaxy model images (check save_results / plotter-interface FITS output — files/adapt_images exists in the consuming stage via save_attributes; what the upstream stage stores decides A vs B).
• (B) Lazy maker + consuming-stage artifact: introduce AdaptImages.from_result(...) returning a lazy object; the dict computes on first access, which happens inside fit() where paths are known — modify_before_fit short-circuits from the consuming stage's own files/adapt_images (loader already exists: autogalaxy/aggregator/agg_util.adapt_images_from). Small public API addition; guide script swaps one call per stage.

Positions mirror the same choice: the solved image-plane positions are a property of the upstream result — persisting them into the upstream stage's files/ as JSON on completion (save_results) lets positions_likelihood_from load instead of re-solving (option A shape); else the consuming-stage/lazy-maker shape (B).

Implementation Steps

1. Inspect persisted artifacts per stage (files/ contents for a completed SLaM stage, using the Astropy Integration: Model fitting #70 test-mode outputs) → pick A or B per component; record the decision on this issue.
2. PyAutoGalaxy: implement the adapt-images disk-first path per the decision; reuse agg_util.adapt_images_from for reconstruction; docstring states the staleness guard (identifier change ⇒ fresh dirs).
3. PyAutoGalaxy tests: completed-fit fixture with/without the artifact; assert load vs compute path taken and identical dict contents.
4. PyAutoLens: persist solved positions on completion; positions_likelihood_from disk-first with identical first-arrival semantics; tests likewise.
5. Validation: instant recipe before/after decompositions (PYAUTO_TEST_MODE=2 PYAUTO_TEST_MODE_SAMPLES=10000); post the numbers here.
6. Ship PyAutoGalaxy PR first, PyAutoLens second (both pending-release); workspace slam_start_here.py untouched unless option B forces the one-call swap (then a minimal linked workspace PR).

Key Files

• PyAutoGalaxy/autogalaxy/analysis/adapt_images/adapt_images.py — galaxy_name_image_dict_via_result_from (the eager recompute)
• PyAutoGalaxy/autogalaxy/analysis/result.py — lazy cached_property layer (option A hook)
• PyAutoGalaxy/autogalaxy/aggregator/agg_util.py — adapt_images_from (existing disk loader)
• PyAutoGalaxy/autogalaxy/analysis/analysis/dataset.py — save_attributes (adapt artifact writer), save_results
• PyAutoLens/autolens/analysis/result.py — positions_likelihood_from and the point-solver path
• autolens_profiling/pipeline_resume/ — read-only validation harness

Testing

• Per-repo unit tests, numpy-only (no JAX in unit tests — repo rule).
• pipeline_resume instant recipe before/after decomposition as the acceptance evidence.

Original Prompt

Click to expand starting prompt

Follow-up implementing the judgment of autolens_profiling#70 (measurement + decomposition there,
2026-07-17 comment): resuming a fully-complete 5-stage SLaM chain costs ~148s of pure overhead on a
laptop-scale dataset, of which ~85% is inter-stage recomputation — positions_likelihood_from (70s)
and adapt-image reconstruction from the upstream result (55s), each paying a fresh JAX compile +
inversion/point-solve per resume. Samples loading is <1%. A checkpoint system was ruled out; these
two targeted fixes were recommended instead.

1. Adapt images disk-first. Every stage already persists the adapt images it consumed into its
   own files/ (@PyAutoGalaxy AnalysisDataset.save_attributes), and
   autogalaxy/aggregator/agg_util.adapt_images_from already reconstructs AdaptImages from that
   artifact. Add a load-from-disk path so a resumed stage composes its analysis from its own saved
   files/adapt_images when present, instead of rebuilding the upstream result's max-LH fit
   (galaxy_name_image_dict_via_result_from). Design care: first arrival at an incomplete stage
   must still compute; only re-arrivals (artifact present) take the shortcut. Opt-out knob if the
   upstream result changed (identifier mismatch invalidates the whole downstream chain anyway, so
   staleness is structurally guarded — state this in the docstring).

2. Persist + reload the positions likelihood. @PyAutoLens positions_likelihood_from re-runs
   the point solver from the upstream tracer on every resume. The resulting object is small and
   serializable — save it into the consuming stage's files/ at composition time and load it when
   present, same first-arrival semantics as (1).

Validation: the pipeline_resume tier in @autolens_profiling (instant recipe:
PYAUTO_TEST_MODE=2 PYAUTO_TEST_MODE_SAMPLES=10000 python3 pipeline_resume/slam_resume.py) gives
before/after resume decompositions in ~3 minutes; target is the ~125s recomputation bucket going to
~0 on re-arrival. Workspace follow-through: slam_start_here.py should exercise the fast path
unchanged (the loaders slot beneath the existing calls, no API change to the guide script) — confirm
and note in the guide only if a flag surfaces.

[Jammy2211]

Step-1 design decision + validation (2026-07-17)

Decision: cache-aside keyed on the upstream result's own paths (neither pure option A nor B from the plan — the inspection found the better seam):

• Result objects carry their own paths, so both eager script-level calls memoize into the upstream stage's files/: galaxy_images_model.fits/galaxy_images_snr.fits (raw pre-clip per-galaxy dict, FITS layout mirroring the aggregator's adapt_images artifact) and multiple_image_positions[_plane_<z>].json. First arrival computes exactly as before and writes; later arrivals load. Zero public-API change, zero script change.
• Zip-cycle trap found during validation: a resumed search's paths.restore() deletes the output dir and re-extracts the .zip, destroying any file written to files/ after completion — the first validation pass showed caches being wiped every resume. Fix: after writing, the cache file is also appended into the search's zip (_append_to_search_zip), making it a permanent part of the completed output.
• The viz-gated fits_adapt_images write is untouched — the cache is its own artifact, unconditional wherever the result has on-disk output (NullPaths results always compute).

Validation (instant pipeline_resume recipe, same session/load, per-stage inter-stage span):

Stage	Resume before	Resume after
source_pix[1] (adapt + positions)	151.0s	0.3s
source_pix[2] (adapt)	130.4s	0.0s
mass_total (adapt + positions)	127.8s	0.1s

Remaining resume floor: imports + stage-1's deliberate check_likelihood_function consistency recompute (~19s incl. first JAX compile), both config-controllable.

Tests: PyAutoGalaxy — FITS round-trip, poisoned-result cache-hit, pathless-result compute (3 new, file green). PyAutoLens — poisoned second-call positions cache-hit (1 new, file green, 14/14).

[Jammy2211]

Library PRs Created — 2026-07-17

• PyAutoGalaxy: feat: cache per-galaxy result images for SLaM resume fast-path (#502) #504 (merge first — PyAutoLens imports _append_to_search_zip from it)
• PyAutoLens: feat: cache solved multiple-image positions for SLaM resume fast-path PyAutoLens#619

Tests: PyAutoGalaxy 985p/0f, PyAutoLens 387p/0f (full suites, task worktree). Shipped through Heart RED on human ack (5 pre-existing reasons, quoted in-session). Six-workspace smoke gate running; results will be posted here. No workspace script changes expected (no public API change).

[Jammy2211]

Smoke Test Results — 2026-07-17

Run against the slam-resume-fastpath worktree libraries (validates PyAutoGalaxy#504 + PyAutoLens#619 downstream).

All 49 smoke tests passed across 6 workspaces.

Workspace	Passed	Failed	Skipped	Total
autofit_workspace	7	0	1	8
autogalaxy_workspace	6	0	0	6
autolens_workspace	7	0	0	7
autolens_workspace_test	12	0	4	16
euclid_strong_lens_modeling_pipeline	6	0	0	6
HowToLens	6	0	0	6

No failures to report. Skipped entries are pre-existing config/build/no_run.yaml exclusions (unrelated to this change).

[Jammy2211]

Smoke results — 2026-07-17 (feature-branch libraries)

Six-workspace smoke run with PyAutoGalaxy#504 + PyAutoLens#619 on the import path — all clean, zero tracebacks:

Workspace	Passed	Failed	Skipped	Total
autofit_workspace	7	0	1	8
autogalaxy_workspace	6	0	0	6
autolens_workspace	7	0	0	7
autolens_workspace_test	12	0	4	16
euclid_strong_lens_modeling_pipeline	6	0	0	6
HowToLens	6	0	0	6

Skips are the configured no_run.yaml entries. The three pre-existing failures Heart recorded on 2026-07-09 did not reproduce. Status caches written for /health status.

No workspace script changes needed (no public API change) — ready for the ordered merge: PyAutoGalaxy#504 first, then PyAutoLens#619.

[Jammy2211]

Shipped — 2026-07-17

PRs

• feat: cache per-galaxy result images for SLaM resume fast-path (#502) #504 (MERGED)
• feat: cache solved multiple-image positions for SLaM resume fast-path PyAutoLens#619 (MERGED)

Summary

SLaM resume fast-path landed: galaxy_name_image_dict_via_result_from and positions_likelihood_from now memoize their expensive products (per-galaxy result images; solved multiple-image positions) into the upstream result's own files/ + zip, loading on every later call. Validated end-to-end via the autolens_profiling pipeline_resume instant recipe: per-stage inter-stage resume cost 151s/130s/128s → 0.3s/0.0s/0.1s. Full suites green (985p/387p); six-workspace smoke 44p/0f.

Session Notes

• Zip-cycle trap: paths.restore() wipes post-completion files/ writes — caches must be appended into the search zip (_append_to_search_zip; shared helper in autogalaxy, candidate for a future public home in PyAutoFit paths).
• The viz-gated fits_adapt_images artifact is unchanged; the cache is its own unconditional artifact.
• Remaining resume floor: imports + stage-1 check_likelihood_function consistency recompute (deliberate, config-gated).
• Closes the arc from autolens_profiling#70 (measurement → judgment → this fix).