fix: run njit serially in forked children of a non-fork-safe threading layer

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 = (

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

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: ...