fix: run numba prange kernels serially in multiprocessing search workers (#676)

autoarray/config/general.yaml

@@ -25,6 +25,6 @@ numba:
   use_numba: true # If True, use numba-accelerated mapper and Sibson interpolation paths when eligible.
   cache: true # If True, cache numba compiled functions on disk.
   nopython: true # If True, require numba compilation without the Python interpreter.
-  parallel: false # If True, enable numba parallel compilation.
+  parallel: false # If True, enable numba parallel compilation. For single-process benchmarking only: it compiles every numba_util.jit kernel with parallel=True, which is unsafe under multiprocessing searches (number_of_cores > 1) on numba's GNU OpenMP threading layer (forked workers are terminated) and oversubscribes the node. The Sibson and direct_conv prange kernels already switch to serial in search workers.
 structures:
   native_binned_only: false           # If True, data structures are only stored in their native and binned format. This is used to reduce memory usage in autocti.

autoarray/inversion/inversion/interferometer_numba/sparse.py

@@ -4,6 +4,7 @@
 from autonerves import cached_property, conf
 
 from autoarray import exc
+from autoarray import numba_util
 from autoarray.dataset.interferometer.dataset import Interferometer
 from autoarray.inversion.inversion.dataset_interface import DatasetInterface
 from autoarray.inversion.inversion.interferometer.sparse import (
@@ -308,8 +309,13 @@ def _numba_parallel() -> bool:
     Whether the parallel (`prange`) kernel is used, read from the same
     `general.yaml -> numba -> parallel` flag the shared `numba_util.jit` decorator reads,
     with the same fallback when no config supplies it.
+
+    Always False in a search worker process
+    (:func:`autoarray.numba_util.parallel_kernel_allowed`), which takes the serial
+    route: numba parallel threads inside a multiprocessing search are never used.
     """
     try:
-        return bool(conf.instance["general"]["numba"]["parallel"])
+        parallel = bool(conf.instance["general"]["numba"]["parallel"])
     except Exception:
         return False
+    return parallel and numba_util.parallel_kernel_allowed()

autoarray/inversion/mesh/interpolator/_sibson_numba.py

@@ -11,6 +11,8 @@
 import numpy as np
 from numba import prange
 
+from autoarray import numba_util
+
 # float64 machine constants as module globals, which numba freezes into the
 # compiled code (and the on-disk cache) as literals.
 EPS = 2.220446049250313e-16
@@ -106,8 +108,7 @@ def _pairwise_sum(values):
     return partials[0]
 
 
-@numba.njit(cache=True, parallel=True, nogil=True)
-def _sibson_from_tables_parallel(
+def _sibson_from_tables(
     query_points,
     seeds,
     fallbacks,
@@ -365,3 +366,14 @@ def _sibson_from_tables_parallel(
         degenerate_out[q] = degenerate
 
     return mappings, sizes, weights, cavity_sizes, overflow_out, degenerate_out
+
+
+# One body, compiled twice: ``prange`` runs as ``range`` under ``parallel=False``, so the two
+# kernels compute identical results. The serial one is a renamed copy so both hit numba's
+# on-disk cache (see ``numba_util.function_copy``).
+_sibson_from_tables_parallel = numba.njit(cache=True, parallel=True, nogil=True)(
+    _sibson_from_tables
+)
+_sibson_from_tables_serial = numba.njit(cache=True, nogil=True)(
+    numba_util.function_copy(_sibson_from_tables, "_sibson_from_tables_serial")
+)

autoarray/inversion/mesh/interpolator/sibson.py

@@ -597,7 +597,7 @@ def _sibson_use_numba():
     return True
 
 
-def sibson_numba_kernel():
+def sibson_numba_kernel(parallel: bool = True):
     """
     Compile (once) the numba kernel evaluating every Sibson query in parallel.
 
@@ -628,22 +628,32 @@ def sibson_numba_kernel():
     on-disk cache still hits across processes -- closures built inside this
     function recompiled in every new process.
 
+    Parameters
+    ----------
+    parallel
+        The ``parallel=True`` (``prange``) kernel, or its ``parallel=False``
+        compilation of the same body (identical results), which search worker
+        processes use; see :func:`autoarray.numba_util.parallel_kernel_allowed`.
+
     Returns
     -------
     The compiled kernel ``(query_points, seeds, fallbacks, points,
     simplices_padded, simplex_neighbors, circumcentres, circumradii_squared,
     circumcircle_valid, max_cavity_triangles, max_neighbors)`` returning the six
     arrays of :func:`sibson_mappings_weights_from_tables`.
     """
-    if "kernel" in _SIBSON_NUMBA_CACHE:
-        return _SIBSON_NUMBA_CACHE["kernel"]
+    if parallel in _SIBSON_NUMBA_CACHE:
+        return _SIBSON_NUMBA_CACHE[parallel]
 
-    from autoarray.inversion.mesh.interpolator._sibson_numba import (
-        _sibson_from_tables_parallel,
-    )
+    from autoarray.inversion.mesh.interpolator import _sibson_numba
 
-    _SIBSON_NUMBA_CACHE["kernel"] = _sibson_from_tables_parallel
-    return _sibson_from_tables_parallel
+    kernel = (
+        _sibson_numba._sibson_from_tables_parallel
+        if parallel
+        else _sibson_numba._sibson_from_tables_serial
+    )
+    _SIBSON_NUMBA_CACHE[parallel] = kernel
+    return kernel
 
 
 def _sibson_numba_eligible(query_points, points):
@@ -665,9 +675,14 @@ def _sibson_mappings_weights_numba(
     max_cavity_triangles,
     max_neighbors,
 ):
-    """NumPy-path Sibson tables from the parallel numba kernel."""
+    """
+    NumPy-path Sibson tables from the numba kernel: the parallel one, or the serial one in a
+    search worker process (:func:`autoarray.numba_util.parallel_kernel_allowed`).
+    """
+    from autoarray import numba_util
+
     circumcentres, circumradii_squared, circumcircle_valid = circumcircles
-    kernel = sibson_numba_kernel()
+    kernel = sibson_numba_kernel(parallel=numba_util.parallel_kernel_allowed())
     return kernel(
         np.ascontiguousarray(query_points, dtype=np.float64),
         np.ascontiguousarray(simplex_indexes, dtype=np.int64),

autoarray/numba_util.py

@@ -1,6 +1,9 @@
 import logging
+import multiprocessing
+import os
 import sys
 import threading
+import types
 
 from autonerves import conf
 
@@ -25,6 +28,69 @@
 _pending = []
 _materialize_lock = threading.Lock()
 
+# The pid of the process that first asked to launch a numba ``parallel=True`` kernel; see
+# ``parallel_kernel_allowed``.
+_parallel_pid = None
+
+
+def parallel_kernel_allowed() -> bool:
+    """
+    Whether this process may launch a numba ``parallel=True`` (``prange``) kernel, or must take
+    a serial route instead: the Sibson kernel's ``parallel=False`` compilation of the same body
+    (identical results), and the interferometer direct_conv curvature's serial kernels.
+
+    False in any Python multiprocessing child (``multiprocessing.parent_process()`` is set, for
+    fork, spawn and forkserver alike -- how PyAutoFit / Nautilus search pools run), and in a
+    process whose pid differs from the one that first called this (an ``os.fork`` child that
+    ``parent_process`` does not see). True otherwise.
+
+    There is never nested numba parallelism inside a search's worker processes, for two
+    reasons:
+
+    * Without TBB, numba's threading layer is GNU OpenMP (``omp``), whose pool registers an
+      after-fork handler that terminates a forked child launching a parallel region after the
+      parent used the pool ("Terminating: fork() called from a process already using GNU
+      OpenMP, this is unsafe."). PyAutoFit evaluates the likelihood once in the parent before
+      forking its workers, so every worker would die.
+    * Each child would start ``NUMBA_NUM_THREADS`` threads (the node's core count by default),
+      so N search processes times N threads oversubscribe the node. Parallelism inside a search
+      comes from its processes.
+
+    The serial variant never initialises a numba threading layer, so the single-process path
+    keeps numba's default layer and parallel performance unchanged. Cheap on repeat calls (it
+    runs every likelihood evaluation).
+    """
+    global _parallel_pid
+
+    if multiprocessing.parent_process() is not None:
+        return False
+
+    pid = os.getpid()
+
+    if _parallel_pid is None:
+        _parallel_pid = pid
+
+    return pid == _parallel_pid
+
+
+def function_copy(func, name):
+    """
+    A copy of the plain Python function ``func`` under the name ``name`` (same code, globals and
+    closure), so one body can be compiled by numba twice -- e.g. ``parallel=True`` and
+    ``parallel=False`` -- with both cached on disk: numba's cache files are named by qualified
+    name and its index key ignores compile flags, so two dispatchers of one function object
+    would overwrite each other's cache entries.
+    """
+    copy = types.FunctionType(
+        func.__code__, func.__globals__, name, func.__defaults__, func.__closure__
+    )
+    copy.__kwdefaults__ = func.__kwdefaults__
+    copy.__module__ = func.__module__
+    copy.__doc__ = func.__doc__
+    prefix, _, _ = func.__qualname__.rpartition(".")
+    copy.__qualname__ = f"{prefix}.{name}" if prefix else name
+    return copy
+
 
 def _materialize_all():
     with _materialize_lock:
@@ -36,6 +102,18 @@ def _materialize_all():
         except ModuleNotFoundError:
             numba = None
 
+        if (
+            numba is not None
+            and any(options["parallel"] for _, options, _, _ in _pending)
+            and not parallel_kernel_allowed()
+        ):
+            logger.warning(
+                "general.yaml -> numba -> parallel is True in a multiprocessing worker: "
+                "numba_util.jit kernels compiled with parallel=True are unsafe under "
+                "multiprocessing searches on numba's GNU OpenMP threading layer. Set it "
+                "False (it is for single-process benchmarking only)."
+            )
+
         while _pending:
             func, options, placeholder, state = _pending.pop()