Skip to content

Guard numba's thread pool against concurrent entry from thread pools - #1797

Draft
cmdupuis3 wants to merge 2 commits into
UXARRAY:mainfrom
cmdupuis3:cmd/parallel-true
Draft

cmdupuis3 wants to merge 2 commits into
UXARRAY:mainfrom
cmdupuis3:cmd/parallel-true

Conversation

@cmdupuis3

@cmdupuis3 cmdupuis3 commented Sep 29, 2026 •

Copy link
Copy Markdown
Collaborator

Closes #1798

Overview

A parallel=True kernel called from a thread pool -- dask's threaded scheduler, typically -- nests numba's pool under it. What that does depends on numba's threading layer:

  • tbb composes; nothing to do.
  • omp starts a full team per calling thread: 179 OS threads from a 12-thread pool.
  • workqueue aborts the process on concurrent entry. That is live today: a chunked neighborhood reduction (the target="parallel" gufunc runs inside apply_ufunc(dask="parallelized")) kills the interpreter. It predates the nogil change; it aborts with the GIL held too.

numba_pool() in uxarray/utils/parallel.py guards each call by layer: a lock under workqueue (on every thread, since a main-thread call racing a dask task aborts too), numba.set_num_threads(1) off the main thread under omp, restored afterwards because the setting is thread-local. The 14 parallel=True kernels move to @parallel_njit, which is njit plus that guard; the neighborhood gufunc wraps its call in it.

The two remedies are per layer on purpose. Capping does not stop the workqueue abort, and capping under the lock would run every pool-thread call one at a time on one thread.

Chunked neighborhood percentile, 64 steps, best of 3:

layer        before            after
tbb          139ms             139ms
omp          179 OS threads    47 OS threads
workqueue    abort             ~1.1s

The workqueue time is the layer, not the lock: main takes 1176ms under dask's sync scheduler.

Test suite: 983 passed, 1 skipped. test_plot_with_features fails identically before and after (matplotlib figure size, unrelated).

PR Checklist

General

  • An issue is created and linked
  • Added appropriate labels (if your uxarray repo permissions allow it)
  • Filled out Overview and Expected Usage (if applicable) sections

Testing & Benchmarking

  • There is adequate test coverage of changes from this PR (add new tests if needed)
  • If this PR could affect performance, ran ASV benchmarks and confirmed they show expected behavior (add a new benchmark if necessary)

Documentation and Examples

  • Docstrings updated with any function changes, and included in all new functions
  • User (public) functions added to docs/api.rst; internal (private) function names start with an underscore (_)

AI Disclosure

AI Usage: Claude Opus 5.5

  • I have tested and take responsibility for all AI-generated content in my PR.

A parallel=True kernel called from a thread pool -- dask's threaded
scheduler, typically -- nests numba's pool under it. What that does
depends on numba's threading layer:

  * tbb composes; nothing to do.
  * omp starts a full team per calling thread: 179 OS threads from a
    12-thread pool.
  * workqueue aborts the process on concurrent entry. That is live today:
    a chunked neighborhood reduction (the target="parallel" gufunc runs
    inside apply_ufunc(dask="parallelized")) kills the interpreter. It
    predates the nogil change; it aborts with the GIL held too.

numba_pool() in uxarray/utils/parallel.py guards each call by layer: a
lock under workqueue (on every thread, since a main-thread call racing a
dask task aborts too), numba.set_num_threads(1) off the main thread
under omp, restored afterwards because the setting is thread-local. The
14 parallel=True kernels move to @parallel_njit, which is njit plus that
guard; the neighborhood gufunc wraps its call in it.

The two remedies are per layer on purpose. Capping does not stop the
workqueue abort, and capping under the lock would run every pool-thread
call one at a time on one thread.

Chunked neighborhood percentile, 64 steps, best of 3:

    layer        before            after
    tbb          139ms             139ms
    omp          179 OS threads    47 OS threads
    workqueue    abort             ~1.1s

The workqueue time is the layer, not the lock: main takes 1176ms under
dask's sync scheduler.

Test suite: 983 passed, 1 skipped. test_plot_with_features fails
identically before and after (matplotlib figure size, unrelated).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@cmdupuis3 cmdupuis3 self-assigned this Sep 29, 2026
@cmdupuis3 cmdupuis3 added the scalability Related to scalability & performance efforts label Sep 29, 2026
Conflict in uxarray/grid/neighbors.py: UXARRAY#1768 replaced the target="parallel"
gufunc with njit kernels, running the serial one on dask blocks and
_reduce_rows_parallel only on in-memory arrays. Took that version, and put
_reduce_rows_parallel under @parallel_njit like the other parallel kernels:
an in-memory reduction called from a user's own thread pool still aborted
under workqueue.

test_parallel drops the neighborhood case, whose kernel no longer exists;
that path now goes through @parallel_njit like the centroid kernel it keeps.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

scalability Related to scalability & performance efforts

Projects

None yet

Development

Successfully merging this pull request may close these issues.

numba threadpool under dask results in thread oversubscription

1 participant