diff --git a/cuda_bindings/cuda/bindings/driver.pyx b/cuda_bindings/cuda/bindings/driver.pyx index 251f5c481cd..d679cd5c964 100644 --- a/cuda_bindings/cuda/bindings/driver.pyx +++ b/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(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(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: diff --git a/cuda_bindings/cuda/bindings/runtime.pyx.in b/cuda_bindings/cuda/bindings/runtime.pyx.in index 284e5f91e13..c95a9efe561 100644 --- a/cuda_bindings/cuda/bindings/runtime.pyx.in +++ b/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(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: diff --git a/cuda_bindings/docs/source/release/12.X.Y-notes.rst b/cuda_bindings/docs/source/release/12.X.Y-notes.rst index f5e5d00a172..4d37acb7a2d 100644 --- a/cuda_bindings/docs/source/release/12.X.Y-notes.rst +++ b/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 ------------- diff --git a/cuda_bindings/tests/test_callback_threads.py b/cuda_bindings/tests/test_callback_threads.py new file mode 100644 index 00000000000..1c0f44f2c8a --- /dev/null +++ b/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