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

Code before:
    cuda_ipc_event_ipc_handle: Optional[bytes] = None
    # The index of the GPU that the tensors are on. This requires that the GPU is
    # assigned by Ray, e.g., using @ray.remote(num_gpus=1).
    ray_gpu_idx: Optional[int] = None
    # The node that the GPU that the tensors are on is on.
    ray_node_id: Optional[str] = None


class CudaIpcTransport(TensorTransportManager):
    def __init__(self, tensor_transport_backend: str):
        pass

    @property
    def tensor_transport_backend(self) -> str:
        return "CUDA_IPC"

    @staticmethod
    def is_one_sided() -> bool:
        return True

    @staticmethod
    def can_abort_transport() -> bool:
        return False

    def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool:
        return torch.cuda.is_available()

    @staticmethod
    def extract_tensor_transport_metadata(
        obj_id: str,
        gpu_object: List["torch.Tensor"],
    ) -> CudaIpcTransportMetadata:

        tensor_meta = []
        cuda_ipc_handles = []
        event_ipc_handle = None
        ray_gpu_idx = None
        ray_node_id = None
        if gpu_object:
            import torch
            from torch.multiprocessing.reductions import reduce_tensor

            # Create an interprocess-shareable CUDA event so that the receiver
            # can wait for the sender's computations to complete.
            event = torch.cuda.Event(interprocess=True)
            torch.cuda.current_stream().record_event(event)

            device = gpu_object[0].device
            ray_gpu_idx = ray.get_gpu_ids()[device.index]
            ray_node_id = ray.get_runtime_context().get_node_id()


Review comment (suggestion):
This will check if cuda is available on the driver, not the actor. Either way, long term i'm not sure if this actor_has_tensor_transport should exist in the form it is now bc we can't have a ray.get on a .remote call

Code after:
    # assigned by Ray, e.g., using @ray.remote(num_gpus=1).
    ray_gpu_idx: Optional[int] = None
    # The node that the GPU that the tensors are on is on.
    ray_node_id: Optional[str] = None


class CudaIpcTransport(TensorTransportManager):
    def __init__(self):
        pass

    @property
    def tensor_transport_backend(self) -> str:
        return "CUDA_IPC"

    @staticmethod
    def is_one_sided() -> bool:
        return True

    @staticmethod
    def can_abort_transport() -> bool:
        return False

    def actor_has_tensor_transport(self, actor: "ray.actor.ActorHandle") -> bool:
        # TODO: Ideally we would check if torch.cuda.is_available() on the actor
        # and if so, return True. But we want to avoid blocking in ray.get() in
        # this method since it gets called before submitting an actor task.
        return True

    def extract_tensor_transport_metadata(
        self,
        obj_id: str,
        gpu_object: List["torch.Tensor"],
    ) -> CudaIpcTransportMetadata:

        tensor_meta = []
        device = None
        cuda_ipc_handles = []
        event_ipc_handle = None
        ray_gpu_idx = None
        ray_node_id = None
        if gpu_object:
            import torch
            from torch.multiprocessing.reductions import reduce_tensor

            device = gpu_object[0].device
            ray_gpu_idx = ray.get_gpu_ids()[device.index]
            ray_node_id = ray.get_runtime_context().get_node_id()

            # Create an interprocess-shareable CUDA event so that the receiver
            # can wait for the sender's computations to complete.
            event = torch.cuda.Event(interprocess=True)