test(bindings): isolate CUDA 12 graph capture by thread

Signed-off-by: Ralf W. Grosse-Kunstleve <rgrossekunst@nvidia.com>

Author: rwgk

cuda_bindings_12/tests/test_cuda.py

@@ -804,7 +804,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

cuda_bindings_12/tests/test_cudart.py

@@ -521,7 +521,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
@@ -1374,7 +1377,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)