You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
cuda.core: make memory pool attributes settable and add trim() - #3236
For DeviceMemoryResource, PinnedMemoryResource and ManagedMemoryResource:
mr.attributes.release_threshold and the three reuse_* policies can be set.
mr.attributes.reserved_mem_high and used_mem_high can be set to 0 to reset the watermark. Any other value raises AttributeError.
New mr.trim(min_bytes_to_keep=0) wraps cuMemPoolTrimTo.
Motivation
attributes already exposes the right fields, but all of them are read-only, and there is no way to trim a pool.
Measuring the peak device memory of each phase of a computation needs used_mem_high and its documented reset (setting it to 0). Stream-ordered allocations do not show up in CuPy's memory hooks, so the pool's watermark is the only source.
A downstream library's build_arena (_memory.py) allocates temporaries with cudaMallocAsync from the device's default pool, and relies on that pool returning unused memory at the next synchronization. DeviceMemoryResource(device) without options raises the default pool's release_threshold to ULLONG_MAX (MP_raise_release_threshold), which changes this for every user of the pool in the process. There is no way to set it back through cuda.core today.
Design
Follows GraphMemoryResourceAttributes, which already has zero-only setters for reserved_mem_high and used_mem_high, with the same AttributeError message.
trim follows GraphMemoryResource.trim(). The optional min_bytes_to_keep maps to cuMemPoolTrimTo's argument, and 0 keeps the meaning of GraphMemoryResource.trim().
release_threshold rejects negative values with ValueError.
The class docstring notes that setting an attribute changes the pool itself, which other resources sharing it (such as a device's default pool) also see.
Setters and trim call the driver like the existing getters, without the capture-mode exchange that MP_raise_release_threshold uses for its implicit write. They are explicit user calls.
Tests
Parametrized over the three pool-backed resources:
Every settable attribute round-trips. Zero-only and read-only attributes reject other writes.
used_mem_high = 0 resets the watermark to the current usage after an allocation is freed.
With release_threshold at the maximum, freed memory stays reserved, trim(min_bytes_to_keep=reserved) keeps it, trim() returns it, and trim() on a closed resource raises.
test_memory.py passes locally (RTX 5070 Ti, driver 615, CUDA 13.4).
release_threshold and the reuse_* policies of a pool-backed memory
resource can now be set, and reserved_mem_high and used_mem_high can be
reset to zero, matching GraphMemoryResourceAttributes. trim() wraps
cuMemPoolTrimTo and takes an optional min_bytes_to_keep.
Memory-pool reuse policies and release thresholds can now be updated. Release thresholds must be nonnegative.
High-watermark settings can be reset to their current values by setting them to zero.
Trim unused pool memory while retaining an optional minimum amount. Without a minimum, trimming releases all unused pool memory.
Reading or updating pool settings and trimming no longer fail or invalidate an in-progress CUDA graph capture in supported capture modes.
The change adds setters for selected memory-pool attributes and adds _MemPool.trim(min_bytes_to_keep=0). Pool attribute operations and trimming preserve the prior capture mode. Tests cover attribute validation, watermark resets, trimming, and operations during capture.
Changes: Memory pool API
Layer / File(s)
Summary
Attribute setters and capture handling cuda_core/cuda/core/_memory/_memory_pool.pyi, cuda_core/cuda/core/_memory/_memory_pool.pyx, cuda_core/tests/test_memory.py, cuda_core/tests/graph/test_graph_memory_resource.py, cuda_core/docs/source/release/1.3.0-notes.rst
Selected reuse attributes, release_threshold, and high-watermark attributes gain setters. Negative thresholds and nonzero high-watermark assignments are rejected. Attribute reads and writes run in relaxed capture mode and restore the prior mode. Tests cover setter behavior, watermark resets, and capture behavior.
Unused pool memory trimming cuda_core/cuda/core/_memory/_memory_pool.pyi, cuda_core/cuda/core/_memory/_memory_pool.pyx, cuda_core/tests/test_memory.py, cuda_core/docs/source/release/1.3.0-notes.rst
_MemPool.trim(min_bytes_to_keep=0) checks that the resource is open and calls cuMemPoolTrimTo in relaxed capture mode. Tests cover retained memory, full trimming, closed resources, and calls during capture.
Possibly related issues
NVIDIA/cuda-python#2362: This change adds setters for existing pool attributes, but does not add the additional attributes or location-based routing requested by the issue.
Possibly related PRs
NVIDIA/cuda-python#3034: This change edits the release-notes page that the PR replaces with generated notes.
Suggested labels:feature
Suggested reviewers:andy-jost
Priority: Normal
Change: Feature
Merge Risk: Low · up to e0e35
The change is mergeable with bounded follow-up: make the capture test clean up after a failure, run dedicated managed-pool tests on supported devices, and clarify the documented setter limits.
Finishing Touches: 💡 1
🛠️ Fix failing CI checks: 💡
Commit to this branch
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
Reviewing files that changed from the base of the PR and between 590a8bf and c7fe116.
Files selected for processing: 4
cuda_core/cuda/core/_memory/_memory_pool.pyi
cuda_core/cuda/core/_memory/_memory_pool.pyx
cuda_core/docs/source/release/1.3.0-notes.rst
cuda_core/tests/test_memory.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 7 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Minor · Functional Correctness · Quick win
suggestion: State the accepted attribute values.
“Take any value” conflicts with the ValueError for a negative release_threshold and the boolean reuse setters. State those limits instead. The CUDA attribute contract also defines the threshold as an unsigned 64-bit value. (docs.nvidia.com)
The reason will be displayed to describe this comment to others. Learn more.
Minor · Maintainability & Code Quality · Quick win
suggestion: Do not skip supported dedicated managed pools.
If a device supports dedicated managed pools but lacks concurrent managed access, skip_if_managed_memory_unsupported skips these tests before MRops() creates the dedicated pool. The supplied create_managed_memory_resource_or_skip helper permits that pool with explicit options. Check the dedicated-pool setup instead, so these tests run on supported devices. As per coding guidelines, “Skip only real setup failures.”
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.
suggestion: Make cleanup independent of gb.end_building(). If a pool call invalidates the capture, end_building() raises. Then gb.close() and mr.close() never run. The owned pool leaks. The next init_cuda test then fails with UnreleasedMemoryResourceWarning, which hides the original failure. Nest the closes in their own try/finally.
Proposed fix
- finally:- # Raises if a pool call invalidated the capture.- gb.end_building()- gb.close()- mr.close()+ finally:+ try:+ # Raises if a pool call invalidated the capture.+ gb.end_building()+ finally:+ gb.close()+ mr.close()
As per coding guidelines: "Keep pool lifetimes inside the test that creates them" and "Release what you construct regardless."
Files skipped from review as they are similar to previous changes: 1
cuda_core/docs/source/release/1.3.0-notes.rst
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 6 remain after this review.
This branch has not been deployed
No deployments
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
cuda.coreEverything related to the cuda.core module
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
github-actions
Bot
added
the
cuda.core
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Summary
For
DeviceMemoryResource,PinnedMemoryResourceandManagedMemoryResource:mr.attributes.release_thresholdand the threereuse_*policies can be set.mr.attributes.reserved_mem_highandused_mem_highcan be set to0to reset the watermark. Any other value raisesAttributeError.mr.trim(min_bytes_to_keep=0)wrapscuMemPoolTrimTo.Motivation
attributesalready exposes the right fields, but all of them are read-only, and there is no way to trim a pool.used_mem_highand its documented reset (setting it to 0). Stream-ordered allocations do not show up in CuPy's memory hooks, so the pool's watermark is the only source.build_arena(_memory.py) allocates temporaries withcudaMallocAsyncfrom the device's default pool, and relies on that pool returning unused memory at the next synchronization.DeviceMemoryResource(device)without options raises the default pool'srelease_thresholdtoULLONG_MAX(MP_raise_release_threshold), which changes this for every user of the pool in the process. There is no way to set it back throughcuda.coretoday.Design
GraphMemoryResourceAttributes, which already has zero-only setters forreserved_mem_highandused_mem_high, with the sameAttributeErrormessage.trimfollowsGraphMemoryResource.trim(). The optionalmin_bytes_to_keepmaps tocuMemPoolTrimTo's argument, and0keeps the meaning ofGraphMemoryResource.trim().release_thresholdrejects negative values withValueError.trimcall the driver like the existing getters, without the capture-mode exchange thatMP_raise_release_thresholduses for its implicit write. They are explicit user calls.Tests
Parametrized over the three pool-backed resources:
used_mem_high = 0resets the watermark to the current usage after an allocation is freed.release_thresholdat the maximum, freed memory stays reserved,trim(min_bytes_to_keep=reserved)keeps it,trim()returns it, andtrim()on a closed resource raises.test_memory.pypasses locally (RTX 5070 Ti, driver 615, CUDA 13.4).