chore: import upstream snapshot with attribution

python/ray/tests/rdt/test_rdt_ipc.py

@@ -0,0 +1,196 @@
+import sys
+
+import pytest
+import torch
+
+import ray
+
+
+@ray.remote(enable_tensor_transport=True)
+class GPUTestActor:
+    def __init__(self):
+        self.tensor = None
+
+    @ray.method(tensor_transport="cuda_ipc")
+    def echo(self, data):
+        self.tensor = data.to("cuda")
+        return self.tensor
+
+    def double(self, data):
+        data.mul_(2)
+        return data
+
+    def wait_tensor_freed(self):
+        rdt_manager = ray.worker.global_worker.rdt_manager
+        ray.experimental.wait_tensor_freed(self.tensor, timeout=10)
+        assert not rdt_manager.rdt_store.has_tensor(self.tensor)
+        return "freed"
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 1}], indirect=True)
+def test_colocated_actors(ray_start_regular):
+    world_size = 2
+    actors = [
+        GPUTestActor.options(num_gpus=0.5, num_cpus=0).remote()
+        for _ in range(world_size)
+    ]
+
+    src_actor, dst_actor = actors[0], actors[1]
+
+    # Create test tensor
+    tensor = torch.tensor([1, 2, 3])
+    rdt_ref = src_actor.echo.remote(tensor)
+
+    # Trigger tensor transfer from src to dst actor
+    ray.get(dst_actor.double.remote(rdt_ref))
+    # Check that the tensor is modified in place, and is reflected on the source actor
+    assert torch.equal(
+        ray.get(rdt_ref, _use_object_store=True),
+        torch.tensor([2, 4, 6], device="cuda"),
+    )
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True)
+def test_different_devices(ray_start_regular):
+    world_size = 2
+    actors = [
+        GPUTestActor.options(num_gpus=1, num_cpus=0).remote() for _ in range(world_size)
+    ]
+
+    src_actor, dst_actor = actors[0], actors[1]
+
+    # Create test tensor
+    tensor = torch.tensor([1, 2, 3])
+    rdt_ref = src_actor.echo.remote(tensor)
+
+    # Trigger tensor transfer from src to dst actor. Since CUDA IPC transport does not
+    # support cross-device tensor transfers, this should raise a ValueError.
+    with pytest.raises(
+        ValueError, match="CUDA IPC transport only supports tensors on the same GPU*"
+    ):
+        ray.get(dst_actor.double.remote(rdt_ref))
+
+
+def test_different_nodes(ray_start_cluster):
+    # Test that inter-node CUDA IPC transfers throw an error.
+    cluster = ray_start_cluster
+    num_nodes = 2
+    num_cpus = 1
+    num_gpus = 1
+    for _ in range(num_nodes):
+        cluster.add_node(num_cpus=num_cpus, num_gpus=num_gpus)
+    ray.init(address=cluster.address)
+
+    world_size = 2
+    actors = [
+        GPUTestActor.options(num_gpus=1, num_cpus=0).remote() for _ in range(world_size)
+    ]
+
+    src_actor, dst_actor = actors[0], actors[1]
+
+    # Create test tensor
+    tensor = torch.tensor([1, 2, 3])
+    rdt_ref = src_actor.echo.remote(tensor)
+
+    # Trigger tensor transfer from src to dst actor. Since CUDA IPC transport does not
+    # support cross-device tensor transfers, this should raise a ValueError.
+    with pytest.raises(
+        ValueError, match="CUDA IPC transport only supports tensors on the same node.*"
+    ):
+        ray.get(dst_actor.double.remote(rdt_ref))
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 1}], indirect=True)
+def test_ref_freed(ray_start_regular):
+    world_size = 2
+    actors = [
+        GPUTestActor.options(num_gpus=0.5, num_cpus=0).remote()
+        for _ in range(world_size)
+    ]
+
+    src_actor, dst_actor = actors[0], actors[1]
+
+    # Create test tensor
+    tensor = torch.tensor([1, 2, 3])
+    rdt_ref = src_actor.echo.remote(tensor)
+
+    # Trigger tensor transfer from src to dst actor
+    res_ref = dst_actor.double.remote(rdt_ref)
+
+    del rdt_ref
+
+    free_res = ray.get(src_actor.wait_tensor_freed.remote())
+    assert free_res == "freed"
+
+    assert torch.equal(
+        ray.get(res_ref, _use_object_store=True),
+        torch.tensor([2, 4, 6], device="cuda"),
+    )
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 1}], indirect=True)
+def test_source_actor_fails_after_transfer(ray_start_regular):
+    world_size = 2
+    actors = [
+        GPUTestActor.options(num_gpus=0.5, num_cpus=0).remote()
+        for _ in range(world_size)
+    ]
+
+    src_actor, dst_actor = actors[0], actors[1]
+
+    # Create test tensor
+    tensor = torch.tensor([1, 2, 3])
+    rdt_ref = src_actor.echo.remote(tensor)
+
+    # Trigger tensor transfer from src to dst actor
+    res_ref = dst_actor.double.remote(rdt_ref)
+    assert torch.equal(
+        ray.get(res_ref, _use_object_store=True),
+        torch.tensor([2, 4, 6], device="cuda"),
+    )
+
+    # Kill the source actor.
+    ray.kill(src_actor)
+    with pytest.raises(ray.exceptions.RayActorError):
+        ray.get(src_actor.wait_tensor_freed.remote())
+
+    # Check that the tensor is still available on the destination actor.
+    assert torch.equal(
+        ray.get(res_ref, _use_object_store=True),
+        torch.tensor([2, 4, 6], device="cuda"),
+    )
+
+
+@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 1}], indirect=True)
+def test_source_actor_fails_before_transfer(ray_start_regular):
+    world_size = 2
+    actors = [
+        GPUTestActor.options(num_gpus=0.5, num_cpus=0).remote()
+        for _ in range(world_size)
+    ]
+
+    src_actor, dst_actor = actors[0], actors[1]
+
+    # Create test tensor
+    tensor = torch.tensor([1, 2, 3])
+    rdt_ref = src_actor.echo.remote(tensor)
+
+    # Wait for object to be created.
+    assert torch.equal(
+        ray.get(rdt_ref, _use_object_store=True),
+        torch.tensor([1, 2, 3], device="cuda"),
+    )
+
+    # Kill the source actor.
+    ray.kill(src_actor)
+    with pytest.raises(ray.exceptions.RayActorError):
+        ray.get(src_actor.wait_tensor_freed.remote())
+
+    # Check that the tensor is still available on the destination actor.
+    with pytest.raises(ray.exceptions.RayTaskError):
+        res_ref = dst_actor.double.remote(rdt_ref)
+        ray.get(res_ref, _use_object_store=True)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))