From 56db9c0fac7317e53948db917969ceb91e63d0c6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 5 Nov 2026 13:15:08 +0000 Subject: [PATCH] Run numba prange kernels serially in multiprocessing search workers (#676) On numba's default GNU OpenMP threading layer, a child forked after the parent launched a parallel region is terminated ("fork() called from a process already using GNU OpenMP"). PyAutoFit evaluates the likelihood in the parent before forking Nautilus workers, so every DelaunayNN NumPy fit with number_of_cores > 1 died. Nested numba threads in each worker would also oversubscribe the node. - numba_util.parallel_kernel_allowed(): False in any multiprocessing child (parent_process() set) or a process whose pid differs from the first caller; function_copy() renames a function so one body can be cached under two compilations (numba's cache index ignores compile flags). - Sibson: one body compiled parallel=True and parallel=False; the NumPy dispatch picks via the predicate (identical outputs, tested). - direct_conv: _numba_parallel() is False in workers, which take the existing serial / symmetric route. - numba_util.jit warns once at materialization when general.yaml numba.parallel is True in a worker; the yaml comment says it is for single-process benchmarking only. - The threading layer and thread count are left untouched, so the single-process path is unchanged. Co-Authored-By: Claude Opus 5.5 --- autoarray/config/general.yaml | 2 +- .../inversion/interferometer_numba/sparse.py | 8 +- .../mesh/interpolator/_sibson_numba.py | 16 +- .../inversion/mesh/interpolator/sibson.py | 35 ++- autoarray/numba_util.py | 78 +++++ test_autoarray/test_numba_fork_safety.py | 267 ++++++++++++++++++ 6 files changed, 392 insertions(+), 14 deletions(-) create mode 100644 test_autoarray/test_numba_fork_safety.py diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index 4f6b67df7..df2bc5502 100644 --- a/autoarray/config/general.yaml +++ b/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. diff --git a/autoarray/inversion/inversion/interferometer_numba/sparse.py b/autoarray/inversion/inversion/interferometer_numba/sparse.py index 05beed804..2d5022fd8 100644 --- a/autoarray/inversion/inversion/interferometer_numba/sparse.py +++ b/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() diff --git a/autoarray/inversion/mesh/interpolator/_sibson_numba.py b/autoarray/inversion/mesh/interpolator/_sibson_numba.py index 2b0a00fd1..450053a44 100644 --- a/autoarray/inversion/mesh/interpolator/_sibson_numba.py +++ b/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") +) diff --git a/autoarray/inversion/mesh/interpolator/sibson.py b/autoarray/inversion/mesh/interpolator/sibson.py index 11f54fa27..2f1789af5 100644 --- a/autoarray/inversion/mesh/interpolator/sibson.py +++ b/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,6 +628,13 @@ 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, @@ -635,15 +642,18 @@ def sibson_numba_kernel(): 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), diff --git a/autoarray/numba_util.py b/autoarray/numba_util.py index 2f424cfc5..e72f37a27 100644 --- a/autoarray/numba_util.py +++ b/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() diff --git a/test_autoarray/test_numba_fork_safety.py b/test_autoarray/test_numba_fork_safety.py new file mode 100644 index 000000000..42b771199 --- /dev/null +++ b/test_autoarray/test_numba_fork_safety.py @@ -0,0 +1,267 @@ +"""No numba parallel threads inside multiprocessing workers (``numba_util.parallel_kernel_allowed``). + +PyAutoFit evaluates the likelihood once in the parent, then forks (or spawns) search workers that +evaluate it again. On numba's ``omp`` threading layer a forked child launching a ``prange`` kernel +after the parent used the pool is terminated, so workers run each kernel's serial compilation. +""" + +import multiprocessing +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import numpy as np +import pytest + +pytest.importorskip("numba") + +from autoarray import numba_util +from autoarray.inversion.mesh.interpolator import sibson +from autoarray.inversion.mesh.interpolator.delaunay import ( + pix_indexes_delaunay_walk_from, + scipy_delaunay_tri_only, +) + +_SCRIPT = textwrap.dedent( + """ + import multiprocessing as mp + import sys + + import numpy as np + import numba + + from autoarray import numba_util + from autoarray.inversion.mesh.interpolator import sibson + from autoarray.inversion.mesh.interpolator.delaunay import ( + pix_indexes_delaunay_walk_from, + scipy_delaunay_tri_only, + ) + + MODE = sys.argv[1] + TIMEOUT = float(sys.argv[2]) + + if MODE == "control": + # Undo the fix in every process: children launch the parallel kernel too. + numba_util.parallel_kernel_allowed = lambda: True + + rng = np.random.default_rng(1) + points = rng.uniform(-1.0, 1.0, size=(60, 2)) + query = rng.uniform(-1.1, 1.1, size=(200, 2)) + simplices, neighbors, vertex_simplex = scipy_delaunay_tri_only(points) + mappings, simplex_indexes = pix_indexes_delaunay_walk_from( + query_points=query, + points=points, + simplices_padded=simplices, + simplex_neighbors=neighbors, + vertex_simplex=vertex_simplex, + xp=np, + return_simplex_indexes=True, + ) + KWARGS = dict( + query_points=query, + points=points, + simplices_padded=simplices, + simplex_neighbors=neighbors, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=mappings[:, 0], + circumcircles=sibson.delaunay_circumcircles_from(points, simplices, xp=np), + max_cavity_triangles=32, + max_neighbors=32, + ) + + + def threading_layer(): + try: + return numba.threading_layer() + except ValueError: # no parallel region launched in this process + return None + + + def work(_): + return sibson._sibson_mappings_weights_numba(**KWARGS), threading_layer() + + + def check(results, parent, fresh_process): + for outputs, layer in results: + # A forked child inherits the parent's record of its layer; a spawned child + # running the serial kernel never initialises one. + if fresh_process: + assert layer is None, layer + for expected, actual in zip(parent, outputs): + np.testing.assert_array_equal(expected, actual) + + + if __name__ == "__main__": + try: + parent = sibson._sibson_mappings_weights_numba(**KWARGS) + except ValueError as exc: # requested threading layer unavailable + print("LAYER_UNAVAILABLE", exc, flush=True) + sys.exit(0) + print("LAYER", threading_layer(), flush=True) + + with mp.get_context("fork").Pool(2) as pool: + check(pool.map_async(work, range(4)).get(timeout=TIMEOUT), parent, False) + print("FORK_OK", flush=True) + + with mp.get_context("spawn").Pool(2) as pool: + check(pool.map_async(work, range(2)).get(timeout=TIMEOUT), parent, True) + print("SPAWN_OK", flush=True) + """ +) + + +def _run_script(tmp_path, mode, timeout, layer="omp"): + script = tmp_path / "fork_sibson.py" + script.write_text(_SCRIPT) + + env = dict(os.environ) + env.pop("NUMBA_NUM_THREADS", None) + env.pop("NUMBA_THREADING_LAYER", None) + if layer is not None: + env["NUMBA_THREADING_LAYER"] = layer + # The subprocess must import this checkout's autoarray, whatever is installed. + package_root = str(Path(numba_util.__file__).resolve().parents[1]) + env["PYTHONPATH"] = os.pathsep.join( + [package_root] + [p for p in env.get("PYTHONPATH", "").split(os.pathsep) if p] + ) + + return subprocess.run( + [sys.executable, str(script), mode, str(timeout)], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + timeout=4 * timeout + 120, + ) + + +needs_fork = pytest.mark.skipif( + "fork" not in multiprocessing.get_all_start_methods(), + reason="needs the fork start method", +) + + +@needs_fork +def test__worker_processes_run_sibson_kernel_after_parent(tmp_path): + """The RAL failure condition: parallel kernel on ``omp`` in the parent, then a fork pool.""" + result = _run_script(tmp_path, mode="dispatch", timeout=60) + if "LAYER_UNAVAILABLE" in result.stdout: + result = _run_script(tmp_path, mode="dispatch", timeout=60, layer=None) + + output = result.stdout + result.stderr + assert result.returncode == 0, output + assert "FORK_OK" in result.stdout, output + assert "SPAWN_OK" in result.stdout, output + + +@needs_fork +def test__control__parallel_kernel_in_forked_children_dies_on_omp(tmp_path): + """With the predicate forced True the same script fails, so the test above is not vacuous.""" + result = _run_script(tmp_path, mode="control", timeout=10) + + if "LAYER_UNAVAILABLE" in result.stdout: + pytest.skip("numba omp threading layer unavailable") + + output = result.stdout + result.stderr + assert "LAYER omp" in result.stdout, output + assert result.returncode != 0, output + assert "FORK_OK" not in result.stdout, output + + +def test__sibson_serial_kernel_equals_parallel_kernel(): + rng = np.random.default_rng(3) + points = rng.uniform(-1.0, 1.0, size=(80, 2)) + query = rng.uniform(-1.2, 1.2, size=(300, 2)) + simplices, neighbors, vertex_simplex = scipy_delaunay_tri_only(points) + mappings, simplex_indexes = pix_indexes_delaunay_walk_from( + query_points=query, + points=points, + simplices_padded=simplices, + simplex_neighbors=neighbors, + vertex_simplex=vertex_simplex, + xp=np, + return_simplex_indexes=True, + ) + centres, radii_squared, valid = sibson.delaunay_circumcircles_from( + points, simplices, xp=np + ) + args = ( + query, + np.asarray(simplex_indexes, dtype=np.int64), + np.asarray(mappings[:, 0], dtype=np.int64), + points, + np.asarray(simplices, dtype=np.int64), + np.asarray(neighbors, dtype=np.int64), + np.asarray(centres, dtype=np.float64), + np.asarray(radii_squared, dtype=np.float64), + np.asarray(valid, dtype=np.bool_), + 32, + 32, + ) + + parallel = sibson.sibson_numba_kernel(parallel=True)(*args) + serial = sibson.sibson_numba_kernel(parallel=False)(*args) + + assert sibson.sibson_numba_kernel(parallel=False) is not sibson.sibson_numba_kernel( + parallel=True + ) + for expected, actual in zip(parallel, serial): + assert expected.dtype == actual.dtype + np.testing.assert_array_equal(expected, actual) + + +@pytest.fixture(name="predicate") +def make_predicate(monkeypatch): + monkeypatch.setattr(numba_util, "_parallel_pid", None) + monkeypatch.setattr(numba_util.multiprocessing, "parent_process", lambda: None) + return monkeypatch + + +def test__predicate__main_process_allowed_and_recorded(predicate): + assert numba_util.parallel_kernel_allowed() is True + assert numba_util.parallel_kernel_allowed() is True + assert numba_util._parallel_pid == os.getpid() + + +def test__predicate__multiprocessing_child_disallowed(predicate): + predicate.setattr(numba_util.multiprocessing, "parent_process", lambda: object()) + + assert numba_util.parallel_kernel_allowed() is False + + +def test__predicate__forked_pid_disallowed(predicate): + assert numba_util.parallel_kernel_allowed() is True + + predicate.setattr(os, "getpid", lambda: -12345) + + assert numba_util.parallel_kernel_allowed() is False + + +def test__function_copy__renames_without_changing_behaviour(): + def body(x, y=2): + return x * y + + copy = numba_util.function_copy(body, "renamed") + + assert copy is not body + assert copy.__name__ == "renamed" + assert copy.__qualname__.endswith(".renamed") + assert copy(3) == body(3) == 6 + + +def test__direct_conv_parallel_route_off_in_worker_processes(monkeypatch): + from autoarray.inversion.inversion.interferometer_numba import sparse + + class Config: + def __getitem__(self, key): + return {"numba": {"parallel": True}} if key == "general" else {} + + monkeypatch.setattr(sparse.conf, "instance", Config()) + + monkeypatch.setattr(numba_util, "parallel_kernel_allowed", lambda: True) + assert sparse._numba_parallel() is True + + monkeypatch.setattr(numba_util, "parallel_kernel_allowed", lambda: False) + assert sparse._numba_parallel() is False