diff --git a/autoarray/numba_util.py b/autoarray/numba_util.py index a0d3c82a..2f424cfc 100644 --- a/autoarray/numba_util.py +++ b/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 diff --git a/autoarray/operators/transformer.py b/autoarray/operators/transformer.py index a609629c..88165911 100644 --- a/autoarray/operators/transformer.py +++ b/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