perf: Skip VRAM estimation, reduce I/O and formatting in test mode 2+

[Jammy2211]

Overview

With PYAUTO_WORKSPACE_SMALL_DATASETS=1 and PYAUTO_DISABLE_CRITICAL_CAUSTICS=1 (shipped in PyAutoLabs/autolens_workspace#34), smoke test scripts still take ~40-50s each under PYAUTOFIT_TEST_MODE=2. Profiling shows the remaining time is dominated by PyAutoFit-level operations: VRAM estimation (16s), pre-fit file I/O (10s), model.info formatting (7s), and result deserialization (5s). Addressing these should bring total runtime under 10s per script.

Plan

• Skip print_vram_use() when PYAUTOFIT_TEST_MODE >= 2 — avoids expensive JAX JIT compilation + memory analysis
• Reduce file I/O in _fit_bypass_test_mode() — skip saving samples, results, and visualization in test mode
• Speed up or skip model.info formatting for large models in test mode
• Speed up result_info_from() — skip expensive sigma quantile calculations in test mode

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)

Work Classification

Library

Branch Survey

Repository	Current Branch	Dirty?
PyAutoFit	main_build	clean

Suggested branch: feature/smoke-test-fast-pyautofit

Implementation Steps

1. Skip VRAM estimation in test mode — In autofit/non_linear/analysis/analysis.py:319, add early return in print_vram_use() when test_mode_level() >= 2. Saves ~16s per script.

2. Reduce file I/O in test mode bypass — In autofit/non_linear/search/abstract_search.py:790, the _fit_bypass_test_mode() method calls save_samples_summary(), save_samples(), save_results(), save_results_combined(), and paths.completed(). Skip save_samples(), save_results(), and save_results_combined() in test mode 2+ (keep save_samples_summary for result access). Saves ~8-10s.

3. Speed up model.info — In autofit/mapper/prior_model/abstract.py:1775, the info property iterates all path/prior tuples, calls find_groups(), and formats via TextFormatter. For MGE models with ~40 gaussians this takes 7s. Options: cache the property, or return a minimal stub when test_mode_level() >= 2.

4. Speed up result_info_from — In autofit/text/text_util.py:53, result_info_from() calls samples_text.summary() twice (sigma=3.0, sigma=1.0). The values_at_sigma() computation is expensive. Skip in test mode 2+.

Key Files

• autofit/non_linear/analysis/analysis.py:319 — print_vram_use()
• autofit/non_linear/search/abstract_search.py:790 — _fit_bypass_test_mode()
• autofit/mapper/prior_model/abstract.py:1775 — info property
• autofit/non_linear/result.py:85 — result.info
• autofit/text/text_util.py:53 — result_info_from()
• autofit/text/samples_text.py:15 — summary()
• autofit/non_linear/test_mode.py — test_mode_level()

Original Prompt

Click to expand starting prompt

Continuation of the smoke test speed-up work. The first phase shipped env vars
for small datasets and disabled critical curves (autolens_workspace#34). This
phase targets the remaining PyAutoFit-level bottlenecks identified by profiling.

With PYAUTO_WORKSPACE_SMALL_DATASETS=1 and PYAUTO_DISABLE_CRITICAL_CAUSTICS=1,
the imaging and interferometer modeling scripts still take ~40-50s each under
PYAUTOFIT_TEST_MODE=2. The remaining time is dominated by:

1. VRAM estimation: ~16s — analysis.print_vram_use() does full JAX JIT
   compilation + memory analysis. Should be skipped when PYAUTOFIT_TEST_MODE >= 2.

2. search.fit pre-fit I/O: ~10s — _fit_bypass_test_mode() writes model.info,
   visualization, saves samples, saves results. Much of this can be skipped in
   test mode 2+.

3. model.info formatting: ~7s — The info property on MGE models with ~40
   gaussians takes 7s to format. Can be cached or stubbed in test mode.

4. Result access: ~5s — max_log_likelihood_instance takes ~5s to
   deserialize. The result_info_from() function calls expensive
   samples_text.summary() twice.

Target: get both scripts under 10s total with all env vars set.

[Jammy2211]

Progress Update — 2026-04-07

PRs Merged

Library PRs:

• perf: Skip VRAM, I/O, model.info in test mode 2+ and add PYAUTO_DISABLE_JAX #1184 — skip VRAM estimation, model.info, result_info, pre/post-fit I/O in test_mode >= 2; add PYAUTO_DISABLE_JAX env var
• feat: PYAUTO_WORKSPACE_SMALL_DATASETS env var for fast smoke tests PyAutoArray#253 — add PYAUTO_WORKSPACE_SMALL_DATASETS env var (cap grids/masks to 15x15)
• feat: add tight_layout wrapper with PYAUTO_FAST_PLOTS env var PyAutoArray#254 — add tight_layout wrapper with PYAUTO_FAST_PLOTS env var (43 call sites across 3 repos)
• perf: skip savefig rendering in PYAUTO_FAST_PLOTS mode PyAutoArray#255 — skip savefig rendering in PYAUTO_FAST_PLOTS mode
• feat: PYAUTO_DISABLE_CRITICAL_CAUSTICS env var to skip curve computation PyAutoGalaxy#325 — add PYAUTO_DISABLE_CRITICAL_CAUSTICS env var
• feat: use tight_layout wrapper for PYAUTO_FAST_PLOTS support PyAutoGalaxy#326 — use tight_layout wrapper
• feat: use tight_layout wrapper for PYAUTO_FAST_PLOTS support PyAutoLens#422 — use tight_layout wrapper

Workspace PRs:

• docs: add PYAUTO_FAST_PLOTS to fast smoke test env vars autolens_workspace#36 — CLAUDE.md docs
• docs: add PYAUTO_FAST_PLOTS to fast smoke test env vars autogalaxy_workspace#12 — CLAUDE.md docs
• docs: add PYAUTO_FAST_PLOTS to fast smoke test env vars autofit_workspace#21 — CLAUDE.md docs

Infrastructure:

• PyAutoBuild release.yml updated with all env vars (direct push to main)
• autofit_workspace run_scripts.sh updated
• Smoke test SKILL.md updated with all env vars

Optimizations Implemented

Optimization	Env Var	Savings	Repo
Cap grids/masks to 15x15 at 0.6"/px	PYAUTO_WORKSPACE_SMALL_DATASETS=1	~90% grid computation	PyAutoArray
Skip critical curve/caustic overlays	PYAUTO_DISABLE_CRITICAL_CAUSTICS=1	~2-5s per plot	PyAutoGalaxy
Force use_jax=False	PYAUTO_DISABLE_JAX=1	Avoids JAX JIT overhead	PyAutoFit
Skip tight_layout() in 43 subplot functions	PYAUTO_FAST_PLOTS=1	~0.5-1s per subplot	PyAutoArray/Galaxy/Lens
Skip savefig rendering	PYAUTO_FAST_PLOTS=1	~2s per subplot with output	PyAutoArray
Skip VRAM estimation	test_mode >= 2	~16s	PyAutoFit
Skip model.info formatting	test_mode >= 2	~7s for MGE models	PyAutoFit
Skip result_info formatting	test_mode >= 2	~5s	PyAutoFit
Skip pre/post-fit file I/O	test_mode >= 2	~10s	PyAutoFit

Results — imaging/simulator.py benchmark

Stage	Time
Original (no env vars)	~100s
+ all env vars (before fast plots)	12.4s
+ PYAUTO_FAST_PLOTS (tight_layout skip)	9.0s
+ PYAUTO_FAST_PLOTS (savefig skip)	6.4s

Remaining time is dominated by Python imports (~4.7s) and numba JIT (~0.7s).

Remaining Work

• Profile remaining smoke test scripts (modeling, interferometer, etc.)
• Import time (~2-6s depending on numba cache) is the irreducible floor per script
• Consider lazy imports as a future optimization if needed

[Jammy2211]

Progress Update — 2026-04-07 (session 2)

PRs Merged (this session)

• perf: skip savefig rendering in PYAUTO_FAST_PLOTS mode PyAutoArray#255 — skip savefig rendering in PYAUTO_FAST_PLOTS mode
• perf: force over_sample_size=2 when PYAUTO_WORKSPACE_SMALL_DATASETS=1 PyAutoArray#257 — force over_sample_size=2 when PYAUTO_WORKSPACE_SMALL_DATASETS=1
• perf: skip radial bins computation when PYAUTO_WORKSPACE_SMALL_DATASETS=1 PyAutoArray#258 — skip radial bins computation when PYAUTO_WORKSPACE_SMALL_DATASETS=1
• feat: add test_mode module for shared test mode utilities PyAutoConf#83 — move test_mode module to autoconf (fix PyAutoArray CI)
• perf: Skip VRAM, I/O, model.info in test mode 2+ and add PYAUTO_DISABLE_JAX #1184 — skip VRAM estimation, model.info, result_info in test_mode >= 2 (was open from previous session)
• refactor: re-export test_mode from autoconf #1185 — re-export test_mode from autoconf
• fix: import is_test_mode from autoconf instead of autofit PyAutoArray#256 — fix hilbert.py import (autofit → autoconf)
• refactor: import is_test_mode from autoconf PyAutoGalaxy#327 — use autoconf.test_mode import
• perf: skip savefig in _save_subplot when PYAUTO_FAST_PLOTS=1 PyAutoGalaxy#328 — skip savefig in _save_subplot
• refactor: import is_test_mode from autoconf PyAutoLens#423 — use autoconf.test_mode import

Results

imaging/simulator.py: ~100s → 3.6s (96% reduction, warm cache)
imaging/modeling.py: ~100s → 19.5s (80% reduction)

• Remaining modeling time: search.fit bypass (7.9s) + subplot data computation (6.1s) + imports (3s)

Remaining bottleneck identified

subplot_fit_imaging spends 4.8s on data computation, not rendering:

• angular_diameter_distance_kpc_z1z2 called 176 times (1.5s) — excessive for a 2-plane lens
• traced_grid_2d_list_from (2.3s) — ray-tracing repeated per subplot panel

Next session will investigate speeding up cosmology/ray-tracing data computation.

[Jammy2211]

Progress Update — 2026-04-07 (session 3)

PRs Merged

• perf: cache cosmology distances and cap MGE gaussians for smoke tests PyAutoGalaxy#329 — cache cosmology distances + cap MGE to 2 gaussians
• perf: cache galaxy_model_image_dict and model_images_of_planes_list PyAutoLens#424 — cache galaxy_model_image_dict and model_images_of_planes_list

Results — imaging/modeling.py

Stage	Time
Original	~100s
+ all env vars (before this session)	50.4s
+ PyAutoFit PR #1184 (VRAM/info skips)	19.5s
+ cached_property on FitImaging	14.6s
+ Cosmology distance cache	13.5s
+ 2-gaussian MGE cap	10.5s

Results — imaging/simulator.py

~100s → 3.6s (warm cache)

What's left

• search.fit test-mode bypass takes 5.9s — dominated by search setup/teardown overhead
• subplot rendering is now at its floor (matplotlib object creation + one physics pass)
• Import time (~2.5s) is irreducible per-script

[Jammy2211]

Progress Update — 2026-04-07 (session 3, final)

PRs Merged (this session)

• perf: skip savefig rendering in PYAUTO_FAST_PLOTS mode PyAutoArray#255 — skip savefig in PYAUTO_FAST_PLOTS
• perf: force over_sample_size=2 when PYAUTO_WORKSPACE_SMALL_DATASETS=1 PyAutoArray#257 — force over_sample_size=2 for small datasets
• perf: skip radial bins computation when PYAUTO_WORKSPACE_SMALL_DATASETS=1 PyAutoArray#258 — skip radial bins computation for small datasets
• perf: skip savefig in _save_subplot when PYAUTO_FAST_PLOTS=1 PyAutoGalaxy#328 — skip savefig in _save_subplot
• perf: cache cosmology distances and cap MGE gaussians for smoke tests PyAutoGalaxy#329 — cache cosmology distances + cap MGE to 2 gaussians
• perf: cache galaxy_model_image_dict and model_images_of_planes_list PyAutoLens#424 — cache galaxy_model_image_dict and model_images_of_planes_list
• fix: PYAUTO_DISABLE_JAX env var was being overridden PyAutoLens#425 — fix: PYAUTO_DISABLE_JAX env var was being overridden (biggest single win)
• feat: add test_mode module for shared test mode utilities PyAutoConf#83 — move test_mode to autoconf
• perf: Skip VRAM, I/O, model.info in test mode 2+ and add PYAUTO_DISABLE_JAX #1184 — skip VRAM/info in test mode (was open, now merged)
• refactor: re-export test_mode from autoconf #1185 — re-export test_mode from autoconf

Final Results

Script	Original	Final	Reduction
imaging/simulator.py	~100s	3.6s	96%
imaging/modeling.py	~100s	5.3s	95%

Key bug found

AnalysisDataset.__init__ had a stale self._use_jax = use_jax that overwrote the PYAUTO_DISABLE_JAX env var check from Analysis.__init__. This caused JAX JIT to run during smoke tests (~3.5s wasted per likelihood eval). Fix saved 5.3s on search.fit alone.

Remaining time breakdown (imaging/modeling.py)

• Imports: 2.6s (irreducible Python module loading)
• Simple model composition: 0.6s
• search.fit: 0.6s (1 NumPy likelihood eval)
• Subplots: 1.0s (matplotlib object creation)
• Everything else: 0.5s

[Jammy2211]

Smoke Test Results — 2026-04-07

All 27 smoke tests passed across 4 workspaces.

Changes shipped this session

Library PRs (already merged):

• PyAutoGalaxy#334 — Fix cosmology distance cache breaking JAX JIT (guard with if xp is np:)
• PyAutoGalaxy#336 — Re-export dataset_util for ag.util.dataset.should_simulate()
• PyAutoLens#430 — Re-export dataset_util for al.util.dataset.should_simulate()
• PyAutoArray#261 — Add should_simulate() utility (auto-deletes stale datasets when PYAUTO_WORKSPACE_SMALL_DATASETS=1)

Workspace PRs (just merged):

• autogalaxy_workspace#14 — Use should_simulate() in smoke-tested scripts
• autolens_workspace#38 — Use should_simulate(), swap real-data smoke tests for auto-simulating equivalents, fix MGE iteration bug in multi/modeling.py

Smoke test skill update:

• JAX likelihood function scripts now run without PYAUTO_DISABLE_JAX=1 or PYAUTO_WORKSPACE_SMALL_DATASETS=1, ensuring the JAX code path is actually exercised

[Jammy2211]

Smoke Test Results — 2026-04-09

Ran after feature/setup-notebook-boilerplate changes (replace pyprojroot boilerplate with autoconf.setup_notebook()).

Workspace	Passed	Failed	Total
autofit_workspace	7	0	7
autogalaxy_workspace	6	0	6
autolens_workspace	7	0	7
autolens_workspace_test	6	1	7

Failures

autolens_workspace_test/database/scrape/general.py

Pre-existing issue — PYAUTO_WORKSPACE_SMALL_DATASETS=1 creates a 15x15 mask but the checked-in dataset is 100x100:

ValueError: operands could not be broadcast together with shapes (100,100) (15,15)

Not related to the setup_notebook changes.

[Jammy2211]

Shipped — 2026-04-09

PRs

Library:

• Add setup_notebook() to replace pyprojroot boilerplate PyAutoConf#85
• Simplify uncomment_jupyter_magic for setup_notebook PyAutoBuild#37
• Remove pyprojroot dependency #1189

Workspace:

• Replace pyprojroot boilerplate with setup_notebook() autofit_workspace#23
• Replace pyprojroot boilerplate with setup_notebook() autogalaxy_workspace#15
• Replace pyprojroot boilerplate with setup_notebook() autolens_workspace#39

Summary

Added autoconf.setup_notebook() — a single-call utility that replaces the 5-line pyprojroot boilerplate in every workspace script. Updated 375 files across 6 repos, simplified the notebook generation pipeline from 5 regex rules to 1, and removed the pyprojroot dependency entirely. All smoke tests pass (26/27, 1 pre-existing failure).

Session Notes

• Used config/ directory as the workspace root marker instead of .git — more specific to PyAuto workspaces and avoids matching wrong repos in nested setups.
• Followed the existing setup_colab pattern in autoconf — same os.chdir + conf.instance.push() approach.
• autolens_workspace_test and autofit_workspace_test still have old pyprojroot boilerplate (commented out, harmless) — could be migrated in a follow-up.