Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions cuda_core/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -290,6 +290,9 @@ def ipc_memory_resource(request, ipc_device):
assert mr.is_ipc_enabled
yield mr
mr.close()
# TODO(seberg): Make sure the `mr` and it's buffers are fully torn down.
# May be unnecessary as `mr.close()` is not parallel with other work.
ipc_device.sync()


@pytest.fixture
Expand Down
4 changes: 4 additions & 0 deletions cuda_core/tests/memory_ipc/test_peer_access.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,8 @@ def test_main(self, ipc_mempool_device_x2, grant_access_in_parent):
assert process.exitcode == 0

buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
dev1.sync()
mr.close()

def child_main(self, mr, buffer):
Expand Down Expand Up @@ -129,4 +131,6 @@ def child_main(self, mr, buffer):
PatternGen(dev0, NBYTES).verify_buffer(buffer, seed=False)

buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
dev1.sync()
mr.close()
5 changes: 5 additions & 0 deletions cuda_core/tests/memory_ipc/test_send_buffers.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ def test_main(self, ipc_device, nmrs):
device = ipc_device
options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
mrs = [DeviceMemoryResource(device, options=options) for _ in range(nmrs)]
buffers = []

try:
# Allocate and fill memory.
Expand All @@ -54,6 +55,10 @@ def test_main(self, ipc_device, nmrs):
pgen.verify_buffer(buffer, seed=True)
buffer.close()
finally:
for buffer in buffers:
buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
device.sync()
for mr in mrs:
mr.close()

Expand Down
18 changes: 15 additions & 3 deletions cuda_core/tests/memory_ipc/test_workerpool.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ def test_main(self, ipc_device, nmrs):
device = ipc_device
options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
mrs = [DeviceMemoryResource(device, options=options) for _ in range(nmrs)]
buffers = []

try:
buffers = [mr.allocate(NBYTES, stream=device.default_stream) for mr, _ in zip(cycle(mrs), range(NTASKS))]
Expand All @@ -45,8 +46,11 @@ def test_main(self, ipc_device, nmrs):
pgen = PatternGen(device, NBYTES)
for buffer in buffers:
pgen.verify_buffer(buffer, seed=True)
buffer.close()
finally:
for buffer in buffers:
buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
device.sync()
for mr in mrs:
mr.close()

Expand Down Expand Up @@ -77,6 +81,7 @@ def test_main(self, ipc_device, nmrs):
device = ipc_device
options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
mrs = [DeviceMemoryResource(device, options=options) for _ in range(nmrs)]
buffers = []

try:
buffers = [mr.allocate(NBYTES, stream=device.default_stream) for mr, _ in zip(cycle(mrs), range(NTASKS))]
Expand All @@ -90,8 +95,11 @@ def test_main(self, ipc_device, nmrs):
pgen = PatternGen(device, NBYTES)
for buffer in buffers:
pgen.verify_buffer(buffer, seed=True)
buffer.close()
finally:
for buffer in buffers:
buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
device.sync()
for mr in mrs:
mr.close()

Expand Down Expand Up @@ -127,6 +135,7 @@ def test_main(self, ipc_device, nmrs):
device = ipc_device
options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
mrs = [DeviceMemoryResource(device, options=options) for _ in range(nmrs)]
buffers = []

try:
buffers = [mr.allocate(NBYTES, stream=device.default_stream) for mr, _ in zip(cycle(mrs), range(NTASKS))]
Expand All @@ -137,8 +146,11 @@ def test_main(self, ipc_device, nmrs):
pgen = PatternGen(device, NBYTES)
for buffer in buffers:
pgen.verify_buffer(buffer, seed=True)
buffer.close()
finally:
for buffer in buffers:
buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
device.sync()
for mr in mrs:
mr.close()

Expand Down
10 changes: 10 additions & 0 deletions cuda_core/tests/test_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -1103,6 +1103,8 @@ def test_device_memory_resource_with_options(init_cuda):
device.sync()
dst_buffer.close()
src_buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
device.sync()


def test_pinned_memory_resource_with_options(init_cuda):
Expand Down Expand Up @@ -1149,6 +1151,8 @@ def test_pinned_memory_resource_with_options(init_cuda):
device.sync()
dst_buffer.close()
src_buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
device.sync()


def test_managed_memory_resource_with_options(init_cuda):
Expand Down Expand Up @@ -1365,6 +1369,8 @@ def test_mempool_ipc_errors(mempool_device):
Buffer.from_ipc_descriptor(mr, handle, stream=device.default_stream)

buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
device.sync()


def test_pinned_mempool_ipc_basic():
Expand Down Expand Up @@ -1405,6 +1411,8 @@ def test_pinned_mempool_ipc_basic():
assert ipc_desc.size == 1024

buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
device.sync()
mr.close()


Expand Down Expand Up @@ -1436,6 +1444,8 @@ def test_pinned_mempool_ipc_errors():
Buffer.from_ipc_descriptor(mr, handle, stream=device.default_stream)

buffer.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
device.sync()
mr.close()


Expand Down
6 changes: 5 additions & 1 deletion cuda_core/tests/test_object_protocols.py
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,11 @@ def sample_ipc_buffer_descriptor(ipc_device):
options = DeviceMemoryResourceOptions(max_size=POOL_SIZE, ipc_enabled=True)
mr = DeviceMemoryResource(ipc_device, options=options)
buf = mr.allocate(64, stream=ipc_device.default_stream)
return buf.ipc_descriptor
descriptor = buf.ipc_descriptor
buf.close()
# TODO(seberg): 2026-06: mr close may be unsafe with incomplete `buf.close()`
ipc_device.sync()
return descriptor


@pytest.fixture
Expand Down
Loading