Skip to content

Commit dc2a46c

Browse files
committed
cuda.core: keep stream captures valid during memory pool attribute calls
The driver refuses pool attribute reads and writes and cuMemPoolTrimTo, and invalidates the capture, while a global-mode capture is active on any thread or a thread-local one on the calling thread. Run them in relaxed capture mode, as MP_raise_release_threshold already did, through helpers shared with it. test_dmr_init_restores_capture_mode used an attribute read as its probe for the restored capture mode. It now reads the thread's mode directly.
1 parent fd26cea commit dc2a46c

5 files changed

Lines changed: 123 additions & 36 deletions

File tree

‎cuda_core/cuda/core/_memory/_memory_pool.pyi‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ class _MemPoolAttributes:
1313
1414
Setting an attribute changes the memory pool itself, which other memory
1515
resources that share the pool, such as a device's default pool, also see.
16+
Reading or setting an attribute does not affect a stream capture in
17+
progress on any thread.
1618
"""
1719
def __init__(self, *args, **kwargs) -> None: ...
1820
def __repr__(self) -> str: ...
@@ -114,7 +116,8 @@ class _MemPool(MemoryResource):
114116
115117
Memory that backs live allocations is not released, nor is memory
116118
freed on a stream until the host has observed the free through a
117-
synchronization.
119+
synchronization. Trimming does not affect a stream capture in progress
120+
on any thread.
118121
119122
Parameters
120123
----------

‎cuda_core/cuda/core/_memory/_memory_pool.pyx‎

Lines changed: 57 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -40,11 +40,46 @@ if TYPE_CHECKING:
4040
from cuda.core.typing import DevicePointerType
4141

4242

43+
# Pool attribute reads and writes and cuMemPoolTrimTo are potentially unsafe
44+
# calls during stream capture: the driver refuses them, and invalidates the
45+
# capture, while a global-mode capture is active on any thread or a
46+
# thread-local one on the calling thread. Capture mode is a per-thread
47+
# property, so the calling thread is switched to relaxed mode around such a
48+
# call and its previous mode is restored afterwards. This has no observable
49+
# effect when nothing is capturing. The call executes immediately rather than
50+
# being recorded into a graph, which is the intent for a process-wide pool.
51+
cdef inline int _enter_relaxed_capture_mode(cydriver.CUstreamCaptureMode* mode) except?-1 nogil:
52+
# On return, mode holds the thread's previous capture mode.
53+
mode[0] = cydriver.CU_STREAM_CAPTURE_MODE_RELAXED
54+
return HANDLE_RETURN(cydriver.cuThreadExchangeStreamCaptureMode(mode))
55+
56+
57+
cdef int _exit_relaxed_capture_mode(
58+
cydriver.CUresult err, cydriver.CUstreamCaptureMode* mode) except?-1 nogil:
59+
"""Restore the thread's capture mode, then raise if the call or the restore failed."""
60+
cdef cydriver.CUresult restore_err = cydriver.cuThreadExchangeStreamCaptureMode(mode)
61+
if err != cydriver.CUDA_SUCCESS:
62+
with gil:
63+
try:
64+
HANDLE_RETURN(err)
65+
except:
66+
if restore_err != cydriver.CUDA_SUCCESS:
67+
# The call's error propagates with the restore failure
68+
# attached as a note (error handling policy).
69+
attach_rollback_failure(
70+
b"cuThreadExchangeStreamCaptureMode", restore_err,
71+
b"failed to restore the thread's stream capture mode")
72+
raise
73+
return HANDLE_RETURN(restore_err)
74+
75+
4376
cdef class _MemPoolAttributes:
4477
"""Provides access to memory pool attributes.
4578
4679
Setting an attribute changes the memory pool itself, which other memory
4780
resources that share the pool, such as a device's default pool, also see.
81+
Reading or setting an attribute does not affect a stream capture in
82+
progress on any thread.
4883
"""
4984
cdef:
5085
MemoryPoolHandle _h_pool
@@ -65,13 +100,21 @@ cdef class _MemPoolAttributes:
65100
)
66101

67102
cdef int _getattribute(self, cydriver.CUmemPool_attribute attr_enum, void* value) except?-1:
103+
cdef cydriver.CUstreamCaptureMode mode
104+
cdef cydriver.CUresult err
68105
with nogil:
69-
HANDLE_RETURN(cydriver.cuMemPoolGetAttribute(as_cu(self._h_pool), attr_enum, value))
106+
_enter_relaxed_capture_mode(&mode)
107+
err = cydriver.cuMemPoolGetAttribute(as_cu(self._h_pool), attr_enum, value)
108+
_exit_relaxed_capture_mode(err, &mode)
70109
return 0
71110

72111
cdef int _setattribute(self, cydriver.CUmemPool_attribute attr_enum, void* value) except?-1:
112+
cdef cydriver.CUstreamCaptureMode mode
113+
cdef cydriver.CUresult err
73114
with nogil:
74-
HANDLE_RETURN(cydriver.cuMemPoolSetAttribute(as_cu(self._h_pool), attr_enum, value))
115+
_enter_relaxed_capture_mode(&mode)
116+
err = cydriver.cuMemPoolSetAttribute(as_cu(self._h_pool), attr_enum, value)
117+
_exit_relaxed_capture_mode(err, &mode)
75118
return 0
76119

77120
cdef int _setflag(self, cydriver.CUmemPool_attribute attr_enum, bint flag) except?-1:
@@ -256,7 +299,8 @@ cdef class _MemPool(MemoryResource):
256299

257300
Memory that backs live allocations is not released, nor is memory
258301
freed on a stream until the host has observed the free through a
259-
synchronization.
302+
synchronization. Trimming does not affect a stream capture in progress
303+
on any thread.
260304

261305
Parameters
262306
----------
@@ -265,8 +309,12 @@ cdef class _MemPool(MemoryResource):
265309
which releases all unused memory.
266310
"""
267311
MP_check_open(self)
312+
cdef cydriver.CUstreamCaptureMode mode
313+
cdef cydriver.CUresult err
268314
with nogil:
269-
HANDLE_RETURN(cydriver.cuMemPoolTrimTo(as_cu(self._h_pool), min_bytes_to_keep))
315+
_enter_relaxed_capture_mode(&mode)
316+
err = cydriver.cuMemPoolTrimTo(as_cu(self._h_pool), min_bytes_to_keep)
317+
_exit_relaxed_capture_mode(err, &mode)
270318

271319
@property
272320
@cython.critical_section
@@ -388,24 +436,16 @@ cdef int MP_raise_release_threshold(_MemPool self) except? -1:
388436
the OS as soon as there are no active suballocations. Setting it to
389437
ULLONG_MAX avoids repeated OS round-trips.
390438

391-
Pool attribute reads and writes are potentially unsafe calls while the
392-
calling thread is inside a global or thread-local stream capture: the
393-
driver refuses them and invalidates the capture. Capture mode is a
394-
per-thread property, so the thread is switched to relaxed mode around the
395-
two calls and its previous mode is restored afterwards. This has no
396-
observable effect when the thread is not capturing. The attribute write
397-
executes immediately rather than being recorded into a graph, which is the
398-
intent for a process-wide pool setting.
439+
The read and the write run in relaxed capture mode, like every pool
440+
attribute access (see ``_enter_relaxed_capture_mode``).
399441
"""
400442
MP_check_open(self)
401443
cdef cydriver.cuuint64_t current_threshold
402444
cdef cydriver.cuuint64_t max_threshold = ULLONG_MAX
403-
cdef cydriver.CUstreamCaptureMode mode = cydriver.CU_STREAM_CAPTURE_MODE_RELAXED
445+
cdef cydriver.CUstreamCaptureMode mode
404446
cdef cydriver.CUresult err
405-
cdef cydriver.CUresult restore_err
406447
with nogil:
407-
# On return, mode holds the thread's previous capture mode.
408-
HANDLE_RETURN(cydriver.cuThreadExchangeStreamCaptureMode(&mode))
448+
_enter_relaxed_capture_mode(&mode)
409449
err = cydriver.cuMemPoolGetAttribute(
410450
as_cu(self._h_pool),
411451
cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
@@ -417,18 +457,7 @@ cdef int MP_raise_release_threshold(_MemPool self) except? -1:
417457
cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
418458
&max_threshold
419459
)
420-
restore_err = cydriver.cuThreadExchangeStreamCaptureMode(&mode)
421-
try:
422-
HANDLE_RETURN(err)
423-
except:
424-
if restore_err != cydriver.CUDA_SUCCESS:
425-
# The attribute error propagates with the restore failure attached
426-
# as a note (error handling policy).
427-
attach_rollback_failure(
428-
b"cuThreadExchangeStreamCaptureMode", restore_err,
429-
b"failed to restore the thread's stream capture mode")
430-
raise
431-
HANDLE_RETURN(restore_err)
460+
_exit_relaxed_capture_mode(err, &mode)
432461
return 0
433462

434463

‎cuda_core/docs/source/release/1.3.0-notes.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,10 @@ New features
8282
value, and ``reserved_mem_high`` and ``used_mem_high`` can be set to zero to
8383
reset the watermark, as for :class:`GraphMemoryResource`. The new
8484
``trim(min_bytes_to_keep=0)`` method returns unused pool memory to the OS.
85+
Reading or setting a pool attribute, and ``trim()``, no longer fail with
86+
``CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED`` and invalidate a stream capture
87+
that is in progress in global mode on any thread, or in thread-local mode
88+
on the calling thread.
8589

8690
- Added :class:`VirtualMemoryBuffer`, the :class:`Buffer` subclass that
8791
:class:`VirtualMemoryResource` returns. Closing it on a capturing stream

‎cuda_core/tests/graph/test_graph_memory_resource.py‎

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import pytest
1010
from helpers.buffers import compare_buffer_to_constant, make_scratch_buffer, set_buffer, thread_unsafe_on_windows
1111

12+
from cuda.bindings import driver
1213
from cuda.core import (
1314
Device,
1415
DeviceMemoryResource,
@@ -404,14 +405,26 @@ def worker():
404405
def test_dmr_init_restores_capture_mode(mempool_device, mode):
405406
"""After construction the thread is back in its previous capture mode.
406407
407-
Had the constructor left the thread in relaxed mode, this pool attribute
408-
read would be accepted. Instead the driver refuses it and invalidates the
409-
capture, which end_building() then ends and reports (#2834).
408+
The constructor switches the thread to relaxed mode around its pool calls
409+
(#2834), and must not leave it there.
410410
"""
411+
412+
def thread_capture_mode():
413+
# The exchange returns the mode it replaces, so exchanging twice reads
414+
# the thread's mode without changing it.
415+
relaxed = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_RELAXED
416+
err, current = driver.cuThreadExchangeStreamCaptureMode(relaxed)
417+
assert err == driver.CUresult.CUDA_SUCCESS
418+
err, _ = driver.cuThreadExchangeStreamCaptureMode(current)
419+
assert err == driver.CUresult.CUDA_SUCCESS
420+
return current
421+
411422
device = mempool_device
412423
gb = device.create_stream().create_graph_builder().begin_building(mode=mode)
413-
dmr = DeviceMemoryResource(device)
414-
with pytest.raises(CUDAError, match="CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED"):
415-
dmr.attributes.release_threshold # noqa: B018 - the read itself is the probe
416-
with pytest.raises(CUDAError, match="CUDA_ERROR_STREAM_CAPTURE_INVALIDATED"):
424+
try:
425+
before = thread_capture_mode()
426+
DeviceMemoryResource(device)
427+
assert thread_capture_mode() == before
428+
finally:
429+
# Raises if construction invalidated the capture.
417430
gb.end_building()

‎cuda_core/tests/test_memory.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

4+
import concurrent.futures
45
import ctypes
56
import gc
67
import multiprocessing as mp
@@ -2865,6 +2866,43 @@ def test_mempool_trim(memory_resource_factory):
28652866
mr.trim()
28662867

28672868

2869+
@pytest.mark.agent_authored(model="claude-opus-5-5")
2870+
@pytest.mark.thread_unsafe(reason="a global-mode capture affects every thread")
2871+
@pytest.mark.parametrize("caller", ["same_thread", "other_thread"])
2872+
@pytest.mark.parametrize("mode", ["global", "thread_local"])
2873+
def test_mempool_calls_during_capture_keep_capture_valid(memory_resource_factory, mode, caller):
2874+
"""Pool attribute reads, writes and trim() run in relaxed capture mode.
2875+
2876+
The driver refuses these calls, and invalidates the capture, while a
2877+
global-mode capture is active on any thread or a thread-local one on the
2878+
calling thread.
2879+
"""
2880+
device = Device()
2881+
device.set_current()
2882+
mr = _create_owned_pool_or_skip(memory_resource_factory, device)
2883+
2884+
def use_pool():
2885+
device.set_current()
2886+
attributes = mr.attributes
2887+
attributes.release_threshold = attributes.release_threshold
2888+
attributes.reuse_allow_opportunistic = attributes.reuse_allow_opportunistic
2889+
attributes.used_mem_high = 0
2890+
mr.trim()
2891+
2892+
gb = device.create_graph_builder().begin_building(mode=mode)
2893+
try:
2894+
if caller == "same_thread":
2895+
use_pool()
2896+
else:
2897+
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
2898+
pool.submit(use_pool).result()
2899+
finally:
2900+
# Raises if a pool call invalidated the capture.
2901+
gb.end_building()
2902+
gb.close()
2903+
mr.close()
2904+
2905+
28682906
def test_mempool_attributes_repr(memory_resource_factory):
28692907
"""Test the repr of memory pool attributes for all memory resource types."""
28702908
MR, MRops = memory_resource_factory

0 commit comments

Comments
 (0)