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.
Pool attribute reads and writes and trim run in relaxed capture mode, through helpers shared with MP_raise_release_threshold. 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. This also fixes the existing getters.
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).
It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.
Reviewing files that changed from the base of the PR and between d9b1aa4 and f7ad557.
Files selected for processing: 1
cuda_core/tests/test_memory.py
Files skipped from review as they are similar to previous changes: 1
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; 11 remain after this review.
📝 Summary
Summary by CodeRabbit
New Features
Memory-pool reuse policies and release thresholds can now be updated. Release thresholds must be nonnegative and take effect at synchronization.
High-watermark settings can now be reset to zero; assigning a nonzero value is not supported.
Trim unused pool memory while retaining an optional minimum amount. By default, trimming releases all eligible unused memory.
Bug Fixes
Reading or updating pool settings and trimming no longer invalidate CUDA graph capture in supported capture modes. Trimming does not release memory backing live allocations or frees not yet observed through synchronization.
The change adds setters for selected memory-pool attributes and adds _MemPool.trim(min_bytes_to_keep=0). Pool attribute operations and trimming use relaxed capture mode and restore the prior mode. Tests cover validation, watermark resets, trimming, and calls during capture.
Selected reuse attributes, release_threshold, and high-watermark attributes gain setters. Nonzero high-watermark assignments are rejected, and threshold inputs are validated. Attribute reads and writes use relaxed capture mode and restore the prior mode. Tests cover setters 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 retaining memory, full trimming, closed resources, and calls during capture.
Assessment against linked issues
Objective
Addressed
Explanation
Make release_threshold and the three reuse_* policies settable; allow reserved_mem_high and used_mem_high to reset to zero. [#3261]
✅
Add trim(min_bytes_to_keep=0) for pool-backed memory resources. [#3261]
✅
Run pool attribute calls and trim() in relaxed capture mode. [#3261]
✅
Suggested reviewers:andy-jost
Priority: Normal
Change: Feature
Merge Risk: Low · up to f7ad5
Mergeable with owner awareness: some supported devices miss dedicated-pool test coverage, and unusual threshold inputs are handled inconsistently. Neither issue appears to block ordinary pool use.
- [ ] 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 · 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.”
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.
The reason will be displayed to describe this comment to others. Learn more.
Thanks, this is a solid addition and the driver side checks out: cuMemPoolSetAttribute accepts exactly the attributes you make settable with the value types you pass, the zero-only reset of the two high-water marks matches what the driver enforces, cuMemPoolTrimTo gets min_bytes_to_keep as documented, every status goes through HANDLE_RETURN, and the relaxed-capture-mode wrapper is right for all three wrappers. The stubs are the hook's output, the tests carry one marker each and cap their pools, and CI is green on the second attempt.
Three things before merge:
cuda_core/docs/source/api_private.rst: _MemPoolAttributes is not listed, so the setter semantics you documented (bytes for release_threshold, the zero-only reset, the note that a write changes the pool for every resource that shares it, the capture note) never render; the resource pages show only "Memory pool attributes." Add _memory._memory_pool._MemPoolAttributes under the cyclass template there, per the API guideline for classes returned but not instantiable, and build the docs once to see them.
_memory_pool.pyx, the release_threshold setter: value: int makes the pinned Cython apply an exact-type check, so numpy.uint64, bool, and int subclasses raise TypeError, while trim(), allocate(), and the rest of the surface accept anything with __index__. Type it as cydriver.cuuint64_t value; that accepts __index__ types and turns a negative or oversized value into OverflowError with Cython's own message, so the explicit < 0 check can go.
The 1.3.0 release note: "take any value" is not what the code does (a non-negative int; a bool), the last sentence is a fix to existing behavior and belongs under "Fixes and enhancements", and the release notes do not spell out driver error codes; say that the getters no longer fail or invalidate a capture.
Smaller points, no blocker: test_mempool_attributes_set only checks that reserved_mem_high = 0 does not raise; assert reserved_mem_high == reserved_mem_current afterwards, as the used_mem_high test does, and the [thread_local, other_thread] cell of the capture test does not exercise the fix. On a mapped IPC resource the setters and trim() fail inside the driver instead of with the clear TypeError that allocate() gives; a sentence in the docstrings would do. One design point to accept consciously: the setters and trim() now run immediately during an active capture on any thread and mutate a process-wide pool, where before the driver refused and invalidated the capture; the class docstring says so, and I agree with it. Last, the first CI attempt's Windows exit-139 crash happened inside the new test_mempool_trim[PinnedMR] under four parallel threads; it passed on the rerun and on the other Windows rows, so I take it as the known Windows crash pattern, but it is worth watching after merge.
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.
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.
Thanks, this is a solid addition and the driver side checks out: cuMemPoolSetAttribute accepts exactly the attributes you make settable with the value types you pass, the zero-only reset of the two high-water marks matches what the driver enforces, cuMemPoolTrimTo gets min_bytes_to_keep as documented, every status goes through HANDLE_RETURN, and the relaxed-capture-mode wrapper is right for all three wrappers. The stubs are the hook's output, the tests carry one marker each and cap their pools, and CI is green on the second attempt.
Three things before merge:
1. `cuda_core/docs/source/api_private.rst`: `_MemPoolAttributes` is not listed, so the setter semantics you documented (bytes for `release_threshold`, the zero-only reset, the note that a write changes the pool for every resource that shares it, the capture note) never render; the resource pages show only "Memory pool attributes." Add `_memory._memory_pool._MemPoolAttributes` under the `cyclass` template there, per the API guideline for classes returned but not instantiable, and build the docs once to see them.
2. `_memory_pool.pyx`, the `release_threshold` setter: `value: int` makes the pinned Cython apply an exact-type check, so `numpy.uint64`, `bool`, and `int` subclasses raise `TypeError`, while `trim()`, `allocate()`, and the rest of the surface accept anything with `__index__`. Type it as `cydriver.cuuint64_t value`; that accepts `__index__` types and turns a negative or oversized value into `OverflowError` with Cython's own message, so the explicit `< 0` check can go.
3. The 1.3.0 release note: "take any value" is not what the code does (a non-negative int; a bool), the last sentence is a fix to existing behavior and belongs under "Fixes and enhancements", and the release notes do not spell out driver error codes; say that the getters no longer fail or invalidate a capture.
Smaller points, no blocker: test_mempool_attributes_set only checks that reserved_mem_high = 0 does not raise; assert reserved_mem_high == reserved_mem_current afterwards, as the used_mem_high test does, and the [thread_local, other_thread] cell of the capture test does not exercise the fix. On a mapped IPC resource the setters and trim() fail inside the driver instead of with the clear TypeError that allocate() gives; a sentence in the docstrings would do. One design point to accept consciously: the setters and trim() now run immediately during an active capture on any thread and mutate a process-wide pool, where before the driver refused and invalidated the capture; the class docstring says so, and I agree with it. Last, the first CI attempt's Windows exit-139 crash happened inside the new test_mempool_trim[PinnedMR] under four parallel threads; it passed on the rerun and on the other Windows rows, so I take it as the known Windows crash pattern, but it is worth watching after merge.
Thanks @Andy-Jost, addressed in the latest commits:
Added _MemPoolAttributes to api_private.rst and linked it from the attributes docstring.
Typed the setter as uint64_t and dropped the < 0 check. I used uint64_t instead of cydriver.cuuint64_t so the generated stub stays value: int. Added a test for bool, an int subclass, and numpy.uint64.
Rewrote the release note and moved the getter fix to "Fixes and enhancements".
The reset test now covers reserved_mem_high, and I dropped the [thread_local, other_thread] capture case. The getters also fail on a mapped IPC pool, so the docstring note covers reads as well.
Reviewing files that changed from the base of the PR and between dc2a46c and d9b1aa4.
Files selected for processing: 5
cuda_core/cuda/core/_memory/_memory_pool.pyi
cuda_core/cuda/core/_memory/_memory_pool.pyx
cuda_core/docs/source/api_private.rst
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; 11 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Thanks, every point is addressed, the optional ones too: _MemPoolAttributes is in api_private.rst and linked from the attributes docstring, the setter takes uint64_t and so accepts NumPy integers and bool and raises OverflowError for negative and oversized values, the release note states the accepted values and moved the getter fix to "Fixes and enhancements", the watermark-reset test checks both high == current, and the inert capture cell is gone. CI is now running. Approving.
On the two open CodeRabbit comments, I'd decline both: the setter now converts integers exactly like trim() and allocate() do, so a stricter operator.index rule belongs to the whole surface rather than to this one setter; and the mapped-pool note is right, since an imported pool has no block pool and the driver rejects every attribute this class exposes, as well as cuMemPoolTrimTo, on it. One naming nit: Cython converts through __int__, not __index__, so test_mempool_release_threshold_accepts_int_like would describe the test better than its current name.
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 modulefeatureNew feature or requestP2Low priority - Nice to have
2 participants
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}
Andy-Jost
left a comment
Closes #3261
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.MP_raise_release_threshold. 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. This also fixes the existing getters.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).