[12.9.x] Manual backport of #3308: lock the global callback map

cuda_bindings/cuda/bindings/driver.pyx

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # This code was automatically generated with version 12.9.0. Do not modify it directly.
-# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=7bd998915de52155f2e27696c9a4d514071b9c721d94e05f0a00c7a07394a961
+# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=52f8f192ef10bb489857ef7efde68c518ee854b6b1922ad9492f0d96fb1eac28
 from typing import Any, Optional
 import cython
 import ctypes
@@ -28650,7 +28650,7 @@ def cuDeviceRegisterAsyncNotification(device, callbackFunc, userData):
     if err != cydriver.CUDA_SUCCESS:
         free(cbData)
     else:
-        m_global._allocated[int(callback)] = cbData
+        m_global._insert(int(callback), cbData)
     _helper_input_void_ptr_free(&cyuserDataHelper)
     if err != cydriver.CUDA_SUCCESS:
         return (_CUresult(err), None)
@@ -28699,8 +28699,7 @@ def cuDeviceUnregisterAsyncNotification(device, callback):
     with nogil:
         err = cydriver.cuDeviceUnregisterAsyncNotification(cydevice, cycallback)
     if err == cydriver.CUDA_SUCCESS:
-        free(m_global._allocated[pcallback])
-        m_global._allocated.erase(<void_ptr>pcallback)
+        free(m_global._take(pcallback))
     return (_CUresult(err),)
 
 @cython.embedsignature(True)
@@ -51019,7 +51018,7 @@ def cuLogsRegisterCallback(callbackFunc, userData):
     if err != cydriver.CUDA_SUCCESS:
         free(cbData)
     else:
-        m_global._allocated[int(callback_out)] = cbData
+        m_global._insert(int(callback_out), cbData)
     _helper_input_void_ptr_free(&cyuserDataHelper)
     if err != cydriver.CUDA_SUCCESS:
         return (_CUresult(err), None)
@@ -51050,8 +51049,7 @@ def cuLogsUnregisterCallback(callback):
     with nogil:
         err = cydriver.cuLogsUnregisterCallback(cycallback)
     if err == cydriver.CUDA_SUCCESS:
-        free(m_global._allocated[pcallback])
-        m_global._allocated.erase(<void_ptr>pcallback)
+        free(m_global._take(pcallback))
     return (_CUresult(err),)
 
 @cython.embedsignature(True)
@@ -52490,6 +52488,19 @@ def cuGraphicsVDPAURegisterOutputSurface(vdpSurface, unsigned int flags):
 
 cdef class cudaBindingsDriverGlobal:
     cdef map[void_ptr, void*] _allocated
+    cdef cython.pymutex _lock
+
+    cdef void _insert(self, void_ptr key, void* value) noexcept:
+        with self._lock:
+            self._allocated[key] = value
+
+    cdef void* _take(self, void_ptr key) noexcept:
+        cdef void* value = NULL
+        with self._lock:
+            if self._allocated.count(key) != 0:
+                value = self._allocated[key]
+                self._allocated.erase(key)
+        return value
 
     def __dealloc__(self):
         for item in self._allocated:

cuda_bindings/cuda/bindings/runtime.pyx.in

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # This code was automatically generated with version 12.9.0. Do not modify it directly.
-# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=e8e09d2b9dc5c50f5e7a6d85f94bd57c0c0776668393328be79deebdefce1c11
+# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=f8e9899c4f71fdba2ecb42e39bdfdbc4fe30e61dbb988cd3d85b64fdba157099
 from typing import Any, Optional
 import cython
 import ctypes
@@ -20167,7 +20167,7 @@ def cudaDeviceRegisterAsyncNotification(int device, callbackFunc, userData):
     if err != cyruntime.cudaSuccess:
         free(cbData)
     else:
-        m_global._allocated[int(callback)] = cbData
+        m_global._insert(int(callback), cbData)
     _helper_input_void_ptr_free(&cyuserDataHelper)
 
     if err != cyruntime.cudaSuccess:
@@ -20212,8 +20212,7 @@ def cudaDeviceUnregisterAsyncNotification(int device, callback):
     with nogil:
         err = cyruntime.cudaDeviceUnregisterAsyncNotification(device, cycallback)
     if err == cyruntime.cudaSuccess:
-        free(m_global._allocated[pcallback])
-        m_global._allocated.erase(<void_ptr>pcallback)
+        free(m_global._take(pcallback))
     return (_cudaError_t(err),)
 {{endif}}
 
@@ -39171,6 +39170,19 @@ def getLocalRuntimeVersion():
 
 cdef class cudaBindingsRuntimeGlobal:
     cdef map[void_ptr, void*] _allocated
+    cdef cython.pymutex _lock
+
+    cdef void _insert(self, void_ptr key, void* value) noexcept:
+        with self._lock:
+            self._allocated[key] = value
+
+    cdef void* _take(self, void_ptr key) noexcept:
+        cdef void* value = NULL
+        with self._lock:
+            if self._allocated.count(key) != 0:
+                value = self._allocated[key]
+                self._allocated.erase(key)
+        return value
 
     def __dealloc__(self):
         for item in self._allocated:

cuda_bindings/docs/source/release/12.X.Y-notes.rst

@@ -33,6 +33,8 @@ Highlights
 Bug fixes
 ---------
 
+* Fixed a crash on free-threaded Python when callbacks (for example ``cuLogsRegisterCallback``) are registered or
+  unregistered from multiple threads at the same time. Access to the internal callback registry is now synchronized.
 
 Miscellaneous
 -------------

cuda_bindings/tests/test_callback_threads.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import threading
+
+import pytest
+
+from cuda.bindings import driver
+
+NUM_THREADS = 16
+NUM_ITERATIONS = 500
+
+
+@pytest.mark.agent_authored(model="claude-sonnet-5-5")
+def test_async_notification_register_unregister_from_many_threads():
+    (err,) = driver.cuInit(0)
+    assert err == driver.CUresult.CUDA_SUCCESS
+    err, device = driver.cuDeviceGet(0)
+    assert err == driver.CUresult.CUDA_SUCCESS
+
+    barrier = threading.Barrier(NUM_THREADS)
+    failures = []
+
+    def worker():
+        try:
+            barrier.wait()
+            for _ in range(NUM_ITERATIONS):
+                err, handle = driver.cuDeviceRegisterAsyncNotification(device, 0, 0)
+                assert err == driver.CUresult.CUDA_SUCCESS
+                (err,) = driver.cuDeviceUnregisterAsyncNotification(device, handle)
+                assert err == driver.CUresult.CUDA_SUCCESS
+        except BaseException as e:
+            failures.append(e)
+
+    threads = [threading.Thread(target=worker) for _ in range(NUM_THREADS)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert not failures, failures