Skip to content

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

New issue
New issue
Closed
Closed
perf: Skip VRAM estimation, reduce I/O and formatting in test mode 2+#1183

Description

@Jammy2211
Jammy2211
opened on Apr 6, 2026

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.

Activity

  1. Jammy2211 commented on Apr 7, 2026

    @Jammy2211
    Jammy2211
    on Apr 7, 2026
    CollaboratorAuthor

    Progress Update — 2026-04-07

    PRs Merged

    Library PRs:

    Workspace PRs:

    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

  2. Jammy2211 commented on Apr 7, 2026

    @Jammy2211
    Jammy2211
    on Apr 7, 2026
    CollaboratorAuthor

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

    PRs Merged (this session)

    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.

  3. Jammy2211 commented on Apr 7, 2026

    @Jammy2211
    Jammy2211
    on Apr 7, 2026
    CollaboratorAuthor

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

    PRs Merged

    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

  4. Jammy2211 commented on Apr 7, 2026

    @Jammy2211
    Jammy2211
    on Apr 7, 2026
    CollaboratorAuthor

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

    PRs Merged (this session)

    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

  5. Jammy2211 commented on Apr 7, 2026

    @Jammy2211
    Jammy2211
    on Apr 7, 2026
    CollaboratorAuthor

    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

  6. Jammy2211 commented on Apr 9, 2026

    @Jammy2211
    Jammy2211
    on Apr 9, 2026
    CollaboratorAuthor

    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.

  7. Jammy2211 commented on Apr 9, 2026

    @Jammy2211
    Jammy2211
    on Apr 9, 2026
    CollaboratorAuthor

    Shipped — 2026-04-09

    PRs

    Library:

    Workspace:

    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.
  8. Jammy2211
    closed this as completedon Apr 10, 2026

  9. Jammy2211 commented on Apr 12, 2026

    @Jammy2211
    Jammy2211
    on Apr 12, 2026
    CollaboratorAuthor

    Smoke Test Results — 2026-04-12

    Workspace Passed Failed Total
    autofit_workspace 7 0 7
    autogalaxy_workspace 3 3 6
    autolens_workspace 7 0 7
    autolens_workspace_test 2 6 8
    Failures (9 total)

    autogalaxy_workspace/imaging/start_here.py

    KeyError: "Could not find any of the following keys in kwargs ('galaxies.galaxy.bulge.profile_list.0.centre_0', ...)"

    Bypass mode sample keys don't match lp_linear profile_list paths when reconstructing max_log_likelihood_instance.

    autogalaxy_workspace/interferometer/start_here.py

    KeyError: "Could not find any of the following keys in kwargs ('galaxies.galaxy.bulge.profile_list.0.centre_0', ...)"

    Same root cause as imaging/start_here.py.

    autogalaxy_workspace/multi/start_here.py

    KeyError: "No path from (('1', 'galaxies', 'galaxy', 'bulge', 'profile_list', '19', 'centre', 'centre_0'), ...) in sample"

    Same profile_list path mismatch in subsamples during graphical model result construction.

    autolens_workspace_test/aggregator/fit_imaging.py

    ValueError: probabilities do not sum to 1

    draw_randomly_via_pdf fails — bypass mode produces sample weights that don't sum to 1.

    autolens_workspace_test/aggregator/fit_interferometer.py

    ValueError: probabilities do not sum to 1

    Same root cause as aggregator/fit_imaging.py.

    autolens_workspace_test/jax_likelihood_functions/imaging/delaunay_mge.py

    AssertionError: delaunay_mge: JAX vmap likelihood mismatch
 x: array([-1.e+99])
 y: array(-536.770444)

    JAX vmap returns -1e99 (penalty) instead of expected likelihood.

    autolens_workspace_test/jax_likelihood_functions/interferometer/rectangular.py

    FileNotFoundError: [Errno 2] No such file or directory: 'dataset/interferometer/simple/uv_wavelengths.fits'

    Missing dataset file — simulator script can't find input FITS.

    autolens_workspace_test/jax_likelihood_functions/point_source/point.py

    AssertionError: point: JAX vmap likelihood mismatch
 x: array([-3.025787])
 y: array(1.313508)

    JAX vmap likelihood doesn't match NumPy reference.

    autolens_workspace_test/jax_likelihood_functions/multi/mge.py

    AssertionError: multi/mge: JAX vmap likelihood mismatch
 x: array([-1.e+99, -1.e+99, -1.e+99])
 y: array(-2174335.96508)

    JAX vmap returns -1e99 (penalty) for all datasets instead of expected likelihoods.

  10. Jammy2211 commented on May 5, 2026

    @Jammy2211
    Jammy2211
    on May 5, 2026
    CollaboratorAuthor

    Smoke Test Results — 2026-05-05

    All 7 smoke tests passed across autolens_workspace.

    Workspace Passed Failed Total
    autolens_workspace 7 0 7

    Scripts tested: imaging/modeling.py, imaging/fit.py, interferometer/modeling.py, group/modeling.py, multi/modeling.py, guides/galaxies.py, guides/modeling/cookbook.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions