diff --git a/src/fast_array_utils/numba/__init__.py b/src/fast_array_utils/numba/__init__.py index 9812a17..dde8e1f 100644 --- a/src/fast_array_utils/numba/__init__.py +++ b/src/fast_array_utils/numba/__init__.py @@ -76,12 +76,11 @@ def _threading_layer(layer_or_category: ThreadingLayer | TheadingCategory, /, pr raise ValueError(msg) # pragma: no cover -def _is_in_unsafe_thread_pool() -> bool: +def _is_on_unsafe_thread() -> bool: import threading - current_thread = threading.current_thread() - # ThreadPoolExecutor threads typically have names like 'ThreadPoolExecutor-0_1' - return current_thread.name.startswith("ThreadPoolExecutor") and threading_layer() not in LAYERS["threadsafe"] + # We deem it unsafe if the caller is not the main thread, and therefore fall back to serial. + return threading.current_thread() is not threading.main_thread() and threading_layer() not in LAYERS["threadsafe"] @overload @@ -112,7 +111,7 @@ def decorator(f: Callable[P, R], /) -> Callable[P, R]: @wraps(f) def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: msg = None - if _is_in_unsafe_thread_pool(): # pragma: no cover + if _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 not _parallel_numba_runtime_is_safe(): msg = ( diff --git a/tests/test_numba.py b/tests/test_numba.py index 503b21c..e95a0c5 100644 --- a/tests/test_numba.py +++ b/tests/test_numba.py @@ -86,21 +86,20 @@ def import_module(name: str, package: str | None = None) -> object: @pytest.mark.parametrize( - ("name", "layer", "expected"), + ("on_main", "layer", "expected"), [ - pytest.param("ThreadPoolExecutor-0_1", "workqueue", True, id="executor-unsafe"), - pytest.param("ThreadPoolExecutor-0_1", "omp", False, id="executor-threadsafe"), - pytest.param("MainThread", "workqueue", False, id="not-executor"), + pytest.param(False, "workqueue", True, id="worker-unsafe"), + pytest.param(False, "omp", False, id="worker-threadsafe"), + pytest.param(True, "workqueue", False, id="main-thread"), ], ) -def test_is_in_unsafe_thread_pool(monkeypatch: pytest.MonkeyPatch, name: str, layer: fa_numba.ThreadingLayer, *, expected: bool) -> None: - def current_thread() -> object: - return type("FakeThread", (), {"name": name})() +def test_is_on_unsafe_thread(monkeypatch: pytest.MonkeyPatch, layer: fa_numba.ThreadingLayer, *, on_main: bool, expected: bool) -> None: + caller = threading.main_thread() if on_main else threading.Thread() - monkeypatch.setattr(threading, "current_thread", current_thread) + monkeypatch.setattr(threading, "current_thread", lambda: caller) monkeypatch.setattr(fa_numba, "threading_layer", lambda: layer) - assert fa_numba._is_in_unsafe_thread_pool() is expected + assert fa_numba._is_on_unsafe_thread() is expected def _set_runtime( @@ -271,7 +270,7 @@ def run(_cmd: list[str], /, **_kwargs: object) -> subprocess.CompletedProcess[st @pytest.mark.parametrize( - ("unsafe_pool", "needs_probe", "probe_safe", "expected", "warning"), + ("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"), @@ -282,7 +281,7 @@ def run(_cmd: list[str], /, **_kwargs: object) -> subprocess.CompletedProcess[st def test_njit_chooses_version( monkeypatch: pytest.MonkeyPatch, *, - unsafe_pool: bool, + unsafe_thread: bool, needs_probe: bool | None, probe_safe: bool | None, expected: bool, @@ -291,7 +290,7 @@ def test_njit_chooses_version( calls: list[bool] = [] _install_fake_njit(monkeypatch, calls) - monkeypatch.setattr(fa_numba, "_is_in_unsafe_thread_pool", lambda: unsafe_pool) + 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")) else: @@ -316,7 +315,7 @@ def test_njit_chooses_version( def test_serial_fallback(monkeypatch: pytest.MonkeyPatch) -> None: values = np.arange(10, dtype=np.float64) - monkeypatch.setattr(fa_numba, "_is_in_unsafe_thread_pool", lambda: False) + monkeypatch.setattr(fa_numba, "_is_on_unsafe_thread", lambda: False) monkeypatch.setattr(probe, "_needs_parallel_runtime_probe", lambda: True) monkeypatch.setattr(probe, "_parallel_numba_runtime_is_safe", lambda: False) wrapped = fa_numba.njit(_sum_prange)