Why does fau need to know about numba internals for thread layer

[selmanozleyen]

I know that if numba exposed

numba.threading_layer_properties()  # {"name": "tbb", "threadsafe": True, "forksafe": True}

we wouldn't need this LAYER/priority mechanism's and usages across the repo. Which is very much prone to bugs and it has a bug. E.g., _threading_layer resolves a category by walking NUMBA_THREADING_LAYER_PRIORITY. numba doesn't do that: in _launch_threads it only consults the priority list in the default branch, and for safe/threadsafe/forksafe it builds its own fixed list, always tbb first https://github.com/numba/numba/blob/0.67.0/numba/np/ufunc/parallel.py#L473-L499 So with a non-default priority the two disagree:

 NUMBA_THREADING_LAYER_PRIORITY="omp workqueue tbb"

 THREADING_LAYER=threadsafe #  we predict omp        numba launches tbb
 THREADING_LAYER=forksafe     # we predict workqueue  numba launches tbb

For example in the forksafe case we predict workqueue, which isn't in LAYERS["threadsafe"], so _is_in_unsafe_thread_pool warns "unsupported threading environment" and drops to serial inside a thread pool, while the layer that actually launched (tbb) is thread-safe.

my proposals

from our side

We already launch a probe subprocess here:

fast-array-utils/src/fast_array_utils/numba/_parallel_runtime.py

Lines 90 to 104 in 57365b8

	def _parallel_numba_runtime_is_safe_cached(key: _ParallelRuntimeProbeKey) -> bool:
	try:
	# The probe command is built from `sys.executable` plus a generated script
	# that only imports modules from a fixed whitelist.
	result = subprocess.run( # noqa: S603
	[key[0], "-c", _parallel_runtime_probe_code(key[3])],
	capture_output=True,
	check=False,
	env=_build_parallel_runtime_probe_env(key),
	text=True,
	timeout=_PARALLEL_RUNTIME_PROBE_TIMEOUT,
	)
	except Exception: # noqa: BLE001
	return False
	return result.returncode == 0 and _PARALLEL_RUNTIME_PROBE_SENTINEL in result.stdout

why don't we just use this to get the information we need here?

from numba

If the subprocess probe is too hacky and we don't want to repeat it we ask numba to expose in an issue?

numba.threading_layer_properties()  # {"name": "tbb", "threadsafe": True, "forksafe": True}

how this would relate to #179

Once knowing the thread layer reliably we can also detect a fork and fall-back to serial in that case. If we also solve #213 then they would also get a speedup on top of this :)

def _is_in_unsafe_fork() -> bool:
    if not _forked: # see #179 for fork detection
        return False                              # not a child → nothing inherited
    layer = _launched_layer()
    return layer is not None and layer not in LAYERS["forksafe"]

Then having an always parallel/serial context manager settings would be its own new feature.

the solution which would remove all LAYER machinary

why not ask numba to fail gracefully? Like try: parallel() except UnsafeParallelError: serial(), would be a very reasonable ask, no? They can predict when their code would fail maybe? Or if it would need to be conservative maybe they would have a strict=True mode for it?

[flying-sheep]

why don't we just use this to get the information we need here?

Yeah, since we cache this now I think it should be fine to just use that!