Repository navigation
refactor: defer heavy imports — import autolens 4.1s → ~1.9s #1505
Description
Overview
import autolens costs ~4.1 s (full example-script stanza ~4.4 s), and profiling shows ~2.5 s of it is heavy packages loaded eagerly but not needed at import: the jax chain (nufftax + blackjax + optax, 1.75 s — 43%, measured by control test), IPython, sqlalchemy, numba (decoration-time), and astropy. All 457 autolens_workspace scripts pay this per-process (~30 min of pure import per full serial smoke sweep), and smoke runs with JAX disabled, so its jax import there is pure waste. This task defers each of these to first use across PyAutoFit, PyAutoArray and PyAutoNerves, targeting import autolens ≤ ~1.9 s with zero public-API change.
Plan
- Defer the jax chain: lazy
af.NSS(PEP 562) in PyAutoFit; lazy nufftax load + batcher patch in PyAutoArray'sTransformerNUFFT. - Defer IPython, sqlalchemy (lazy
saproxy + annotation fixes), and the eagerautofit.databasepulls in PyAutoFit. - Defer numba to first call via a lazy-compile wrapper in
autoarray.numba_util.jit(fixes all 29 decorated functions centrally). - Defer astropy in
autonerves.fitsable; dedupe the workspace-version warning (currently fires 3× per process). - Verify: import-time absence proof, full suites in all three repos, NSS / NUFFT / aggregator behaviour runs, targeted workspace smoke incl.
# ENV: jaxscripts. - Scoped out: matplotlib deferral (enters via the autolens/autogalaxy visualizer chain; every workspace script imports
apltanyway, so no smoke/user win — optional follow-up for bare-library import only).
Detailed implementation plan
Affected Repositories
- PyAutoFit (primary)
- PyAutoArray
- PyAutoNerves
Branch Survey
| Repository | Current Branch | Dirty? |
|---|---|---|
| ./PyAutoFit | main | clean |
| ./PyAutoArray | main | clean |
| ./PyAutoNerves | main | clean |
Suggested branch: feature/lazy-heavy-imports
Note: PyAutoFit is also claimed by stored-sample-reconstruction-guard (#1486, library half shipped as PR#1504) — deliberate file-disjoint override, recorded in Mind active.md. The version-stamp-sync-guards branch also touches autofit/__init__.py (version-stamp lines); coordinate at merge if both are in flight.
Implementation Steps (PyAutoFit, ~0.9 s)
autofit/__init__.py— remove eager imports ofNSS(L98),GridSearchAggregator(L15),Aggregator(L37),Query(L42); add PEP 562 module__getattr__resolving these four lazily, caching inglobals()(pattern:autolens/__init__.py:165-181). NSS stays fully supported — lazy, never removed.autofit/non_linear/fitness.py— movefrom IPython.display import clear_output(L3) into the single use site (~L439, quick-update branch).autofit/database/sqlalchemy_.py— makesa/declarativelazy proxies:__getattr__imports sqlalchemy on first attribute access and caches; fall back to the existingMockSQlAlchemyon ImportError.- Add
from __future__ import annotationsto modules evaluatingsa.orm.Sessionin signatures at def time (Py3.12):non_linear/settings.py,search/mcmc/abstract_mcmc.py,search/nest/abstract_nest.py,search/nest/nss/search.py,search/nest/nautilus/search.py, plus a full per-file sweep of everysqlalchemy_importer. Verify no eager__annotations__introspection in these modules. autofit/non_linear/paths/database.py— deferFitandInfoimports to function-local at their use sites (~L112, L236-243, L345).autofit/non_linear/grid/grid_search/result_builder.py—from __future__ import annotations+ movefrom autofit.database import Prior(annotation-only) underif TYPE_CHECKING.
Implementation Steps (PyAutoArray, ~0.75 s)
autoarray/operators/transformer.py— replace module-level nufftax try-import (L26-29) and version-gated_patch_nufftax_batchers()call (L90-94) with an idempotent_load_nufftax()(import, 0.6.x check, batcher patch — a global JAX registration that must precede the first traced call) invoked at the top ofTransformerNUFFT.__init__(existing_nufftax is Noneguard at ~L646). All other_nufftaxuses are inTransformerNUFFTmethods.autoarray/numba_util.py—jit()returns a wrapper that on first call imports numba (falling back to the plain function onModuleNotFoundError), compiles with the same nopython/cache/parallel/fastmath options, caches, delegates. Sweep for attribute access on decorated functions (.py_funcetc.) first.
Implementation Steps (PyAutoNerves, ~0.17 s + noise)
autonerves/fitsable.py— remove module-level astropy try-import (L12-15, keep TYPE_CHECKING block); lazy_fits()accessor used at the six runtime sites (~L63, 86, 88, 90, 210, 238).autonerves/workspace.py— module-level_warned_messagesset; the threewarnings.warnsites (~L227, 237, 257) skip byte-identical repeats, soimport autolenswarns once, not three times. The direct-call tests use distincttmp_pathroots — unaffected.
Key Files
PyAutoFit/autofit/__init__.py— lazy NSS/Aggregator/Query/GridSearchAggregatorPyAutoFit/autofit/database/sqlalchemy_.py— lazysaproxyPyAutoArray/autoarray/operators/transformer.py— lazy nufftax + batcher patchPyAutoArray/autoarray/numba_util.py— lazy-compile jit wrapperPyAutoNerves/autonerves/fitsable.py,PyAutoNerves/autonerves/workspace.py
Verification
- Timing: 3× cold
import autolensbefore/after; target ≤ ~1.9 s (matplotlib retained by design). - Absence proof:
python -X importtime -c "import autolens"shows no jax/jaxlib/blackjax/optax/nufftax/IPython/sqlalchemy/numba/astropy. - Behaviour:
af.NSS/af.Aggregator/af.Query/af.GridSearchAggregatorresolve; runautofit_workspace_test/scripts/searches/NSS.py, an interferometerTransformerNUFFTscript, a database/aggregator flow. - Full pytest suites in PyAutoFit, PyAutoArray, PyAutoNerves; downstream autogalaxy/autolens suites via ship_library;
unittest-nojaxleg stays green. - Targeted workspace smoke incl.
# ENV: jaxscripts before PRs open.
Follow-up (not this task)
- Brain RefactorDecision flags pyautonerves as "unwitnessed" — its witness map lacks
test_autonerves; fix the map. - Optional: matplotlib deferral for bare-library import (unit-test startup), via the autolens/autogalaxy visualizer chain.
Original Prompt
Click to expand starting prompt
Can you do a profiling of how long import times on autolens_workspace
example scripts and performance in general is, and assess if we can speed
it up as it is important for running smoke tests and general user
performance being fast
Full measured profile in PyAutoMind/active/import_time_lazy_heavy_imports.md.
Activity
Jammy2211 commentedon Aug 19, 2026
Verification results (worktree
feature/lazy-heavy-imports)Headline:
import autolens4.07 s → ~2.1 s (full script stanza ~4.4 s → ~2.0 s).python -X importtimeproves jax, jaxlib, blackjax, optax, nufftax, IPython, sqlalchemy, numba and astropy are all absent from the bare import tree. The autonerves workspace-version warning now fires once per process (was 3×).Unit suites (all green)
blackjax.ns;_HAS_NSSFalse on canonical main too)Workspace smoke (
pyauto-heart smoke --root <worktree>, all six workspaces)The 4 autolens_workspace_test failures are pre-existing, not from this branch — a control run against canonical
main(same harness, same env) fails the identical four scripts:imaging/jax_likelihood/lp.py— vmap result is ~half the hardcoded pin; matches the in-flight positions-LH penalty-accumulation fix (orphan worktree commitfb1aefe0b"sum PositionsLH penalties instead of doubling the last entry" — the pin was calibrated against the doubling behaviour).interferometer/jax_likelihood/rectangular.py— misses its pin by 0.0124% against rtol 1e-4 (stale-pin drift).interferometer/jax_likelihood/mge.py,multi_dataset/jax_likelihood/mge.py— same pin-drift family.The
imaging/subhalo_recovery.pytimeout (300 s cap) occurred only under the six-workspace sweep's load: an uncontended worktree rerun passes it at 252.1 s (canonical: 252.5 s — no slowdown from this branch) and reproduces exactly the canonical result, 20/24 with the same four pre-existing failures.Behaviour checks
af.NSS/af.Aggregator/af.Query/af.GridSearchAggregator/af.dbresolve via PEP 562 and construct.db.open_database("sqlite://")round-trip works (sqlalchemy loads on first use).TransformerNUFFTnumpy/jax visibilities agree exactly;transform_mapping_matrixandimage_fromwork; a pickled instance in a state-reset module (multiprocessing-worker simulation) works.Drawermodel fit runs through the annotated search plumbing.from autofit import *users exist; sphinx autosummary resolves lazy names via getattr.Jammy2211 commentedon Aug 19, 2026
Library PRs created (all
pending-release, all gates green)Heart readiness at ship: GREEN, score 100. Six-workspace smoke against the branch: identical to canonical
main(the four autolens_workspace_test jax_likelihood failures are control-proven pre-existing — see the verification comment above).Result:
import autolens4.07 s → ~2.1 s; full example-script stanza ~4.4 s → ~2.0 s; jax/jaxlib/blackjax/optax/nufftax/IPython/sqlalchemy/numba/astropy all absent from the bare import tree; workspace-version warning deduped to once per process.Workspace impact: None — all five PRs declare "None — internal changes only"; no workspace migration or new demos needed (option iii). Merge is the remaining human decision.
Follow-ups noted (not this task): Brain witness map lacks
test_autonerves; optional matplotlib deferral for bare-library import; the pre-existing jax_likelihood pin failures belong to the positions-LH penalty-accumulation arc + stale pins.Merge pull request #1506 from PyAutoLabs/feature/lazy-heavy-imports
Jammy2211 commentedon Aug 19, 2026
Shipped — all five PRs merged 2026-08-19
PyAutoFit#1506, PyAutoArray#451, PyAutoNerves#151, PyAutoGalaxy#575, PyAutoLens#705 — all 17 CI checks green (incl.
unittest-nojax), remote branches deleted.Final measured result (warm cache; earlier ~2.1 s readings were polluted by the fresh worktree's one-time matplotlib font-cache build and first-run bytecode compilation):
import autolens: 4.07 s → 1.2–1.3 s (~3.2× faster) — beats the ≤1.9 s acceptance target.apltanyway), the PyAuto libraries' own module tree ~0.38 s.Closing. Follow-ups tracked as Mind drafts: pre-existing autolens_workspace_test jax_likelihood pin failures (positions-LH penalty arc + stale pins), Brain witness map missing
test_autonerves.refactor: defer blackjax/jax, sqlalchemy and IPython imports to first…