Repository navigation
Rework memory reclamation API; reduce memory usage - #3118
Merged
Merged
Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #3118 +/- ##
==========================================
- Coverage 16.39% 10.17% -6.22%
==========================================
Files 123 122 -1
Lines 9686 9304 -382
==========================================
- Hits 1588 947 -641
- Misses 8098 8357 +259 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
The previous 2 GiB/worker estimate ignored CUDA context and library overhead plus peak allocations inside individual tests, letting the cap float up to 64+ workers on large systems. A 4 GiB budget accounts for ~1 GiB baseline (context + loaded libraries) and leaves room for multi-GiB peak test allocations. Also reset the device before querying free memory so the budget reflects true capacity, and `@info`-log the computed budget (device, free memory, gpu_jobs, cpu_threads, cpu_free) so users can see why the cap landed where it did. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two related changes to reclaim GPU memory that was effectively leaking:
1. Refactor cuSOLVER state to use an object-bound finalizer
The fat cusolverDnHandle (and analogous sparse/mg wrappers) held a
`workspace_gpu::CuVector{UInt8}` that can grow to hundreds of MiB. Its
finalizer was attached to `current_task()`, which on test workers never
dies — so `CUDA.reclaim()` could empty idle handle caches but couldn't
touch the pinned workspace. Probing showed ~200 MiB of pool memory per
cuSOLVER-using worker sitting there permanently.
Wrap the raw handles in `mutable struct` and attach the finalizer to
the wrapper instead. Add a `pre_reclaim_hooks` mechanism in CUDACore
so libraries can drop their TLS state on reclaim; `reclaim()` then
runs these hooks + GC + regular hooks, so finalizers fire and the
underlying buffers are released before the pool is trimmed. cuSOLVER
registers a hook that clears its dense/sparse/mg TLS entries.
2. Cap the test-worker memory pool release threshold at 64 MiB
The default `MEMPOOL_ATTR_RELEASE_THRESHOLD` is `typemax(UInt64)` —
freed stream-ordered buffers stay cached indefinitely, inflating NVML
readings. Set a 64 MiB cap on test workers so released buffers go back
to the driver quickly. Trades a per-call pool-alloc cost for peak GPU
RSS tracking the actual live working set.
Impact on library tests: big wins where workspace was hoarded (e.g.
cutensornet/contractions 1508 → 642 MiB), small/neutral elsewhere.
19813 library tests pass cleanly with no errors or warnings.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Same pattern as the prior cuSOLVER commit: wrap the raw handle in a mutable struct with an object-bound finalizer, and register a `pre_reclaim_hooks` entry that clears the library's TLS state. When CUDA.reclaim() runs, the fat-handle wrapper becomes unreferenced, GC collects it, its finalizer returns the raw handle to the idle cache, and the subsequent HandleCache hook destroys the idle handle properly. Memory savings for cuBLAS/cuSPARSE are smaller than cuSOLVER (no cached workspace CuVector; just the ~handle), but this removes the last place library state stays pinned across the worker's lifetime, and makes the lifecycle consistent across libraries. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replace the two untyped hook arrays (pre_reclaim_hooks, reclaim_hooks) and the three copies of the phase ladder with a single Reclaimable abstract type, a ReclaimLevel enum, and one reclaim(level) entry point. Libraries register their HandleCache and TaskLocalCache instances from __init__ and implement drop!/purge! via dispatch; CUDACore invokes them with per-hook error isolation, modelled on Base.atexit_hooks. retry_reclaim walks ReclaimLevel one rung at a time. Convert cuDNN, cuRAND, cuTENSOR, cuTensorNet, and cuStateVec from the task-bound finalizer pattern (finalizer(current_task(), …)) to the object-bound fat-handle pattern used by cuBLAS/cuSPARSE/cuSOLVER. The old task-bound variant only released resources once the owning task was garbage-collected, so those libraries previously leaked memory under pressure; they now participate in RECLAIM_DROP_STATE. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Inspired by PyTorch's `FreeMemoryCallback` (see
c10/cuda/CUDACachingAllocator.cpp), which invokes registered callbacks
before any expensive allocation work. `RECLAIM_PURGE_IDLE` sits at the
bottom of the ladder and runs `purge!` on every `Reclaimable` without
touching sync, GC, or pool trim. If a library happens to be holding
idle cached handles from dead tasks, we get them back immediately; if
not, it's a cheap no-op and we escalate.
The existing `RECLAIM_PURGE_CACHES` rung stays in place at its former
position: after `RECLAIM_GC_{MINOR,FULL}` have run, handle wrappers
from dead tasks get finalized and their raw handles end up in the
caches. That rung now catches those new entries before we'd have to
escalate all the way to `RECLAIM_DROP_STATE` (which also clears the
current task's TLS).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Collapse 8 ReclaimLevel rungs into 4 (PURGE/SYNC/GC/DROP); each level is a complete reclaim at that aggressiveness, so reclaim(level) runs the matching step directly instead of walking the whole ladder. - foreach_reclaimable now snapshots the registry under the lock and runs callbacks outside it (Base.atexit-style), so a slow or re-entrant callback can't pin other threads. - Drop the closure allocation in task_dict. - Add a note in the ReclaimLevel docstring on why DROP is safe (descriptors/plans are context-bound, user-held wrappers protect their raw handles) and that callbacks must not switch the active device. - Guard the test pool-cap setup with stream_ordered(dev). - Fix the stale "4 GiB" worker-budget comment. - Expand the reclaim testset to cover every level and verify retry_reclaim's escalation behavior. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Rework GPU memory reclamation so libraries can release resources under memory pressure instead of pinning them for the lifetime of the owning task. Probing showed CUDA.jl test workers were sitting on ~200 MiB of cuSOLVER workspace each, plus smaller per-library overhead, even after
CUDA.reclaim(), because finalizers were attached tocurrent_task()and the worker tasks never died.What changes
Object-bound finalizers. Library handles (cuBLAS, cuDNN, cuRAND, cuSOLVER, cuSPARSE, cuStateVec, cuTENSOR, cuTensorNet) move from
finalizer(current_task(), …)to a mutableHandlewrapper with its own finalizer. When TLS state is cleared, the wrapper becomes GC-eligible and its finalizer returns the raw handle to the idle cache.Typed reclaim registry. Replace the two untyped
pre_reclaim_hooks/reclaim_hooksglobals with aReclaimableabstract type, dispatching ondrop!(release live state) andpurge!(destroy idle resources).HandleCacheand a newTaskLocalCache{K,V}areReclaimable; libraries register their instances from__init__.New
reclaimAPI.reclaim(level::ReclaimLevel)with a four-rung ladder, from cheapest to most aggressive:RECLAIM_PURGEHandleCaches (no GC, no sync)RECLAIM_SYNCRECLAIM_GCRECLAIM_DROPretry_reclaimwalks the ladder one rung at a time on alloc failure. Breaking: oldreclaim(sz::Int)returning bytes is gone —reclaim(level)returnsnothing, default isRECLAIM_DROP.Test infra. Tighter per-worker budget (1 GiB, was 2 GiB) so we cap parallelism at a realistic baseline + peak. Cap the test pool's
MEMPOOL_ATTR_RELEASE_THRESHOLDto 256 MiB so freed buffers go back to the driver and NVML readings track actual live working set rather than running max.Impact
Tests where workspace was hoarded see large drops (cuTensorNet contractions: 1508 → 642 MiB peak); other libraries see smaller wins. All library tests pass cleanly.