[core][rdt] ray.get launches concurrent one-sided transfers for multiple ObjectRefs by stephanie-wang · Pull Request #61773 · ray-project/ray

stephanie-wang · 2026-03-16T21:00:20Z

Description

Instead of fetching objects one at a time during ray.get, we launch a fetch request per ObjectRef, then wait for all of them to complete. This should allow overlapping network transfers for requests of multiple ObjectRefs at a time. Note that this PR only supports this for ray.get, not for transfer of task arguments.

The PR refactors the TensorTransportManager to add new methods, fetch_multiple_transfers and wait_fetch_complete. The default implementation for these calls the synchronous recv_multiple_transfers during fetch and simply returns the tensors during wait_fetch_complete. Backends that support asynchronous fetching can override the new methods to allow concurrent transfers (see NixlTensorTransport for an example).

This PR also adds timeout support for ray.get on RDT objects and unit tests for RDTManager.

Related issues

Closes #61453.

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

gemini-code-assist

Code Review

This pull request introduces concurrent fetching for RDT objects in ray.get by refactoring the tensor transport mechanism to support an asynchronous fetch/wait pattern. The changes are well-structured, introducing a new FetchRequest to manage state for asynchronous operations and providing default synchronous implementations for backward compatibility. The NixlTensorTransport is updated to leverage this new asynchronous pattern, and the RDTManager and Worker are modified to pipeline multiple fetch requests. The implementation appears correct and aligns with the goal of improving performance for ray.get with multiple RDT objects. I have one suggestion regarding a potentially unused class that could cause confusion.

gemini-code-assist · 2026-03-16T21:04:56Z

+@dataclass
+class TransferMetadata:
+    """Base class for in-flight tensor transfer state.
+
+    This class holds the minimal state needed to track an async transfer.
+    Backend-specific implementations should extend this class with additional fields.
+
+    Args:
+        tensors: The tensors being transferred.
+    """
+
+    tensors: List[Any]


The TransferMetadata dataclass appears to be unused in this pull request. The NixlFetchRequest, which holds the in-flight transfer state, inherits from FetchRequest, not this class. Additionally, there is another TransferMetadata NamedTuple defined in rdt_manager.py, which could lead to confusion. If this class is not intended for use, consider removing it to improve clarity.

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

dayshah

It feels a little weird to have the ray.get path and the implicit get path diverge so much. But I guess eventually we'll unify them a bit if we ever get around to having rdt more integrated into the dependency waiter

dayshah · 2026-03-17T19:47:13Z

+
+
+class _PipelineCheckingTransport(TensorTransportManager):
+    """Fake one-sided transport that records the order of fetch/wait calls.


oh nice, we've been needing these fake transports

Now we can with the custom transport API :)

dayshah · 2026-03-17T19:48:52Z

+    manager = _build_manager([obj_id], backend=_TWO_SIDED_BACKEND_NAME)
+
+    with pytest.raises(ValueError, match="use_object_store=True"):
+        manager.get_rdt_objects([obj_id], use_object_store=False)


I feel like this test already exists / should exist in the gloo tests where ray.get-ting should fail

Hmm actually I could not find such an existing test. I figured it's nice to have a unit test for the behavior here since it can run much faster. Actually let me update the match string though.

dayshah · 2026-03-17T20:14:25Z

+                self._wait_fetch(object_id, fetch_request)
+            except Exception as e:
+                if trigger_exception is None:
+                    trigger_exception = e


why wait for all the fetches before raising, i think this is different than the normal ray.get semantic where if any of the obj id's error it's immediately raised without waiting for other obj's

Ah yeah the problem is that _wait_fetch also does some cleanup :( Right now it's only relevant for NIXL, though, since it's cleaning up memory registrations.

Another option is to update it so that FetchRequest.del does the cleanup.

dayshah · 2026-03-17T20:16:10Z

+                )
+            else:
+                result[object_id] = rdt_store.wait_and_get_object(
+                    object_id, timeout=ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS


the timeouts here could probably be 0 right, otherwise something went wrong in the fetch?

We use this codepath for both ray.get (where the timeout here should be 0) and the implicit get (where the timeout here is real, and we should actually wait). This was an issue before this PR too, I added a TODO.

dayshah · 2026-03-17T20:16:53Z

+                    trigger_exception = e
+
+        if trigger_exception is not None:
+            raise trigger_exception


the other objects won't get popped out if one of the obj's raises an exception?

Ah good catch, thanks.

dayshah · 2026-03-17T20:22:27Z

+        # the remaining fetches.
+        for object_id, fetch_request in fetch_requests.items():
+            try:
+                self._wait_fetch(object_id, fetch_request)


wait fetch won't abide by the timeout?

Yeah it didn't before either. Let me see if I can update this.

stephanie-wang · 2026-03-17T22:28:05Z

It feels a little weird to have the ray.get path and the implicit get path diverge so much. But I guess eventually we'll unify them a bit if we ever get around to having rdt more integrated into the dependency waiter

Yeah I realized this too, but the goal here was just to get this in now and we can later update the implicit get path.

Co-authored-by: Dhyey Shah <dhyey2019@gmail.com> Signed-off-by: Stephanie Wang <smwang@cs.washington.edu>

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

cursor

Cursor Bugbot has reviewed your changes and found 1 potential issue.

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

cursor · 2026-05-12T18:07:54Z

+                        owner_address="",
+                        call_site="",
+                    )
+        return result


Fetched RDT objects not stored for repeated ray.get

Medium Severity

Objects fetched via _trigger_fetch and _wait_fetch in fetch_and_get_rdt_objects are returned but never added to the rdt_store. The old _fetch_object code called rdt_store.add_object(obj_id, tensors) after a successful fetch, allowing subsequent ray.get calls on the same ObjectRef to find the object in phase 1's store check. Without this caching, every ray.get call on the same RDT ObjectRef triggers a new network transfer, which may fail if the source has already garbage-collected the object.

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

…ple ObjectRefs (ray-project#61773) Instead of fetching objects one at a time during `ray.get`, we launch a fetch request per ObjectRef, then wait for all of them to complete. This should allow overlapping network transfers for requests of multiple ObjectRefs at a time. Note that this PR only supports this for `ray.get`, not for transfer of task arguments. The PR refactors the TensorTransportManager to add new methods, `fetch_multiple_transfers` and `wait_fetch_complete`. The default implementation for these calls the synchronous `recv_multiple_transfers` during fetch and simply returns the tensors during `wait_fetch_complete`. Backends that support asynchronous fetching can override the new methods to allow concurrent transfers (see NixlTensorTransport for an example). This PR also adds timeout support for ray.get on RDT objects and unit tests for RDTManager. Closes ray-project#61453. --------- Signed-off-by: Stephanie wang <smwang@cs.washington.edu> Co-authored-by: Dhyey Shah <dhyey2019@gmail.com> Signed-off-by: anindyam1969 <amukherjee@kinetica.com>

stephanie-wang added 6 commits March 10, 2026 16:22

split nixl transfer into fetch then wait

d9baca8

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

Only add transfer if doesn't already exist

d766145

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

generic transfer metadata

bfe6fbd

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

fetch returns context for request

d0a7405

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

update

7ecdc89

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

comment

a14212e

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

stephanie-wang requested a review from a team as a code owner March 16, 2026 21:00

gemini-code-assist Bot reviewed Mar 16, 2026

View reviewed changes

cursor Bot reviewed Mar 16, 2026

View reviewed changes

Comment thread python/ray/experimental/rdt/rdt_manager.py Outdated

Comment thread python/ray/experimental/rdt/rdt_manager.py

Comment thread python/ray/experimental/rdt/tensor_transport_manager.py Outdated

stephanie-wang added 3 commits March 16, 2026 15:52

update

518b12b

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

async object store fetch

d0b2bfe

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

unit test

7ddc1cf

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

cursor Bot reviewed Mar 17, 2026

View reviewed changes

Comment thread python/ray/_private/worker.py

Comment thread python/ray/experimental/rdt/rdt_manager.py Outdated

ray-gardener Bot added the core Issues that should be addressed in Ray Core label Mar 17, 2026

stephanie-wang added 2 commits March 17, 2026 11:14

TODO

98d6f2d

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

fix

cbba0ea

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

cursor Bot reviewed Mar 17, 2026

View reviewed changes

Comment thread python/ray/experimental/rdt/rdt_manager.py

Comment thread python/ray/experimental/rdt/rdt_manager.py Outdated

dayshah reviewed Mar 17, 2026

View reviewed changes

Update python/ray/experimental/rdt/rdt_manager.py

20757cd

Co-authored-by: Dhyey Shah <dhyey2019@gmail.com> Signed-off-by: Stephanie Wang <smwang@cs.washington.edu>

cursor Bot reviewed Mar 17, 2026

View reviewed changes

Comment thread python/ray/experimental/rdt/rdt_manager.py Outdated

stephanie-wang added the @author-action-required The PR author is responsible for the next step. Remove tag to send back to the reviewer. label Mar 18, 2026

stephanie-wang added 6 commits March 24, 2026 10:29

update

b7b23c7

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

update

3924469

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

get existing objects

c2aa744

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

Fetch objects before deserialize

5091bb3

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

Add timeout

bff5940

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

cleanup

a7b4ca6

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

stephanie-wang removed the @author-action-required The PR author is responsible for the next step. Remove tag to send back to the reviewer. label Mar 24, 2026

cursor Bot reviewed Mar 24, 2026

View reviewed changes

Comment thread python/ray/experimental/rdt/rdt_store.py

Comment thread python/ray/tests/rdt/test_rdt_manager.py

Comment thread python/ray/tests/rdt/test_rdt_manager.py

Comment thread python/ray/experimental/rdt/rdt_manager.py

stephanie-wang requested review from a team, SongGuyang, WangTaoTheTonic, kfstorm and raulchen as code owners April 27, 2026 22:40

github-actions Bot added unstale A PR that has been marked unstale. It will not get marked stale again if this label is on it. and removed stale The issue is stale. It will be closed within 7 days unless there are further conversation labels Apr 28, 2026

stephanie-wang added 2 commits April 28, 2026 15:53

Merge remote-tracking branch 'upstream/master' into async-transfers

78ed178

update

0730c2a

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

stephanie-wang force-pushed the async-transfers branch from d9115be to 0730c2a Compare April 28, 2026 22:53

stephanie-wang enabled auto-merge (squash) April 28, 2026 22:54

github-actions Bot added the go add ONLY when ready to merge, run all tests label Apr 28, 2026

cursor Bot reviewed Apr 28, 2026

View reviewed changes

Comment thread python/ray/experimental/rdt/rdt_manager.py

Comment thread python/ray/experimental/rdt/rdt_manager.py

aslonnie removed request for a team and aslonnie April 29, 2026 02:51

stephanie-wang added 2 commits April 29, 2026 20:34

update

0405224

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

Merge remote-tracking branch 'upstream/master' into async-transfers

9d05836

github-actions Bot disabled auto-merge April 30, 2026 03:34

cursor Bot reviewed Apr 30, 2026

View reviewed changes

Comment thread python/ray/experimental/rdt/nixl_tensor_transport.py

stephanie-wang added 2 commits April 29, 2026 21:03

requirements?

20eca1c

Signed-off-by: Stephanie wang <smwang@cs.washington.edu>

Merge remote-tracking branch 'upstream/master' into async-transfers

e5e651e

cursor Bot reviewed May 5, 2026

View reviewed changes

Comment thread python/ray/experimental/rdt/rdt_manager.py

stephanie-wang force-pushed the async-transfers branch from ff9a483 to e5e651e Compare May 12, 2026 18:00

cursor Bot reviewed May 12, 2026

View reviewed changes

edoakes merged commit f77ec0b into ray-project:master May 13, 2026
4 of 6 checks passed



		class _PipelineCheckingTransport(TensorTransportManager):
		"""Fake one-sided transport that records the order of fetch/wait calls.

Conversation

stephanie-wang commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related issues

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Mar 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dayshah left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

stephanie-wang commented Mar 17, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

cursor Bot May 12, 2026

Choose a reason for hiding this comment

Fetched RDT objects not stored for repeated ray.get

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

stephanie-wang commented Mar 16, 2026 •

edited

Loading