chore: import upstream snapshot with attribution

python/ray/tests/rdt/test_rdt_custom.py

@@ -0,0 +1,156 @@
+import multiprocessing.shared_memory as shm
+import pickle
+import sys
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import numpy
+import pytest
+
+import ray
+from ray.experimental import (
+    CommunicatorMetadata,
+    TensorTransportManager,
+    TensorTransportMetadata,
+    register_tensor_transport,
+)
+
+
+@dataclass
+class ShmTransportMetadata(TensorTransportMetadata):
+    shm_name: Optional[str] = None
+    shm_size: Optional[int] = None
+
+
+@dataclass
+class ShmCommunicatorMetadata(CommunicatorMetadata):
+    pass
+
+
+class SharedMemoryTransport(TensorTransportManager):
+    def __init__(self):
+        self.shared_memory_objects: Dict[str, shm.SharedMemory] = {}
+
+    def tensor_transport_backend(self) -> str:
+        return "shared_memory"
+
+    @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 True
+
+    def extract_tensor_transport_metadata(
+        self,
+        obj_id: str,
+        rdt_object: List[numpy.ndarray],
+    ) -> TensorTransportMetadata:
+
+        tensor_meta = []
+        if rdt_object:
+            for tensor in rdt_object:
+                tensor_meta.append((tensor.shape, tensor.dtype))
+
+        serialized_rdt_object = pickle.dumps(rdt_object)
+        size = len(serialized_rdt_object)
+        # Shm name can't be as long as the obj_id, so we truncate it.
+        name = obj_id[:20]
+        shm_obj = shm.SharedMemory(name=name, create=True, size=size)
+        shm_obj.buf[:size] = serialized_rdt_object
+        self.shared_memory_objects[obj_id] = shm_obj
+
+        return ShmTransportMetadata(
+            tensor_meta=tensor_meta, tensor_device="cpu", shm_name=name, shm_size=size
+        )
+
+    def get_communicator_metadata(
+        self,
+        src_actor: "ray.actor.ActorHandle",
+        dst_actor: "ray.actor.ActorHandle",
+        backend: Optional[str] = None,
+    ) -> CommunicatorMetadata:
+        return ShmCommunicatorMetadata()
+
+    def recv_multiple_tensors(
+        self,
+        obj_id: str,
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+        target_buffers: Optional[List[Any]] = None,
+    ):
+        shm_name = tensor_transport_metadata.shm_name
+        size = tensor_transport_metadata.shm_size
+        shm_block = shm.SharedMemory(name=shm_name)
+        recv_tensors = pickle.loads(shm_block.buf[:size])
+        shm_block.close()
+        return recv_tensors
+
+    def send_multiple_tensors(
+        self,
+        tensors: List[numpy.ndarray],
+        tensor_transport_metadata: TensorTransportMetadata,
+        communicator_metadata: CommunicatorMetadata,
+    ):
+        pass
+
+    def garbage_collect(
+        self,
+        obj_id: str,
+        tensor_transport_meta: TensorTransportMetadata,
+        tensors: List[numpy.ndarray],
+    ):
+        self.shared_memory_objects[obj_id].close()
+        self.shared_memory_objects[obj_id].unlink()
+        del self.shared_memory_objects[obj_id]
+
+    def abort_transport(
+        self,
+        obj_id: str,
+        communicator_metadata: CommunicatorMetadata,
+    ):
+        pass
+
+
+def test_register_and_use_custom_transport(ray_start_regular):
+    register_tensor_transport(
+        "shared_memory", ["cpu"], SharedMemoryTransport, numpy.ndarray
+    )
+
+    @ray.remote
+    class Actor:
+        @ray.method(tensor_transport="shared_memory")
+        def echo(self, data):
+            return data
+
+        def non_rdt_echo(self, data):
+            return data
+
+        def sum(self, data):
+            return data.sum().item()
+
+    # Classes defined in test files get pickled by ref. So we need to
+    # explicitly pickle the transport class in this module by value.
+    # Note that this doesn't happen if you define the transport class on the
+    # driver, something with pytest convinces cloudpickle to pickle by ref.
+    from ray import cloudpickle
+
+    cloudpickle.register_pickle_by_value(sys.modules[SharedMemoryTransport.__module__])
+
+    actors = [Actor.remote() for _ in range(2)]
+    ref = actors[0].echo.remote(numpy.array([1, 2, 3]))
+    result = actors[1].sum.remote(ref)
+    assert ray.get(result) == 6
+
+    # Test that non-rdt methods that return the data type still work.
+    ref = actors[0].non_rdt_echo.remote(numpy.array([1, 2, 3]))
+    result = actors[1].sum.remote(ref)
+    assert ray.get(result) == 6
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))