diff --git a/src/fast_array_utils/numba/__init__.py b/src/fast_array_utils/numba/__init__.py index fd5c1dd..bf31ec1 100644 --- a/src/fast_array_utils/numba/__init__.py +++ b/src/fast_array_utils/numba/__init__.py @@ -11,6 +11,7 @@ from __future__ import annotations +import os import sys import warnings from functools import update_wrapper, wraps @@ -48,6 +49,26 @@ def _is_on_unsafe_thread() -> bool: return threading.current_thread() is not threading.main_thread() and _parallel_numba_runtime_layer() not in LAYERS["threadsafe"] +_inherited_layer: ThreadingLayer | None = None + + +def _after_fork_in_child() -> None: + global _inherited_layer # noqa: PLW0603 + try: + _inherited_layer = None if (numba := sys.modules.get("numba")) is None else numba.threading_layer() + except ValueError: # no parallel code ran before the fork + _inherited_layer = None + + +if hasattr(os, "register_at_fork"): # Unix only; only sees `os.fork`, and only if this module was imported before + os.register_at_fork(after_in_child=_after_fork_in_child) + + +def _is_in_unsafe_fork() -> bool: + # e.g. GNU OpenMP terminates forked children + return _inherited_layer is not None and _inherited_layer not in LAYERS["forksafe"] + + @overload def njit[**P, R](fn: Callable[P, R], /, *, nogil: bool = True) -> Callable[P, R]: ... @overload @@ -57,6 +78,9 @@ def njit[**P, R](fn: Callable[P, R] | None = None, /, *, nogil: bool = True) -> On call, this function dispatches to a parallel or serial numba function, depending on the current threading environment. + In a process forked after numba launched a threading layer that isn’t fork-safe + (GNU OpenMP on Linux), it runs the serial function. + TBB is only fork-safe when forked from the thread that launched it, which isn’t detected. Parameters ---------- @@ -83,7 +107,13 @@ def decorator(f: Callable[P, R], /) -> Callable[P, R]: @wraps(f) def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: msg = None - if _is_on_unsafe_thread(): # pragma: no cover + if _is_in_unsafe_fork(): + msg = ( + f"Detected a fork after numba launched the {_inherited_layer} threading layer, which isn’t fork-safe here. " + f"Running {f.__name__} in serial mode. To avoid this fallback, use the `spawn` or `forkserver` start method, " + "or if you can’t control how the process is forked, set `NUMBA_THREADING_LAYER=forksafe` before running parallel code." + ) + elif _is_on_unsafe_thread(): # pragma: no cover msg = f"Detected unsupported threading environment. Trying to run {f.__name__} in serial mode. In case of problems, install `tbb`." elif _needs_parallel_runtime_probe() and _parallel_numba_runtime_layer() is None: msg = ( diff --git a/tests/test_numba.py b/tests/test_numba.py index b71ce2f..dcf8b03 100644 --- a/tests/test_numba.py +++ b/tests/test_numba.py @@ -4,7 +4,9 @@ from __future__ import annotations import importlib +import os import subprocess +import sys import threading import warnings from typing import TYPE_CHECKING @@ -61,6 +63,29 @@ def test_is_on_unsafe_thread(monkeypatch: pytest.MonkeyPatch, layer: fa_numba.Th assert fa_numba._is_on_unsafe_thread() is expected +@pytest.mark.parametrize( + ("layer", "expected"), + [ + pytest.param(None, False, id="nothing-launched"), + pytest.param("workqueue", False, id="workqueue"), + pytest.param("omp", sys.platform == "linux", id="omp"), + pytest.param("tbb", False, id="tbb"), + ], +) +def test_is_in_unsafe_fork(monkeypatch: pytest.MonkeyPatch, layer: fa_numba.ThreadingLayer | None, *, expected: bool) -> None: + def threading_layer() -> str: + if layer is None: + raise ValueError + return layer + + monkeypatch.setattr(numba, "threading_layer", threading_layer) + monkeypatch.setattr(fa_numba, "_inherited_layer", None) + fa_numba._after_fork_in_child() + monkeypatch.setattr(numba, "threading_layer", lambda: "omp") # a layer launched in the child isn’t inherited + + assert fa_numba._is_in_unsafe_fork() is expected + + def _set_runtime( monkeypatch: pytest.MonkeyPatch, *, @@ -234,17 +259,19 @@ def run(_cmd: list[str], /, **_kwargs: object) -> subprocess.CompletedProcess[st @pytest.mark.parametrize( - ("unsafe_thread", "needs_probe", "probe_safe", "expected", "warning"), + ("unsafe_fork", "unsafe_thread", "needs_probe", "probe_safe", "expected", "warning"), [ - pytest.param(True, None, None, False, "unsupported threading environment", id="thread-pool"), - pytest.param(False, True, False, False, "unsupported numba parallel runtime", id="probe-fails"), - pytest.param(False, True, True, True, None, id="probe-passes"), - pytest.param(False, False, None, True, None, id="no-probe"), + pytest.param(True, True, None, None, False, "fork after numba launched", id="unsafe-fork"), + pytest.param(False, True, None, None, False, "unsupported threading environment", id="thread-pool"), + pytest.param(False, False, True, False, False, "unsupported numba parallel runtime", id="probe-fails"), + pytest.param(False, False, True, True, True, None, id="probe-passes"), + pytest.param(False, False, False, None, True, None, id="no-probe"), ], ) def test_njit_chooses_version( monkeypatch: pytest.MonkeyPatch, *, + unsafe_fork: bool, unsafe_thread: bool, needs_probe: bool | None, probe_safe: bool | None, @@ -254,6 +281,7 @@ def test_njit_chooses_version( calls: list[bool] = [] _install_fake_njit(monkeypatch, calls) + monkeypatch.setattr(fa_numba, "_is_in_unsafe_fork", lambda: unsafe_fork) monkeypatch.setattr(fa_numba, "_is_on_unsafe_thread", lambda: unsafe_thread) if needs_probe is None: monkeypatch.setattr(probe, "_needs_parallel_runtime_probe", lambda: pytest.fail("probe should not be consulted")) @@ -293,3 +321,29 @@ def test_serial_fallback(monkeypatch: pytest.MonkeyPatch) -> None: def test_njit_nogil(monkeypatch: pytest.MonkeyPatch) -> None: _install_fake_njit(monkeypatch, [], expected_nogil=False) fa_numba.njit(nogil=False)(_return_true) + + +_FORK_SCRIPT = """ +import os +import numpy as np +import scipy.sparse as sp +from fast_array_utils.conv import to_dense + +x = sp.random(100, 100, density=0.1, format="csr") +to_dense(x) # launches the threading layer +if (pid := os.fork()) == 0: + os._exit(0 if np.array_equal(to_dense(x), x.toarray()) else 1) +os._exit(os.waitstatus_to_exitcode(os.waitpid(pid, 0)[1])) +""" + + +@pytest.mark.skipif(sys.platform != "linux", reason="GNU OpenMP only terminates forked children on Linux") +def test_fork_after_omp() -> None: + pytest.importorskip("scipy") + importlib.import_module("numba.np.ufunc.omppool") # fail instead of skipping if the omp layer is missing + result = subprocess.run( # noqa: S603 + [sys.executable, "-c", _FORK_SCRIPT], env={**os.environ, "NUMBA_THREADING_LAYER": "omp"}, capture_output=True, text=True, check=False + ) + assert result.returncode == 0, result.stderr + assert "fork after numba launched the omp threading layer" in result.stderr + assert "NUMBA_THREADING_LAYER=forksafe" in result.stderr diff --git a/typings/numba/__init__.pyi b/typings/numba/__init__.pyi index 6eb44a6..06e9945 100644 --- a/typings/numba/__init__.pyi +++ b/typings/numba/__init__.pyi @@ -2,6 +2,8 @@ from collections.abc import Iterable from typing import SupportsIndex, overload +from fast_array_utils.numba import ThreadingLayer + from .core import config as config from .core.decorators import njit as njit @@ -10,3 +12,4 @@ def prange(stop: SupportsIndex, /) -> Iterable[int]: ... @overload def prange(start: SupportsIndex, stop: SupportsIndex, step: SupportsIndex = ..., /) -> Iterable[int]: ... def get_num_threads() -> int: ... +def threading_layer() -> ThreadingLayer: ...