Merge pull request #451 from PyAutoLabs/feature/lazy-heavy-imports

refactor: defer nufftax/jax and numba imports to first use

Author: Jammy2211

autoarray/numba_util.py

@@ -1,4 +1,6 @@
 import logging
+import sys
+import threading
 
 from autonerves import conf
 
@@ -15,24 +17,61 @@
     parallel = False
 
 
-def jit(nopython=nopython, cache=cache, parallel=parallel, fastmath=False):
+# Decorated functions are queued here and only handed to numba on the first
+# call of any of them, keeping ``import numba`` off the library import path.
+# Materialization converts every queued function at once and rebinds the
+# defining module's global, because numba's nopython mode must resolve
+# cross-calls between decorated functions to real dispatchers at compile time.
+_pending = []
+_materialize_lock = threading.Lock()
 
-    def wrapper(func):
 
-        try:
+def _materialize_all():
+    with _materialize_lock:
+        if not _pending:
+            return
 
+        try:
             import numba
+        except ModuleNotFoundError:
+            numba = None
 
-            return numba.jit(
-                func,
-                nopython=nopython,
-                cache=cache,
-                parallel=parallel,
-                fastmath=fastmath,
-            )
+        while _pending:
+            func, options, placeholder, state = _pending.pop()
 
-        except ModuleNotFoundError:
+            if numba is None:
+                target = func
+            else:
+                target = numba.jit(func, **options)
+
+            state["target"] = target
+
+            module = sys.modules.get(func.__module__)
+            if module is not None and getattr(module, func.__name__, None) is placeholder:
+                setattr(module, func.__name__, target)
+
+
+def jit(nopython=nopython, cache=cache, parallel=parallel, fastmath=False):
+    options = dict(
+        nopython=nopython,
+        cache=cache,
+        parallel=parallel,
+        fastmath=fastmath,
+    )
+
+    def wrapper(func):
+        import functools
+
+        state = {"target": None}
+
+        @functools.wraps(func)
+        def lazy(*args, **kwargs):
+            if state["target"] is None:
+                _materialize_all()
+            return state["target"](*args, **kwargs)
+
+        _pending.append((func, options, lazy, state))
 
-            return func
+        return lazy
 
     return wrapper

autoarray/operators/transformer.py

@@ -23,10 +23,29 @@ class NUFFTPlaceholder:
 from autoarray.operators import transformer_util
 
 
-try:
-    import nufftax as _nufftax
-except ModuleNotFoundError:
-    _nufftax = None
+# nufftax pulls in jax at import (~0.7s), which sessions that never touch an
+# interferometer transformer should not pay for — deferred to _load_nufftax(),
+# called from TransformerNUFFT's entry points (not only __init__, because
+# unpickled instances in multiprocessing workers never re-run __init__).
+_nufftax = None
+_nufftax_loaded = False
+
+
+def _load_nufftax():
+    global _nufftax, _nufftax_loaded
+    if _nufftax_loaded:
+        return _nufftax
+    _nufftax_loaded = True
+    try:
+        import nufftax
+    except ModuleNotFoundError:
+        return None
+    _nufftax = nufftax
+    _version = tuple(int(v) for v in _nufftax.__version__.split(".")[:2])
+    # Only the 0.6.x series both has the primitives module and needs the shim.
+    if (0, 6) <= _version < (0, 7):
+        _patch_nufftax_batchers()
+    return _nufftax
 
 
 def _patch_nufftax_batchers():
@@ -87,13 +106,6 @@ def batcher(args, dims, **kwargs):
         )
 
 
-if _nufftax is not None:
-    _version = tuple(int(v) for v in _nufftax.__version__.split(".")[:2])
-    # Only the 0.6.x series both has the primitives module and needs the shim.
-    if (0, 6) <= _version < (0, 7):
-        _patch_nufftax_batchers()
-
-
 def pynufft_exception():
     raise ModuleNotFoundError(
         "\n--------------------\n"
@@ -643,7 +655,7 @@ def __init__(
         """
         from astropy import units
 
-        if _nufftax is None:
+        if _load_nufftax() is None:
             nufftax_exception()
 
         if chunk_size is not None and chunk_size <= 0:
@@ -685,6 +697,8 @@ def _forward_native(self, image_native_2d, xp=np):
         fixed-size chunks via ``jax.lax.scan`` (JAX path) or a Python loop
         (numpy path) — caps the nufftax gather-buffer allocation per call.
         """
+        _load_nufftax()
+
         K = int(self._x.shape[0])
 
         if xp.__name__.startswith("jax"):
@@ -789,6 +803,8 @@ def image_from(
         the sparse-operator dirty image is scale-consistent across all three
         transformers.
         """
+        _load_nufftax()
+
         n_y, n_x = self.real_space_mask.shape_native
         n_modes = (n_x, n_y)  # nufftax wants (n1, n2) = (N_x, N_y)
         K = int(self._x.shape[0])
@@ -860,6 +876,8 @@ def transform_mapping_matrix(self, mapping_matrix, xp=np):
         ``n_src`` separate NUFFT invocations and blow up the JIT graph
         for pixelization-heavy fits (notably double-source-plane).
         """
+        _load_nufftax()
+
         n_src = mapping_matrix.shape[1]
         rows, cols = self.real_space_mask.slim_to_native_tuple
         n_y, n_x = self.real_space_mask.shape_native