Repository navigation
Conversation
Buffers of a VirtualMemoryResource whose handle_type is "posix_fd" can now be exported and imported like pool-backed buffers. Buffer.ipc_descriptor exports every physical allocation that backs a VirtualMemoryBuffer as a file descriptor, Buffer.from_ipc_descriptor imports the descriptor with a VirtualMemoryResource of the receiving process, and a buffer or its descriptor can be sent through multiprocessing directly. A grown buffer imports as one contiguous range with the same byte layout. The import reuses MemAllocationHandle: import_mem_allocation_handle wraps cuMemImportFromShareableHandle in the same box and cuMemRelease deleter as create_mem_allocation_handle, with the importer's access descriptors. VirtualMemoryIPCBufferDescriptor carries one IPCAllocationHandle and one size per allocation and owns the file descriptors; multiprocessing duplicates them into the receiving process. VirtualMemoryResource pickles as (device, options), which lets a buffer pickle as (resource, descriptor). Fabric handles are not shared yet: no available test system reports fabric support or exposes an IMEX channel, so that path could not be verified. Part of NVIDIA#2980. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Andy-Jost
added
P0
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test 2231df7 |
The CI wheel build compiles with -Werror and -Wsign-compare, and the VMM import path compared a size_t count with the Py_ssize_t that len() returns. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
/ok to test 682db0f |
|
Under pytest-run-parallel the same test body runs in several threads. The alias test checks mapping state by address, and a freed address can be reused by another thread's allocation; the descriptor test checks a closed file descriptor number, which another thread can reuse. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
/ok to test 4ed2c03 |
Address the review of the VirtualMemoryResource IPC path. Import: check every imported allocation's location against the resource's before mapping (the same device, or the same host location type) and raise a ValueError that names both; reject a descriptor with a size but no allocations; attach notes to driver failures that say which chunk failed and why a peer-controlled size fails at the map. The imported buffer no longer keeps the descriptor: the driver copies the shared allocation into a handle of its own and keeps no reference to the file descriptor, so the descriptors live only as long as the caller holds the descriptor object. Export: an imported buffer, and a buffer grown from one, cannot be exported again because the driver exports only allocations created with the requested handle type; say so instead of surfacing INVALID_VALUE. modify_allocation() rejects a config whose handle_type differs from the resource's, so every chunk of a buffer is exportable the same way. __reduce__ preserves subclasses. Tests: parametrize the memory_ipc suite over VirtualMemoryResource (VirtualMR), with the pool-only scenarios guarded or skipped; PatternGen works on the first `size` bytes of a larger buffer; the fd-leak harness asserts on the success path; new VMM-only tests for many-chunk ranges, a second device with and without peers, descriptor validation and rollback, re-export, empty buffers, and a DLPack consumer. Docs: file descriptor and memory lifetime, multiprocessing-only transport, trust boundary, device rule, re-export limit; VirtualMemoryIPCBufferDescriptor in api_private; the from_ipc_descriptor stub has its mr type again. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
/ok to test ef9b712 |
…M IPC Apply the decisions on the VirtualMemoryResource IPC design. No exporter-side cache. VirtualMemoryBuffer.ipc_descriptor exports on every access and returns a new descriptor that owns its file descriptors; they close when the descriptor is released. Pickling a buffer for multiprocessing builds a transient descriptor, so a sender holds no file descriptors between sends. A spawned Process is the one exception: it receives the file descriptors when it is created, after pickling, so the transient descriptor stays on the spawning Popen until the Process object is released. A descriptor pins the physical memory while it exists, and nothing else does. No re-export of imported memory. Every physical allocation records whether cuMemImportFromShareableHandle produced it (mem_allocation_is_imported), and modify_allocation reuses the allocation handles of its input, so the mark reaches every alias, in-place grow, and move. Exporting such a buffer, directly or through pickling, raises a RuntimeError before any driver call that says the buffer contains memory imported from another process, that it cannot be exported again, and that the caller should forward the descriptor it imported from or copy into a buffer it owns. A Queue feeder thread reports the same error and drops the item. Raw-fd transport. VirtualMemoryIPCBufferDescriptor exposes handle_type, sizes, handles, and fds, and from_fds() rebuilds a descriptor in the receiving process from file descriptors received another way (the integers are duplicated; IPCAllocationHandle objects are shared). The descriptor's own pickling stays multiprocessing-only. No ipc_enabled option: handle_type is the switch and is_ipc_enabled the query, and the docs say so. Tests: file descriptor counts return to baseline on both sides, export under a lowered RLIMIT_NOFILE, the re-export errors for imported and derived buffers (direct, Pipe, Process arguments, Queue feeder thread), and raw-fd round trips through os.dup and a Unix socket with SCM_RIGHTS. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
/ok to test b804e18 |
…warding Write down in VMM_DESIGN.md the opt-in design that would let an imported VirtualMemoryResource buffer be forwarded (keep the received handles with the imported buffer and duplicate them on export), why the option would belong on the resource, and its cost, so that it does not have to be re-derived. Not implemented; nothing needs it yet. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
The module-level import of the POSIX-only resource module made pytest fail to collect tests/memory_ipc/test_vmm_ipc.py on Windows, so the Windows rows ran no cuda.core tests at all. Every test in the module is skipped on Windows by its fixture; only the import has to move. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
/ok to test |
…emory Take ownership of each exported or duplicated file descriptor before appending its handle, so a failed append cannot close the descriptor twice. Validate imported chunk sizes against the minimum granularity that cuMemMap requires rather than the resource's granularity option. Name both causes in the note on a failed import: a handle that is not a CUDA allocation of the expected type, or an owning device that this process cannot use. Report is_ipc_enabled False for host-located resources: host NUMA virtual memory has no verified allocation path, so it is not shared yet; the docs and the design note say so. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
/ok to test |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Summary
Priority: ➖ Normal Change: Feature Merge Risk: ⚪ Minimal · up to
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
ℹ️ Review info
⚙️ Run configuration
- Configuration used: Repository: NVIDIA/cuda-python/.coderabbit.yaml
- Review profile: CHILL
- Plan: Enterprise
- Run ID:
bfea5288-7073-4618-b8d0-848bfb4dc3ff
📒 Files selected for processing (26)
cuda_core/cuda/core/_cpp/rt/VMM_DESIGN.mdcuda_core/cuda/core/_cpp/rt/api.hppcuda_core/cuda/core/_cpp/rt/driver_api.hppcuda_core/cuda/core/_cpp/rt/virtual_memory.cppcuda_core/cuda/core/_memory/_buffer.pyicuda_core/cuda/core/_memory/_buffer.pyxcuda_core/cuda/core/_memory/_ipc.pxdcuda_core/cuda/core/_memory/_ipc.pyicuda_core/cuda/core/_memory/_ipc.pyxcuda_core/cuda/core/_memory/_virtual_memory_resource.pyicuda_core/cuda/core/_memory/_virtual_memory_resource.pyxcuda_core/cuda/core/_rt.pxdcuda_core/cuda/core/_rt.pyxcuda_core/docs/source/api_private.rstcuda_core/docs/source/release/1.3.0-notes.rstcuda_core/tests/conftest.pycuda_core/tests/helpers/buffers.pycuda_core/tests/memory_ipc/test_errors.pycuda_core/tests/memory_ipc/test_event_ipc.pycuda_core/tests/memory_ipc/test_leaks.pycuda_core/tests/memory_ipc/test_memory_ipc.pycuda_core/tests/memory_ipc/test_send_buffers.pycuda_core/tests/memory_ipc/test_serialize.pycuda_core/tests/memory_ipc/test_vmm_ipc.pycuda_core/tests/memory_ipc/test_workerpool.pycuda_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.
rparolin
left a comment
rparolin
left a comment
There was a problem hiding this comment.
lgtm. I'd consider breaking larger PRs like this into smaller more reviewable chunks in future though.
| if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemImportFromShareableHandle, &handle, os_handle, handle_type))) { | ||
| return {}; | ||
| } | ||
| return wrap_mem_allocation(handle, size, std::move(access), true); |
There was a problem hiding this comment.
Nit: I'd consider moving the construction of access inline with the wrap_mem_allocation call so its already an r-value and you don't need to explicitly move it. The only down side is you are moving the std::vector construction into the GILReleaseGuard scope. However you would be avoiding constructing the access std::vector in the error case.
There was a problem hiding this comment.
Constructing the vector inline at the call site would move it after cuMemImportFromShareableHandle meaning it would need rollback logic to handle std::bad_alloc. I think it's simpler to leave it alone.
juenglin
left a comment
juenglin
left a comment
There was a problem hiding this comment.
1. buf.ipc_descriptor.fds returns closed file descriptors
ipc_descriptor is a property, and on a VirtualMemoryBuffer every access builds a new descriptor that closes its file descriptors when it is released. In this line the temporary is released when the statement finishes, so fds holds numbers that are already closed:
fds = buf.ipc_descriptor.fdsThe next open() can reuse one of those numbers, and sending fds over a Unix socket then hands the receiver an unrelated file. The same line is safe for pool-backed buffers, which cache their descriptor. cuda_core/AGENTS.md also says a property should be cheap and idempotent; this one calls the driver and opens file descriptors. The base class defines it as a property, so it has to stay one. An explicit export method next to it and a warning in the docstring would make the lifetime visible.
2. A pickled buffer is imported with the sender's device numbering
VirtualMemoryResource.__reduce__ rebuilds the resource in the receiver as VirtualMemoryResource(sender_device_index, sender_options). Device indices are per-process, for example under a per-worker CUDA_VISIBLE_DEVICES. Two problems follow:
- The owning device's index may name a different GPU. The location check catches this, so
queue.put(buf)fails withValueErrorinstead of mapping the wrong device. peersindices are not checked. A peer index that names a different GPU in the receiver grants access to that GPU, and the intended peer only faults later.
Could the resource pickle devices by UUID, both the owner and the peers, and map them back in the receiver? At minimum, the docs should say that the pickle path uses the sender's options and device indices, and that Buffer.from_ipc_descriptor(my_resource, desc) is the way to use the receiver's own.
3. The Process workaround depends on private multiprocessing internals
VirtualMemoryBuffer.__reduce__ calls multiprocessing.context.get_spawning_popen(), which is not a public API, and stores the descriptor as an attribute on the Popen object. A future Python release can break that. Python 3.14 also changes the default start method on Linux to forkserver, so the tests should cover it. There is a lifetime cost as well: the parent keeps the memory pinned until the Process object is garbage-collected, and many programs keep a list of finished Process objects until they exit.
4. modify_allocation now rejects a different handle_type
_check_config now raises ValueError when config.handle_type differs from the resource's. Previously a grow could create its new chunk with another handle type, such as None; now that raises. The 1.3.0 behavior-change note still says only the location must be kept. This restriction belongs in that paragraph.
5. Fabric support would change published behavior
The main API is ready for fabric: ipc_descriptor, from_ipc_descriptor, handle_type, and is_ipc_enabled. A few places, though, treat "a shareable handle" as "a POSIX file descriptor". Adding fabric would change what they promise:
from_fds()takeshandle_type: it accepts"fabric"today and returns a descriptor full of file descriptors that claims to hold fabric handles. That call would have to start raising, so it's better to drop the parameter or allow only"posix_fd". A laterfrom_fabric_handles()would then be purely additive.handlestype: it is annotated astuple[IPCAllocationHandle, ...], andIPCAllocationHandlewraps a file descriptor. A 64-byte fabric handle doesn't fit, so the annotation would have to widen. Typing the property as broadly as its docstring now would avoid that.- General rules in the docs: "plain
pickleraises" and "a descriptor pins the memory while it exists" are stated for every handle type. Fabric handles are plain bytes, and I don't believe they hold a reference to the memory, though I haven't confirmed that. Scoping both rules toposix_fdnow keeps them from becoming false later. - The importer-device check: it compares the owning GPU's index with the resource's GPU index. Across nodes, those indices can't be compared. Either the docs should limit this rule to
posix_fd, or fabric will need its own rule.
| """The memory belongs to the exporter's device; the importer maps it for that device and may add peers.""" | ||
|
|
||
| @pytest.mark.agent_authored(model="claude-fable-5-1") | ||
| def test_main(self, vmm_ipc_device_x2): |
There was a problem hiding this comment.
Local testing has this test flaking about 25% of the time. Could use a followup or a flaky marker
$ for i in $(seq 1 12); do python -m pytest \
"tests/memory_ipc/test_vmm_ipc.py::TestImportOnSecondDevice::test_main" \
-q -p no:randomly -p no:flaky; done
pass=9 fail=3
There was a problem hiding this comment.
Two data points. On a two-GPU H100 node (driver 615.71.09) this test passed 12 of 12 times in fresh processes, and the memory_ipc directory passed in random and in file order. On a two-GPU H200 NVL node (driver 595.58.03) it fails every time at the peer_read assertion with zeros and garbage instead of the pattern, and on that node the pre-existing pool test test_peer_access.py::TestBufferPeerAccessAfterImport and the simpleP2P sample fail the same way, so peer reads on that node return wrong data regardless of VMM. Does test_peer_access.py pass reliably on your machine? If it flakes too, this is a peer-access property of the host rather than this test. Your GPU model, driver version, and the failing assertion would settle it.
| # The resource holds no driver object that other processes must share; | ||
| # the receiving process rebuilds an equivalent one from the device and | ||
| # the options. This is what lets a Buffer pickle as (resource, descriptor). | ||
| return type(self), (self.device.device_id if self.device is not None else 0, self.config) |
There was a problem hiding this comment.
Passing the whole config across the process boundary here, including the peers list, allows a receiving process to violate CUDA_VISIBLE_DEVICES. We should use a UUID or perhaps straight up refuse to pickle a config that names peers.
This repro requires a multi GPU machine, but demonstrates how this allows the receiving process to write into GPU memory that shouldn't be visible to it:
Details
import ctypes
import multiprocessing as mp
import os
import sys
import traceback
def uuids():
from cuda.core import Device
out = {}
for i in range(8):
try:
out[i] = str(Device(i).uuid)[:8]
except Exception:
break
return out
def child(q_in, q_out):
try:
from cuda.bindings import driver
from cuda.core import Device
Device(0).set_current()
seen = uuids()
buf = q_in.get() # unpickles as (resource, descriptor)
peers = tuple(buf.memory_resource.config.peers)
# What access did the rebuilt resource actually grant, per device?
flags = {}
for ordinal in seen:
loc = driver.CUmemLocation()
loc.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
loc.id = ordinal
err, f = driver.cuMemGetAccess(loc, driver.CUdeviceptr(int(buf.handle)))
flags[seen[ordinal]] = int(f) if err == driver.CUresult.CUDA_SUCCESS else f"<{err}>"
# Write from the child's device 1 -- whatever physical GPU that now is.
Device(1).set_current()
(err,) = driver.cuMemsetD32(driver.CUdeviceptr(int(buf.handle)), 0xDEADBEEF, 8)
driver.cuCtxSynchronize()
q_out.put(("ok", {
"child devices": seen,
"peers list from the pickle": peers,
"access flags by GPU uuid": flags,
"wrote from GPU": seen.get(1),
"write result": str(err),
}))
except BaseException:
q_out.put(("err", traceback.format_exc()))
def main():
from cuda.bindings import driver
from cuda.core import Device, VirtualMemoryResource, VirtualMemoryResourceOptions
print(f"parent CUDA_VISIBLE_DEVICES = {os.environ.get('CUDA_VISIBLE_DEVICES')!r}")
dev = Device(0)
dev.set_current()
seen = uuids()
print(f"parent devices = {seen}")
mr = VirtualMemoryResource(
dev, config=VirtualMemoryResourceOptions(handle_type="posix_fd", peers=[1])
)
buf = mr.allocate(1)
driver.cuMemsetD32(driver.CUdeviceptr(int(buf.handle)), 0x11111111, 8)
driver.cuCtxSynchronize()
def read():
host = (ctypes.c_uint32 * 4)()
driver.cuMemcpyDtoH(ctypes.addressof(host), driver.CUdeviceptr(int(buf.handle)), 16)
return [hex(x) for x in host]
owner, intended = seen.get(0), seen.get(1)
print(f"owner GPU = {owner}")
print(f"peers=[1] means GPU = {intended} <- the only GPU authorized")
print(f"memory before = {read()}")
os.environ["CUDA_VISIBLE_DEVICES"] = "0,2" # the child sees a different ordering
ctx = mp.get_context("spawn")
q_in, q_out = ctx.Queue(), ctx.Queue()
p = ctx.Process(target=child, args=(q_in, q_out))
p.start()
q_in.put(buf) # pickles the buffer -> (resource, descriptor)
try:
status, payload = q_out.get(timeout=120)
except Exception as exc:
status, payload = "timeout", repr(exc)
p.join(timeout=30)
if p.is_alive():
p.terminate()
p.join(timeout=10)
if status != "ok":
print(f"child {status}:\n{payload}")
return 1
print()
for key, value in payload.items():
print(f"child {key:26} = {value}")
after = read()
wrote = payload["wrote from GPU"]
print()
print(f"memory after = {after}")
print(f"GPU that wrote = {wrote}")
print(f"GPU the sender authorized = {intended}")
print(f"same GPU? = {wrote == intended}")
buf.close()
return 0
if __name__ == "__main__":
sys.exit(main())There was a problem hiding this comment.
Agreed that pickling device indices is wrong across processes with different CUDA_VISIBLE_DEVICES. The same applies to DeviceMemoryResource.__reduce__, which pickles the device index too, so I filed #3097 to fix both paths together (UUIDs, or no peers in the pickled config).
One note on the framing: the receiver can only name GPUs it can see, and it can import the same descriptor with any peers through Buffer.from_ipc_descriptor, so the sender cannot constrain the receiver's peers once it hands over the handle. The bug is that peers=[1] means a different GPU in the receiver, not that a hidden GPU becomes reachable.
… no test skips - from_fds() takes no handle_type; a file descriptor can only carry posix_fd - The descriptor, resource, and release-note text scope the pin, pickle, and same-device rules to POSIX file descriptors, and the fds docstrings state the descriptor's lifetime rule - The 1.3.0 notes record that a per-call modify_allocation config must keep the resource's handle_type - VMM_DESIGN.md records the fabric direction: a descriptor class of its own - TestManyChunkBuffer and TestReexport no longer skip when the driver refuses the decoy's address hint; the grow relocates in that case too Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
# Conflicts: # cuda_core/tests/memory_ipc/test_errors.py
… in the VMM IPC tests init_cuda no longer runs gc.collect() at every teardown (NVIDIA#3053), so objects left by earlier tests (child pipes, queues, descriptors) can be collected in the middle of a test that counts process-wide file descriptors, and the count then drops below its baseline. Collect before taking the baseline. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
| if (CUDA_SUCCESS != (err = DRIVER_CALL(cuMemImportFromShareableHandle, &handle, os_handle, handle_type))) { | ||
| return {}; | ||
| } | ||
| return wrap_mem_allocation(handle, size, std::move(access), true); |
There was a problem hiding this comment.
Constructing the vector inline at the call site would move it after cuMemImportFromShareableHandle meaning it would need rollback logic to handle std::bad_alloc. I think it's simpler to leave it alone.
| """The memory belongs to the exporter's device; the importer maps it for that device and may add peers.""" | ||
|
|
||
| @pytest.mark.agent_authored(model="claude-fable-5-1") | ||
| def test_main(self, vmm_ipc_device_x2): |
There was a problem hiding this comment.
Two data points. On a two-GPU H100 node (driver 615.71.09) this test passed 12 of 12 times in fresh processes, and the memory_ipc directory passed in random and in file order. On a two-GPU H200 NVL node (driver 595.58.03) it fails every time at the peer_read assertion with zeros and garbage instead of the pattern, and on that node the pre-existing pool test test_peer_access.py::TestBufferPeerAccessAfterImport and the simpleP2P sample fail the same way, so peer reads on that node return wrong data regardless of VMM. Does test_peer_access.py pass reliably on your machine? If it flakes too, this is a peer-access property of the host rather than this test. Your GPU model, driver version, and the failing assertion would settle it.
| # The resource holds no driver object that other processes must share; | ||
| # the receiving process rebuilds an equivalent one from the device and | ||
| # the options. This is what lets a Buffer pickle as (resource, descriptor). | ||
| return type(self), (self.device.device_id if self.device is not None else 0, self.config) |
There was a problem hiding this comment.
Agreed that pickling device indices is wrong across processes with different CUDA_VISIBLE_DEVICES. The same applies to DeviceMemoryResource.__reduce__, which pickles the device index too, so I filed #3097 to fix both paths together (UUIDs, or no peers in the pickled config).
One note on the framing: the receiver can only name GPUs it can see, and it can import the same descriptor with any peers through Buffer.from_ipc_descriptor, so the sender cannot constrain the receiver's peers once it hands over the handle. The bug is that peers=[1] means a different GPU in the receiver, not that a hidden GPU becomes reachable.
| ) | ||
| # The handle type is part of the resource's sharing contract, like the | ||
| # location: every chunk of a buffer must be exportable the same way. | ||
| if cfg.handle_type != self.config.handle_type: |
There was a problem hiding this comment.
4.
modify_allocationnow rejects a differenthandle_type
_check_confignow raisesValueErrorwhenconfig.handle_typediffers from the resource's. [...] The 1.3.0 behavior-change note still says only the location must be kept. This restriction belongs in that paragraph.
Added to the modify_allocation paragraph in the 1.3.0 notes in 85f1d6d.
| return self | ||
|
|
||
| @classmethod | ||
| def from_fds( |
There was a problem hiding this comment.
5. Fabric support would change published behavior
from_fds()takeshandle_type: it accepts"fabric"today and returns a descriptor full of file descriptors that claims to hold fabric handles. That call would have to start raising, so it's better to drop the parameter or allow only"posix_fd".
Done in 85f1d6d: from_fds takes no handle_type and builds a posix_fd descriptor by construction. A fabric handle is 64 bytes of plain data that pickles directly, so no from_fabric_handles is needed later either.
| return self._chunk_sizes | ||
|
|
||
| @property | ||
| def handles(self) -> tuple[IPCAllocationHandle, ...]: |
There was a problem hiding this comment.
handlestype: it is annotated astuple[IPCAllocationHandle, ...], andIPCAllocationHandlewraps a file descriptor. A 64-byte fabric handle doesn't fit, so the annotation would have to widen.
Keeping the precise type. Fabric handles differ in lifetime (no close, no pin), transport (plain pickle), and importer rule, so fabric sharing gets a descriptor class of its own, and this class keeps IPCAllocationHandles. VMM_DESIGN.md records that direction in 85f1d6d.
| :class:`IPCAllocationHandle` objects: the descriptor owns them and closes | ||
| them when it is released. | ||
|
|
||
| A descriptor pins the physical memory while it exists, in every process |
There was a problem hiding this comment.
General rules in the docs: "plain
pickleraises" and "a descriptor pins the memory while it exists" are stated for every handle type. Fabric handles are plain bytes [...] Scoping both rules toposix_fdnow keeps them from becoming false later.
Scoped in 85f1d6d: the descriptor docstring states that it holds POSIX file descriptors and that its rules describe them, and the resource docstring says posix_fd is the only handle type shared today.
| ): | ||
| raise ValueError( | ||
| f"chunk {i + 1} of the descriptor is memory on {_describe_location(imported_prop.location)}, " | ||
| f"but this resource maps memory on {_describe_location(prop.location)}; import the descriptor " |
There was a problem hiding this comment.
The importer-device check: it compares the owning GPU's index with the resource's GPU index. Across nodes, those indices can't be compared. Either the docs should limit this rule to
posix_fd, or fabric will need its own rule.
The same-device rule is now stated under the file-descriptor rules in the resource docstring (85f1d6d). Fabric gets its own importer rule with its own descriptor class; VMM_DESIGN.md records it.
|
Since we didn't do a formal design cycle on this feature, here are a few key points behind the design that might not be obvious from the diff. One source of complexity is a difference in how the POSIX fds used for IPC transport must be handled. A VMM buffers imported via IPC cannot be re-exported to a third process. The driver won't export an imported allocation, and the received fds that could stand in for it are closed after import for the same reason. In the future we could add an opt-in to Plain More broadly, fd-based transport only has to be good enough. POSIX fds are the older VMM handle type: they work on most drivers but consume system resources and are awkward to use. The next change on this front will add fabric handles, which are plain bytes that travel through ordinary The rules themselves are written up in |
This branch has not been deployed
4 participants
Summary
VirtualMemoryResourcebuffers can now be shared with other processes the way pool-backed buffers are.Buffer.ipc_descriptoron aVirtualMemoryBufferexports every physical allocation that backs the buffer as a POSIX file descriptor,Buffer.from_ipc_descriptorimports the descriptor with aVirtualMemoryResourceof the receiving process for the device that owns the memory, and a buffer or its descriptor can be sent throughmultiprocessingdirectly.VirtualMemoryResource.is_ipc_enabledisTruewhen the resource'shandle_typeis"posix_fd"on Linux and the memory lives on a device;handle_typeis the switch andis_ipc_enabledthe query. Host-located virtual memory is not shared yet. A grown buffer imports as one contiguous range with the same byte layout.Changes
_rt:import_mem_allocation_handle()wrapscuMemImportFromShareableHandlein the same box andcuMemReleasedeleter ascreate_mem_allocation_handle(), with the importer's access descriptors; the box records whether the allocation was imported (mem_allocation_is_imported()), and the mark travels with the allocation into every range that maps it.cuMemImportFromShareableHandlejoins the driver function table._ipc:VirtualMemoryIPCBufferDescriptor, anIPCBufferDescriptorsubclass that carries oneIPCAllocationHandleand one size per physical allocation, in address order. It exposeshandle_type,sizes,handles, andfds, andfrom_fds()rebuilds a descriptor in the receiving process from file descriptors received through a transport of your own (integers are duplicated, so the caller keeps ownership of its own;IPCAllocationHandleobjects are shared). The descriptor's own pickling staysmultiprocessing-only, which duplicates the file descriptors into the receiver; plainpickleraises.Buffer.from_ipc_descriptordispatches on the descriptor's class and rejects a mismatched resource withTypeError.VirtualMemoryResource:is_ipc_enabled; the import path, which imports every allocation with the resource's handle type, checks that each one lives on the resource's device, reserves the total, maps the allocations in order with the resource's access options for its device and peers, and records the deallocation stream asallocate()does;__reduce__, so a buffer pickles as (resource, descriptor).VirtualMemoryBuffer.ipc_descriptorexports on every access and returns a new descriptor; nothing is cached on the buffer. Amodify_allocation()config must keep the resource'shandle_type.ValueErrorbefore any driver call. Driver failures during export, import, or map carry a note that names the allocation and, for the map, says that a size which does not match the exported allocation fails there.VMM_DESIGN.md, docstrings,VirtualMemoryIPCBufferDescriptorin the private API reference, and a 1.3.0 release note.Rules
Queueor sent to aPoolpins the memory in the sender until the receiver has unpickled it. Pickling a buffer builds a transient descriptor, so a sender holds no file descriptors between sends, and the importer does not keep the descriptor it imported from. A buffer passed as aProcessargument is the one exception: a spawned child receives the file descriptors when it is created, after pickling, so that descriptor lives on thePopenuntil theProcessobject is released.peersoption. A resource for another device gets aValueErrorthat names both devices.ipc_descriptorand pickling on an imported buffer, or on any buffermodify_allocation()derived from one (alias, in-place grow, or move), raise aRuntimeErrorbefore any driver call that says so and points to the two alternatives: forward the descriptor it was imported from, or copy into a buffer you own. AQueuefeeder thread reports the same error and drops the item; the queue stays usable. This matches the contract PyTorch applies to received CUDA tensors.Tests
The existing
memory_ipcsuite runs againstVirtualMemoryResourcethrough a thirdipc_memory_resourceparameter (VirtualMR), with the scenarios that need the pool half of the protocol (allocation handles, the registry,uuid, plain pickle of a buffer, re-export of an imported buffer) guarded or skipped for it.test_vmm_ipc.pycovers what is specific to virtual memory: a three-chunk range placed by a relocating grow, import on a second device with and without peer access, the descriptor validation and rollback paths, the re-export errors for imported and derived buffers (direct,Pipe,Processarguments, and theQueuefeeder thread), file descriptor counts returning to baseline on both sides, export under a loweredRLIMIT_NOFILE, raw-fd round trips throughos.dupand a Unix socket withSCM_RIGHTS, empty buffers, a DLPack consumer, same-process aliasing, and file descriptor ownership.Fabric handles
The path for
handle_type="fabric"would differ only in the payload: a 64-byteCUmemFabricHandleper allocation instead of a file descriptor. It is left out of this PR because no available test system reportsCU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTEDor exposes an IMEX channel, so it could not be verified.is_ipc_enabledreportsFalsefor"fabric"and the docs say so.Related Work
Part of #2980. Builds on the
_rthandle layer that #2917 gaveVirtualMemoryResource.🤖 Generated with Claude Code