Skip to content

Commit 4490ecb

Browse files
Andy-Jostclaude
Andy-Jost
and
claude
authored
cuda.core: fix the peer-access cleanup before pool destroy (#2871)
* cuda.core: skip the pool's own device when clearing peer access at destroy clear_mempool_peer_access() named every device, including the pool's own, in one cuMemPoolSetAccess call. The driver validates the whole map before applying any entry and rejects a non-read-write entry for the owning device, so the call failed every time and the nvbug 5698116 workaround never took effect. Leave the owning device out and skip the call when no other device exists. Host pools and imported pools pass -1 and keep the previous behavior. Adds a two-GPU test that grants peer access, closes the pool, and checks that a pool created on the recycled handle does not report it. Issue #2807 Co-Authored-By: Claude Fable 5.1 <[email protected]> * cuda.core: revoke pool peer access one device at a time The driver validates a cuMemPoolSetAccess request as a whole, so one request that lists every device is rejected if any listed device cannot be changed: the owning device, a device without memory-map support, or devices of more than one kind. Issue one request per non-owner device so that a rejection blocks only that device. This also revokes access for imported pools, whose owner is unknown and was listed. Issue #2807 Co-Authored-By: Claude Fable 5.1 <[email protected]> --------- Co-authored-by: Claude Fable 5.1 <[email protected]>
1 parent 381128d commit 4490ecb

3 files changed

Lines changed: 64 additions & 12 deletions

File tree

‎cuda_core/cuda/core/_cpp/rt/memory.cpp‎

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
#include <mutex>
1515
#include <stdexcept>
1616
#include <utility>
17-
#include <vector>
1817
#ifndef _WIN32
1918
#include <unistd.h>
2019
#endif
@@ -35,32 +34,42 @@ struct MemoryPoolBox {
3534

3635
// Helper to clear peer access before destroying a memory pool.
3736
// Works around nvbug 5698116: recycled pool handles inherit peer access state.
37+
// The driver validates a request as a whole and rejects it if any entry cannot
38+
// be applied: the pool's own device, a device without memory-map support, or
39+
// devices of more than one kind in one request. Each device is therefore
40+
// revoked with its own request, and the owning device (-1 if the pool has
41+
// none) is not requested at all.
3842
// Must be noexcept since it's called from a shared_ptr deleter.
39-
static void clear_mempool_peer_access(CUmemoryPool pool) noexcept {
43+
static void clear_mempool_peer_access(CUmemoryPool pool, int owner_device) noexcept {
4044
try {
4145
int device_count = 0;
4246
if (p_cuDeviceGetCount(&device_count) != CUDA_SUCCESS || device_count <= 0) {
4347
return;
4448
}
4549

46-
std::vector<CUmemAccessDesc> clear_access(device_count);
50+
CUmemAccessDesc revoke{};
51+
revoke.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
52+
revoke.flags = CU_MEM_ACCESS_FLAGS_PROT_NONE;
4753
for (int i = 0; i < device_count; ++i) {
48-
clear_access[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
49-
clear_access[i].location.id = i;
50-
clear_access[i].flags = CU_MEM_ACCESS_FLAGS_PROT_NONE;
54+
if (i == owner_device) {
55+
continue;
56+
}
57+
revoke.location.id = i;
58+
p_cuMemPoolSetAccess(pool, &revoke, 1); // Best effort
5159
}
52-
p_cuMemPoolSetAccess(pool, clear_access.data(), device_count); // Best effort
5360
} catch (...) {
5461
// Swallow exceptions - this is best-effort cleanup in destructor context
5562
}
5663
}
5764

58-
static MemoryPoolHandle wrap_mempool_owned(CUmemoryPool pool) {
65+
// owner_device is the ordinal of a device-located pool, or -1 when there is no
66+
// device to skip (host pools) or the location is unknown (imported pools).
67+
static MemoryPoolHandle wrap_mempool_owned(CUmemoryPool pool, int owner_device) {
5968
auto box = std::shared_ptr<const MemoryPoolBox>(
6069
new MemoryPoolBox{pool},
61-
[](const MemoryPoolBox* b) {
70+
[owner_device](const MemoryPoolBox* b) {
6271
GILReleaseGuard gil;
63-
clear_mempool_peer_access(b->resource);
72+
clear_mempool_peer_access(b->resource, owner_device);
6473
pw_cuMemPoolDestroy(b->resource);
6574
delete b;
6675
}
@@ -74,7 +83,8 @@ MemoryPoolHandle create_mempool_handle(const CUmemPoolProps& props) {
7483
if (CUDA_SUCCESS != (err = p_cuMemPoolCreate(&pool, &props))) {
7584
return {};
7685
}
77-
return wrap_mempool_owned(pool);
86+
int owner_device = props.location.type == CU_MEM_LOCATION_TYPE_DEVICE ? props.location.id : -1;
87+
return wrap_mempool_owned(pool, owner_device);
7888
}
7989

8090
MemoryPoolHandle create_mempool_handle_ref(CUmemoryPool pool) {
@@ -98,7 +108,7 @@ MemoryPoolHandle create_mempool_handle_ipc(int fd, CUmemAllocationHandleType han
98108
if (CUDA_SUCCESS != (err = p_cuMemPoolImportFromShareableHandle(&pool, handle_ptr, handle_type, 0))) {
99109
return {};
100110
}
101-
return wrap_mempool_owned(pool);
111+
return wrap_mempool_owned(pool, -1);
102112
}
103113

104114
// ============================================================================

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,3 +120,12 @@ Fixes and enhancements
120120

121121
- :attr:`Stream.device` and related queries on a stream whose context is not
122122
current now restore the caller's context even when the device query fails.
123+
124+
- Closing a :class:`DeviceMemoryResource` now revokes the peer access it
125+
granted before the driver destroys the pool. The revocation used to be one
126+
request that named every device, including the pool's own; the driver
127+
rejects such a request as a whole, so it never took effect and a pool
128+
created later could inherit that access through a recycled handle. Each
129+
device is now revoked with its own request, so no single device can block
130+
the others.
131+
(`#2807 <https://github.com/NVIDIA/cuda-python/issues/2807>`__)

‎cuda_core/tests/test_memory.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2481,6 +2481,39 @@ def test_dmr_mempool_get_access_peer(mempool_device_x2):
24812481
assert DMR_mempool_get_access(mr, peer.device_id) == ""
24822482

24832483

2484+
@pytest.mark.thread_unsafe(reason="depends on the driver handing back the pool handle that was just destroyed")
2485+
@pytest.mark.agent_authored(model="claude-fable-5-1")
2486+
def test_closed_pool_peer_access_not_inherited_by_recycled_handle(mempool_device_x2):
2487+
"""Peer access granted to a closed pool does not survive into a pool that reuses its handle."""
2488+
from cuda.core._memory._device_memory_resource import DMR_mempool_get_access
2489+
2490+
dev, peer = mempool_device_x2
2491+
options = DeviceMemoryResourceOptions(max_size=POOL_SIZE)
2492+
mr = DeviceMemoryResource(dev, options)
2493+
mr.peer_accessible_by = [peer]
2494+
assert DMR_mempool_get_access(mr, peer.device_id) == "rw"
2495+
old_handle = int(mr.handle)
2496+
mr.close()
2497+
2498+
# The driver usually hands the freed handle straight back. Keep the misses
2499+
# alive so that a retry cannot land on the same address twice.
2500+
pools = []
2501+
try:
2502+
for _ in range(8):
2503+
recycled = DeviceMemoryResource(dev, options)
2504+
pools.append(recycled)
2505+
if int(recycled.handle) == old_handle:
2506+
break
2507+
else:
2508+
pytest.skip("the driver did not recycle the destroyed pool handle")
2509+
# Unless the peer access is revoked before the pool is destroyed, the
2510+
# recycled handle still reports it (nvbug 5698116).
2511+
assert DMR_mempool_get_access(recycled, peer.device_id) == ""
2512+
finally:
2513+
for pool in pools:
2514+
pool.close()
2515+
2516+
24842517
def test_dmr_peer_accessible_by_setter_empty(mempool_device):
24852518
"""Assigning an empty peer-access set to a fresh owned pool is a no-op."""
24862519
# max_size caps VA to dodge Windows MCDM OOM

0 commit comments

Comments
 (0)