[core][rdt] Rework passing transport through for bring your own transport (#59216)

Signed-off-by: dayshah <dhyey2019@gmail.com>

Author: dayshah

doc/source/ray-core/doc_code/direct_transport_gloo.py

@@ -140,7 +140,7 @@ def sum(self, tensor: torch.Tensor):
     ray.get(tensor)
 
 assert (
-    "Currently ray.get() only supports OBJECT_STORE and NIXL tensor transport, got TensorTransportEnum.GLOO, please specify the correct tensor transport in ray.get()"
+    "Trying to use two-sided tensor transport: GLOO for ray.get. This is only supported for one-sided transports such as NIXL or the OBJECT_STORE."
     in str(e.value)
 )
 

python/ray/_private/custom_types.py

@@ -124,9 +124,7 @@
 # See `common.proto` for more details.
 class TensorTransportEnum(Enum):
     OBJECT_STORE = TensorTransport.Value("OBJECT_STORE")
-    NCCL = TensorTransport.Value("NCCL")
-    GLOO = TensorTransport.Value("GLOO")
-    NIXL = TensorTransport.Value("NIXL")
+    DIRECT_TRANSPORT = TensorTransport.Value("DIRECT_TRANSPORT")
 
     @classmethod
     def from_str(cls, name: str) -> "TensorTransportEnum":

python/ray/_private/worker.py

@@ -801,7 +801,7 @@ def put_object(
         value: Any,
         owner_address: Optional[str] = None,
         _is_experimental_channel: bool = False,
-        _tensor_transport: str = "object_store",
+        _tensor_transport: str = TensorTransportEnum.OBJECT_STORE.name,
     ):
         """Put value in the local object store.
 
@@ -835,18 +835,15 @@ def put_object(
                 "ray.ObjectRef in a list and call 'put' on it."
             )
         tensors = None
-        tensor_transport: TensorTransportEnum = TensorTransportEnum.from_str(
-            _tensor_transport
+        from ray.experimental.gpu_object_manager.util import (
+            normalize_and_validate_tensor_transport,
+            validate_one_sided,
         )
-        if tensor_transport not in [
-            TensorTransportEnum.OBJECT_STORE,
-            TensorTransportEnum.NIXL,
-        ]:
-            raise ValueError(
-                "Currently, Ray Direct Transport only supports 'object_store' and 'nixl' for tensor transport in ray.put()."
-            )
+
+        tensor_transport = normalize_and_validate_tensor_transport(_tensor_transport)
+        validate_one_sided(tensor_transport, "ray.put")
         try:
-            if tensor_transport != TensorTransportEnum.OBJECT_STORE:
+            if tensor_transport != TensorTransportEnum.OBJECT_STORE.name:
                 (
                     serialized_value,
                     tensors,
@@ -867,19 +864,24 @@ def put_object(
         # object. Instead, clients will keep the object pinned.
         pin_object = not _is_experimental_channel
 
+        tensor_transport_enum = TensorTransportEnum.OBJECT_STORE
+        if tensor_transport != TensorTransportEnum.OBJECT_STORE.name:
+            tensor_transport_enum = TensorTransportEnum.DIRECT_TRANSPORT
+
         # This *must* be the first place that we construct this python
         # ObjectRef because an entry with 0 local references is created when
         # the object is Put() in the core worker, expecting that this python
         # reference will be created. If another reference is created and
         # removed before this one, it will corrupt the state in the
         # reference counter.
+
         ret = self.core_worker.put_object(
             serialized_value,
             pin_object=pin_object,
             owner_address=owner_address,
             inline_small_object=True,
             _is_experimental_channel=_is_experimental_channel,
-            tensor_transport_val=tensor_transport.value,
+            tensor_transport_val=tensor_transport_enum.value,
         )
         if tensors:
             self.gpu_object_manager.put_object(ret, tensor_transport, tensors)
@@ -896,43 +898,26 @@ def deserialize_objects(
         self,
         serialized_objects,
         object_refs,
-        tensor_transport_hint: Optional[TensorTransportEnum] = None,
+        tensor_transport_hint: Optional[str] = None,
     ):
         gpu_objects: Dict[str, List["torch.Tensor"]] = {}
         for obj_ref, (_, _, tensor_transport) in zip(object_refs, serialized_objects):
-            # TODO: Here tensor_transport_hint is set by the user in ray.get(), tensor_transport is set
-            # in serialize_objects by ray.method(tensor_transport="xxx"), and obj_ref.tensor_transport()
-            # is set by ray.put(). We may clean up this logic in the future.
             if (
                 tensor_transport is None
                 or tensor_transport == TensorTransportEnum.OBJECT_STORE
-            ) and (
-                obj_ref is None
-                or obj_ref.tensor_transport() == TensorTransportEnum.OBJECT_STORE.value
             ):
                 # The object is not a gpu object, so we cannot use other external transport to
                 # fetch it.
                 continue
 
-            # If the object is a gpu object, we can choose to use the object store or other external
-            # transport to fetch it. The `tensor_transport_hint` has the highest priority, then the
-            # tensor_transport in obj_ref.tensor_transport(), then the tensor_transport in serialize_objects,
-            # then the default value `OBJECT_STORE`.
-            chosen_tensor_transport = (
-                tensor_transport_hint
-                or (
-                    TensorTransportEnum(obj_ref.tensor_transport()) if obj_ref else None
-                )
-                or tensor_transport
-                or TensorTransportEnum.OBJECT_STORE
-            )
-
             object_id = obj_ref.hex()
             if object_id not in gpu_objects:
                 # If using a non-object store transport, then tensors will be sent
                 # out-of-band. Get them before deserializing the object store data.
+                # The user can choose OBJECT_STORE as the hint to fetch the RDT object
+                # through the object store.
                 gpu_objects[object_id] = self.gpu_object_manager.get_gpu_object(
-                    object_id, tensor_transport=chosen_tensor_transport
+                    object_id, tensor_transport=tensor_transport_hint
                 )
 
         # Function actor manager or the import thread may call pickle.loads
@@ -983,16 +968,6 @@ def get_objects(
                     f"Attempting to call `get` on the value {object_ref}, "
                     "which is not an ray.ObjectRef."
                 )
-        tensor_transport: TensorTransportEnum = (
-            TensorTransportEnum.from_str(_tensor_transport)
-            if _tensor_transport is not None
-            else None
-        )
-        assert tensor_transport in [
-            TensorTransportEnum.OBJECT_STORE,
-            TensorTransportEnum.NIXL,
-            None,
-        ], "Currently, RDT only supports 'object_store' and 'nixl' for tensor transport in ray.get()."
         timeout_ms = (
             int(timeout * 1000) if timeout is not None and timeout != -1 else -1
         )
@@ -1004,7 +979,7 @@ def get_objects(
         )
 
         debugger_breakpoint = b""
-        for data, metadata, _ in serialized_objects:
+        for _, metadata, _ in serialized_objects:
             if metadata:
                 metadata_fields = metadata.split(b",")
                 if len(metadata_fields) >= 2 and metadata_fields[1].startswith(
@@ -1016,8 +991,17 @@ def get_objects(
         if skip_deserialization:
             return None, debugger_breakpoint
 
+        if _tensor_transport is not None:
+            from ray.experimental.gpu_object_manager.util import (
+                normalize_and_validate_tensor_transport,
+            )
+
+            _tensor_transport = normalize_and_validate_tensor_transport(
+                _tensor_transport
+            )
+
         values = self.deserialize_objects(
-            serialized_objects, object_refs, tensor_transport_hint=tensor_transport
+            serialized_objects, object_refs, tensor_transport_hint=_tensor_transport
         )
         if not return_exceptions:
             # Raise exceptions instead of returning them to the user.

python/ray/_raylet.pyx

@@ -2874,7 +2874,7 @@ cdef class CoreWorker:
                 c_object_ids, timeout_ms, results)
         check_status(op_status)
 
-        return RayObjectsToSerializedRayObjects(results)
+        return RayObjectsToSerializedRayObjects(results, object_refs)
 
     def get_if_local(self, object_refs):
         """Get objects from local plasma store directly
@@ -2886,7 +2886,7 @@ cdef class CoreWorker:
             check_status(
                 CCoreWorkerProcess.GetCoreWorker().GetIfLocal(
                     c_object_ids, &results))
-        return RayObjectsToSerializedRayObjects(results)
+        return RayObjectsToSerializedRayObjects(results, object_refs)
 
     def object_exists(self, ObjectRef object_ref, memory_store_only=False):
         cdef:

python/ray/actor.py

@@ -473,9 +473,14 @@ def annotate_method(method: Callable[_P, _Ret]):
         if "enable_task_events" in kwargs and kwargs["enable_task_events"] is not None:
             method.__ray_enable_task_events__ = kwargs["enable_task_events"]
         if "tensor_transport" in kwargs:
-            method.__ray_tensor_transport__ = TensorTransportEnum.from_str(
-                kwargs["tensor_transport"]
+            tensor_transport = kwargs["tensor_transport"]
+            from ray.experimental.gpu_object_manager.util import (
+                normalize_and_validate_tensor_transport,
             )
+
+            tensor_transport = normalize_and_validate_tensor_transport(tensor_transport)
+            method.__ray_tensor_transport__ = tensor_transport
+
         return method
 
     # Check if decorator is called without parentheses (args[0] would be the function)
@@ -521,7 +526,7 @@ def __init__(
         enable_task_events: bool,
         decorator: Optional[Any] = None,
         signature: Optional[List[inspect.Parameter]] = None,
-        tensor_transport: Optional[TensorTransportEnum] = None,
+        tensor_transport: Optional[str] = None,
     ):
         """Initialize an _ActorMethodMetadata.
 
@@ -599,7 +604,7 @@ def __init__(
         enable_task_events: bool,
         decorator=None,
         signature: Optional[List[inspect.Parameter]] = None,
-        tensor_transport: Optional[TensorTransportEnum] = None,
+        tensor_transport: Optional[str] = None,
     ):
         """Initialize an ActorMethod.
 
@@ -649,10 +654,10 @@ def __init__(
         # and return the resulting ObjectRefs.
         self._decorator = decorator
 
-        # If the task call doesn't specify a tensor transport option, use `_tensor_transport`
+        # If the task call doesn't specify a tensor transport option, use `OBJECT_STORE`
         # as the default transport for this actor method.
         if tensor_transport is None:
-            tensor_transport = TensorTransportEnum.OBJECT_STORE
+            tensor_transport = TensorTransportEnum.OBJECT_STORE.name
         self._tensor_transport = tensor_transport
 
     def __call__(self, *args, **kwargs):
@@ -695,7 +700,12 @@ def options(self, **options):
 
         tensor_transport = options.get("tensor_transport", None)
         if tensor_transport is not None:
-            options["tensor_transport"] = TensorTransportEnum.from_str(tensor_transport)
+            from ray.experimental.gpu_object_manager.util import (
+                normalize_and_validate_tensor_transport,
+            )
+
+            tensor_transport = normalize_and_validate_tensor_transport(tensor_transport)
+            options["tensor_transport"] = tensor_transport
 
         class FuncWrapper:
             def remote(self, *args, **kwargs):
@@ -800,7 +810,7 @@ def _remote(
         concurrency_group=None,
         _generator_backpressure_num_objects=None,
         enable_task_events=None,
-        tensor_transport: Optional[TensorTransportEnum] = None,
+        tensor_transport: Optional[str] = None,
     ):
         if num_returns is None:
             num_returns = self._num_returns
@@ -820,23 +830,23 @@ def _remote(
         if tensor_transport is None:
             tensor_transport = self._tensor_transport
 
-        if tensor_transport != TensorTransportEnum.OBJECT_STORE:
+        if tensor_transport != TensorTransportEnum.OBJECT_STORE.name:
             if num_returns != 1:
                 raise ValueError(
-                    f"Currently, methods with tensor_transport={tensor_transport.name} only support 1 return value. "
+                    f"Currently, methods with tensor_transport={tensor_transport} only support 1 return value. "
                     "Please make sure the actor method is decorated with `@ray.method(num_returns=1)` (the default)."
                 )
             if not self._actor._ray_enable_tensor_transport:
                 raise ValueError(
-                    f'Currently, methods with .options(tensor_transport="{tensor_transport.name}") are not supported when enable_tensor_transport=False. '
+                    f'Currently, methods with .options(tensor_transport="{tensor_transport}") are not supported when enable_tensor_transport=False. '
                     "Please set @ray.remote(enable_tensor_transport=True) on the actor class definition."
                 )
             gpu_object_manager = ray._private.worker.global_worker.gpu_object_manager
             if not gpu_object_manager.actor_has_tensor_transport(
                 self._actor, tensor_transport
             ):
                 raise ValueError(
-                    f'{self._actor} does not have tensor transport {tensor_transport.name} available. If using a collective-based transport ("nccl" or "gloo"), please create a communicator with '
+                    f'{self._actor} does not have tensor transport {tensor_transport} available. If using a collective-based transport ("nccl" or "gloo"), please create a communicator with '
                     "`ray.experimental.collective.create_collective_group` "
                     "before calling actor tasks with non-default tensor_transport."
                 )
@@ -876,7 +886,7 @@ def invocation(args, kwargs):
             invocation = self._decorator(invocation)
 
         object_refs = invocation(args, kwargs)
-        if tensor_transport != TensorTransportEnum.OBJECT_STORE:
+        if tensor_transport != TensorTransportEnum.OBJECT_STORE.name:
             # Currently, we only support transfer tensor out-of-band when
             # num_returns is 1.
             assert isinstance(object_refs, ObjectRef)
@@ -979,14 +989,16 @@ def create(
         self.enable_task_events = {}
         self.generator_backpressure_num_objects = {}
         self.concurrency_group_for_methods = {}
-        self.method_name_to_tensor_transport: Dict[str, TensorTransportEnum] = {}
+        self.method_name_to_tensor_transport: Dict[str, str] = {}
 
         # Check whether any actor methods specify a non-default tensor transport.
         self.has_tensor_transport_methods = any(
             getattr(
-                method, "__ray_tensor_transport__", TensorTransportEnum.OBJECT_STORE
+                method,
+                "__ray_tensor_transport__",
+                TensorTransportEnum.OBJECT_STORE.name,
             )
-            != TensorTransportEnum.OBJECT_STORE
+            != TensorTransportEnum.OBJECT_STORE.name
             for _, method in actor_methods
         )
 
@@ -1941,7 +1953,7 @@ def __init__(
         method_generator_backpressure_num_objects: Dict[str, int],
         method_enable_task_events: Dict[str, bool],
         enable_tensor_transport: bool,
-        method_name_to_tensor_transport: Dict[str, TensorTransportEnum],
+        method_name_to_tensor_transport: Dict[str, str],
         actor_method_cpus: int,
         actor_creation_function_descriptor,
         cluster_and_job,
@@ -1968,7 +1980,7 @@ def __init__(
                 this actor. If True, then methods can be called with
                 .options(tensor_transport=...) to specify a non-default tensor
                 transport.
-            method_name_to_tensor_transport: Dictionary mapping method names to their tensor transport settings.
+            method_name_to_tensor_transport: Dictionary mapping method names to their tensor transport type.
             actor_method_cpus: The number of CPUs required by actor methods.
             actor_creation_function_descriptor: The function descriptor for actor creation.
             cluster_and_job: The cluster and job information.
@@ -2079,7 +2091,7 @@ def _actor_method_call(
         concurrency_group_name: Optional[str] = None,
         generator_backpressure_num_objects: Optional[int] = None,
         enable_task_events: Optional[bool] = None,
-        tensor_transport: Optional[TensorTransportEnum] = None,
+        tensor_transport: Optional[str] = None,
     ):
         """Method execution stub for an actor handle.
 
@@ -2157,6 +2169,12 @@ def _actor_method_call(
         if generator_backpressure_num_objects is None:
             generator_backpressure_num_objects = -1
 
+        tensor_transport_enum = TensorTransportEnum.OBJECT_STORE
+        if (
+            tensor_transport is not None
+            and tensor_transport != TensorTransportEnum.OBJECT_STORE.name
+        ):
+            tensor_transport_enum = TensorTransportEnum.DIRECT_TRANSPORT
         object_refs = worker.core_worker.submit_actor_task(
             self._ray_actor_language,
             self._ray_actor_id,
@@ -2171,7 +2189,7 @@ def _actor_method_call(
             concurrency_group_name if concurrency_group_name is not None else b"",
             generator_backpressure_num_objects,
             enable_task_events,
-            tensor_transport.value,
+            tensor_transport_enum.value,
         )
 
         if num_returns == STREAMING_GENERATOR_RETURN: