Skip to content

Keep a shared thread pool alive when the thread count is reset - #24770

Open
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:threadpool-keep-shared-pools
Open

shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:threadpool-keep-shared-pools

Conversation

@shoumikhin

Copy link
Copy Markdown
Contributor

What is wrong today

From Python, _unsafe_reset_threadpool(n) is the only way to set the CPU thread count. If it runs after a method that uses XNNPACK is loaded, the next runs can crash or hang:

method = Runtime.get().load_program(pte.buffer).load_method("forward")
_unsafe_reset_threadpool(16)
for _ in range(20):
    method.execute([x])  # segfault, abort, or hang

On Linux x86_64 with an 8-thread default pool, a reset to 16 after loading crashed in 10 of 12 runs, and a reset to 4 also in 10 of 12. The Android bindings reset the pool each time a Module is created, so a second Module can hit it too.

Why it happens

The reset destroys the current pthreadpool and creates a new one. But XNNPACK got that pthreadpool from get_pthreadpool() when the method was loaded, and it keeps using it. So the loaded method runs on a freed pool.

Resizing the pool in place is not safe either. XNNPACK sizes its per-thread buffers by the thread count when it plans a method, and it does not plan again when the count changes.

What this change does

A reset no longer destroys a pthreadpool that get_pthreadpool() has handed out. It stays alive with its size, next to the new one:

  • XNNPACK methods loaded before the reset keep running on their old pool.
  • Methods loaded after the reset, and ExecuTorch's own kernels, use the new pool.
  • A reset back to a kept size reuses that pool, so at most one pool per size is kept. _unsafe_destroy_threadpool() frees kept pools too.
  • A pool that was never handed out is destroyed as before, so the usual "reset, then load" order costs nothing extra.

The header and the Python docstring now explain this.

Tests

  • New ThreadPoolTest.ResetKeepsHandedOutPthreadpoolUsable: after a reset, the held pool still runs work on 2 threads, and a reset back to its size reuses it. It fails without the fix and passes with it.
  • Ran extension_threadpool_test, the native CPU engine tests, the Python tests for the pybindings and XNNPACK, and the LLM custom op tests that call the reset, on Linux x86_64. Tests that failed there fail on main too.
  • The script above now passes after resets to 16 and to 4. With 20 resets between 2 thread counts and a new method loaded after each, all 21 methods stay correct and the process gains no threads from round to round.

Fixes #24765

`_unsafe_reset_threadpool()` destroyed the current pthreadpool and created
a new one. Delegates such as XNNPACK get that pthreadpool from
`get_pthreadpool()` when a method is loaded and keep using it, so after a
reset the loaded method ran on a freed pool and crashed or hung. From
Python the reset is the only way to set the thread count, so calling it
after `load_method` was enough to hit this.

Now a pthreadpool that `get_pthreadpool()` has handed out stays alive
with its size when a reset replaces it. XNNPACK methods loaded before the
reset keep running on it, and methods loaded after it use the new pool. A
reset back to a kept size reuses that pool, so at most one pool per size
is kept, and `_unsafe_destroy_threadpool()` frees the kept pools too. A
pool that was never handed out is destroyed as before, so a reset before
loading costs nothing extra. Resizing a pool in place is not an option,
because XNNPACK sizes per-thread buffers by the thread count when it
plans a method.

Fixes pytorch#24765
@shoumikhin shoumikhin added the release notes: runtime Changes related to the core runtime which loads the program methods, initializes delegates, and runs label Nov 7, 2026
@shoumikhin shoumikhin added the release notes: runtime Changes related to the core runtime which loads the program methods, initializes delegates, and runs label Nov 7, 2026
Copilot AI balanced review requested due to automatic review settings November 7, 2026 05:16
@pytorch-bot

pytorch-bot Bot commented Nov 7, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24770

Note: Links to docs will display an error until the docs builds have been completed.

❌ You can merge normally! (2 Unrelated Failures), 3 Unclassified Failures

As of commit ce80ef8 with merge base 75af78c (image):

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

FLAKY - The following job failed but was likely due to flakiness present on trunk:

BROKEN TRUNK - The following job failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 7, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

@shoumikhin

Copy link
Copy Markdown
Contributor Author

@claude /pr-review

@claude

claude Bot commented Nov 7, 2026 •

Copy link
Copy Markdown

Claude finished @shoumikhin's task in 25m 5s —— View job


This stops XNNPACK methods from running on a freed pool, but get_pthreadpool() now keeps every pool it returns alive until the process exits, Python and Android can't free them, the new docs don't say what happens to methods loaded before a reset, and the #24765 crash has no automated test, so this needs discussion.

Detailed review

PR Review: #24770

Summary

A pthreadpool that get_pthreadpool() has returned now stays alive across _unsafe_reset_threadpool(), and a later reset to its size reuses it. The problems: kept pools are never released, the docs on already-loaded methods are incomplete, and the reported crash has no end-to-end test. Not checked: how this interacts with the Meta-internal UseNThreadsThreadPoolGuard, which is not in the OSS tree.

API Design

  • get_pthreadpool() pins every pool it returns, and nothing releases them (extension/threadpool/threadpool.cpp:222-223, :77-79).
    • Sticky flag: the getter sets threadpool_shared_, and nothing ever clears it. A pool counts as handed out once anyone has read it, not only while someone holds it. It is then kept until _unsafe_destroy_threadpool() or process exit, even after every method that used it is gone.
    • Short-lived callers pin too: backends/cpu/runtime/KernelProvider.cpp:19 calls get_pthreadpool() on every CPU-backend init and execute (CPUBackend.cpp:99,223). Keeping the pool does not keep that backend's old methods running, because CPUPlan::execute rejects any method whose pool or thread count changed (CPUPlan.cpp:150-154).
    • No release from Python or Android: _unsafe_destroy_threadpool() is the only release. It has no Python binding (pybindings.cpp:5709-5723) and no caller anywhere in the repo, so in Python and Android, the two callers this PR targets, kept pools live until exit.
    • Trigger: a benchmark resets to each n in 1..16, loads a fresh program and method, runs it once and drops it. Before this PR that was already safe and kept one pool alive. Now it keeps 16 pools, about 120 worker threads.
    • Android: every distinct numThreads passed to a Module (extension/android/jni/jni_layer.cpp:387) leaves one pool behind.
    • PR description: the "gains no threads" check only alternated between 2 sizes. "Reset, then load costs nothing extra" holds only if nothing has called get_pthreadpool() since the previous reset.
    • Suggestion: tie a pool's lifetime to whoever holds it. For example, XNNPACK takes a std::shared_ptr<pthreadpool> from a new accessor, and XNNExecutor keeps it as long as its runtime, as it already does with its other shared_ptr members (backends/xnnpack/runtime/XNNExecutor.h:47-51). retired_threadpools_ would then hold weak_ptrs, used only to reuse a pool by size. If the sticky flag stays, the get_pthreadpool() doc (threadpool.h:131-137) should at least say that calling it keeps the current pool alive until the process exits.
  • The new docs don't say what happens to methods loaded before a reset (extension/pybindings/pybindings.pyi:467-469, extension/threadpool/threadpool.h:72-75).
    • Two pools in one method: both docs suggest the new count reaches such a method only after a reload. That holds only for its XNNPACK partitions. Its ExecuTorch kernels use parallel_for → ThreadPool::run() (extension/threadpool/thread_parallel.cpp:47,87), so they move to the new pool immediately. One method then runs on two pools with two thread counts, and _threadpool_get_thread_count() reports a count its XNNPACK partitions don't use.
    • backends/cpu: methods loaded before the reset don't keep running at all. They fail with "CPU execution properties changed; reload the static method" (CPUPlan.cpp:150-154), as CPUEngineTest.RejectsExecutionAfterThreadPoolReset (backends/native/test/runtime/cpu/test_cpu_engine.cpp:1046) asserts.
    • FAQ: docs/source/using-executorch-faqs.md:46 still says to reset "prior to loading or running a model". With this PR, a reset between load and run is silently not applied to XNNPACK.
    • No log: keeping a pool logs nothing (threadpool.cpp:77-79). The No safe way to set the CPU thread count from Python: _unsafe_reset_threadpool after load_method crashes or hangs XNNPACK #24765 pattern no longer crashes, but the new count now applies only partly, with no sign of it.
    • Fix: say in both docstrings that methods loaded before the reset should be reloaded, fix the FAQ line, and add one ET_LOG in the branch that keeps the pool. Fix this →

Testing

  • The reported crash has no automated test.
    • End-to-end: no test loads an XNNPACK method, resets, then executes. extension/pybindings/test/test_pybindings.py already lowers programs with XnnpackPartitioner, so the repro script from the description fits there.
    • Unshared pools: no test checks that a pool get_pthreadpool() never returned is destroyed rather than kept.
    • Destroy: the new retired_threadpools_.clear() in _unsafe_destroy_threadpool() (threadpool.cpp:105) is untested.
    • Fix this →
  • The new test skips restoring the thread count when an assertion fails.
    • Problem: extension/threadpool/test/threadpool_test.cpp:325 is a bare call, and its result is not checked. A failing ASSERT_* at :279, :283, :288, :290 or :321 returns before it.
    • Effect: the singleton stays at 2 or 3 threads for later tests in the binary: TestNoThreadPoolGuard.*, and thread_parallel_test.cpp, which test/CMakeLists.txt:22 builds into the same binary.
    • Fix: use an RAII restorer like RestoreThreadCount in backends/native/test/runtime/cpu/test_cpu_engine.cpp:1059-1066. Fix this →

Recommendation

Needs Discussion

The blocker is a maintainer decision on how kept pools are owned. The PR's way is a sticky flag that any get_pthreadpool() call sets, which keeps pools alive for the life of the process, with no way to free them from Python or Android. The alternative is reference counting by the code that actually holds each pool. The doc and test items are small.


  • Gather context (diff, changed files, checklist)
  • Run three independent reviewers
  • Consolidate findings
  • Fact-check findings
  • Post review

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. release notes: runtime Changes related to the core runtime which loads the program methods, initializes delegates, and runs

Projects

None yet

Development

Successfully merging this pull request may close these issues.

No safe way to set the CPU thread count from Python: _unsafe_reset_threadpool after load_method crashes or hangs XNNPACK

2 participants