From 8e3dd5b1fa51aef001e8d31f6212755112fa5914 Mon Sep 17 00:00:00 2001 From: Jinfeng Date: Fri, 30 Oct 2026 23:35:12 +0000 Subject: [PATCH 1/2] Lock the callback bookkeeping map in cuda.bindings Under free-threaded Python, concurrent register/unregister of callbacks (cuLogs*, *AsyncNotification) corrupts the shared m_global._allocated std::map. Guard it with a cython.pymutex. driver.pyx and runtime.pyx are regenerated from cybind, not hand-edited. Closes #3282 --- cuda_bindings_13/cuda/bindings/driver.pyx | 25 ++++++++++++++++------ cuda_bindings_13/cuda/bindings/runtime.pyx | 20 +++++++++++++---- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/cuda_bindings_13/cuda/bindings/driver.pyx b/cuda_bindings_13/cuda/bindings/driver.pyx index eb037ffb4ed..ef2feb43e43 100644 --- a/cuda_bindings_13/cuda/bindings/driver.pyx +++ b/cuda_bindings_13/cuda/bindings/driver.pyx @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # This code was automatically generated with version 13.4.1. Do not modify it directly. -# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=348b02f0a76a3a271fbb11535f02d660d811d3c21a877e46bfc070bbabf43e76 +# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=8932ee959b6a2f1e69316f813cb8e8492c8e112094812dc1310c49970f00779a from typing import Any, Optional import cython import ctypes @@ -30899,7 +30899,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) @@ -30948,8 +30948,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) @@ -55743,7 +55742,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) @@ -55774,8 +55773,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) @@ -57375,6 +57373,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_13/cuda/bindings/runtime.pyx b/cuda_bindings_13/cuda/bindings/runtime.pyx index 4c6be9ee02d..a0486b038cf 100644 --- a/cuda_bindings_13/cuda/bindings/runtime.pyx +++ b/cuda_bindings_13/cuda/bindings/runtime.pyx @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # This code was automatically generated with version 13.4.1. Do not modify it directly. -# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=ea95eb12334cbd4c1bf3fc8e890fb6f9493063b12faea853bb3d8c6288934182 +# CYTHON-BINDINGS-GENERATED-DO-NOT-MODIFY-THIS-FILE: format=1; content-sha256=4cc5d08de16bc7cac781ba6e469ffe16b0d8e30beea5ec6f5e579f0fae24e834 from typing import Any, Optional import cython import ctypes @@ -21228,7 +21228,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: @@ -21270,8 +21270,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),) @cython.embedsignature(True) @@ -40262,6 +40261,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: From 1ceee4566320e376370c0f89b6b264d2debde834 Mon Sep 17 00:00:00 2001 From: Jinfeng Date: Fri, 30 Oct 2026 23:48:18 +0000 Subject: [PATCH 2/2] add a test case to increase probability triggering contention --- .../tests/test_callback_threads.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) create mode 100644 cuda_bindings_13/tests/test_callback_threads.py diff --git a/cuda_bindings_13/tests/test_callback_threads.py b/cuda_bindings_13/tests/test_callback_threads.py new file mode 100644 index 00000000000..0b445402c55 --- /dev/null +++ b/cuda_bindings_13/tests/test_callback_threads.py @@ -0,0 +1,45 @@ +# 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 +from cuda_python_test_helpers import driver_version_less_than + +NUM_THREADS = 16 +NUM_ITERATIONS = 500 + + +@pytest.mark.agent_authored(model="claude-sonnet-5-5") +@pytest.mark.skipif( + driver_version_less_than(12090), + reason="cuLogs* functions were introduced in CUDA 12.9", +) +@pytest.mark.thread_unsafe(reason="starts its own threads and needs them to overlap") +def test_culogs_register_unregister_from_many_threads(): + (err,) = driver.cuInit(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.cuLogsRegisterCallback(0, None) + assert err == driver.CUresult.CUDA_SUCCESS + (err,) = driver.cuLogsUnregisterCallback(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