From e55a371f0a948acdcea0d3bde03fc47ae07bbc25 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 15 Oct 2026 11:01:28 -0700 Subject: [PATCH 1/2] test(cuda.bindings): reuse worker-thread contexts under pytest-run-parallel 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 --- cuda_bindings/tests/conftest.py | 87 ++++++++++++++++++++++++++++----- 1 file changed, 74 insertions(+), 13 deletions(-) diff --git a/cuda_bindings/tests/conftest.py b/cuda_bindings/tests/conftest.py index fada7d95601..a782e21796f 100644 --- a/cuda_bindings/tests/conftest.py +++ b/cuda_bindings/tests/conftest.py @@ -6,6 +6,7 @@ import inspect import pathlib import sys +import threading from contextlib import contextmanager import pytest @@ -37,19 +38,72 @@ def pytest_configure(config): config.pluginmanager.register(_CudaBindingsParallelPlugin(), name="_cuda_bindings_parallel_plugin") -@contextmanager -def _thread_context(): - # Context setting up `device` and `ctx` for individual threads on - # pytest-run-parallel - err, device = cuda.cuDeviceGet(0) - assert err == cuda.CUresult.CUDA_SUCCESS - err, ctx = cuda.cuCtxCreate(None, 0, device) - assert err == cuda.CUresult.CUDA_SUCCESS - try: - yield device, ctx - finally: - (err,) = cuda.cuCtxDestroy(ctx) +class _ThreadContextPool: + """CUDA contexts for pytest-run-parallel worker threads, reused across tests. + + pytest-run-parallel starts new threads for every test, and each thread + needs a current context. Creating and destroying a context per thread per + test serializes in the driver: with four threads it cost about 0.75 s per + test and made the free-threaded bindings run more than ten times slower. + + The pool gives each running thread a context of its own, as before, and + keeps it for the next test. The plugin empties the pool at the end of each + test module, so a context lives no longer than the module-scoped ``ctx`` + fixture does in a run without threads. A context that reports an error at + the end of a test is destroyed rather than reused. + """ + + def __init__(self): + self._lock = threading.Lock() + self._idle = [] + + @staticmethod + def _create(): + err, device = cuda.cuDeviceGet(0) + assert err == cuda.CUresult.CUDA_SUCCESS + # cuCtxCreate makes the new context current; pop it so that context() + # can push and pop it like a reused one. + err, ctx = cuda.cuCtxCreate(None, 0, device) assert err == cuda.CUresult.CUDA_SUCCESS + err, _ = cuda.cuCtxPopCurrent() + assert err == cuda.CUresult.CUDA_SUCCESS + return device, ctx + + @contextmanager + def context(self): + with self._lock: + entry = self._idle.pop() if self._idle else None + if entry is None: + entry = self._create() + device, ctx = entry + (err,) = cuda.cuCtxPushCurrent(ctx) + assert err == cuda.CUresult.CUDA_SUCCESS + try: + yield device, ctx + finally: + (sync_err,) = cuda.cuCtxSynchronize() + # Pop rather than unset, so that the context that was current + # before the test (on the main thread, for tests that are not run + # in parallel) is current again. + err, _ = cuda.cuCtxPopCurrent() + assert err == cuda.CUresult.CUDA_SUCCESS + if sync_err == cuda.CUresult.CUDA_SUCCESS: + with self._lock: + self._idle.append(entry) + else: + (err,) = cuda.cuCtxDestroy(ctx) + assert err == cuda.CUresult.CUDA_SUCCESS + + def clear(self): + """Destroy the idle contexts. Call only when no test is running.""" + with self._lock: + entries, self._idle = self._idle, [] + for _, ctx in entries: + (err,) = cuda.cuCtxDestroy(ctx) + assert err == cuda.CUresult.CUDA_SUCCESS + + +_thread_contexts = _ThreadContextPool() def _wrap_worker_cuda_test(func): @@ -63,7 +117,7 @@ def _wrap_worker_cuda_test(func): @functools.wraps(func) def wrapper(*args, **kwargs): kwargs = dict(kwargs) # copy before mutating - with _thread_context() as (device, ctx): + with _thread_contexts.context() as (device, ctx): # device is None when reusing an existing context (defensive path); # keep whatever the fixture provided in kwargs as-is. if wants_device and device is not None: @@ -103,6 +157,13 @@ def pytest_collection_modifyitems(self, config, items): if _item_needs_thread_ctx(item): item.obj = _wrap_worker_cuda_test(item.obj) + @pytest.hookimpl(trylast=True) + def pytest_runtest_teardown(self, item, nextitem): + # The worker threads of this test have joined. Destroy the pooled + # contexts when the next test is in another module, or there is none. + if nextitem is None or getattr(nextitem, "module", None) is not getattr(item, "module", None): + _thread_contexts.clear() + @pytest.fixture(scope="module") def cuda_driver(): From 5a0ef2e7f2ee4259d30a253aa9944f3f36fba609 Mon Sep 17 00:00:00 2001 From: Andy Jost Date: Thu, 15 Oct 2026 12:24:15 -0700 Subject: [PATCH 2/2] test(cuda.bindings): capture thread-locally in tests run on parallel 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 --- cuda_bindings/tests/conftest.py | 16 +++++++--------- cuda_bindings/tests/test_cuda.py | 5 ++++- cuda_bindings/tests/test_cudart.py | 10 ++++++++-- 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/cuda_bindings/tests/conftest.py b/cuda_bindings/tests/conftest.py index a782e21796f..926248450ef 100644 --- a/cuda_bindings/tests/conftest.py +++ b/cuda_bindings/tests/conftest.py @@ -49,8 +49,11 @@ class _ThreadContextPool: The pool gives each running thread a context of its own, as before, and keeps it for the next test. The plugin empties the pool at the end of each test module, so a context lives no longer than the module-scoped ``ctx`` - fixture does in a run without threads. A context that reports an error at - the end of a test is destroyed rather than reused. + fixture does in a run without threads. + + The worker threads of a test now run at the same time. Do not add a + synchronizing call between tests here: while another thread captures a + stream in global mode, such a call fails and breaks that capture. """ def __init__(self): @@ -81,18 +84,13 @@ def context(self): try: yield device, ctx finally: - (sync_err,) = cuda.cuCtxSynchronize() # Pop rather than unset, so that the context that was current # before the test (on the main thread, for tests that are not run # in parallel) is current again. err, _ = cuda.cuCtxPopCurrent() assert err == cuda.CUresult.CUDA_SUCCESS - if sync_err == cuda.CUresult.CUDA_SUCCESS: - with self._lock: - self._idle.append(entry) - else: - (err,) = cuda.cuCtxDestroy(ctx) - assert err == cuda.CUresult.CUDA_SUCCESS + with self._lock: + self._idle.append(entry) def clear(self): """Destroy the idle contexts. Call only when no test is running.""" diff --git a/cuda_bindings/tests/test_cuda.py b/cuda_bindings/tests/test_cuda.py index 369190c6b64..d85724fe49a 100644 --- a/cuda_bindings/tests/test_cuda.py +++ b/cuda_bindings/tests/test_cuda.py @@ -706,7 +706,10 @@ def host_fn(userData): err, stream = cuda.cuStreamCreate(0) assert err == cuda.CUresult.CUDA_SUCCESS try: - (err,) = cuda.cuStreamBeginCapture(stream, cuda.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_GLOBAL) + # Thread-local mode: pytest-run-parallel runs this test in several threads + # at once, and in global mode a call such as cuGraphInstantiate in one + # thread fails while another thread is capturing. + (err,) = cuda.cuStreamBeginCapture(stream, cuda.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_THREAD_LOCAL) assert err == cuda.CUresult.CUDA_SUCCESS (err,) = cuda.cuLaunchHostFunc(stream, callback, ctypes.addressof(counter)) assert err == cuda.CUresult.CUDA_SUCCESS diff --git a/cuda_bindings/tests/test_cudart.py b/cuda_bindings/tests/test_cudart.py index d0be4166c0d..b16391210ed 100644 --- a/cuda_bindings/tests/test_cudart.py +++ b/cuda_bindings/tests/test_cudart.py @@ -616,7 +616,10 @@ def test_cudart_cudaStreamGetCaptureInfo(): assert status == cudart.cudaStreamCaptureStatus.cudaStreamCaptureStatusNone # start capture - (err,) = cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal) + # Thread-local mode: pytest-run-parallel runs this test in several threads + # at once, and in global mode a call such as cudaGraphInstantiate in one + # thread fails while another thread is capturing. + (err,) = cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal) assertSuccess(err) # validate that stream is capturing now @@ -1469,7 +1472,10 @@ def host_fn(userData): err, stream = cudart.cudaStreamCreate() assertSuccess(err) try: - (err,) = cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal) + # Thread-local mode: pytest-run-parallel runs this test in several threads + # at once, and in global mode a call such as cudaGraphInstantiate in one + # thread fails while another thread is capturing. + (err,) = cudart.cudaStreamBeginCapture(stream, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal) assertSuccess(err) (err,) = cudart.cudaLaunchHostFunc(stream, callback, ctypes.addressof(counter)) assertSuccess(err)