[core][rdt] Implement receive side memory pool

[Sparks0219]

Following up on #61573 and implementing the receiver side memory pool. Unlike the sender side where once the UDF finishes we store the tensors into the memory pool, on the receiver side once NIXL writes the transferred tensor into the memory pool the user could be actively using it. Hence using a python weakref to detect when the tensor goes out of scope and the corresponding memory block can be returned to the pool. Added some unit and integration tests for it.

[gemini-code-assist]

Code Review

This pull request implements receiver-side memory pool management for NIXL RDT optimization, enabling incoming tensors to be read directly into a pre-registered pool and automatically freed via weakref finalizers. The review feedback highlights several critical issues: a deduplication collision and OOM risk when forwarding existing pool-backed views, potential alignment errors during tensor viewing due to unaligned block offsets, crashes when handling empty (0-sized) tensors, and a memory leak where pool blocks are not immediately freed if a transfer fails.

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[github-actions]

This pull request has been automatically closed because there has been no more activity in the 14 days
since being marked stale.

Please feel free to reopen or open a new pull request if you'd still like this to be addressed.

Again, you can always ask for help on our discussion forum or Ray's public slack channel.

Thanks again for your contribution!

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

^{Reviewed by Cursor Bugbot for commit 0e32523. Configure here.}

[cursor]

Padding breaks exact-fit sender pools

Medium Severity

_allocate_memory_blocks always rounds each request up to BLOCK_ALIGNMENT before carving free space. That shared path is also used by sender-side allocate_for_tensors, so pools sized to the sum of tensor storage bytes can now raise NixlOutOfMemoryError for workloads that previously fit—for example two 12-byte float32 tensors in a 24-byte pool, where the first allocation consumes 16 bytes and only 8 remain.

Additional Locations (1)

• python/ray/experimental/rdt/nixl_memory_pool.py#L367-L369

 

^{Reviewed by Cursor Bugbot for commit 0e32523. Configure here.}

[cursor]

Alignment too small for wide dtypes

Low Severity

BLOCK_ALIGNMENT is 8, but Tensor.view requires the byte offset to be divisible by the target dtype’s element size. Dtypes such as complex128 need 16-byte alignment, so a block starting at offset 8 or 24 makes _view_for_block raise when building the typed pool view.

Additional Locations (1)

• python/ray/experimental/rdt/nixl_memory_pool.py#L294-L297

 

^{Reviewed by Cursor Bugbot for commit 0e32523. Configure here.}

[stephanie-wang]

I think this last condition is actually not necessary, since you have to do a second copy from the memory pool device to the final target device anyway.

[stephanie-wang]

Hmm I feel like there should be a cleaner design where you don't need to pass in an additional pool_blocks arg, and you just pass pool_blocks as the tensors argument here instead, and do the clone right before returning the tensors to the caller. Also wondering if the self._remove_tensor_descs(tensors) line is actually valid here since it seems to get called on the copy target instead of the pool blocks that were registered. What do you think?

[Sparks0219]

The reason we passed pool_blocks as an additional field is that on the receive path we don't maintain a mapping from tensors to blocks unlike on the sender side in _allocated_blocks. This is because the pool blocks aren't ref counted and repeatedly used since we copy out of them, so we just pass it to this function for the case where the transfer fails.

We also don't call remove_tensor_descs for the receive side pool tensors since they don't need to be ref counted at all, its set to false.

I'm realizing it's getting a bit confusing to differentiate between the sender and receiver side pool paths now since they go through different paths 😅. What I could do is have just temporarily populate _allocated_blocks which is what the send side uses to track tensor -> pool blocks, and just clean them up as well when we copy out.