[core][rdt] Register your own transport at runtime for RDT

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

Author: dayshah

doc/source/ray-core/api/direct-transport.rst

@@ -34,4 +34,5 @@ Advanced APIs
     :nosignatures:
     :toctree: doc/
 
-    ray.experimental.wait_tensor_freed
+    ray.experimental.wait_tensor_freed
+    ray.experimental.register_tensor_transport

doc/source/ray-core/direct-transport.rst

@@ -251,6 +251,22 @@ You can also use NIXL to retrieve the result from references created by :func:`r
    :start-after: __nixl_put__and_get_start__
    :end-before: __nixl_put__and_get_end__
 
+
+Registering a new tensor transport
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Ray allows users to register new tensor transports for use in RDT at runtime. To register a new tensor transport, use the :func:`ray.experimental.register_tensor_transport <ray.experimental.register_tensor_transport>` function.
+To implement a new tensor transport, you need to implement the abstract interface defined in :class:`ray.experimental.gpu_object_manager.tensor_transport_manager.TensorTransportManager`.
+Then you can simply give `register_tensor_transport` the transport name, devices, and the class that implements `TensorTransportManager`.
+NIXL, NCCL, and GLOO are registered through this API as well, see ``nixl_tensor_transport.py`` for a reference example.
+
+.. code-block:: python
+
+   from ray.experimental.gpu_object_manager import register_tensor_transport
+
+   register_tensor_transport("NIXL", ["cuda", "cpu"], NixlTensorTransport)
+
+
 Summary
 -------
 

python/ray/experimental/__init__.py

@@ -1,5 +1,9 @@
 from ray.experimental.dynamic_resources import set_resource
-from ray.experimental.gpu_object_manager import GPUObjectManager, wait_tensor_freed
+from ray.experimental.gpu_object_manager import (
+    GPUObjectManager,
+    register_tensor_transport,
+    wait_tensor_freed,
+)
 from ray.experimental.locations import get_local_object_locations, get_object_locations
 
 __all__ = [
@@ -8,4 +12,5 @@
     "set_resource",
     "GPUObjectManager",
     "wait_tensor_freed",
+    "register_tensor_transport",
 ]

python/ray/experimental/gpu_object_manager/__init__.py

@@ -2,5 +2,6 @@
     GPUObjectManager,
     wait_tensor_freed,
 )
+from ray.experimental.gpu_object_manager.util import register_tensor_transport
 
-__all__ = ["GPUObjectManager", "wait_tensor_freed"]
+__all__ = ["GPUObjectManager", "wait_tensor_freed", "register_tensor_transport"]

python/ray/experimental/gpu_object_manager/util.py

@@ -1,5 +1,5 @@
 import threading
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Dict, List
 
 from ray._private.custom_types import TensorTransportEnum
 from ray.experimental.gpu_object_manager.collective_tensor_transport import (
@@ -11,31 +11,58 @@
 from ray.experimental.gpu_object_manager.tensor_transport_manager import (
     TensorTransportManager,
 )
+from ray.util.annotations import PublicAPI
 
 if TYPE_CHECKING:
     import torch
 
 
 # Class definitions for transport managers
-transport_manager_classes: dict[str, TensorTransportManager] = {
-    "NIXL": NixlTensorTransport,
-    "GLOO": CollectiveTensorTransport,
-    "NCCL": CollectiveTensorTransport,
-}
-
-transport_devices = {
-    "NIXL": ["cuda", "cpu"],
-    "GLOO": ["cpu"],
-    "NCCL": ["cuda"],
-}
+transport_manager_classes: Dict[str, type[TensorTransportManager]] = {}
 
+transport_devices: Dict[str, List[str]] = {}
 
 # Singleton instances of transport managers
-transport_managers = {}
+transport_managers: Dict[str, TensorTransportManager] = {}
 
 transport_managers_lock = threading.Lock()
 
 
+@PublicAPI(stability="alpha")
+def register_tensor_transport(
+    transport_name: str,
+    devices: List[str],
+    transport_manager_class: type[TensorTransportManager],
+):
+    """
+    Register a new tensor transport for use in Ray.
+
+    Args:
+        transport_name: The name of the transport protocol.
+        devices: List of device types supported by this transport (e.g., ["cuda", "cpu"]).
+        transport_manager_class: A class that implements TensorTransportManager.
+
+    Raises:
+        ValueError: If transport_manager_class is not a class or does not subclass TensorTransportManager.
+    """
+    global transport_manager_classes
+    global transport_devices
+
+    if not issubclass(transport_manager_class, TensorTransportManager):
+        raise ValueError(
+            f"transport_manager_class {transport_manager_class.__name__} must be a subclass of TensorTransportManager."
+        )
+
+    transport_name = transport_name.upper()
+    transport_manager_classes[transport_name] = transport_manager_class
+    transport_devices[transport_name] = devices
+
+
+register_tensor_transport("NIXL", ["cuda", "cpu"], NixlTensorTransport)
+register_tensor_transport("GLOO", ["cpu"], CollectiveTensorTransport)
+register_tensor_transport("NCCL", ["cuda"], CollectiveTensorTransport)
+
+
 def get_tensor_transport_manager(
     transport_name: str,
 ) -> "TensorTransportManager":