diff --git a/cuda_core/cuda/core/_memory/_memory_pool.pyi b/cuda_core/cuda/core/_memory/_memory_pool.pyi index 364f901e3a8..18ece4234fb 100644 --- a/cuda_core/cuda/core/_memory/_memory_pool.pyi +++ b/cuda_core/cuda/core/_memory/_memory_pool.pyi @@ -9,33 +9,63 @@ from cuda.core.typing import DevicePointerType class _MemPoolAttributes: - """Provides access to memory pool attributes.""" + """Provides access to memory pool attributes. + + Setting an attribute changes the memory pool itself, which other memory + resources that share the pool, such as a device's default pool, also see. + Reading or setting an attribute does not affect a stream capture in + progress on any thread. The driver rejects reading or setting an + attribute of a memory resource mapped from another process + (``is_mapped``), which raises ``CUDAError``. + """ def __init__(self, *args, **kwargs) -> None: ... def __repr__(self) -> str: ... @property def reuse_follow_event_dependencies(self) -> bool: """Allow memory to be reused when there are event dependencies between streams.""" + @reuse_follow_event_dependencies.setter + def reuse_follow_event_dependencies(self, value: bool) -> None: ... @property def reuse_allow_opportunistic(self) -> bool: """Allow reuse of completed frees without dependencies.""" + @reuse_allow_opportunistic.setter + def reuse_allow_opportunistic(self, value: bool) -> None: ... @property def reuse_allow_internal_dependencies(self) -> bool: """Allow insertion of new stream dependencies for memory reuse.""" + @reuse_allow_internal_dependencies.setter + def reuse_allow_internal_dependencies(self, value: bool) -> None: ... @property def release_threshold(self) -> int: - """Amount of reserved memory to hold before OS release.""" + """ + Amount of reserved memory, in bytes, to hold before OS release. When + more than this is reserved, the pool returns unused memory to the OS + at the next synchronization of an event, stream or device. + """ + @release_threshold.setter + def release_threshold(self, value: int) -> None: ... @property def reserved_mem_current(self) -> int: """Current amount of backing memory allocated.""" @property def reserved_mem_high(self) -> int: - """High watermark of backing memory allocated.""" + """ + High watermark of backing memory allocated. It can be set to zero to + reset it to the current value. + """ + @reserved_mem_high.setter + def reserved_mem_high(self, value: int) -> None: ... @property def used_mem_current(self) -> int: """Current amount of memory in use.""" @property def used_mem_high(self) -> int: - """High watermark of memory in use.""" + """ + High watermark of memory in use. It can be set to zero to reset it to + the current value. + """ + @used_mem_high.setter + def used_mem_high(self, value: int) -> None: ... class _MemPool(MemoryResource): def __init__(self) -> None: ... @@ -83,9 +113,24 @@ class _MemPool(MemoryResource): asynchronously. Must be passed explicitly; pass ``device.default_stream`` to use the default stream. """ + def trim(self, min_bytes_to_keep: int=0) -> None: + """Release unused memory held by the memory pool back to the OS. + + Memory that backs live allocations is not released, nor is memory + freed on a stream until the host has observed the free through a + synchronization. Trimming does not affect a stream capture in progress + on any thread. The driver rejects trimming a memory resource mapped + from another process (``is_mapped``), which raises ``CUDAError``. + + Parameters + ---------- + min_bytes_to_keep : int, optional + Reserved memory, in bytes, that the pool keeps. Defaults to 0, + which releases all unused memory. + """ @property def attributes(self) -> _MemPoolAttributes: - """Memory pool attributes.""" + """:obj:`~_memory._memory_pool._MemPoolAttributes` object with the attributes of the memory pool.""" @property def handle(self) -> object: """Handle to the underlying memory pool.""" diff --git a/cuda_core/cuda/core/_memory/_memory_pool.pyx b/cuda_core/cuda/core/_memory/_memory_pool.pyx index e9c7912f8fc..7d071f62238 100644 --- a/cuda_core/cuda/core/_memory/_memory_pool.pyx +++ b/cuda_core/cuda/core/_memory/_memory_pool.pyx @@ -6,7 +6,7 @@ from __future__ import annotations cimport cython from libc.limits cimport ULLONG_MAX -from libc.stdint cimport uintptr_t +from libc.stdint cimport uint64_t, uintptr_t from libc.string cimport memset from cuda.bindings cimport cydriver @@ -40,8 +40,49 @@ if TYPE_CHECKING: from cuda.core.typing import DevicePointerType +# Pool attribute reads and writes and cuMemPoolTrimTo are potentially unsafe +# calls during stream capture: the driver refuses them, and invalidates the +# capture, while a global-mode capture is active on any thread or a +# thread-local one on the calling thread. Capture mode is a per-thread +# property, so the calling thread is switched to relaxed mode around such a +# call and its previous mode is restored afterwards. This has no observable +# effect when nothing is capturing. The call executes immediately rather than +# being recorded into a graph, which is the intent for a process-wide pool. +cdef inline int _enter_relaxed_capture_mode(cydriver.CUstreamCaptureMode* mode) except?-1 nogil: + # On return, mode holds the thread's previous capture mode. + mode[0] = cydriver.CU_STREAM_CAPTURE_MODE_RELAXED + return HANDLE_RETURN(cydriver.cuThreadExchangeStreamCaptureMode(mode)) + + +cdef int _exit_relaxed_capture_mode( + cydriver.CUresult err, cydriver.CUstreamCaptureMode* mode) except?-1 nogil: + """Restore the thread's capture mode, then raise if the call or the restore failed.""" + cdef cydriver.CUresult restore_err = cydriver.cuThreadExchangeStreamCaptureMode(mode) + if err != cydriver.CUDA_SUCCESS: + with gil: + try: + HANDLE_RETURN(err) + except: + if restore_err != cydriver.CUDA_SUCCESS: + # The call's error propagates with the restore failure + # attached as a note (error handling policy). + attach_rollback_failure( + b"cuThreadExchangeStreamCaptureMode", restore_err, + b"failed to restore the thread's stream capture mode") + raise + return HANDLE_RETURN(restore_err) + + cdef class _MemPoolAttributes: - """Provides access to memory pool attributes.""" + """Provides access to memory pool attributes. + + Setting an attribute changes the memory pool itself, which other memory + resources that share the pool, such as a device's default pool, also see. + Reading or setting an attribute does not affect a stream capture in + progress on any thread. The driver rejects reading or setting an + attribute of a memory resource mapped from another process + (``is_mapped``), which raises ``CUDAError``. + """ cdef: MemoryPoolHandle _h_pool @@ -61,10 +102,33 @@ cdef class _MemPoolAttributes: ) cdef int _getattribute(self, cydriver.CUmemPool_attribute attr_enum, void* value) except?-1: + cdef cydriver.CUstreamCaptureMode mode + cdef cydriver.CUresult err with nogil: - HANDLE_RETURN(cydriver.cuMemPoolGetAttribute(as_cu(self._h_pool), attr_enum, value)) + _enter_relaxed_capture_mode(&mode) + err = cydriver.cuMemPoolGetAttribute(as_cu(self._h_pool), attr_enum, value) + _exit_relaxed_capture_mode(err, &mode) return 0 + cdef int _setattribute(self, cydriver.CUmemPool_attribute attr_enum, void* value) except?-1: + cdef cydriver.CUstreamCaptureMode mode + cdef cydriver.CUresult err + with nogil: + _enter_relaxed_capture_mode(&mode) + err = cydriver.cuMemPoolSetAttribute(as_cu(self._h_pool), attr_enum, value) + _exit_relaxed_capture_mode(err, &mode) + return 0 + + cdef int _setflag(self, cydriver.CUmemPool_attribute attr_enum, bint flag) except?-1: + cdef int value = flag + return self._setattribute(attr_enum, &value) + + cdef int _reset_watermark(self, cydriver.CUmemPool_attribute attr_enum, str name, value) except?-1: + if value != 0: + raise AttributeError(f"Attribute '{name}' may only be set to zero (got {value}).") + cdef cydriver.cuuint64_t zero = 0 + return self._setattribute(attr_enum, &zero) + @property def reuse_follow_event_dependencies(self) -> bool: """Allow memory to be reused when there are event dependencies between streams.""" @@ -72,6 +136,10 @@ cdef class _MemPoolAttributes: self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES, &value) return bool(value) + @reuse_follow_event_dependencies.setter + def reuse_follow_event_dependencies(self, value: bool) -> None: + self._setflag(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES, value) + @property def reuse_allow_opportunistic(self) -> bool: """Allow reuse of completed frees without dependencies.""" @@ -79,6 +147,10 @@ cdef class _MemPoolAttributes: self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC, &value) return bool(value) + @reuse_allow_opportunistic.setter + def reuse_allow_opportunistic(self, value: bool) -> None: + self._setflag(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC, value) + @property def reuse_allow_internal_dependencies(self) -> bool: """Allow insertion of new stream dependencies for memory reuse.""" @@ -86,13 +158,25 @@ cdef class _MemPoolAttributes: self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES, &value) return bool(value) + @reuse_allow_internal_dependencies.setter + def reuse_allow_internal_dependencies(self, value: bool) -> None: + self._setflag(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES, value) + @property def release_threshold(self) -> int: - """Amount of reserved memory to hold before OS release.""" + """ + Amount of reserved memory, in bytes, to hold before OS release. When + more than this is reserved, the pool returns unused memory to the OS + at the next synchronization of an event, stream or device. + """ cdef cydriver.cuuint64_t value self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, &value) return int(value) + @release_threshold.setter + def release_threshold(self, uint64_t value) -> None: + self._setattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, &value) + @property def reserved_mem_current(self) -> int: """Current amount of backing memory allocated.""" @@ -102,11 +186,19 @@ cdef class _MemPoolAttributes: @property def reserved_mem_high(self) -> int: - """High watermark of backing memory allocated.""" + """ + High watermark of backing memory allocated. It can be set to zero to + reset it to the current value. + """ cdef cydriver.cuuint64_t value self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH, &value) return int(value) + @reserved_mem_high.setter + def reserved_mem_high(self, value: int) -> None: + self._reset_watermark( + cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH, "reserved_mem_high", value) + @property def used_mem_current(self) -> int: """Current amount of memory in use.""" @@ -116,11 +208,19 @@ cdef class _MemPoolAttributes: @property def used_mem_high(self) -> int: - """High watermark of memory in use.""" + """ + High watermark of memory in use. It can be set to zero to reset it to + the current value. + """ cdef cydriver.cuuint64_t value self._getattribute(cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_HIGH, &value) return int(value) + @used_mem_high.setter + def used_mem_high(self, value: int) -> None: + self._reset_watermark( + cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_USED_MEM_HIGH, "used_mem_high", value) + cdef class _MemPool(MemoryResource): @@ -193,10 +293,33 @@ cdef class _MemPool(MemoryResource): cdef Stream s = Stream_accept(stream) _MP_deallocate(self, ptr, size, s) + def trim(self, size_t min_bytes_to_keep=0) -> None: + """Release unused memory held by the memory pool back to the OS. + + Memory that backs live allocations is not released, nor is memory + freed on a stream until the host has observed the free through a + synchronization. Trimming does not affect a stream capture in progress + on any thread. The driver rejects trimming a memory resource mapped + from another process (``is_mapped``), which raises ``CUDAError``. + + Parameters + ---------- + min_bytes_to_keep : int, optional + Reserved memory, in bytes, that the pool keeps. Defaults to 0, + which releases all unused memory. + """ + MP_check_open(self) + cdef cydriver.CUstreamCaptureMode mode + cdef cydriver.CUresult err + with nogil: + _enter_relaxed_capture_mode(&mode) + err = cydriver.cuMemPoolTrimTo(as_cu(self._h_pool), min_bytes_to_keep) + _exit_relaxed_capture_mode(err, &mode) + @property @cython.critical_section def attributes(self) -> _MemPoolAttributes: - """Memory pool attributes.""" + """:obj:`~_memory._memory_pool._MemPoolAttributes` object with the attributes of the memory pool.""" MP_check_open(self) cdef _MemPoolAttributes attributes if self._attributes is None: @@ -313,24 +436,16 @@ cdef int MP_raise_release_threshold(_MemPool self) except? -1: the OS as soon as there are no active suballocations. Setting it to ULLONG_MAX avoids repeated OS round-trips. - Pool attribute reads and writes are potentially unsafe calls while the - calling thread is inside a global or thread-local stream capture: the - driver refuses them and invalidates the capture. Capture mode is a - per-thread property, so the thread is switched to relaxed mode around the - two calls and its previous mode is restored afterwards. This has no - observable effect when the thread is not capturing. The attribute write - executes immediately rather than being recorded into a graph, which is the - intent for a process-wide pool setting. + The read and the write run in relaxed capture mode, like every pool + attribute access (see ``_enter_relaxed_capture_mode``). """ MP_check_open(self) cdef cydriver.cuuint64_t current_threshold cdef cydriver.cuuint64_t max_threshold = ULLONG_MAX - cdef cydriver.CUstreamCaptureMode mode = cydriver.CU_STREAM_CAPTURE_MODE_RELAXED + cdef cydriver.CUstreamCaptureMode mode cdef cydriver.CUresult err - cdef cydriver.CUresult restore_err with nogil: - # On return, mode holds the thread's previous capture mode. - HANDLE_RETURN(cydriver.cuThreadExchangeStreamCaptureMode(&mode)) + _enter_relaxed_capture_mode(&mode) err = cydriver.cuMemPoolGetAttribute( as_cu(self._h_pool), cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, @@ -342,18 +457,7 @@ cdef int MP_raise_release_threshold(_MemPool self) except? -1: cydriver.CUmemPool_attribute.CU_MEMPOOL_ATTR_RELEASE_THRESHOLD, &max_threshold ) - restore_err = cydriver.cuThreadExchangeStreamCaptureMode(&mode) - try: - HANDLE_RETURN(err) - except: - if restore_err != cydriver.CUDA_SUCCESS: - # The attribute error propagates with the restore failure attached - # as a note (error handling policy). - attach_rollback_failure( - b"cuThreadExchangeStreamCaptureMode", restore_err, - b"failed to restore the thread's stream capture mode") - raise - HANDLE_RETURN(restore_err) + _exit_relaxed_capture_mode(err, &mode) return 0 diff --git a/cuda_core/docs/source/api_private.rst b/cuda_core/docs/source/api_private.rst index bc8c511b356..ccec08ce5af 100644 --- a/cuda_core/docs/source/api_private.rst +++ b/cuda_core/docs/source/api_private.rst @@ -50,6 +50,7 @@ CUDA runtime _memory._ipc.IPCBufferDescriptor _memory._ipc.VirtualMemoryFdBufferDescriptor _memory._managed_buffer.AccessedBySetProxy + _memory._memory_pool._MemPoolAttributes CUDA graphs diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index 5d6521d30a8..aad3df03963 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -76,6 +76,17 @@ New features top-level ``cuda.core`` namespace, alongside :class:`CUDAWarning`. Previously they were only available from a private module. +- The memory pool attributes of :class:`DeviceMemoryResource`, + :class:`PinnedMemoryResource` and :class:`ManagedMemoryResource` can now be + set: ``release_threshold`` takes a non-negative integer, the three + ``reuse_*`` policies take a bool, and ``reserved_mem_high`` and + ``used_mem_high`` can be set to zero to reset the watermark, as for + :class:`GraphMemoryResource`. The new :meth:`DeviceMemoryResource.trim` + method, also on the other two resources, returns unused pool memory to the + OS. Setting an attribute and ``trim()`` take effect immediately, even during + a stream capture, and leave the capture valid. + (`#3261 `__) + - Added :class:`VirtualMemoryBuffer`, the :class:`Buffer` subclass that :class:`VirtualMemoryResource` returns. Closing it on a capturing stream other than a default stream raises, because virtual memory deallocation @@ -155,6 +166,13 @@ Fixes and enhancements new pool with options during such a capture is still refused by the driver. (`#2834 `__) +- Reading the :attr:`~DeviceMemoryResource.attributes` of a + :class:`DeviceMemoryResource`, :class:`PinnedMemoryResource` or + :class:`ManagedMemoryResource` no longer fails, or invalidates a stream + capture, while the calling thread is inside a ``thread_local`` capture or + any thread is inside a ``global`` one. + (`#3236 `__) + - Ending, closing, or collecting a top-level :class:`~graph.GraphBuilder` whose capture was invalidated, or left unjoined by a :meth:`~graph.GraphBuilder.join` that failed partway, no longer crashes. diff --git a/cuda_core/tests/graph/test_graph_memory_resource.py b/cuda_core/tests/graph/test_graph_memory_resource.py index 19aac877a4a..faf56c061a3 100644 --- a/cuda_core/tests/graph/test_graph_memory_resource.py +++ b/cuda_core/tests/graph/test_graph_memory_resource.py @@ -9,6 +9,7 @@ import pytest from helpers.buffers import compare_buffer_to_constant, make_scratch_buffer, set_buffer, thread_unsafe_on_windows +from cuda.bindings import driver from cuda.core import ( Device, DeviceMemoryResource, @@ -404,14 +405,26 @@ def worker(): def test_dmr_init_restores_capture_mode(mempool_device, mode): """After construction the thread is back in its previous capture mode. - Had the constructor left the thread in relaxed mode, this pool attribute - read would be accepted. Instead the driver refuses it and invalidates the - capture, which end_building() then ends and reports (#2834). + The constructor switches the thread to relaxed mode around its pool calls + (#2834), and must not leave it there. """ + + def thread_capture_mode(): + # The exchange returns the mode it replaces, so exchanging twice reads + # the thread's mode without changing it. + relaxed = driver.CUstreamCaptureMode.CU_STREAM_CAPTURE_MODE_RELAXED + err, current = driver.cuThreadExchangeStreamCaptureMode(relaxed) + assert err == driver.CUresult.CUDA_SUCCESS + err, _ = driver.cuThreadExchangeStreamCaptureMode(current) + assert err == driver.CUresult.CUDA_SUCCESS + return current + device = mempool_device gb = device.create_stream().create_graph_builder().begin_building(mode=mode) - dmr = DeviceMemoryResource(device) - with pytest.raises(CUDAError, match="CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED"): - dmr.attributes.release_threshold # noqa: B018 - the read itself is the probe - with pytest.raises(CUDAError, match="CUDA_ERROR_STREAM_CAPTURE_INVALIDATED"): + try: + before = thread_capture_mode() + DeviceMemoryResource(device) + assert thread_capture_mode() == before + finally: + # Raises if construction invalidated the capture. gb.end_building() diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 6451ae0ba40..9cb505ee2cb 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import concurrent.futures import ctypes import gc import multiprocessing as mp @@ -2775,6 +2776,168 @@ def test_mempool_attributes(ipc_enabled, memory_resource_factory, property_name, assert value >= current_value, f"{property_name} should be >= {current_prop}" +def _create_owned_pool_or_skip(memory_resource_factory, device): + MR, MRops = memory_resource_factory + if MR is DeviceMemoryResource: + if not device.properties.memory_pools_supported: + pytest.skip("Device does not support mempool operations") + return MR(device, options=MRops(max_size=POOL_SIZE)) + if MR is PinnedMemoryResource: + skip_if_pinned_memory_unsupported(device) + return MR(options=MRops(max_size=POOL_SIZE)) + skip_if_managed_memory_unsupported(device) + return create_managed_memory_resource_or_skip(options=MRops()) + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_mempool_attributes_set(memory_resource_factory): + device = Device() + device.set_current() + mr = _create_owned_pool_or_skip(memory_resource_factory, device) + attributes = mr.attributes + try: + for name in ( + "reuse_follow_event_dependencies", + "reuse_allow_opportunistic", + "reuse_allow_internal_dependencies", + ): + for flag in (False, True): + setattr(attributes, name, flag) + assert getattr(attributes, name) is flag + + attributes.release_threshold = 1 << 20 + assert attributes.release_threshold == 1 << 20 + attributes.release_threshold = 0 + assert attributes.release_threshold == 0 + for value in (-1, 2**64): + with pytest.raises(OverflowError): + attributes.release_threshold = value + + for name in ("reserved_mem_high", "used_mem_high"): + with pytest.raises(AttributeError, match=f"'{name}' may only be set to zero"): + setattr(attributes, name, 1) + for name in ("reserved_mem_current", "used_mem_current"): + with pytest.raises(AttributeError): + setattr(attributes, name, 0) + finally: + mr.close() + + +class _IntSubclass(int): + pass + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +@pytest.mark.parametrize( + "make_value", + [ + pytest.param(lambda: True, id="bool"), + pytest.param(lambda: _IntSubclass(4096), id="int-subclass"), + pytest.param(lambda: np.uint64(4096), id="np-uint64"), + ], +) +def test_mempool_release_threshold_accepts_index_types(memory_resource_factory, make_value): + device = Device() + device.set_current() + mr = _create_owned_pool_or_skip(memory_resource_factory, device) + try: + value = make_value() + mr.attributes.release_threshold = value + assert mr.attributes.release_threshold == int(value) + finally: + mr.close() + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +@pytest.mark.parametrize("kind", ["reserved", "used"]) +def test_mempool_watermark_reset(memory_resource_factory, kind): + device = Device() + device.set_current() + mr = _create_owned_pool_or_skip(memory_resource_factory, device) + high, current = f"{kind}_mem_high", f"{kind}_mem_current" + try: + buffer = mr.allocate(POOL_SIZE // 2, stream=device.default_stream) + buffer.close() + device.sync() + # Return the freed memory so that the reserved watermark is above the + # current value too. + mr.trim() + assert getattr(mr.attributes, high) >= POOL_SIZE // 2 + assert getattr(mr.attributes, current) < POOL_SIZE // 2 + + setattr(mr.attributes, high, 0) + assert getattr(mr.attributes, high) == getattr(mr.attributes, current) + finally: + mr.close() + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_mempool_trim(memory_resource_factory): + device = Device() + device.set_current() + mr = _create_owned_pool_or_skip(memory_resource_factory, device) + try: + # Keep freed memory reserved so that only trim() returns it. + mr.attributes.release_threshold = 2**64 - 1 + buffer = mr.allocate(POOL_SIZE // 2, stream=device.default_stream) + buffer.close() + device.sync() + reserved = mr.attributes.reserved_mem_current + assert reserved > 0 + + mr.trim(min_bytes_to_keep=reserved) + assert mr.attributes.reserved_mem_current == reserved + + mr.trim() + assert mr.attributes.reserved_mem_current == 0 + finally: + mr.close() + + with pytest.raises(RuntimeError, match="has been closed"): + mr.trim() + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +@pytest.mark.thread_unsafe(reason="a global-mode capture affects every thread") +@pytest.mark.parametrize( + ("mode", "caller"), + # A thread-local capture does not restrict other threads, so pairing it + # with other_thread would pass without relaxed mode. + [("global", "same_thread"), ("global", "other_thread"), ("thread_local", "same_thread")], +) +def test_mempool_calls_during_capture_keep_capture_valid(memory_resource_factory, mode, caller): + """Pool attribute reads, writes and trim() run in relaxed capture mode. + + The driver refuses these calls, and invalidates the capture, while a + global-mode capture is active on any thread or a thread-local one on the + calling thread. + """ + device = Device() + device.set_current() + mr = _create_owned_pool_or_skip(memory_resource_factory, device) + + def use_pool(): + device.set_current() + attributes = mr.attributes + attributes.release_threshold = attributes.release_threshold + attributes.reuse_allow_opportunistic = attributes.reuse_allow_opportunistic + attributes.used_mem_high = 0 + mr.trim() + + gb = device.create_graph_builder().begin_building(mode=mode) + try: + if caller == "same_thread": + use_pool() + else: + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + pool.submit(use_pool).result() + finally: + # Raises if a pool call invalidated the capture. + gb.end_building() + gb.close() + mr.close() + + def test_mempool_attributes_repr(memory_resource_factory): """Test the repr of memory pool attributes for all memory resource types.""" MR, MRops = memory_resource_factory