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
Contributor
|
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. |
Contributor
Author
|
/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>
Contributor
Author
|
/ok to test 682db0f |
Contributor
|
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>
Contributor
Author
|
/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>
Contributor
Author
|
/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>
Contributor
Author
|
/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>
Contributor
Author
|
/ok to test |
This branch has not been deployed
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
1 participant
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.
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;handle_typeis the switch andis_ipc_enabledthe query. 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 where the resource allocates (the same device, or the same host location type), 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