wehub-resource-sync
1122 lines
37 KiB
Python
1122 lines
37 KiB
Python
import logging
|
|
import random
|
|
import re
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import ray
|
|
from ray._common.test_utils import SignalActor, wait_for_condition
|
|
from ray.experimental.collective import create_collective_group
|
|
from ray.experimental.rdt.collective_tensor_transport import (
|
|
CollectiveTransportMetadata,
|
|
)
|
|
|
|
# tensordict is not supported on macos ci, so we skip the tests
|
|
support_tensordict = sys.platform != "darwin"
|
|
|
|
if support_tensordict:
|
|
from tensordict import TensorDict
|
|
|
|
|
|
# TODO: check whether concurrency groups are created correctly if
|
|
# enable_tensor_transport is True or if any methods are decorated with
|
|
# @ray.method(tensor_transport=...). Check that specifying
|
|
# .options(tensor_transport=...) fails if enable_tensor_transport is False.
|
|
@ray.remote
|
|
class GPUTestActor:
|
|
@ray.method(tensor_transport="gloo")
|
|
def echo(self, data):
|
|
return data
|
|
|
|
def add(self, a, b):
|
|
return a + b
|
|
|
|
def double(self, data):
|
|
if isinstance(data, list):
|
|
return [self.double(d) for d in data]
|
|
if support_tensordict and isinstance(data, TensorDict):
|
|
return data.apply(lambda x: x * 2)
|
|
return data * 2
|
|
|
|
def increment(self, data):
|
|
data += 1
|
|
return data
|
|
|
|
def get_out_of_band_tensors(self, obj_id: str, timeout=None):
|
|
rdt_store = ray._private.worker.global_worker.rdt_manager.rdt_store
|
|
if timeout is None:
|
|
timeout = 0
|
|
return rdt_store.wait_and_get_object(obj_id, timeout)
|
|
|
|
def get_num_rdt_objects(self):
|
|
rdt_manager = ray._private.worker.global_worker.rdt_manager
|
|
return rdt_manager.rdt_store.get_num_objects()
|
|
|
|
def fail(self, error_message):
|
|
raise Exception(error_message)
|
|
|
|
|
|
@ray.remote
|
|
class ErrorActor:
|
|
@ray.method(tensor_transport="gloo")
|
|
def send(self, tensor):
|
|
return tensor
|
|
|
|
def recv(self, tensor):
|
|
return tensor
|
|
|
|
def clear_rdt_store(self):
|
|
rdt_store = ray._private.worker.global_worker.rdt_manager.rdt_store
|
|
|
|
with rdt_store._lock:
|
|
assert len(rdt_store._rdt_store) > 0
|
|
rdt_store._rdt_store.clear()
|
|
|
|
@ray.method(concurrency_group="_ray_system")
|
|
def block_background_thread(self):
|
|
time.sleep(100)
|
|
|
|
def block_main_thread(self):
|
|
time.sleep(100)
|
|
|
|
|
|
@pytest.mark.parametrize("data_size_bytes", [100])
|
|
def test_gc_rdt_object(ray_start_regular, data_size_bytes):
|
|
"""
|
|
For small data, GPU objects are inlined, but the actual data lives
|
|
on the remote actor. Therefore, if we decrement the reference count
|
|
upon inlining, we may cause the tensors on the sender actor to be
|
|
freed before transferring to the receiver actor.
|
|
|
|
# TODO(kevin85421): Add a test for large CPU data that is not inlined
|
|
# after https://github.com/ray-project/ray/issues/54281 is fixed.
|
|
"""
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
small_tensor = torch.randn((1,))
|
|
cpu_data = b"1" * data_size_bytes
|
|
data = [small_tensor, cpu_data]
|
|
sender = actors[0]
|
|
receiver = actors[1]
|
|
|
|
ref1 = sender.echo.remote(data)
|
|
ref2 = receiver.double.remote(ref1)
|
|
ref3 = receiver.double.remote(ref1)
|
|
|
|
result = ray.get(ref2)
|
|
assert result[0] == pytest.approx(small_tensor * 2)
|
|
assert result[1] == cpu_data * 2
|
|
result = ray.get(ref3)
|
|
assert result[0] == pytest.approx(small_tensor * 2)
|
|
assert result[1] == cpu_data * 2
|
|
|
|
wait_for_condition(
|
|
lambda: ray.get(receiver.get_num_rdt_objects.remote()) == 0,
|
|
timeout=10,
|
|
retry_interval_ms=100,
|
|
)
|
|
|
|
del ref1
|
|
|
|
wait_for_condition(
|
|
lambda: ray.get(sender.get_num_rdt_objects.remote()) == 0,
|
|
timeout=10,
|
|
retry_interval_ms=100,
|
|
)
|
|
|
|
|
|
def test_gc_rdt_metadata(ray_start_regular):
|
|
actors = [GPUTestActor.remote() for _ in range(2)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
tensor = torch.randn((100, 100))
|
|
ref = actors[0].echo.remote(tensor)
|
|
rdt_ref_id = ref.hex()
|
|
rdt_manager = ray._private.worker.global_worker.rdt_manager
|
|
assert rdt_manager.is_managed_object(rdt_ref_id)
|
|
ray.get(actors[1].double.remote(ref))
|
|
del ref
|
|
|
|
wait_for_condition(
|
|
lambda: not rdt_manager.is_managed_object(rdt_ref_id),
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("data_size_bytes", [100])
|
|
def test_gc_del_ref_before_recv_finish(ray_start_regular, data_size_bytes):
|
|
"""
|
|
This test deletes the ObjectRef of the GPU object before calling
|
|
`ray.get` to ensure the receiver finishes receiving the GPU object.
|
|
"""
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
small_tensor = torch.randn((1,))
|
|
cpu_data = b"1" * data_size_bytes
|
|
data = [small_tensor, cpu_data]
|
|
sender = actors[0]
|
|
receiver = actors[1]
|
|
|
|
ref1 = sender.echo.remote(data)
|
|
ref2 = receiver.double.remote(ref1)
|
|
|
|
del ref1
|
|
|
|
result = ray.get(ref2)
|
|
assert result[0] == pytest.approx(small_tensor * 2)
|
|
assert result[1] == cpu_data * 2
|
|
|
|
wait_for_condition(
|
|
lambda: ray.get(receiver.get_num_rdt_objects.remote()) == 0,
|
|
timeout=10,
|
|
retry_interval_ms=100,
|
|
)
|
|
wait_for_condition(
|
|
lambda: ray.get(sender.get_num_rdt_objects.remote()) == 0,
|
|
timeout=10,
|
|
retry_interval_ms=100,
|
|
)
|
|
|
|
|
|
def test_gc_intra_actor_rdt_object(ray_start_regular):
|
|
"""
|
|
This test checks that passes a GPU object ref to the same actor multiple times.
|
|
"""
|
|
actor = GPUTestActor.remote()
|
|
create_collective_group([actor], backend="gloo")
|
|
|
|
small_tensor = torch.randn((1,))
|
|
|
|
ref = actor.echo.remote(small_tensor)
|
|
result = actor.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
result = actor.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
del ref
|
|
|
|
wait_for_condition(
|
|
lambda: ray.get(actor.get_num_rdt_objects.remote()) == 0,
|
|
timeout=10,
|
|
retry_interval_ms=100,
|
|
)
|
|
|
|
|
|
def test_gc_pass_ref_to_same_and_different_actors(ray_start_regular):
|
|
"""
|
|
This test checks that passes a GPU object ref to the same actor and a different actor.
|
|
"""
|
|
actor1 = GPUTestActor.remote()
|
|
actor2 = GPUTestActor.remote()
|
|
create_collective_group([actor1, actor2], backend="gloo")
|
|
|
|
small_tensor = torch.randn((1,))
|
|
|
|
ref = actor1.echo.remote(small_tensor)
|
|
result1 = actor1.double.remote(ref)
|
|
result2 = actor2.double.remote(ref)
|
|
assert ray.get(result1) == pytest.approx(small_tensor * 2)
|
|
assert ray.get(result2) == pytest.approx(small_tensor * 2)
|
|
|
|
wait_for_condition(
|
|
lambda: ray.get(actor2.get_num_rdt_objects.remote()) == 0,
|
|
timeout=10,
|
|
retry_interval_ms=100,
|
|
)
|
|
|
|
del ref
|
|
|
|
wait_for_condition(
|
|
lambda: ray.get(actor1.get_num_rdt_objects.remote()) == 0,
|
|
timeout=10,
|
|
retry_interval_ms=100,
|
|
)
|
|
|
|
|
|
def test_p2p(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
small_tensor = torch.randn((1,))
|
|
sender = actors[0]
|
|
receiver = actors[1]
|
|
|
|
ref = sender.echo.remote(small_tensor)
|
|
result = receiver.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
medium_tensor = torch.randn((500, 500))
|
|
ref = sender.echo.remote(medium_tensor)
|
|
result = receiver.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx(medium_tensor * 2)
|
|
|
|
|
|
def test_p2p_errors_before_group_creation(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
|
|
small_tensor = torch.randn((1,))
|
|
sender = actors[0]
|
|
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="Actor.* does not have tensor transport GLOO available.*",
|
|
):
|
|
sender.echo.remote(small_tensor)
|
|
|
|
|
|
@pytest.mark.parametrize("has_tensor_transport_method", [True, False])
|
|
def test_p2p_blocking(ray_start_regular, has_tensor_transport_method):
|
|
"""Test that p2p transfers still work when sender is blocked in another
|
|
task. This should work whether the actor has (a) a tensor transport method
|
|
(a method decorated with @ray.method(tensor_transport=...)) or (b) an actor-level decorator
|
|
@ray.remote(enable_tensor_transport=True)."""
|
|
|
|
class _GPUTestActor:
|
|
def double(self, data):
|
|
if isinstance(data, list):
|
|
return [self.double(d) for d in data]
|
|
if support_tensordict and isinstance(data, TensorDict):
|
|
return data.apply(lambda x: x * 2)
|
|
return data * 2
|
|
|
|
def infinite_sleep(self, signal):
|
|
signal.send.remote()
|
|
while True:
|
|
time.sleep(0.1)
|
|
|
|
if has_tensor_transport_method:
|
|
# Test tensor transport annotation via ray.method.
|
|
@ray.remote
|
|
class GPUTestActor(_GPUTestActor):
|
|
@ray.method(tensor_transport="gloo")
|
|
def echo(self, data):
|
|
return data
|
|
|
|
else:
|
|
# Test tensor transport annotation via ray.remote.
|
|
@ray.remote(enable_tensor_transport=True)
|
|
class GPUTestActor(_GPUTestActor):
|
|
def echo(self, data):
|
|
return data
|
|
|
|
sender, receiver = GPUTestActor.remote(), GPUTestActor.remote()
|
|
signal = SignalActor.remote()
|
|
create_collective_group([sender, receiver], backend="gloo")
|
|
tensor = torch.randn((500, 500))
|
|
# If the actor does not have a tensor transport method declared, declare it
|
|
# dynamically using .options().
|
|
sender_fn = (
|
|
sender.echo
|
|
if has_tensor_transport_method
|
|
else sender.echo.options(tensor_transport="gloo")
|
|
)
|
|
ref = sender_fn.remote(tensor)
|
|
|
|
# Start a blocking task on the sender actor.
|
|
sender.infinite_sleep.remote(signal)
|
|
ray.get(signal.wait.remote(), timeout=10)
|
|
|
|
# Ensure that others can still receive the object.
|
|
result = receiver.double.remote(ref)
|
|
result = ray.get(result, timeout=10)
|
|
assert result == pytest.approx(tensor * 2)
|
|
|
|
|
|
def test_p2p_with_cpu_data(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
sender = actors[0]
|
|
receiver = actors[1]
|
|
|
|
cpu_data = 123
|
|
ref = sender.echo.remote(cpu_data)
|
|
result = receiver.double.remote(ref)
|
|
assert ray.get(result) == cpu_data * 2
|
|
|
|
|
|
def test_send_same_ref_to_same_actor_task_multiple_times(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
small_tensor = torch.randn((1,))
|
|
sender = actors[0]
|
|
receiver = actors[1]
|
|
|
|
ref = sender.echo.remote(small_tensor)
|
|
result = receiver.add.remote(ref, ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
wait_for_condition(
|
|
lambda: ray.get(receiver.get_num_rdt_objects.remote()) == 0,
|
|
timeout=10,
|
|
retry_interval_ms=100,
|
|
)
|
|
|
|
|
|
def test_send_same_ref_to_same_actor_multiple_times(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
small_tensor = torch.randn((1,))
|
|
sender = actors[0]
|
|
receiver = actors[1]
|
|
|
|
ref = sender.echo.remote(small_tensor)
|
|
result = receiver.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
result = receiver.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
|
|
def test_intra_rdt_tensor_transfer(ray_start_regular):
|
|
actor = GPUTestActor.remote()
|
|
create_collective_group([actor], backend="gloo")
|
|
|
|
small_tensor = torch.randn((1,))
|
|
|
|
# Intra-actor communication for pure GPU tensors
|
|
ref = actor.echo.remote(small_tensor)
|
|
result = actor.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
# Intra-actor communication for mixed CPU and GPU data
|
|
cpu_data = random.randint(0, 100)
|
|
data = [small_tensor, cpu_data]
|
|
ref = actor.echo.remote(data)
|
|
result = actor.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx([small_tensor * 2, cpu_data * 2])
|
|
|
|
# Intra-actor communication for multiple GPU tensors
|
|
tensor1 = torch.randn((1,))
|
|
tensor2 = torch.randn((2,))
|
|
data = [tensor1, tensor2, cpu_data]
|
|
ref = actor.echo.remote(data)
|
|
result = actor.double.remote(ref)
|
|
result = ray.get(result)
|
|
|
|
assert result[0] == pytest.approx(tensor1 * 2)
|
|
assert result[1] == pytest.approx(tensor2 * 2)
|
|
assert result[2] == cpu_data * 2
|
|
|
|
|
|
def test_send_same_ref_multiple_times_intra_actor(ray_start_regular):
|
|
actor = GPUTestActor.remote()
|
|
create_collective_group([actor], backend="gloo")
|
|
|
|
small_tensor = torch.randn((1,))
|
|
|
|
ref = actor.echo.remote(small_tensor)
|
|
result = actor.add.remote(ref, ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
|
|
def test_mix_cpu_gpu_data(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
tensor = torch.randn((1,))
|
|
cpu_data = random.randint(0, 100)
|
|
|
|
data = [tensor, cpu_data]
|
|
|
|
sender, receiver = actors[0], actors[1]
|
|
ref = sender.echo.remote(data)
|
|
ref = receiver.double.remote(ref)
|
|
result = ray.get(ref)
|
|
|
|
assert result[0] == pytest.approx(tensor * 2)
|
|
assert result[1] == cpu_data * 2
|
|
|
|
|
|
def test_object_in_plasma(ray_start_regular):
|
|
"""
|
|
This test uses a CPU object that is large enough to be stored
|
|
in plasma instead of being inlined in the gRPC message.
|
|
"""
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
tensor = torch.randn((1,))
|
|
cpu_data = b"1" * 1000 * 1000
|
|
data = [tensor, cpu_data]
|
|
|
|
sender, receiver = actors[0], actors[1]
|
|
ref = sender.echo.remote(data)
|
|
ref = receiver.double.remote(ref)
|
|
result = ray.get(ref)
|
|
|
|
assert result[0] == pytest.approx(tensor * 2)
|
|
assert result[1] == cpu_data * 2
|
|
|
|
|
|
def test_multiple_tensors(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
tensor1 = torch.randn((1,))
|
|
tensor2 = torch.randn((2,))
|
|
if support_tensordict:
|
|
td1 = TensorDict(
|
|
{"action1": torch.randn((2,)), "reward1": torch.randn((2,))}, batch_size=[2]
|
|
)
|
|
td2 = TensorDict(
|
|
{"action2": torch.randn((2,)), "reward2": torch.randn((2,))}, batch_size=[2]
|
|
)
|
|
else:
|
|
td1 = 0
|
|
td2 = 0
|
|
cpu_data = random.randint(0, 100)
|
|
data = [tensor1, tensor2, cpu_data, td1, td2]
|
|
|
|
sender, receiver = actors[0], actors[1]
|
|
ref = sender.echo.remote(data)
|
|
ref = receiver.double.remote(ref)
|
|
result = ray.get(ref)
|
|
|
|
assert result[0] == pytest.approx(tensor1 * 2)
|
|
assert result[1] == pytest.approx(tensor2 * 2)
|
|
assert result[2] == cpu_data * 2
|
|
if support_tensordict:
|
|
assert result[3]["action1"] == pytest.approx(td1["action1"] * 2)
|
|
assert result[3]["reward1"] == pytest.approx(td1["reward1"] * 2)
|
|
assert result[4]["action2"] == pytest.approx(td2["action2"] * 2)
|
|
assert result[4]["reward2"] == pytest.approx(td2["reward2"] * 2)
|
|
|
|
|
|
def test_trigger_out_of_band_tensor_transfer(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
src_actor, dst_actor = actors[0], actors[1]
|
|
|
|
tensor = torch.tensor([1, 2, 3])
|
|
rdt_ref = src_actor.echo.remote(tensor)
|
|
rdt_ref_id = rdt_ref.hex()
|
|
|
|
# Check src_actor has the GPU object
|
|
ret_val_src = ray.get(src_actor.get_out_of_band_tensors.remote(rdt_ref_id))
|
|
assert ret_val_src is not None
|
|
assert len(ret_val_src) == 1
|
|
assert torch.equal(ret_val_src[0], tensor)
|
|
|
|
rdt_manager = ray._private.worker.global_worker.rdt_manager
|
|
rdt_manager.add_rdt_ref(rdt_ref, src_actor, "GLOO")
|
|
|
|
# Trigger out-of-band tensor transfer from src_actor to dst_actor.
|
|
task_args = (rdt_ref,)
|
|
rdt_manager.queue_or_trigger_out_of_band_tensor_transfer(dst_actor, task_args)
|
|
|
|
rdt_manager.set_tensor_transport_metadata_and_trigger_queued_operations(
|
|
rdt_ref_id,
|
|
CollectiveTransportMetadata(
|
|
tensor_meta=[(tensor.shape, tensor.dtype)],
|
|
tensor_device=tensor.device.type,
|
|
),
|
|
)
|
|
|
|
# Check dst_actor has the GPU object
|
|
ret_val_dst = ray.get(
|
|
dst_actor.get_out_of_band_tensors.remote(rdt_ref_id, timeout=10)
|
|
)
|
|
assert ret_val_dst is not None
|
|
assert len(ret_val_dst) == 1
|
|
assert torch.equal(ret_val_dst[0], tensor)
|
|
|
|
|
|
def test_fetch_rdt_object_to_driver(ray_start_regular):
|
|
actor = GPUTestActor.remote()
|
|
create_collective_group([actor], backend="gloo")
|
|
|
|
tensor1 = torch.tensor([1, 2, 3])
|
|
tensor2 = torch.tensor([4, 5, 6])
|
|
|
|
# Case 1: Single tensor
|
|
ref = actor.echo.remote(tensor1)
|
|
assert torch.equal(ray.get(ref, _use_object_store=True), tensor1)
|
|
|
|
# Case 2: Multiple tensors
|
|
ref = actor.echo.remote([tensor1, tensor2])
|
|
result = ray.get(ref, _use_object_store=True)
|
|
assert torch.equal(result[0], tensor1)
|
|
assert torch.equal(result[1], tensor2)
|
|
|
|
# Case 3: Mixed CPU and GPU data
|
|
data = [tensor1, tensor2, 7]
|
|
ref = actor.echo.remote(data)
|
|
result = ray.get(ref, _use_object_store=True)
|
|
assert torch.equal(result[0], tensor1)
|
|
assert torch.equal(result[1], tensor2)
|
|
assert result[2] == 7
|
|
|
|
|
|
def test_invalid_tensor_transport(ray_start_regular):
|
|
with pytest.raises(ValueError, match="Invalid tensor transport"):
|
|
|
|
@ray.remote
|
|
class InvalidActor:
|
|
@ray.method(tensor_transport="invalid")
|
|
def echo(self, data):
|
|
return data
|
|
|
|
actor = GPUTestActor.remote()
|
|
with pytest.raises(ValueError, match="Invalid tensor transport"):
|
|
actor.double.options(tensor_transport="invalid").remote(torch.randn((1,)))
|
|
|
|
with pytest.raises(ValueError, match="Invalid tensor transport"):
|
|
ray.put(torch.randn((1,)), _tensor_transport="invalid")
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not support_tensordict,
|
|
reason="tensordict is not supported on this platform",
|
|
)
|
|
def test_tensordict_transfer(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
td = TensorDict(
|
|
{"action": torch.randn((2,)), "reward": torch.randn((2,))}, batch_size=[2]
|
|
)
|
|
sender, receiver = actors[0], actors[1]
|
|
ref = sender.echo.remote(td)
|
|
result = receiver.double.remote(ref)
|
|
td_result = ray.get(result)
|
|
|
|
assert td_result["action"] == pytest.approx(td["action"] * 2)
|
|
assert td_result["reward"] == pytest.approx(td["reward"] * 2)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not support_tensordict,
|
|
reason="tensordict is not supported on this platform",
|
|
)
|
|
def test_nested_tensordict(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
inner_td = TensorDict(
|
|
{"action": torch.randn((2,)), "reward": torch.randn((2,))}, batch_size=[2]
|
|
)
|
|
outer_td = TensorDict(
|
|
{"inner_td": inner_td, "test": torch.randn((2,))}, batch_size=[2]
|
|
)
|
|
sender = actors[0]
|
|
receiver = actors[1]
|
|
rdt_ref = sender.echo.remote(outer_td)
|
|
ret_val_src = ray.get(receiver.double.remote(rdt_ref))
|
|
assert ret_val_src is not None
|
|
assert torch.equal(ret_val_src["inner_td"]["action"], inner_td["action"] * 2)
|
|
assert torch.equal(ret_val_src["inner_td"]["reward"], inner_td["reward"] * 2)
|
|
assert torch.equal(ret_val_src["test"], outer_td["test"] * 2)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not support_tensordict,
|
|
reason="tensordict is not supported on this platform",
|
|
)
|
|
def test_tensor_extracted_from_tensordict_in_rdt_store(ray_start_regular):
|
|
actor = GPUTestActor.remote()
|
|
create_collective_group([actor], backend="gloo")
|
|
|
|
td = TensorDict(
|
|
{"action": torch.randn((2,)), "reward": torch.randn((2,))}, batch_size=[2]
|
|
).to("cpu")
|
|
rdt_ref = actor.echo.remote(td)
|
|
|
|
# Since the tensor is extracted from the tensordict, the `ret_val_src` will be a list of tensors
|
|
# instead of a tensordict.
|
|
ret_val_src = ray.get(actor.get_out_of_band_tensors.remote(rdt_ref.hex()))
|
|
assert ret_val_src is not None
|
|
assert len(ret_val_src) == 2
|
|
assert torch.equal(ret_val_src[0], td["action"])
|
|
assert torch.equal(ret_val_src[1], td["reward"])
|
|
|
|
|
|
@pytest.mark.parametrize("enable_tensor_transport", [True, False])
|
|
def test_dynamic_tensor_transport_via_options(
|
|
ray_start_regular, enable_tensor_transport
|
|
):
|
|
"""Test that tensor_transport can be set dynamically via .options() at call
|
|
time, if enable_tensor_transport is set to True in @ray.remote."""
|
|
|
|
class TestActor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def normal_method(self):
|
|
return "normal"
|
|
|
|
def tensor_method(self):
|
|
return torch.randn(5, 5)
|
|
|
|
def double(self, data):
|
|
return data * 2
|
|
|
|
if enable_tensor_transport:
|
|
TestActor = ray.remote(enable_tensor_transport=True)(TestActor)
|
|
else:
|
|
TestActor = ray.remote(TestActor)
|
|
|
|
# Create actor without any tensor_transport decorators
|
|
sender = TestActor.remote()
|
|
receiver = TestActor.remote()
|
|
create_collective_group([sender, receiver], backend="gloo")
|
|
|
|
# Test normal method call
|
|
result = ray.get(sender.normal_method.remote())
|
|
assert result == "normal"
|
|
|
|
# Test method call with tensor_transport specified via .options()
|
|
if enable_tensor_transport:
|
|
# If enable_tensor_transport is set to True, then it's okay to use
|
|
# dynamic tensor_transport.
|
|
ref = sender.tensor_method.options(tensor_transport="gloo").remote()
|
|
tensor = ray.get(ref, _use_object_store=True)
|
|
result = ray.get(receiver.double.remote(ref), _use_object_store=True)
|
|
assert result == pytest.approx(tensor * 2)
|
|
else:
|
|
# If enable_tensor_transport is not set, then user cannot use
|
|
# dynamic tensor_transport.
|
|
with pytest.raises(
|
|
ValueError,
|
|
match='Currently, methods with .options\\(tensor_transport="GLOO"\\) are not supported when enable_tensor_transport=False. Please set @ray.remote\\(enable_tensor_transport=True\\) on the actor class definition.',
|
|
):
|
|
ref = sender.tensor_method.options(tensor_transport="gloo").remote()
|
|
|
|
|
|
def test_app_error_inter_actor(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
src_actor, dst_actor = actors[0], actors[1]
|
|
|
|
# Make sure the receiver can receive an exception from the sender.
|
|
ref = src_actor.fail.options(tensor_transport="gloo").remote("test_app_error")
|
|
with pytest.raises(Exception, match="test_app_error"):
|
|
ray.get(dst_actor.double.remote(ref))
|
|
|
|
# Make sure the sender and receiver do not hang.
|
|
small_tensor = torch.randn((1,))
|
|
ref = src_actor.echo.remote(small_tensor)
|
|
result = dst_actor.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
|
|
def test_app_error_intra_actor(ray_start_regular):
|
|
actor = GPUTestActor.remote()
|
|
create_collective_group([actor], backend="gloo")
|
|
|
|
# Make sure the receiver can receive an exception from the sender.
|
|
ref = actor.fail.options(tensor_transport="gloo").remote("test_app_error")
|
|
with pytest.raises(Exception, match="test_app_error"):
|
|
ray.get(actor.double.remote(ref))
|
|
|
|
# Make sure the sender and receiver do not hang.
|
|
small_tensor = torch.randn((1,))
|
|
ref = actor.echo.remote(small_tensor)
|
|
result = actor.double.remote(ref)
|
|
assert ray.get(result) == pytest.approx(small_tensor * 2)
|
|
|
|
|
|
def test_app_error_fetch_to_driver(ray_start_regular):
|
|
actor = GPUTestActor.remote()
|
|
create_collective_group([actor], backend="gloo")
|
|
|
|
ref = actor.fail.options(tensor_transport="gloo").remote("test_app_error")
|
|
with pytest.raises(Exception, match="test_app_error"):
|
|
ray.get(ref, _use_object_store=True)
|
|
|
|
# Make sure the driver can receive an exception from the actor.
|
|
small_tensor = torch.tensor([1, 2, 3])
|
|
ref = actor.echo.remote(small_tensor)
|
|
assert torch.equal(ray.get(ref, _use_object_store=True), small_tensor)
|
|
|
|
|
|
@ray.remote
|
|
class FailingRDTActor:
|
|
def __init__(self):
|
|
self.attempts = 0
|
|
|
|
@ray.method(
|
|
tensor_transport="gloo", max_task_retries=1, retry_exceptions=[ValueError]
|
|
)
|
|
def fail_first_attempt(self):
|
|
self.attempts += 1
|
|
if self.attempts == 1:
|
|
raise ValueError("first-attempt failure")
|
|
return torch.tensor([1, 2, 3])
|
|
|
|
@ray.method(
|
|
tensor_transport="gloo", max_task_retries=1, retry_exceptions=[ValueError]
|
|
)
|
|
def rdt_obj_always_fails(self):
|
|
self.attempts += 1
|
|
raise ValueError("permanent failure")
|
|
|
|
def consume(self, tensor):
|
|
return tensor
|
|
|
|
def get_num_rdt_objects(self):
|
|
return ray._private.worker.global_worker.rdt_manager.rdt_store.get_num_objects()
|
|
|
|
|
|
def test_rdt_retry_then_succeeds(ray_start_regular):
|
|
"""
|
|
Retryable exception on first attempt, success on second
|
|
Only one entry should be in the RDTStore.
|
|
"""
|
|
sender = FailingRDTActor.remote()
|
|
receiver = FailingRDTActor.remote()
|
|
create_collective_group([sender, receiver], backend="gloo")
|
|
|
|
ref = sender.fail_first_attempt.remote()
|
|
result = ray.get(receiver.consume.remote(ref))
|
|
assert torch.equal(result, torch.tensor([1, 2, 3]))
|
|
|
|
# Sender should hold one primary entry for this ref
|
|
assert ray.get(sender.get_num_rdt_objects.remote()) == 1
|
|
|
|
|
|
def test_rdt_retry_fetch_through_obj_store(ray_start_regular):
|
|
"""
|
|
Retryable exception on first attempt, successful fetch to driver on second
|
|
"""
|
|
sender = FailingRDTActor.remote()
|
|
create_collective_group([sender], backend="gloo")
|
|
|
|
ref = sender.fail_first_attempt.remote()
|
|
assert torch.equal(ray.get(ref, _use_object_store=True), torch.tensor([1, 2, 3]))
|
|
|
|
|
|
def test_rdt_retries_exhausted_raises(ray_start_regular):
|
|
"""
|
|
When all retries fail, the user's exception must propagate to the
|
|
consumer via the CPU path (no direct_transport_metadata is set on the
|
|
final reply, so the consumer sees the error when deserializing the arg).
|
|
"""
|
|
sender = FailingRDTActor.remote()
|
|
receiver = FailingRDTActor.remote()
|
|
create_collective_group([sender, receiver], backend="gloo")
|
|
|
|
ref = sender.rdt_obj_always_fails.remote()
|
|
with pytest.raises(Exception, match="permanent failure"):
|
|
ray.get(receiver.consume.remote(ref))
|
|
|
|
|
|
def test_write_after_save(ray_start_regular):
|
|
"""Check that an actor can safely write to a tensor after saving it to its
|
|
local state by calling `ray.experimental.wait_tensor_freed`."""
|
|
|
|
@ray.remote(enable_tensor_transport=True)
|
|
class GPUTestActor:
|
|
@ray.method(tensor_transport="gloo")
|
|
def save(self, data: torch.Tensor):
|
|
# Save the tensor to the actor's local state.
|
|
self.data = data
|
|
return data
|
|
|
|
def receive(self, data: torch.Tensor):
|
|
return data
|
|
|
|
def increment_saved(self):
|
|
ray.experimental.wait_tensor_freed(self.data)
|
|
# Write to the saved tensor.
|
|
self.data += 1
|
|
return self.data
|
|
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
medium_tensor = torch.randn((500, 500))
|
|
sender, receiver = actors
|
|
ref = sender.save.remote(medium_tensor)
|
|
# Sender writes to the GPU object while Ray sends the object to a receiver
|
|
# task in the background.
|
|
tensor1 = sender.increment_saved.remote()
|
|
tensor2 = receiver.receive.remote(ref)
|
|
|
|
# The sender task should not have returned yet because the ObjectRef is
|
|
# still in scope.
|
|
with pytest.raises(ray.exceptions.GetTimeoutError):
|
|
ray.get(tensor1, timeout=1)
|
|
|
|
del ref
|
|
# Check that Ray completed the transfer of the original tensor before the
|
|
# sender writes to it.
|
|
assert torch.allclose(ray.get(tensor1), medium_tensor + 1)
|
|
assert torch.allclose(ray.get(tensor2), medium_tensor)
|
|
|
|
|
|
def test_wait_tensor_freed(ray_start_regular):
|
|
"""Unit test for ray.experimental.wait_tensor_freed. Check that the call
|
|
returns when the tensor has been freed from the GPU object store."""
|
|
rdt_store = ray.worker.global_worker.rdt_manager.rdt_store
|
|
obj_id = "random_id"
|
|
tensor = torch.randn((1,))
|
|
rdt_store.add_object_primary(obj_id, [tensor], "GLOO")
|
|
|
|
assert rdt_store.has_object(obj_id)
|
|
with pytest.raises(TimeoutError):
|
|
ray.experimental.wait_tensor_freed(tensor, timeout=1)
|
|
assert rdt_store.has_object(obj_id)
|
|
|
|
# Simulate garbage collection in a background thread.
|
|
def gc():
|
|
time.sleep(0.1)
|
|
rdt_store.pop_object(obj_id)
|
|
|
|
gc_thread = threading.Thread(target=gc)
|
|
gc_thread.start()
|
|
# Now the wait_tensor_freed call should be able to return.
|
|
ray.experimental.wait_tensor_freed(tensor)
|
|
gc_thread.join()
|
|
assert not rdt_store.has_object(obj_id)
|
|
|
|
|
|
def test_wait_tensor_freed_double_tensor(ray_start_regular):
|
|
"""Unit test for ray.experimental.wait_tensor_freed when multiple objects
|
|
contain the same tensor."""
|
|
rdt_store = ray.worker.global_worker.rdt_manager.rdt_store
|
|
obj_id1 = "random_id1"
|
|
obj_id2 = "random_id2"
|
|
tensor = torch.randn((1,))
|
|
rdt_store.add_object_primary(obj_id1, [tensor], "GLOO")
|
|
rdt_store.add_object_primary(obj_id2, [tensor], "GLOO")
|
|
|
|
assert rdt_store.has_object(obj_id1)
|
|
assert rdt_store.has_object(obj_id2)
|
|
with pytest.raises(TimeoutError):
|
|
ray.experimental.wait_tensor_freed(tensor, timeout=1)
|
|
assert rdt_store.has_object(obj_id1)
|
|
assert rdt_store.has_object(obj_id2)
|
|
|
|
# Simulate garbage collection in a background thread.
|
|
def gc(obj_id):
|
|
time.sleep(0.1)
|
|
rdt_store.pop_object(obj_id)
|
|
|
|
# Free one object. Tensor should still be stored.
|
|
gc_thread = threading.Thread(target=gc, args=(obj_id1,))
|
|
gc_thread.start()
|
|
with pytest.raises(TimeoutError):
|
|
ray.experimental.wait_tensor_freed(tensor, timeout=1)
|
|
gc_thread.join()
|
|
assert not rdt_store.has_object(obj_id1)
|
|
|
|
# Free the other object. Now the wait_tensor_freed call should be able to
|
|
# return.
|
|
gc_thread = threading.Thread(target=gc, args=(obj_id2,))
|
|
gc_thread.start()
|
|
ray.experimental.wait_tensor_freed(tensor)
|
|
gc_thread.join()
|
|
assert not rdt_store.has_object(obj_id2)
|
|
|
|
|
|
def test_send_back_and_dst_warning(ray_start_regular):
|
|
# Test warning when object is sent back to the src actor and to dst actors
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
|
|
src_actor, dst_actor = actors[0], actors[1]
|
|
|
|
tensor = torch.tensor([1, 2, 3])
|
|
|
|
warning_message = r"RDT ObjectRef\(.+\)"
|
|
|
|
with pytest.warns(UserWarning, match=warning_message):
|
|
t = src_actor.echo.remote(tensor)
|
|
t1 = src_actor.echo.remote(t) # Sent back to the source actor
|
|
t2 = dst_actor.echo.remote(t) # Also sent to another actor
|
|
ray.get([t1, t2], _use_object_store=True)
|
|
|
|
# Second transmission of ObjectRef `t` to `dst_actor` should not trigger a warning
|
|
# Verify no `pytest.warns` context is used here because no warning should be raised
|
|
t3 = dst_actor.echo.remote(t)
|
|
ray.get(t3, _use_object_store=True)
|
|
|
|
|
|
def test_duplicate_objectref_transfer(ray_start_regular):
|
|
world_size = 2
|
|
actors = [GPUTestActor.remote() for _ in range(world_size)]
|
|
create_collective_group(actors, backend="gloo")
|
|
actor0, actor1 = actors[0], actors[1]
|
|
|
|
small_tensor = torch.randn((1,))
|
|
|
|
# Store the original value for comparison
|
|
original_value = small_tensor
|
|
|
|
ref = actor0.echo.remote(small_tensor)
|
|
|
|
# Pass the same ref to actor1 twice
|
|
result1 = actor1.increment.remote(ref)
|
|
result2 = actor1.increment.remote(ref)
|
|
|
|
# Both should return original_value + 1 because each increment task should receive the same object value.
|
|
val1 = ray.get(result1)
|
|
val2 = ray.get(result2)
|
|
|
|
# Check for correctness
|
|
assert val1 == pytest.approx(
|
|
original_value + 1
|
|
), f"Result1 incorrect: got {val1}, expected {original_value + 1}"
|
|
assert val2 == pytest.approx(
|
|
original_value + 1
|
|
), f"Result2 incorrect: got {val2}, expected {original_value + 1}"
|
|
|
|
# Additional check: results should be equal (both got clean copies)
|
|
assert val1 == pytest.approx(
|
|
val2
|
|
), f"Results differ: result1={val1}, result2={val2}"
|
|
|
|
|
|
def test_transfer_from_not_actor_creator(ray_start_regular):
|
|
@ray.remote
|
|
class Actor:
|
|
@ray.method(tensor_transport="gloo")
|
|
def create(self):
|
|
return torch.tensor([1, 2, 3])
|
|
|
|
def consume(self, obj):
|
|
return obj
|
|
|
|
def do_transfer(self, a1, a2):
|
|
create_collective_group([a1, a2], backend="torch_gloo")
|
|
return ray.get(a1.consume.remote(a2.create.remote()))
|
|
|
|
actor = [Actor.remote() for _ in range(3)]
|
|
assert ray.get(actor[2].do_transfer.remote(actor[0], actor[1])) == pytest.approx(
|
|
torch.tensor([1, 2, 3])
|
|
)
|
|
|
|
|
|
def test_send_fails(ray_start_regular):
|
|
actors = [ErrorActor.remote() for _ in range(2)]
|
|
create_collective_group(actors, backend="torch_gloo")
|
|
|
|
# The gpu object will be gone when we trigger the transfer
|
|
# so the send will error out
|
|
rdt_ref = actors[0].send.remote(torch.randn((100, 100)))
|
|
ray.get(actors[0].clear_rdt_store.remote())
|
|
result_ref = actors[1].recv.remote(rdt_ref)
|
|
|
|
with pytest.raises(ray.exceptions.ActorDiedError):
|
|
ray.get(result_ref)
|
|
|
|
|
|
def test_send_actor_dies_before_creating(ray_start_regular):
|
|
actors = [ErrorActor.remote() for _ in range(2)]
|
|
create_collective_group(actors, backend="torch_gloo")
|
|
|
|
# Block the main thread so the object doesn't get created before the kill
|
|
actors[0].block_main_thread.remote()
|
|
gpu_obj_ref = actors[0].send.remote(torch.randn(100, 100))
|
|
result_ref = actors[1].recv.remote(gpu_obj_ref)
|
|
ray.kill(actors[0])
|
|
|
|
with pytest.raises(ray.exceptions.ActorDiedError):
|
|
ray.get(result_ref)
|
|
|
|
|
|
def test_send_actor_dies_before_sending(ray_start_regular):
|
|
actors = [ErrorActor.remote() for _ in range(2)]
|
|
create_collective_group(actors, backend="torch_gloo")
|
|
|
|
rdt_ref = actors[0].send.remote(torch.randn(100, 100))
|
|
# Wait for the object to actually be created on the sender
|
|
ray.wait([rdt_ref])
|
|
# Block the background thread before triggering the transfer
|
|
# so the send doesn't happen before the actor is killed
|
|
actors[0].block_background_thread.remote()
|
|
result_ref = actors[1].recv.remote(rdt_ref)
|
|
ray.kill(actors[0])
|
|
|
|
with pytest.raises(ray.exceptions.ActorDiedError):
|
|
ray.get(result_ref)
|
|
|
|
|
|
def test_recv_actor_dies(ray_start_regular, caplog, propagate_logs):
|
|
actors = [ErrorActor.remote() for _ in range(2)]
|
|
create_collective_group(actors, backend="torch_gloo")
|
|
|
|
# Do a transfer with the receiver's background thread blocked,
|
|
# so the recv doesn't happen before the actor is killed
|
|
rdt_ref = actors[0].send.remote(torch.randn((100, 100)))
|
|
actors[1].block_background_thread.remote()
|
|
result_ref = actors[1].recv.remote(rdt_ref)
|
|
ray.kill(actors[1])
|
|
|
|
def check_logs():
|
|
records = caplog.records
|
|
return any(
|
|
record.levelno == logging.ERROR
|
|
and re.search(r"RDT transfer with.*failed", record.message)
|
|
for record in records
|
|
) and any(
|
|
record.levelno == logging.ERROR
|
|
and "Destroyed collective group" in record.message
|
|
for record in records
|
|
)
|
|
|
|
wait_for_condition(check_logs)
|
|
|
|
with pytest.raises(ray.exceptions.ActorDiedError):
|
|
ray.get(result_ref)
|
|
with pytest.raises(ray.exceptions.ActorDiedError):
|
|
ray.get(actors[0].recv.remote(1))
|
|
|
|
|
|
@pytest.mark.skip(
|
|
"Lineage Reconstruction currently results in a check failure with RDT"
|
|
)
|
|
def test_rdt_lineage_reconstruction(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=0)
|
|
ray.init(address=cluster.address)
|
|
cluster.add_node(num_cpus=1)
|
|
worker_to_kill = cluster.add_node(num_cpus=1, resources={"to_restart": 1})
|
|
|
|
@ray.remote(max_restarts=1, max_task_retries=1, resources={"to_restart": 1})
|
|
class RecvRestartableActor:
|
|
def recv(self, obj):
|
|
return obj
|
|
|
|
send_actor = GPUTestActor.remote()
|
|
recv_actor = RecvRestartableActor.remote()
|
|
create_collective_group([send_actor, recv_actor], backend="gloo")
|
|
|
|
one_mb_tensor = torch.randn((1024 * 1024,))
|
|
ref = recv_actor.recv.remote(send_actor.echo.remote(one_mb_tensor))
|
|
ray.wait([ref], fetch_local=False)
|
|
cluster.remove_node(worker_to_kill, allow_graceful=False)
|
|
cluster.add_node(num_cpus=1, resources={"to_restart": 1})
|
|
assert ray.get(ref).nbytes >= (1024 * 1024)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|