Repository: ray-project/ray
File: python/ray/experimental/gpu_object_manager/util.py (Python)

Code before:
DEFAULT_TRANSPORTS = ["NIXL", "GLOO", "NCCL", "CUDA_IPC", "JAX"]

register_tensor_transport("NIXL", ["cuda", "cpu"], NixlTensorTransport)
register_tensor_transport("GLOO", ["cpu"], GLOOTensorTransport)
register_tensor_transport("NCCL", ["cuda"], NCCLTensorTransport)
register_tensor_transport("CUDA_IPC", ["cuda"], CudaIpcTransport)

try:
    import jax

    register_tensor_transport(
        "JAX", ["tpu", "cpu", "cuda"], JaxTensorTransport, jax.Array
    )
except ImportError:
    pass


def get_transport_data_type(tensor_transport: str) -> type:
    if tensor_transport not in transport_manager_info:
        raise ValueError(f"Unsupported tensor transport protocol: {tensor_transport}")

    data_type = transport_manager_info[tensor_transport].data_type
    if data_type is None:
        import torch

        return torch.Tensor
    return data_type


def get_tensor_transport_manager(
    transport_name: str,
) -> "TensorTransportManager":
    """Get the tensor transport manager for the given tensor transport protocol.

    Args:
        transport_name: The tensor transport protocol to use for the GPU object.

    Returns:
        TensorTransportManager: The tensor transport manager for the given tensor transport protocol.
    """
    global transport_manager_info
    global transport_managers
    global transport_managers_lock

    with transport_managers_lock:
        if transport_name in transport_managers:
            return transport_managers[transport_name]

        if transport_name not in transport_manager_info:
            raise ValueError(f"Unsupported tensor transport protocol: {transport_name}")


Review comment (suggestion):
Do you want to just make the current tensor transport managers explicitly specify torch.Tensor? It's not very intuitive why torch.Tensor should be the default, and this way the existing tensor transports can test the custom data type code.

Code after:
_default_transports_registered = False


def _ensure_default_transports_registered():
    global _default_transports_registered
    with transport_managers_lock:
        if _default_transports_registered:
            return
        _default_transports_registered = True
        try:
            import torch

            register_tensor_transport(
                "NIXL", ["cuda", "cpu"], NixlTensorTransport, torch.Tensor
            )
            register_tensor_transport(
                "GLOO", ["cpu"], GLOOTensorTransport, torch.Tensor
            )
            register_tensor_transport(
                "NCCL", ["cuda"], NCCLTensorTransport, torch.Tensor
            )
            register_tensor_transport(
                "CUDA_IPC", ["cuda"], CudaIpcTransport, torch.Tensor
            )
        except ImportError:
            pass


def get_transport_data_type(tensor_transport: str) -> type:
    _ensure_default_transports_registered()
    if tensor_transport not in transport_manager_info:
        raise ValueError(f"Unsupported tensor transport protocol: {tensor_transport}")

    return transport_manager_info[tensor_transport].data_type


def get_tensor_transport_manager(
    transport_name: str,
) -> "TensorTransportManager":
    """Get the tensor transport manager for the given tensor transport protocol.

    Args:
        transport_name: The tensor transport protocol to use for the GPU object.

    Returns:
        TensorTransportManager: The tensor transport manager for the given tensor transport protocol.
    """
    global transport_manager_info
    global transport_managers
    global transport_managers_lock