Repository navigation
test(cuda.bindings): reuse worker-thread contexts under pytest-run-parallel - #3122
Conversation
…rallel The pytest-run-parallel plugin in the bindings conftest made each worker thread create and destroy a CUDA context around every test. The driver serializes context creation, so with four threads every test paid about 0.75 s. On free-threaded rows that run the bindings tests, the step took 36 to 39 minutes and the jobs hit the 60-minute limit. Keep one context per running thread, but reuse it. A pool hands each worker an idle context or creates one; the worker pushes it, runs the test, synchronizes, and pops it, so the previous current context is restored. A context that reports an error at the sync is destroyed, not reused. The plugin destroys the pooled contexts at each module boundary, so a context lives no longer than the module-scoped ctx fixture does in a run without threads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Andy-Jost
added
P0
📝 Summary
Priority: ➖ Normal Change: Bug fix Merge Risk: 🔵 Low · up to
|
|
…threads With pooled contexts the worker threads of a test start together; before, serialized context creation kept them apart. Three tests begin a stream capture in global mode, and in that mode a call such as cuGraphInstantiate in one thread fails while another thread is capturing. The host-function graph tests failed this way on the free-threaded rows. Capture in thread-local mode instead; what the tests check does not depend on it. Drop the synchronize at the end of each pooled test for the same reason: it could break a global capture running in another thread. Pooled contexts now behave like the module-scoped ctx fixture, which does not synchronize between tests either. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
ℹ️ Review info
⚙️ Run configuration
- Configuration used: Repository: NVIDIA/cuda-python/.coderabbit.yaml
- Review profile: CHILL
- Plan: Enterprise
- Run ID:
8dd0495a-1801-469a-8d80-52f4691873de
📒 Files selected for processing: 3
cuda_bindings/tests/conftest.pycuda_bindings/tests/test_cuda.pycuda_bindings/tests/test_cudart.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
| with self._lock: | ||
| self._idle.append(entry) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🔎 Supported by static analysis
🏁 Script executed:
sed -n '1,185p' cuda_bindings/tests/conftest.py
sed -n '600,655p' cuda_bindings/tests/test_cudart.py
sed -n '1460,1505p' cuda_bindings/tests/test_cudart.py
sed -n '695,735p' cuda_bindings/tests/test_cuda.pyRepository: NVIDIA/cuda-python
Length of output: 12896
🏁 Script executed:
set -o pipefail
printf '%s\n' '--- capture call sites ---'
rg -n -F --glob 'cuda_bindings/tests/test_cudart.py' -e 'cudaStreamBeginCapture' -e 'cudaStreamEndCapture' -e 'cudaStreamCaptureMode' .
rg -n -F --glob 'cuda_bindings/tests/test_cuda.py' -e 'cuStreamBeginCapture' -e 'cuStreamEndCapture' -e 'CU_STREAM_CAPTURE_MODE' .
printf '%s\n' '--- wrapper and teardown references ---'
rg -n -F --glob 'cuda_bindings/tests/*.py' -e '_thread_contexts' -e '_wrap_worker_cuda_test' -e 'pytest_runtest_teardown' -e 'pytest_runtest_call' -e 'pytest_runtest_makereport' .
printf '%s\n' '--- relevant complete blocks ---'
python3 - <<'PY'
from pathlib import Path
for name in ['cuda_bindings/tests/test_cudart.py', 'cuda_bindings/tests/test_cuda.py']:
p=Path(name)
lines=p.read_text().splitlines()
print(f'### {name}')
hits=[i for i,l in enumerate(lines) if 'StreamBeginCapture' in l or 'StreamEndCapture' in l]
ranges=[]
for i in hits:
start=max(0,i-18); end=min(len(lines),i+35)
if not ranges or start>ranges[-1][1]:
ranges.append([start,end])
else:
ranges[-1][1]=max(ranges[-1][1],end)
for start,end in ranges:
print(f'--- lines {start+1}-{end} ---')
for n in range(start,end):
print(f'{n+1:5}: {lines[n]}')
PYRepository: NVIDIA/cuda-python
Length of output: 10683
🌐 Web query:
CUDA Runtime API cudaStreamEndCapture same thread cudaStreamCaptureModeThreadLocal Relaxed
💡 Result:
`cudaStreamCaptureModeThreadLocal` does **not** relax the same-thread requirement. It only changes which threads are subject to capture’s restrictions on potentially unsafe API calls.
NVIDIA’s Runtime API says `cudaStreamEndCapture()` must be called on the thread that began capture whenever the mode is **not** `cudaStreamCaptureModeRelaxed`. So with `ThreadLocal`, begin and end capture on the same host thread. With `Relaxed`, ending on a different thread is permitted by that rule—but capture must still end on the same stream it began on, and other capture-conflict restrictions still apply. ([docs.nvidia.com](https://docs.nvidia.com/cuda/archive/13.4.2/cuda-runtime-api/cuda_runtime_api/group__CUDART__STREAM.html?utm_source=openai))
Citations:
- 1: https://docs.nvidia.com/cuda/archive/13.4.2/cuda-runtime-api/cuda_runtime_api/group__CUDART__STREAM.html?utm_source=openai
important: Discard the context when a worker test raises.
_ThreadContextPool.context() returns the popped context to _idle even when the test raises. Pytest clears the pool only at a module boundary or session end. A failed ThreadLocal capture can therefore leave active capture state and its stream in a context that a later worker reuses. CUDA requires a non-relaxed capture to end on its starting thread.
Destroy the context after an exception instead of adding it to _idle. This is a bounded test-isolation and resource-lifetime issue, not proof that every later test will fail.
mdboom
left a comment
mdboom
left a comment
There was a problem hiding this comment.
LGTM. Only some really niche things found in robo-review.
This appears to be an issue introduced with #3061 and uncovered by this change. When capturing graphs in global mode, while one thread is capturing, a sync or a graph instantiate in any other thread fails and breaks the capture. Under the old conftest, each thread first created a context, which takes about 0.15 s and is serialized, so the four test bodies never overlapped. |
3 participants
Summary
Every free-threaded nightly row on the newest CUDA has hit the 60-minute job limit since the limit went back to 60 minutes in #3019. The cuda.bindings test step causes it. On these rows it takes 36 to 39 minutes, against under 3 minutes on the GIL build of the same row.
The pytest-run-parallel plugin in
cuda_bindings/tests/conftest.pywraps each test that uses thedevicefixture. The module-scoped autousectxfixture depends ondevice, so the plugin wraps every test. The wrapper then has each worker thread callcuCtxCreateandcuCtxDestroyaround the test. With--parallel-threads=4, that is four context creations per test, and the driver serializes them.This PR keeps one context per running thread but reuses it:
ctxfixture does in a run without threads.The worker threads of a test now start together. Before, the serialized context creation kept them about 0.15 s apart. Three tests begin a stream capture in global mode, and in that mode a call such as
cuGraphInstantiatein one thread fails while another thread is capturing. The two host-function graph tests from #3061 failed this way on the first CI run of this PR, and one crashed on Windows. These three tests now capture in thread-local mode. What they check does not depend on the mode.Runs without
--parallel-threadsdo not register the plugin and are unchanged.Evidence
From the nightly of 2026-10-13, comparing the Python 3.15t and 3.15 rows on linux-64, CUDA 13.4.2, L4:
The 3.14t rows on 12.9.1 and 13.0.2 pass because they test against released bindings and skip this step. The H100 x2 row spent 54 minutes in this step and never reached the cuda.core tests. The cost dates from #2194, which enabled parallel threads in CI. The longer nightly limit from #2802 hid it until #3019.
Testing
A/B run of the cuda.bindings suite on one H100 with
--parallel-threads=4, Python 3.14 (GIL build) and cuda-bindings 13.4.1. Both arms used the same test tree and random seed; onlyconftest.pydiffered.--count=5The two single-run arms fail the same 25 items. They come from that machine's environment: libnvvm, the cuFile setup, and a build-hook test that needs a source checkout. The free-threaded rows of this PR's CI will confirm the result on free-threaded Python.
Free-threaded run on one H100 with Python 3.14.7t and the cuda.bindings wheel that this PR's CI built,
--parallel-threads=4:--count=30cuStreamSynchronizeMain and this PR fail the same 28 items, all from that machine's environment: missing libnvfatbin, setuptools, and pyglet, libnvvm, and the cuFile setup. The last arm reproduces the Windows crash. It confirms that the global-mode tests race once the threads start together.
Checked the pool logic with a fake driver module that keeps a context stack per thread. It ran under pytest 9.1.0 and pytest-run-parallel 0.10.0 with four threads, with and without
--count=5and random order. The checks: each thread sees its own context as current; the main thread's context is current before and after every test; no context is left alive at session end. A deliberate break, skipping the pop, fails the checks.pre-commit run --all-filespasses.Expected effect: the free-threaded 13.4.2 L4 job should drop from over 60 minutes to about 40. This is an estimate from the CI step timings and the A/B run above.
🤖 Generated with Claude Code