Repository: ray-project/ray
Pull request: [core][gpu-objects] Support `ray.get` on the driver process for GPU objects
<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

If an object ID exists in `gpu_object_ref`, it means that the driver is attempting to fetch the GPU objects. In that case, tensors will not be transferred out of band. Instead, we need to transfer the tensors via object store.

## Related issue number

Closes #51272 

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for https://docs.ray.io/en/master/.
    - [ ] I've added any new APIs to the API Reference. For example, if I added a
           method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(


Problem statement:
[core][gpu-objects] Driver tries to get the data from in-actor store
### Description

The driver is not allowed to be in an NCCL group in Ray GPU objects. Hence, if the driver wants to retrieve data from the in-actor store, we can move the data from the in-actor store to the object store so that the driver can access it.

### Use case

_No response_

Hints from issue discussion:

Commit message: update

Signed-off-by: kaihsun <kaihsun@anyscale.com>

Patch to review:
diff --git a/python/ray/_private/gpu_object_manager.py b/python/ray/_private/gpu_object_manager.py
index f1c3857767b2..0761f55fe3fd 100644
--- a/python/ray/_private/gpu_object_manager.py
+++ b/python/ray/_private/gpu_object_manager.py
@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Tuple
 
+import ray
 from ray._private.custom_types import TensorTransportEnum
 from ray._raylet import ObjectRef
 from ray.actor import ActorHandle
@@ -50,6 +51,9 @@ def __init__(self):
         # The entries in this dictionary are 1:1 with ObjectRefs created by this process with a tensor_transport hint and that are currently in scope.
         self.gpu_object_refs: Dict[ObjectRef, GPUObjectMeta] = {}
 
+        # A dictionary that maps from an object ID to the corresponding ObjectRef.
+        self.gpu_object_id_to_obj_ref: Dict[str, ObjectRef] = {}
+
     def has_gpu_object(self, obj_id: str) -> bool:
         return obj_id in self.gpu_object_store
 
@@ -79,6 +83,19 @@ def __ray_get_tensor_meta__(self, obj_id: str):
 
         return src_actor.__ray_call__.remote(__ray_get_tensor_meta__, obj_id)
 
+    def is_driver(self, obj_id: str) -> bool:
+        """
+        Check if the GPU object is managed by this process.
+
+        Args:
+            obj_id: The object ID of the GPU object.
+
+        Returns:
+            True if the current process is the driver process coordinating the data transfer
+            of this GPU object.
+        """
+        return obj_id in self.gpu_object_id_to_obj_ref
+
     def add_gpu_object_ref(
         self,
         obj_ref: ObjectRef,
@@ -104,11 +121,13 @@ def add_gpu_object_ref(
             tensor_transport_backend=tensor_transport_backend,
             tensor_meta=tensor_meta,
         )
+        self.gpu_object_id_to_obj_ref[obj_ref.hex()] = obj_ref
 
     # TODO(kevin85421): Call this function to remove the `obj_ref` from the `gpu_object_refs` dictionary
     # to allow garbage collection of the object.
     def remove_gpu_object_ref(self, obj_ref: ObjectRef):
         del self.gpu_object_refs[obj_ref]
+        del self.gpu_object_id_to_obj_ref[obj_ref.hex()]
 
     def _get_gpu_object_ref(self, obj_ref: ObjectRef) -> Optional[GPUObjectMeta]:
         return self.gpu_object_refs[obj_ref]
@@ -141,6 +160,33 @@ def _recv_gpu_object(
             util.__ray_recv__, communicator_name, obj_id, src_rank, tensor_meta
         )
 
+    def fetch_gpu_object(self, obj_id: str):
+        """
+        Fetches the GPU object from the source actor's GPU object store via the object store
+        instead of out-of-band tensor transfer and stores the tensors in `gpu_object_store`.
+
+        This is useful when the current process is the driver process that coordinates the data
+        transfer of this GPU object. For example, when the driver process calls `ray.get(gpu_object_ref)`,
+        the tensors will not be transferred out-of-band automatically. We need to fetch the tensors from
+        the corresponding actor's GPU object store via the object store API.
+
+        Args:
+            obj_id: The object ID of the GPU object.
+
+        Returns:
+            None
+        """
+        if obj_id in self.gpu_object_store:
+            return
+
+        obj_ref = self.gpu_object_id_to_obj_ref[obj_id]
+        gpu_object_meta = self.gpu_object_refs[obj_ref]
+        src_actor = gpu_object_meta.src_actor
+        tensors = ray.get(
+            src_actor.__ray_call__.remote(util.__ray_fetch_gpu_object__, obj_id)
+        )
+        self.gpu_object_store[obj_id] = tensors
+
     def trigger_out_of_band_tensor_transfer(
         self, dst_actor: ActorHandle, task_args: Tuple[Any, ...]
     ):
diff --git a/python/ray/_private/gpu_object_manager_util.py b/python/ray/_private/gpu_object_manager_util.py
index e273e386ec38..73774b02988e 100644
--- a/python/ray/_private/gpu_object_manager_util.py
+++ b/python/ray/_private/gpu_object_manager_util.py
@@ -81,3 +81,17 @@ def __ray_recv__(
         collective.recv(tensor, src_rank, group_name=communicator_name)
         tensors.append(tensor)
     gpu_object_manager.add_gpu_object(obj_id, tensors)
+
+
+def __ray_fetch_gpu_object__(self, obj_id: str):
+    """Helper function that runs on the src actor to fetch tensors from the GPU object store via the object store."""
+    gpu_object_manager = global_worker.gpu_object_manager
+    assert gpu_object_manager.has_gpu_object(
+        obj_id
+    ), f"obj_id={obj_id} not found in GPU object store"
+    tensors = gpu_object_manager.get_gpu_object(obj_id)
+    # TODO(kevin85421): The current garbage collection implementation for the
+    # in-actor object store is naive. We garbage collect each object after it
+    # is consumed once.
+    gpu_object_manager.remove_gpu_object(obj_id)
+    return tensors
diff --git a/python/ray/_private/serialization.py b/python/ray/_private/serialization.py
index 2ea4a0f89e30..637e332e8770 100644
--- a/python/ray/_private/serialization.py
+++ b/python/ray/_private/serialization.py
@@ -277,8 +277,12 @@ def _deserialize_pickle5_data(
         worker = ray._private.worker.global_worker
 
         gpu_object_manager = worker.gpu_object_manager
-        enable_gpu_objects = gpu_object_manager.has_gpu_object(object_id)
+        enable_gpu_objects = gpu_object_manager.has_gpu_object(
+            object_id
+        ) or gpu_object_manager.is_driver(object_id)
         if enable_gpu_objects:
+            if gpu_object_manager.is_driver(object_id):
+                gpu_object_manager.fetch_gpu_object(object_id)
             tensors = gpu_object_manager.get_gpu_object(object_id)
             ctx.reset_out_of_band_tensors(tensors)
             # TODO(kevin85421): The current garbage collection implementation for the in-actor object store
diff --git a/python/ray/tests/test_gpu_objects_gloo.py b/python/ray/tests/test_gpu_objects_gloo.py
index 495f2e82cd4f..c023eb4d0b67 100644
--- a/python/ray/tests/test_gpu_objects_gloo.py
+++ b/python/ray/tests/test_gpu_objects_gloo.py
@@ -149,6 +149,32 @@ def test_trigger_out_of_band_tensor_transfer(ray_start_regular):
     assert torch.equal(ret_val_dst[0], tensor)
 
 
+def test_fetch_gpu_object_to_driver(ray_start_regular):
+    actor = GPUTestActor.remote()
+    create_collective_group([actor], backend="torch_gloo")
+
+    tensor1 = torch.tensor([1, 2, 3])
+    tensor2 = torch.tensor([4, 5, 6])
+
+    # Case 1: Single tensor
+    ref = actor.echo.remote(tensor1)
+    assert torch.equal(ray.get(ref), tensor1)
+
+    # Case 2: Multiple tensors
+    ref = actor.echo.remote([tensor1, tensor2])
+    result = ray.get(ref)
+    assert torch.equal(result[0], tensor1)
+    assert torch.equal(result[1], tensor2)
+
+    # Case 3: Mixed CPU and GPU data
+    data = [tensor1, tensor2, 7]
+    ref = actor.echo.remote(data)
+    result = ray.get(ref)
+    assert torch.equal(result[0], tensor1)
+    assert torch.equal(result[1], tensor2)
+    assert result[2] == 7
+
+
 def test_invalid_tensor_transport(ray_start_regular):
     with pytest.raises(ValueError, match="Invalid tensor transport"):
 


Code review comments:
diff_hunk: @@ -79,6 +83,19 @@ def __ray_get_tensor_meta__(self, obj_id: str):
 
         return src_actor.__ray_call__.remote(__ray_get_tensor_meta__, obj_id)
 
+    def is_driver(self, obj_id: str) -> bool:
path: python/ray/_private/gpu_object_manager.py
text: @user1:
The function name does not match the docstring.

How about is_managed_object, or something along those lines? You can rename self.gpu_object_refs to self.managed_gpu_object_refs, etc.

@author:
Updated: https://github.com/ray-project/ray/pull/53902/commits/cc90f1f891b0b798f82948401c9d584a25d47d9a

diff_hunk: @@ -50,6 +51,9 @@ def __init__(self):
         # The entries in this dictionary are 1:1 with ObjectRefs created by this process with a tensor_transport hint and that are currently in scope.
         self.gpu_object_refs: Dict[ObjectRef, GPUObjectMeta] = {}
 
+        # A dictionary that maps from an object ID to the corresponding ObjectRef.
+        self.gpu_object_id_to_obj_ref: Dict[str, ObjectRef] = {}
path: python/ray/_private/gpu_object_manager.py
text: @user1:
Let's not introduce this map as it seems unnecessarily complex.

It would be better to standardize all of the internal data structures to use either hex string or ObjectID (which can be created from a hex string) as the key, and then I don't think we need the additional map.

@author:
Updated.