chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,102 @@
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self, init_value, fail_after=None, sys_exit=False):
|
||||
self.i = init_value
|
||||
self.fail_after = fail_after
|
||||
self.sys_exit = sys_exit
|
||||
|
||||
self.count = 0
|
||||
|
||||
def _fail_if_needed(self):
|
||||
if self.fail_after and self.count > self.fail_after:
|
||||
# Randomize the failures to better cover multi actor scenarios.
|
||||
if random.random() > 0.5:
|
||||
if self.sys_exit:
|
||||
os._exit(1)
|
||||
else:
|
||||
raise RuntimeError("injected fault")
|
||||
|
||||
def inc(self, x):
|
||||
self.i += x
|
||||
self.count += 1
|
||||
self._fail_if_needed()
|
||||
return self.i
|
||||
|
||||
def double_and_inc(self, x):
|
||||
self.i *= 2
|
||||
self.i += x
|
||||
return self.i
|
||||
|
||||
def echo(self, x):
|
||||
self.count += 1
|
||||
self._fail_if_needed()
|
||||
return x
|
||||
|
||||
def append_to(self, lst):
|
||||
lst.append(self.i)
|
||||
return lst
|
||||
|
||||
def inc_two(self, x, y):
|
||||
self.i += x
|
||||
self.i += y
|
||||
return self.i
|
||||
|
||||
def sleep(self, x):
|
||||
time.sleep(x)
|
||||
return x
|
||||
|
||||
@ray.method(num_returns=2)
|
||||
def return_two(self, x):
|
||||
return x, x + 1
|
||||
|
||||
def read_input(self, x):
|
||||
return x
|
||||
|
||||
@ray.method(num_returns=2)
|
||||
def inc_and_return_two(self, x):
|
||||
self.i += x
|
||||
return self.i, self.i + 1
|
||||
|
||||
@ray.method(num_returns=1)
|
||||
def return_two_as_one(self, x):
|
||||
return x, x + 1
|
||||
|
||||
@ray.method(num_returns=2)
|
||||
def return_two_from_three(self, x):
|
||||
return x, x + 1, x + 2
|
||||
|
||||
@ray.method(num_returns=2)
|
||||
def return_two_but_raise_exception(self, x):
|
||||
raise RuntimeError
|
||||
return 1, 2
|
||||
|
||||
def get_events(self):
|
||||
return getattr(self, "__ray_cgraph_events", [])
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Collector:
|
||||
def __init__(self):
|
||||
self.results = []
|
||||
|
||||
def collect(self, x):
|
||||
self.results.append(x)
|
||||
return self.results
|
||||
|
||||
def collect_two(self, x, y):
|
||||
self.results.append(x)
|
||||
self.results.append(y)
|
||||
return self.results
|
||||
|
||||
def collect_three(self, x, y, z):
|
||||
self.results.append(x)
|
||||
self.results.append(y)
|
||||
self.results.append(z)
|
||||
return self.results
|
||||
@@ -0,0 +1,497 @@
|
||||
# coding: utf-8
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
import ray.experimental.collective as collective
|
||||
from ray.dag import InputNode, MultiOutputNode
|
||||
from ray.experimental.channel import CPUCommunicator
|
||||
from ray.experimental.collective.conftest import (
|
||||
AbstractNcclGroup,
|
||||
CPUTorchTensorWorker,
|
||||
check_nccl_group_init,
|
||||
check_nccl_group_teardown,
|
||||
)
|
||||
from ray.experimental.util.types import ReduceOp
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import cupy as cp
|
||||
import torch
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if sys.platform != "linux" and sys.platform != "darwin":
|
||||
pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
|
||||
|
||||
|
||||
class MockCommunicator(CPUCommunicator):
|
||||
"""
|
||||
Use a mock communicator to test the actor schedules.
|
||||
"""
|
||||
|
||||
def __init__(self, world_size: int, actor_handles: List["ray.actor.ActorHandle"]):
|
||||
self._world_size = world_size
|
||||
self._actor_handles = actor_handles
|
||||
|
||||
def send(self, value: "torch.Tensor", peer_rank: int) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def recv(
|
||||
self,
|
||||
shape: Tuple[int],
|
||||
dtype: "torch.dtype",
|
||||
peer_rank: int,
|
||||
allocator: Optional[
|
||||
Callable[[Tuple[int], "torch.dtype"], "torch.Tensor"]
|
||||
] = None,
|
||||
) -> "torch.Tensor":
|
||||
raise NotImplementedError
|
||||
|
||||
def allgather(
|
||||
self,
|
||||
send_buf: "torch.Tensor",
|
||||
recv_buf: "torch.Tensor",
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def allreduce(
|
||||
self,
|
||||
send_buf: "torch.Tensor",
|
||||
recv_buf: "torch.Tensor",
|
||||
op: ReduceOp,
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def reducescatter(
|
||||
self,
|
||||
send_buf: "torch.Tensor",
|
||||
recv_buf: "torch.Tensor",
|
||||
op: ReduceOp,
|
||||
) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def recv_stream(self) -> Optional["cp.cuda.ExternalStream"]:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def send_stream(self) -> Optional["cp.cuda.ExternalStream"]:
|
||||
raise NotImplementedError
|
||||
|
||||
def destroy(self) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
@ray.remote
|
||||
class DDPWorker:
|
||||
def __init__(self):
|
||||
return
|
||||
|
||||
def backward(self, _):
|
||||
return 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
|
||||
def test_all_reduce_duplicate_actors(ray_start_regular):
|
||||
"""
|
||||
Test an error is thrown when two input nodes from the same actor bind to
|
||||
an all-reduce.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options()
|
||||
worker = actor_cls.remote()
|
||||
|
||||
with InputNode() as inp:
|
||||
computes = [worker.return_tensor.bind(inp) for _ in range(2)]
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Expected unique actor handles, but found duplicate actor handles from input nodes",
|
||||
):
|
||||
collective.allreduce.bind(computes)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
|
||||
def test_all_reduce_custom_comm_wrong_actors(ray_start_regular):
|
||||
"""
|
||||
Test an error is thrown when an all-reduce binds to a custom NCCL group and
|
||||
a wrong set of actors.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options()
|
||||
|
||||
num_workers = 2
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
nccl_group = AbstractNcclGroup([workers[0]])
|
||||
with InputNode() as inp:
|
||||
computes = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Expected actor handles to match the custom communicator group",
|
||||
):
|
||||
collective.allreduce.bind(computes, transport=nccl_group)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
|
||||
def test_all_reduce_bind_list_of_nodes_duplicate_nodes(ray_start_regular):
|
||||
"""
|
||||
Test an error is thrown when an all-reduce binds to lists of nodes
|
||||
that are duplicated.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options()
|
||||
|
||||
num_workers = 2
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
nccl_group = AbstractNcclGroup([workers[0]])
|
||||
with InputNode() as inp:
|
||||
computes_0 = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
computes_1 = [workers[0].return_tensor.bind(inp) for _ in range(2)]
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Expected unique actor handles at list at index",
|
||||
):
|
||||
collective.allreduce.bind([computes_0, computes_1], transport=nccl_group)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
|
||||
def test_all_reduce_bind_list_of_nodes_unequal_number_of_nodes(ray_start_regular):
|
||||
"""
|
||||
Test an error is thrown when an all-reduce binds to lists of nodes
|
||||
of different number of nodes across actors.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options()
|
||||
|
||||
num_workers = 2
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
nccl_group = AbstractNcclGroup([workers[0]])
|
||||
with InputNode() as inp:
|
||||
computes_0 = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
computes_1 = [worker.return_tensor.bind(inp) for worker in workers[1:]]
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Expected all input lists to have the same number of nodes",
|
||||
):
|
||||
collective.allreduce.bind([computes_0, computes_1], transport=nccl_group)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
|
||||
def test_all_reduce_bind_list_of_nodes_different_actors(ray_start_regular):
|
||||
"""
|
||||
Test an error is thrown when an all-reduce binds to a list of nodes
|
||||
from different set of actors.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options()
|
||||
|
||||
num_workers = 3
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
nccl_group = AbstractNcclGroup([workers[0]])
|
||||
with InputNode() as inp:
|
||||
computes_0 = [worker.return_tensor.bind(inp) for worker in workers[:2]]
|
||||
computes_1 = [worker.return_tensor.bind(inp) for worker in workers[1:]]
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Expected all input lists to have the same set of actor handles",
|
||||
):
|
||||
collective.allreduce.bind([computes_0, computes_1], transport=nccl_group)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
|
||||
def test_all_reduce_bind_list_of_nodes_different_dtypes(ray_start_regular):
|
||||
"""
|
||||
Test an error is thrown when an all-reduce binds to a list of nodes
|
||||
that execute with tensors of different dtypes.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options()
|
||||
|
||||
num_workers = 3
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
comm = MockCommunicator(num_workers, workers)
|
||||
with InputNode() as inp:
|
||||
computes_0 = [worker.return_tensor.bind(inp[0], inp[1]) for worker in workers]
|
||||
computes_1 = [worker.return_tensor.bind(inp[0], inp[2]) for worker in workers]
|
||||
collectives = collective.allreduce.bind(
|
||||
[computes_0, computes_1], transport=comm
|
||||
)
|
||||
recvs = [
|
||||
worker.recv_tensors.bind(*collective)
|
||||
for worker, collective in zip(workers, collectives)
|
||||
]
|
||||
dag = MultiOutputNode(recvs)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Expected all input tensors to have the same dtype",
|
||||
):
|
||||
import torch
|
||||
|
||||
ray.get(compiled_dag.execute(1, torch.float16, torch.float32))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
|
||||
)
|
||||
def test_comm_all_reduces(ray_start_regular, monkeypatch):
|
||||
"""
|
||||
Test different communicators are used for different all-reduce calls of
|
||||
different sets of actors.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
|
||||
|
||||
num_workers = 2
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
with InputNode() as inp:
|
||||
computes = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
# There are two all-reduces, each on one actor.
|
||||
collectives = [collective.allreduce.bind([compute]) for compute in computes]
|
||||
# collective[0] is the only CollectiveOutputNode for each all-reduce.
|
||||
dag = MultiOutputNode([collective[0] for collective in collectives])
|
||||
|
||||
compiled_dag, mock_nccl_group_set = check_nccl_group_init(
|
||||
monkeypatch,
|
||||
dag,
|
||||
{
|
||||
(frozenset([workers[0]]), None),
|
||||
(frozenset([workers[1]]), None),
|
||||
},
|
||||
)
|
||||
|
||||
check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
|
||||
)
|
||||
def test_comm_deduplicate_all_reduces(ray_start_regular, monkeypatch):
|
||||
"""
|
||||
Test communicators are deduplicated when all-reduces are called on the same
|
||||
group of actors more than once.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
|
||||
|
||||
num_workers = 2
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
with InputNode() as inp:
|
||||
tensors = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
collectives = collective.allreduce.bind(tensors)
|
||||
collectives = collective.allreduce.bind(collectives)
|
||||
dag = MultiOutputNode(collectives)
|
||||
|
||||
compiled_dag, mock_nccl_group_set = check_nccl_group_init(
|
||||
monkeypatch,
|
||||
dag,
|
||||
{(frozenset(workers), None)},
|
||||
)
|
||||
|
||||
check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
|
||||
)
|
||||
def test_comm_deduplicate_p2p_and_collective(ray_start_regular, monkeypatch):
|
||||
"""
|
||||
Test communicators are deduplicated when the collective and the P2P are on
|
||||
the same set of actors.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
|
||||
|
||||
num_workers = 2
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
with InputNode() as inp:
|
||||
computes = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
collectives = collective.allreduce.bind(computes)
|
||||
recvs = [
|
||||
# Each of the 2 workers receives from the other.
|
||||
workers[0].recv.bind(
|
||||
collectives[1].with_tensor_transport(transport="nccl")
|
||||
),
|
||||
workers[1].recv.bind(
|
||||
collectives[0].with_tensor_transport(transport="nccl")
|
||||
),
|
||||
]
|
||||
dag = MultiOutputNode(recvs)
|
||||
|
||||
compiled_dag, mock_nccl_group_set = check_nccl_group_init(
|
||||
monkeypatch,
|
||||
dag,
|
||||
{(frozenset(workers), None)},
|
||||
)
|
||||
|
||||
check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
|
||||
|
||||
with InputNode() as inp:
|
||||
computes = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
collectives = collective.allreduce.bind(computes)
|
||||
# Sender is workers[0] and receiver is workers[1].
|
||||
dag = workers[1].recv.bind(
|
||||
collectives[0].with_tensor_transport(transport="nccl")
|
||||
)
|
||||
dag = MultiOutputNode([dag, collectives[1]])
|
||||
|
||||
compiled_dag, mock_nccl_group_set = check_nccl_group_init(
|
||||
monkeypatch,
|
||||
dag,
|
||||
{(frozenset(workers), None)},
|
||||
)
|
||||
|
||||
check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
|
||||
)
|
||||
def test_custom_comm(ray_start_regular, monkeypatch):
|
||||
"""
|
||||
Test a custom GPU communicator is used when specified and a default
|
||||
communicator is used otherwise.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
|
||||
|
||||
num_workers = 2
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
comm = AbstractNcclGroup(workers)
|
||||
with InputNode() as inp:
|
||||
computes = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
collectives = collective.allreduce.bind(computes, transport=comm)
|
||||
collectives = collective.allreduce.bind(collectives)
|
||||
dag = workers[0].recv.bind(
|
||||
collectives[1].with_tensor_transport(transport="nccl")
|
||||
)
|
||||
dag = MultiOutputNode([dag, collectives[0]])
|
||||
|
||||
compiled_dag, mock_nccl_group_set = check_nccl_group_init(
|
||||
monkeypatch,
|
||||
dag,
|
||||
{
|
||||
(frozenset(workers), comm),
|
||||
(frozenset(workers), None),
|
||||
},
|
||||
)
|
||||
|
||||
check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
|
||||
|
||||
comm = AbstractNcclGroup(workers)
|
||||
with InputNode() as inp:
|
||||
computes = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
collectives = collective.allreduce.bind(computes)
|
||||
collectives = collective.allreduce.bind(collectives)
|
||||
dag = workers[0].recv.bind(collectives[1].with_tensor_transport(transport=comm))
|
||||
dag = MultiOutputNode([dag, collectives[0]])
|
||||
|
||||
compiled_dag, mock_nccl_group_set = check_nccl_group_init(
|
||||
monkeypatch,
|
||||
dag,
|
||||
{
|
||||
(frozenset(workers), comm),
|
||||
(frozenset(workers), None),
|
||||
},
|
||||
)
|
||||
|
||||
check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_regular", [{"num_cpus": 4, "num_gpus": 4}], indirect=True
|
||||
)
|
||||
def test_custom_comm_init_teardown(ray_start_regular, monkeypatch):
|
||||
"""
|
||||
Test custom NCCL groups are properly initialized and destroyed.
|
||||
1. Test when multiple type hints have the same `transport=custom_nccl_group`,
|
||||
the `custom_nccl_group` is initialized only once.
|
||||
2. Test all initialized NCCL groups are destroyed during teardown.
|
||||
"""
|
||||
actor_cls = CPUTorchTensorWorker.options(num_cpus=0, num_gpus=1)
|
||||
|
||||
num_workers = 2
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
|
||||
comm = AbstractNcclGroup(workers)
|
||||
|
||||
with InputNode() as inp:
|
||||
tensors = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
allreduce = collective.allreduce.bind(tensors, transport=comm)
|
||||
dag = workers[0].recv.bind(allreduce[1].with_tensor_transport(transport=comm))
|
||||
dag = MultiOutputNode([dag, allreduce[0]])
|
||||
|
||||
compiled_dag, mock_nccl_group_set = check_nccl_group_init(
|
||||
monkeypatch,
|
||||
dag,
|
||||
{(frozenset(workers), comm)},
|
||||
)
|
||||
|
||||
check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
|
||||
|
||||
comm_1 = AbstractNcclGroup(workers)
|
||||
comm_2 = AbstractNcclGroup(workers)
|
||||
comm_3 = AbstractNcclGroup(workers)
|
||||
|
||||
with InputNode() as inp:
|
||||
tensors = [worker.return_tensor.bind(inp) for worker in workers]
|
||||
allreduce1 = collective.allreduce.bind(tensors, transport=comm_1)
|
||||
allreduce2 = collective.allreduce.bind(allreduce1, transport=comm_2)
|
||||
dag = workers[0].recv.bind(
|
||||
allreduce2[1].with_tensor_transport(transport=comm_3)
|
||||
)
|
||||
dag = MultiOutputNode([dag, allreduce2[0]])
|
||||
|
||||
compiled_dag, mock_nccl_group_set = check_nccl_group_init(
|
||||
monkeypatch,
|
||||
dag,
|
||||
{
|
||||
(frozenset(workers), comm_1),
|
||||
(frozenset(workers), comm_2),
|
||||
(frozenset(workers), comm_3),
|
||||
},
|
||||
)
|
||||
|
||||
check_nccl_group_teardown(monkeypatch, compiled_dag, mock_nccl_group_set)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 4}], indirect=True)
|
||||
@pytest.mark.parametrize("num_workers", [2, 4])
|
||||
def test_exec_schedules_ddp(ray_start_regular, num_workers):
|
||||
"""
|
||||
Test the execution schedules for the DDP strategy. Each worker should have
|
||||
identical schedules.
|
||||
"""
|
||||
actor_cls = DDPWorker.options(num_cpus=1)
|
||||
workers = [actor_cls.remote() for _ in range(num_workers)]
|
||||
comm = MockCommunicator(num_workers, workers)
|
||||
|
||||
outputs = []
|
||||
with InputNode() as inp:
|
||||
grads = [worker.backward.bind(inp) for worker in workers]
|
||||
grads_reduced = collective.allreduce.bind(grads, transport=comm)
|
||||
outputs.extend(grads_reduced)
|
||||
grads = [worker.backward.bind(grad) for worker, grad in zip(workers, grads)]
|
||||
grads_reduced = collective.allreduce.bind(grads, transport=comm)
|
||||
outputs.extend(grads_reduced)
|
||||
dag = MultiOutputNode(outputs)
|
||||
|
||||
compiled_dag = dag.experimental_compile(_default_communicator=comm)
|
||||
actor_to_execution_schedule = list(
|
||||
compiled_dag.actor_to_execution_schedule.values()
|
||||
)
|
||||
expected_schedule = actor_to_execution_schedule[0]
|
||||
for schedule in actor_to_execution_schedule[1:]:
|
||||
assert schedule == expected_schedule
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if os.environ.get("PARALLEL_CI"):
|
||||
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
||||
else:
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,361 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
import ray.cluster_utils
|
||||
import ray.experimental.collective as collective
|
||||
from ray.dag import InputNode
|
||||
from ray.dag.output_node import MultiOutputNode
|
||||
from ray.exceptions import RayChannelError, RayTaskError
|
||||
from ray.experimental.channel.cpu_communicator import CPUCommunicator
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
@ray.remote
|
||||
class CPUTorchTensorWorker:
|
||||
def __init__(self):
|
||||
self.device = torch.device(type="cpu")
|
||||
|
||||
def send(self, shape, dtype, value: int, send_tensor=True):
|
||||
if not send_tensor:
|
||||
return 1
|
||||
return torch.ones(shape, dtype=dtype) * value
|
||||
|
||||
def send_dict(self, entries):
|
||||
results = {}
|
||||
for key, entry in entries.items():
|
||||
value, shape, dtype = entry
|
||||
results[key] = torch.ones(shape, dtype=dtype) * value
|
||||
return results
|
||||
|
||||
def send_or_raise(self, shape, dtype, value: int, raise_exception=False):
|
||||
if raise_exception:
|
||||
raise RuntimeError()
|
||||
return torch.ones(shape, dtype=dtype) * value
|
||||
|
||||
def recv(self, tensor):
|
||||
assert tensor.device == self.device
|
||||
return (tensor[0].item(), tensor.shape, tensor.dtype)
|
||||
|
||||
def recv_dict(self, tensor_dict):
|
||||
vals = {}
|
||||
for i, tensor in tensor_dict.items():
|
||||
assert tensor.device == self.device
|
||||
vals[i] = self.recv(tensor)
|
||||
return vals
|
||||
|
||||
def compute_with_tuple_args(self, args, i: int):
|
||||
shape, dtype, value = args[i]
|
||||
tensor = torch.ones(shape, dtype=dtype) * value
|
||||
return tensor
|
||||
|
||||
def recv_tensor(self, tensor):
|
||||
assert tensor.device == self.device
|
||||
return tensor
|
||||
|
||||
def return_tensor(self, size: int) -> torch.Tensor:
|
||||
return torch.ones(size)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 0,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_p2p_basic(ray_start_cluster):
|
||||
sender = CPUTorchTensorWorker.remote()
|
||||
receiver = CPUTorchTensorWorker.remote()
|
||||
|
||||
cpu_group = CPUCommunicator(2, [sender, receiver])
|
||||
|
||||
shape = (10,)
|
||||
dtype = torch.float16
|
||||
|
||||
with InputNode() as inp:
|
||||
dag = sender.send.bind(inp.shape, inp.dtype, inp[0])
|
||||
dag = dag.with_tensor_transport(transport=cpu_group)
|
||||
dag = receiver.recv.bind(dag)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
for i in range(3):
|
||||
ref = compiled_dag.execute(i, shape=shape, dtype=dtype)
|
||||
assert ray.get(ref) == (i, shape, dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 0,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_allreduce_basic(ray_start_cluster):
|
||||
num_workers = 2
|
||||
workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers)]
|
||||
|
||||
cpu_group = CPUCommunicator(num_workers, workers)
|
||||
|
||||
with InputNode() as inp:
|
||||
computes = [
|
||||
worker.compute_with_tuple_args.bind(inp, i)
|
||||
for i, worker in enumerate(workers)
|
||||
]
|
||||
collectives = collective.allreduce.bind(computes, transport=cpu_group)
|
||||
recvs = [
|
||||
worker.recv.bind(collective)
|
||||
for worker, collective in zip(workers, collectives)
|
||||
]
|
||||
dag = MultiOutputNode(recvs)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
for i in range(3):
|
||||
i += 1
|
||||
shape = (i * 10,)
|
||||
dtype = torch.float16
|
||||
ref = compiled_dag.execute(
|
||||
[(shape, dtype, i + idx) for idx in range(num_workers)]
|
||||
)
|
||||
result = ray.get(ref)
|
||||
reduced_val = sum(i + idx for idx in range(num_workers))
|
||||
assert result == [(reduced_val, shape, dtype) for _ in workers]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 0,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_allreduce_get_partial(ray_start_cluster):
|
||||
num_workers = 2
|
||||
workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers)]
|
||||
|
||||
cpu_group = CPUCommunicator(num_workers, workers)
|
||||
|
||||
shape = (10,)
|
||||
dtype = torch.float16
|
||||
|
||||
with InputNode() as inp:
|
||||
computes = [
|
||||
worker.compute_with_tuple_args.bind(inp, i)
|
||||
for i, worker in enumerate(workers)
|
||||
]
|
||||
collectives = collective.allreduce.bind(computes, transport=cpu_group)
|
||||
recv = workers[0].recv.bind(collectives[0])
|
||||
tensor = workers[1].recv_tensor.bind(collectives[0])
|
||||
dag = MultiOutputNode([recv, tensor, collectives[1]])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
for i in range(3):
|
||||
ref = compiled_dag.execute(
|
||||
[(shape, dtype, i + idx + 1) for idx in range(num_workers)]
|
||||
)
|
||||
result = ray.get(ref)
|
||||
metadata, tensor, _ = result
|
||||
reduced_val = sum(i + idx + 1 for idx in range(num_workers))
|
||||
assert metadata == (reduced_val, shape, dtype)
|
||||
expected_tensor_val = torch.ones(shape, dtype=dtype) * reduced_val
|
||||
assert torch.equal(tensor, expected_tensor_val)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 0,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_allreduce_wrong_shape(ray_start_cluster):
|
||||
"""
|
||||
Test an error is thrown when the tensors in an all-reduce have different shapes.
|
||||
"""
|
||||
num_workers = 2
|
||||
workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers)]
|
||||
|
||||
cpu_group = CPUCommunicator(num_workers, workers)
|
||||
|
||||
dtype = torch.float16
|
||||
|
||||
with InputNode() as inp:
|
||||
computes = [
|
||||
worker.compute_with_tuple_args.bind(inp, i)
|
||||
for i, worker in enumerate(workers)
|
||||
]
|
||||
collectives = collective.allreduce.bind(computes, transport=cpu_group)
|
||||
recvs = [
|
||||
worker.recv.bind(collective)
|
||||
for worker, collective in zip(workers, collectives)
|
||||
]
|
||||
dag = MultiOutputNode(recvs)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
ref = compiled_dag.execute([((20,), dtype, idx + 1) for idx in range(num_workers)])
|
||||
reduced_val = (1 + num_workers) * num_workers / 2
|
||||
assert ray.get(ref) == [(reduced_val, (20,), dtype) for _ in range(num_workers)]
|
||||
|
||||
ref = compiled_dag.execute(
|
||||
[((10 * (idx + 1),), dtype, idx + 1) for idx in range(num_workers)]
|
||||
)
|
||||
# Execution hangs because of shape mismatch and a task error is raised.
|
||||
with pytest.raises(RayTaskError):
|
||||
ray.get(ref)
|
||||
|
||||
# Since we have buffered channels, the execution should not error, but the
|
||||
# get should error, as the dag should no longer work after the application-
|
||||
# level exception.
|
||||
ref = compiled_dag.execute([((20,), dtype, 1) for _ in workers])
|
||||
with pytest.raises(RayChannelError):
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 0,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_allreduce_scheduling(ray_start_cluster):
|
||||
"""
|
||||
Test scheduling avoids potential deadlocks that arise from all-reduce operations.
|
||||
|
||||
inp --> x(0) --> +------------+
|
||||
| | all-reduce |
|
||||
--> y(1) --> +------------+
|
||||
|
|
||||
--> t(0) --> recv(1)
|
||||
|
||||
In the above graph, x, y, t are tensors, and the numbers inside parentheses
|
||||
identify the actors. If actor 1 launches an all-reduce with tensor y while
|
||||
actor 0 starts sending t, then actor 1 waits for actor 0 to join the all-reduce
|
||||
while actor 1 waits for actor 0 to receive t.
|
||||
"""
|
||||
num_workers = 2
|
||||
workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers)]
|
||||
|
||||
cpu_group = CPUCommunicator(num_workers, workers)
|
||||
|
||||
shape = (10,)
|
||||
dtype = torch.float16
|
||||
|
||||
with InputNode() as inp:
|
||||
# Tensors in the all-reduce.
|
||||
x = workers[0].send.bind(shape, dtype, inp)
|
||||
y = workers[1].send.bind(shape, dtype, inp)
|
||||
|
||||
# Tensor to be sent from workers[0] to workers[1].
|
||||
t = workers[0].send.bind(shape, dtype, inp)
|
||||
t = t.with_tensor_transport(transport=cpu_group)
|
||||
|
||||
collectives = collective.allreduce.bind([x, y], transport=cpu_group)
|
||||
recv = workers[1].recv.bind(t)
|
||||
dag = MultiOutputNode([collectives[0], collectives[1], recv])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
value = 10
|
||||
ref = compiled_dag.execute(value)
|
||||
result = ray.get(ref)
|
||||
reduced_value = value * 2
|
||||
expected_tensor_val = torch.ones(shape, dtype=dtype) * reduced_value
|
||||
assert torch.equal(result[0], expected_tensor_val)
|
||||
assert torch.equal(result[1], expected_tensor_val)
|
||||
assert result[2] == (value, shape, dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 0,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_allreduce_duplicate_actors(ray_start_cluster):
|
||||
"""
|
||||
Test an error is thrown when two input nodes from the same actor bind to
|
||||
an all-reduce.
|
||||
"""
|
||||
num_workers = 2
|
||||
worker = CPUTorchTensorWorker.remote()
|
||||
|
||||
cpu_group = CPUCommunicator(num_workers, [worker, worker])
|
||||
|
||||
with InputNode() as inp:
|
||||
computes = [worker.return_tensor.bind(inp) for _ in range(2)]
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"Expected unique actor handles, but found duplicate actor handles "
|
||||
"from input nodes"
|
||||
),
|
||||
):
|
||||
collective.allreduce.bind(computes, transport=cpu_group)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 0,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_allreduce_wrong_actors(ray_start_cluster):
|
||||
"""
|
||||
Test an error is thrown when an all-reduce binds to a wrong set of actors.
|
||||
"""
|
||||
num_workers = 2
|
||||
workers = [CPUTorchTensorWorker.remote() for _ in range(num_workers * 2)]
|
||||
|
||||
cpu_group = CPUCommunicator(num_workers, workers[:2])
|
||||
|
||||
with InputNode() as inp:
|
||||
computes = [worker.return_tensor.bind(inp) for worker in workers[2:]]
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Expected actor handles to match the custom communicator group",
|
||||
):
|
||||
collective.allreduce.bind(computes, transport=cpu_group)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if os.environ.get("PARALLEL_CI"):
|
||||
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
||||
else:
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,814 @@
|
||||
# coding: utf-8
|
||||
import copy
|
||||
import logging
|
||||
import pickle
|
||||
import re
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
import ray._private
|
||||
import ray.cluster_utils
|
||||
from ray._common.test_utils import SignalActor
|
||||
from ray._common.utils import (
|
||||
get_or_create_event_loop,
|
||||
)
|
||||
from ray._private.test_utils import (
|
||||
run_string_as_driver_nonblocking,
|
||||
wait_for_pid_to_exit,
|
||||
)
|
||||
from ray.dag import DAGContext, InputNode, MultiOutputNode
|
||||
from ray.dag.tests.experimental.actor_defs import Actor
|
||||
from ray.exceptions import ActorDiedError, RayChannelError, RayChannelTimeoutError
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skipif(
|
||||
sys.platform != "linux" and sys.platform != "darwin",
|
||||
reason="Requires Linux or MacOS",
|
||||
),
|
||||
pytest.mark.timeout(500),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temporary_change_timeout(request):
|
||||
ctx = DAGContext.get_current()
|
||||
original = ctx.submit_timeout
|
||||
ctx.submit_timeout = request.param
|
||||
yield ctx.submit_timeout
|
||||
ctx.submit_timeout = original
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def zero_teardown_timeout(request):
|
||||
ctx = DAGContext.get_current()
|
||||
original = ctx.teardown_timeout
|
||||
ctx.teardown_timeout = 0
|
||||
yield ctx.teardown_timeout
|
||||
ctx.teardown_timeout = original
|
||||
|
||||
|
||||
def test_kwargs_not_supported(ray_start_regular):
|
||||
a = Actor.remote(0)
|
||||
|
||||
# Binding InputNode as kwarg is not supported.
|
||||
with InputNode() as i:
|
||||
dag = a.inc_two.bind(x=i, y=1)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"Compiled DAG currently does not support binding to other DAG "
|
||||
"nodes as kwargs",
|
||||
):
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Binding another DAG node as kwarg is not supported.
|
||||
with InputNode() as i:
|
||||
dag = a.inc.bind(i)
|
||||
dag = a.inc_two.bind(x=dag, y=1)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"Compiled DAG currently does not support binding to other DAG "
|
||||
"nodes as kwargs",
|
||||
):
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Binding normal Python value as a kwarg is supported.
|
||||
with InputNode() as i:
|
||||
dag = a.inc_two.bind(i, y=1)
|
||||
compiled_dag = dag.experimental_compile()
|
||||
assert ray.get(compiled_dag.execute(2)) == 3
|
||||
|
||||
|
||||
def test_dag_exception_basic(ray_start_regular, capsys):
|
||||
# Test application throwing exceptions with a single task.
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as inp:
|
||||
dag = a.inc.bind(inp)
|
||||
|
||||
# Can throw an error.
|
||||
compiled_dag = dag.experimental_compile()
|
||||
ref = compiled_dag.execute("hello")
|
||||
with pytest.raises(TypeError) as exc_info:
|
||||
ray.get(ref)
|
||||
# Traceback should match the original actor class definition.
|
||||
assert "self.i += x" in str(exc_info.value)
|
||||
|
||||
# Can throw an error multiple times.
|
||||
ref = compiled_dag.execute("hello")
|
||||
with pytest.raises(TypeError) as exc_info:
|
||||
ray.get(ref)
|
||||
# Traceback should match the original actor class definition.
|
||||
assert "self.i += x" in str(exc_info.value)
|
||||
|
||||
# Can use the DAG after exceptions are thrown.
|
||||
assert ray.get(compiled_dag.execute(1)) == 1
|
||||
|
||||
|
||||
def test_dag_exception_chained(ray_start_regular, capsys):
|
||||
# Test application throwing exceptions with a task that depends on another
|
||||
# task.
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as inp:
|
||||
dag = a.inc.bind(inp)
|
||||
dag = a.inc.bind(dag)
|
||||
|
||||
# Can throw an error.
|
||||
compiled_dag = dag.experimental_compile()
|
||||
ref = compiled_dag.execute("hello")
|
||||
with pytest.raises(TypeError) as exc_info:
|
||||
ray.get(ref)
|
||||
# Traceback should match the original actor class definition.
|
||||
assert "self.i += x" in str(exc_info.value)
|
||||
|
||||
# Can throw an error multiple times.
|
||||
ref = compiled_dag.execute("hello")
|
||||
with pytest.raises(TypeError) as exc_info:
|
||||
ray.get(ref)
|
||||
# Traceback should match the original actor class definition.
|
||||
assert "self.i += x" in str(exc_info.value)
|
||||
|
||||
# Can use the DAG after exceptions are thrown.
|
||||
assert ray.get(compiled_dag.execute(1)) == 2
|
||||
|
||||
|
||||
def test_dag_exception_multi_output(ray_start_regular, capsys):
|
||||
# Test application throwing exceptions with a DAG with multiple outputs.
|
||||
a = Actor.remote(0)
|
||||
b = Actor.remote(0)
|
||||
with InputNode() as inp:
|
||||
dag = MultiOutputNode([a.inc.bind(inp), b.inc.bind(inp)])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Verify that fetching each output individually raises the error.
|
||||
refs = compiled_dag.execute("hello")
|
||||
for ref in refs:
|
||||
with pytest.raises(TypeError) as exc_info:
|
||||
ray.get(ref)
|
||||
# Traceback should match the original actor class definition.
|
||||
assert "self.i += x" in str(exc_info.value)
|
||||
|
||||
# Verify that another bad input exhibits the same behavior.
|
||||
refs = compiled_dag.execute("hello")
|
||||
for ref in refs:
|
||||
with pytest.raises(TypeError) as exc_info:
|
||||
ray.get(ref)
|
||||
# Traceback should match the original actor class definition.
|
||||
assert "self.i += x" in str(exc_info.value)
|
||||
|
||||
# Verify that the DAG can be used after the errors.
|
||||
assert ray.get(compiled_dag.execute(1)) == [1, 1]
|
||||
|
||||
|
||||
def test_dag_errors(ray_start_regular):
|
||||
a = Actor.remote(0)
|
||||
dag = a.inc.bind(1)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="No InputNode found in the DAG: when traversing upwards, "
|
||||
"no upstream node was found for",
|
||||
):
|
||||
dag.experimental_compile()
|
||||
|
||||
a2 = Actor.remote(0)
|
||||
with InputNode() as inp:
|
||||
dag = MultiOutputNode([a.inc.bind(inp), a2.inc.bind(1)])
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Compiled DAGs require each task to take a ray.dag.InputNode or "
|
||||
"at least one other DAGNode as an input",
|
||||
):
|
||||
dag.experimental_compile()
|
||||
|
||||
@ray.remote
|
||||
def f(x):
|
||||
return x
|
||||
|
||||
with InputNode() as inp:
|
||||
dag = f.bind(inp)
|
||||
with pytest.raises(
|
||||
NotImplementedError,
|
||||
match="Compiled DAGs currently only support actor method nodes",
|
||||
):
|
||||
dag.experimental_compile()
|
||||
|
||||
with InputNode() as inp:
|
||||
dag = a.inc.bind(inp)
|
||||
compiled_dag = dag.experimental_compile()
|
||||
ref = compiled_dag.execute(1)
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match=(
|
||||
re.escape(
|
||||
"wait() expected a list of ray.ObjectRef or ray.ObjectRefGenerator, "
|
||||
"got <class 'ray.experimental.compiled_dag_ref.CompiledDAGRef'>"
|
||||
)
|
||||
),
|
||||
):
|
||||
ray.wait(ref)
|
||||
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match=(
|
||||
re.escape(
|
||||
"wait() expected a list of ray.ObjectRef or ray.ObjectRefGenerator, "
|
||||
"got list containing "
|
||||
"<class 'ray.experimental.compiled_dag_ref.CompiledDAGRef'>"
|
||||
)
|
||||
),
|
||||
):
|
||||
ray.wait([ref])
|
||||
|
||||
with pytest.raises(TypeError, match=r".*was found to be non-serializable.*"):
|
||||
ray.put([ref])
|
||||
|
||||
with pytest.raises(ValueError, match="CompiledDAGRef cannot be copied."):
|
||||
copy.copy(ref)
|
||||
|
||||
with pytest.raises(ValueError, match="CompiledDAGRef cannot be deep copied."):
|
||||
copy.deepcopy(ref)
|
||||
|
||||
with pytest.raises(ValueError, match="CompiledDAGRef cannot be pickled."):
|
||||
pickle.dumps(ref)
|
||||
|
||||
with pytest.raises(
|
||||
TypeError, match="CompiledDAGRef cannot be used as Ray task/actor argument."
|
||||
):
|
||||
f.remote(ref)
|
||||
|
||||
with pytest.raises(
|
||||
TypeError, match="CompiledDAGRef cannot be used as Ray task/actor argument."
|
||||
):
|
||||
a2.inc.remote(ref)
|
||||
|
||||
result = ray.get(ref)
|
||||
assert result == 1
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
r"ray.get\(\) can only be called once "
|
||||
r"on a CompiledDAGRef, and it was already called."
|
||||
),
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
def test_get_timeout(ray_start_regular, zero_teardown_timeout):
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as inp:
|
||||
dag = a.sleep.bind(inp)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
ref = compiled_dag.execute(5)
|
||||
|
||||
timed_out = False
|
||||
epsilon = 0.1 # Allow for some slack in the timeout checking
|
||||
try:
|
||||
start_time = time.monotonic()
|
||||
ray.get(ref, timeout=1)
|
||||
except RayChannelTimeoutError:
|
||||
duration = time.monotonic() - start_time
|
||||
assert duration > 1 - epsilon
|
||||
assert duration < 1 + epsilon
|
||||
timed_out = True
|
||||
assert timed_out
|
||||
|
||||
compiled_dag.teardown(kill_actors=True)
|
||||
|
||||
|
||||
def test_buffered_get_timeout(ray_start_regular, zero_teardown_timeout):
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as inp:
|
||||
dag = a.sleep.bind(inp)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
# The tasks will execute in order and sleep 1s, 1s, then 0s, respectively.
|
||||
refs = [
|
||||
compiled_dag.execute(1),
|
||||
compiled_dag.execute(1),
|
||||
compiled_dag.execute(0),
|
||||
]
|
||||
|
||||
with pytest.raises(RayChannelTimeoutError):
|
||||
# The final task takes <1s on its own, but because it's queued behind the
|
||||
# other two that take 1s each, this should time out.
|
||||
ray.get(refs[-1], timeout=1)
|
||||
|
||||
compiled_dag.teardown(kill_actors=True)
|
||||
|
||||
|
||||
def test_get_with_zero_timeout(ray_start_regular):
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self, signal_actor):
|
||||
self.signal_actor = signal_actor
|
||||
|
||||
def send(self, x):
|
||||
self.signal_actor.send.remote()
|
||||
return x
|
||||
|
||||
signal_actor = SignalActor.remote()
|
||||
a = Actor.remote(signal_actor)
|
||||
with InputNode() as inp:
|
||||
dag = a.send.bind(inp)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
ref = compiled_dag.execute(1)
|
||||
# Give enough time for DAG execution result to be ready
|
||||
ray.get(signal_actor.wait.remote())
|
||||
time.sleep(0.1)
|
||||
# Use timeout=0 to either get result immediately or raise an exception
|
||||
result = ray.get(ref, timeout=0)
|
||||
assert result == 1
|
||||
|
||||
|
||||
class TestDAGExceptionCompileMultipleTimes:
|
||||
@pytest.mark.parametrize("use_multi_output_node", [False, True])
|
||||
def test_compile_twice_fails(self, ray_start_regular, use_multi_output_node: bool):
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
if use_multi_output_node:
|
||||
dag = MultiOutputNode([a.echo.bind(i)])
|
||||
else:
|
||||
dag = a.echo.bind(i)
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Trying to compile again should fail.
|
||||
expected_err = (
|
||||
"It is not allowed to call `experimental_compile` on the same DAG "
|
||||
"object multiple times no matter whether `teardown` is called or not. "
|
||||
"Please reuse the existing compiled DAG or create a new one."
|
||||
)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=expected_err,
|
||||
):
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Even if we teardown the DAG, trying to compile again should still fail.
|
||||
compiled_dag.teardown()
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=expected_err,
|
||||
):
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
def test_compile_twice_with_different_nodes(self, ray_start_regular):
|
||||
a = Actor.remote(0)
|
||||
b = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
branch1 = a.echo.bind(i)
|
||||
branch2 = b.echo.bind(i)
|
||||
dag = MultiOutputNode([branch1, branch2])
|
||||
compiled_dag = dag.experimental_compile()
|
||||
compiled_dag.teardown()
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="The DAG was compiled more than once. The following two "
|
||||
"nodes call `experimental_compile`: ",
|
||||
):
|
||||
branch2.experimental_compile()
|
||||
|
||||
|
||||
def test_exceed_max_buffered_results(ray_start_regular):
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
dag = a.inc.bind(i)
|
||||
|
||||
compiled_dag = dag.experimental_compile(_max_buffered_results=1)
|
||||
|
||||
refs = []
|
||||
for i in range(2):
|
||||
ref = compiled_dag.execute(1)
|
||||
# Hold the refs to avoid get() being called on the ref
|
||||
# when it goes out of scope
|
||||
refs.append(ref)
|
||||
|
||||
# ray.get() on the 2nd ref fails because the DAG cannot buffer 2 results.
|
||||
with pytest.raises(
|
||||
ray.exceptions.RayCgraphCapacityExceeded,
|
||||
match=(
|
||||
"The compiled graph can't have more than 1 buffered results, "
|
||||
r"and you currently have 1 buffered results. Call `ray.get\(\)` on "
|
||||
r"CompiledDAGRef's \(or await on CompiledDAGFuture's\) to retrieve "
|
||||
"results, or increase `_max_buffered_results` if buffering is "
|
||||
"desired, note that this will increase driver memory usage."
|
||||
),
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
def test_exceed_max_buffered_results_multi_output(ray_start_regular):
|
||||
a = Actor.remote(0)
|
||||
b = Actor.remote(0)
|
||||
with InputNode() as inp:
|
||||
dag = MultiOutputNode([a.inc.bind(inp), b.inc.bind(inp)])
|
||||
|
||||
compiled_dag = dag.experimental_compile(_max_buffered_results=1)
|
||||
|
||||
refs = []
|
||||
for _ in range(2):
|
||||
ref = compiled_dag.execute(1)
|
||||
# Hold the refs to avoid get() being called on the ref
|
||||
# when it goes out of scope
|
||||
refs.append(ref)
|
||||
|
||||
# If there are results not fetched from an execution, that execution
|
||||
# still counts towards the number of buffered results.
|
||||
ray.get(refs[0][0])
|
||||
|
||||
# ray.get() on the 2nd ref fails because the DAG cannot buffer 2 results.
|
||||
with pytest.raises(
|
||||
ray.exceptions.RayCgraphCapacityExceeded,
|
||||
match=(
|
||||
"The compiled graph can't have more than 1 buffered results, "
|
||||
r"and you currently have 1 buffered results. Call `ray.get\(\)` on "
|
||||
r"CompiledDAGRef's \(or await on CompiledDAGFuture's\) to retrieve "
|
||||
"results, or increase `_max_buffered_results` if buffering is "
|
||||
"desired, note that this will increase driver memory usage."
|
||||
),
|
||||
):
|
||||
ray.get(ref[0])
|
||||
|
||||
|
||||
def test_dag_fault_tolerance_chain(ray_start_regular):
|
||||
actors = [
|
||||
Actor.remote(0, fail_after=10 if i == 0 else None, sys_exit=False)
|
||||
for i in range(4)
|
||||
]
|
||||
with InputNode() as i:
|
||||
dag = i
|
||||
for a in actors:
|
||||
dag = a.echo.bind(dag)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
for i in range(9):
|
||||
ref = compiled_dag.execute(i)
|
||||
results = ray.get(ref)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
for i in range(9):
|
||||
ref = compiled_dag.execute(i)
|
||||
results = ray.get(ref)
|
||||
assert results == i
|
||||
|
||||
compiled_dag.teardown()
|
||||
|
||||
# All actors are still alive.
|
||||
ray.get([actor.sleep.remote(0) for actor in actors])
|
||||
|
||||
# Remaining actors can be reused.
|
||||
actors.pop(0)
|
||||
with InputNode() as i:
|
||||
dag = i
|
||||
for a in actors:
|
||||
dag = a.echo.bind(dag)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
for i in range(10):
|
||||
ref = compiled_dag.execute(i)
|
||||
results = ray.get(ref)
|
||||
assert results == i
|
||||
|
||||
|
||||
def test_dag_fault_tolerance(ray_start_regular):
|
||||
actors = [
|
||||
Actor.remote(0, fail_after=10 if i == 0 else None, sys_exit=False)
|
||||
for i in range(4)
|
||||
]
|
||||
with InputNode() as i:
|
||||
out = [a.inc.bind(i) for a in actors]
|
||||
dag = MultiOutputNode(out)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
for i in range(9):
|
||||
refs = compiled_dag.execute(1)
|
||||
assert ray.get(refs) == [i + 1] * len(actors)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
for i in range(9, 20):
|
||||
refs = compiled_dag.execute(1)
|
||||
assert ray.get(refs) == [i + 1] * len(actors)
|
||||
|
||||
compiled_dag.teardown()
|
||||
|
||||
# All actors are still alive.
|
||||
ray.get([actor.sleep.remote(0) for actor in actors])
|
||||
|
||||
# Remaining actors can be reused.
|
||||
actors.pop(0)
|
||||
with InputNode() as i:
|
||||
out = [a.inc.bind(i) for a in actors]
|
||||
dag = MultiOutputNode(out)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
for i in range(10):
|
||||
ray.get(compiled_dag.execute(1))
|
||||
|
||||
|
||||
def test_dag_fault_tolerance_sys_exit(ray_start_regular):
|
||||
actors = [
|
||||
Actor.remote(0, fail_after=10 if i == 0 else None, sys_exit=True)
|
||||
for i in range(4)
|
||||
]
|
||||
with InputNode() as i:
|
||||
out = [a.inc.bind(i) for a in actors]
|
||||
dag = MultiOutputNode(out)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
for i in range(9):
|
||||
refs = compiled_dag.execute(1)
|
||||
assert ray.get(refs) == [i + 1] * len(actors)
|
||||
|
||||
with pytest.raises(
|
||||
ActorDiedError, match="The actor died unexpectedly before finishing this task."
|
||||
):
|
||||
for i in range(9):
|
||||
refs = compiled_dag.execute(1)
|
||||
ray.get(refs)
|
||||
|
||||
# Remaining actors are still alive.
|
||||
with pytest.raises(ray.exceptions.RayActorError):
|
||||
ray.get(actors[0].echo.remote("hello"))
|
||||
actors.pop(0)
|
||||
ray.get([actor.echo.remote("hello") for actor in actors])
|
||||
|
||||
# Remaining actors can be reused.
|
||||
with InputNode() as i:
|
||||
out = [a.inc.bind(i) for a in actors]
|
||||
dag = MultiOutputNode(out)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
for i in range(10):
|
||||
refs = compiled_dag.execute(1)
|
||||
ray.get(refs)
|
||||
|
||||
|
||||
def test_dag_teardown_while_running(ray_start_regular):
|
||||
a = Actor.remote(0)
|
||||
|
||||
with InputNode() as inp:
|
||||
dag = a.sleep.bind(inp)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
ref = compiled_dag.execute(3) # 3-second slow task running async
|
||||
compiled_dag.teardown()
|
||||
try:
|
||||
ray.get(ref) # Sanity check the channel doesn't block.
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Check we can still use the actor after first DAG teardown.
|
||||
with InputNode() as inp:
|
||||
dag = a.sleep.bind(inp)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
ref = compiled_dag.execute(0.1)
|
||||
result = ray.get(ref)
|
||||
assert result == 0.1
|
||||
|
||||
|
||||
def test_asyncio_exceptions(ray_start_regular):
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
dag = a.inc.bind(i)
|
||||
|
||||
loop = get_or_create_event_loop()
|
||||
compiled_dag = dag.experimental_compile(enable_asyncio=True)
|
||||
|
||||
async def main():
|
||||
fut = await compiled_dag.execute_async(1)
|
||||
result = await fut
|
||||
assert result == 1
|
||||
|
||||
fut = await compiled_dag.execute_async("hello")
|
||||
with pytest.raises(TypeError) as exc_info:
|
||||
await fut
|
||||
# Traceback should match the original actor class definition.
|
||||
assert "self.i += x" in str(exc_info.value)
|
||||
|
||||
# Can throw an error multiple times.
|
||||
fut = await compiled_dag.execute_async("hello")
|
||||
with pytest.raises(TypeError) as exc_info:
|
||||
await fut
|
||||
# Traceback should match the original actor class definition.
|
||||
assert "self.i += x" in str(exc_info.value)
|
||||
|
||||
# Can use the DAG after exceptions are thrown.
|
||||
fut = await compiled_dag.execute_async(1)
|
||||
result = await fut
|
||||
assert result == 2
|
||||
|
||||
loop.run_until_complete(main())
|
||||
|
||||
|
||||
def test_channel_read_after_close(ray_start_regular):
|
||||
# Tests that read to a channel after Compiled Graph teardown raises a
|
||||
# RayChannelError exception as the channel is closed (see issue #46284).
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def foo(self, arg):
|
||||
return arg
|
||||
|
||||
a = Actor.remote()
|
||||
with InputNode() as inp:
|
||||
dag = a.foo.bind(inp)
|
||||
|
||||
dag = dag.experimental_compile()
|
||||
ref = dag.execute(1)
|
||||
dag.teardown()
|
||||
|
||||
with pytest.raises(RayChannelError, match="Channel closed."):
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
def test_channel_write_after_close(ray_start_regular):
|
||||
# Tests that write to a channel after Compiled Graph teardown raises a
|
||||
# RayChannelError exception as the channel is closed.
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def foo(self, arg):
|
||||
return arg
|
||||
|
||||
a = Actor.remote()
|
||||
with InputNode() as inp:
|
||||
dag = a.foo.bind(inp)
|
||||
|
||||
dag = dag.experimental_compile()
|
||||
dag.teardown()
|
||||
|
||||
with pytest.raises(RayChannelError, match="Channel closed."):
|
||||
dag.execute(1)
|
||||
|
||||
|
||||
def test_multi_arg_exception(shutdown_only):
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
o1, o2 = a.return_two_but_raise_exception.bind(i)
|
||||
dag = MultiOutputNode([o1, o2])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
for _ in range(3):
|
||||
x, y = compiled_dag.execute(1)
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get(x)
|
||||
with pytest.raises(RuntimeError):
|
||||
ray.get(y)
|
||||
|
||||
|
||||
def test_multi_arg_exception_async(shutdown_only):
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
o1, o2 = a.return_two_but_raise_exception.bind(i)
|
||||
dag = MultiOutputNode([o1, o2])
|
||||
|
||||
compiled_dag = dag.experimental_compile(enable_asyncio=True)
|
||||
|
||||
async def main():
|
||||
for _ in range(3):
|
||||
x, y = await compiled_dag.execute_async(1)
|
||||
with pytest.raises(RuntimeError):
|
||||
await x
|
||||
with pytest.raises(RuntimeError):
|
||||
await y
|
||||
|
||||
loop = get_or_create_event_loop()
|
||||
loop.run_until_complete(main())
|
||||
|
||||
|
||||
def test_signature_mismatch(shutdown_only):
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def w(self, x):
|
||||
return 1
|
||||
|
||||
def f(self, x, *, y):
|
||||
pass
|
||||
|
||||
def g(self, x, y, z=1):
|
||||
pass
|
||||
|
||||
worker = Worker.remote()
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match=(
|
||||
r"got an unexpected keyword argument 'y'\. The function `w` has a "
|
||||
r"signature `\(x\)`, but the given arguments to `bind` doesn't match\. "
|
||||
r".*args:.*kwargs:.*"
|
||||
),
|
||||
):
|
||||
with InputNode() as inp:
|
||||
_ = worker.w.bind(inp, y=inp)
|
||||
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match=(
|
||||
r"too many positional arguments\. The function `w` has a signature "
|
||||
r"`\(x\)`, but the given arguments to `bind` doesn't match\. "
|
||||
r"args:.*kwargs:.*"
|
||||
),
|
||||
):
|
||||
with InputNode() as inp:
|
||||
_ = worker.w.bind(inp, inp)
|
||||
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
# Starting from Python 3.12, the error message includes "keyword-only."
|
||||
# Therefore, we need to match both "required keyword-only argument" and
|
||||
# "required argument."
|
||||
match=(
|
||||
r"missing a required (keyword-only )?argument: 'y'\. "
|
||||
r"The function `f` has a signature `\(x, \*, y\)`, "
|
||||
r"but the given arguments to `bind` doesn't match\. "
|
||||
r"args:.*kwargs:.*"
|
||||
),
|
||||
):
|
||||
with InputNode() as inp:
|
||||
_ = worker.f.bind(inp)
|
||||
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match=(
|
||||
r"missing a required argument: 'y'\. The function `g` has a signature "
|
||||
r"`\(x, y, z=1\)`, but the given arguments to `bind` doesn't match\. "
|
||||
r"args:.*kwargs:.*"
|
||||
),
|
||||
):
|
||||
with InputNode() as inp:
|
||||
_ = worker.g.bind(inp)
|
||||
|
||||
|
||||
def test_missing_input_node():
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def f(self, input):
|
||||
return input
|
||||
|
||||
def add(self, a, b):
|
||||
return a + b
|
||||
|
||||
actor = Actor.remote()
|
||||
|
||||
with ray.dag.InputNode() as dag_input:
|
||||
input0, input1, input2 = dag_input[0], dag_input[1], dag_input[2]
|
||||
_ = actor.f.bind(input1)
|
||||
dag = actor.add.bind(input0, input2)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Compiled Graph expects input to be accessed "
|
||||
"using all of attributes 0, 1, 2, "
|
||||
"but 1 is unused. "
|
||||
"Ensure all input attributes are used and contribute "
|
||||
"to the computation of the Compiled Graph output.",
|
||||
):
|
||||
dag.experimental_compile()
|
||||
|
||||
|
||||
def test_sigint_get_dagref(ray_start_cluster):
|
||||
driver_script = """
|
||||
import ray
|
||||
from ray.dag import InputNode
|
||||
import time
|
||||
|
||||
ray.init()
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def sleep(self, x):
|
||||
time.sleep(x)
|
||||
|
||||
a = Actor.remote()
|
||||
with InputNode() as inp:
|
||||
dag = a.sleep.bind(inp)
|
||||
compiled_dag = dag.experimental_compile()
|
||||
ref = compiled_dag.execute(100)
|
||||
print("executing", flush=True)
|
||||
ray.get(ref)
|
||||
"""
|
||||
driver_proc = run_string_as_driver_nonblocking(
|
||||
driver_script, env={"RAY_CGRAPH_teardown_timeout": "0"}
|
||||
)
|
||||
# wait for graph execution to start
|
||||
assert driver_proc.stdout.readline() == b"executing\n"
|
||||
driver_proc.send_signal(signal.SIGINT) # ctrl+c
|
||||
# teardown will kill actors after timeout
|
||||
wait_for_pid_to_exit(driver_proc.pid, 10)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,712 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pydot
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.dag import InputNode, MultiOutputNode
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cleanup_files():
|
||||
"""Clean up files generated during the test."""
|
||||
|
||||
def _cleanup_files(filename: str):
|
||||
for ext in ["", ".png", ".pdf", ".jpeg", ".dot"]:
|
||||
file_path = filename + ext
|
||||
if os.path.exists(file_path):
|
||||
os.remove(file_path)
|
||||
|
||||
return _cleanup_files
|
||||
|
||||
|
||||
def test_visualize_basic(ray_start_regular, cleanup_files):
|
||||
"""
|
||||
Expect output or dot_source:
|
||||
MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
|
||||
0 -> 1 [label=SharedMemoryType]
|
||||
1 -> 2 [label=SharedMemoryType]
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def echo(self, x):
|
||||
return x
|
||||
|
||||
actor = Actor.remote()
|
||||
|
||||
with InputNode() as i:
|
||||
dag = actor.echo.bind(i)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Call the visualize method
|
||||
dot_source = compiled_dag.visualize()
|
||||
|
||||
graphs = pydot.graph_from_dot_data(dot_source)
|
||||
graph = graphs[0]
|
||||
|
||||
node_names = {node.get_name() for node in graph.get_nodes()}
|
||||
edge_pairs = {
|
||||
(edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
|
||||
}
|
||||
|
||||
expected_nodes = {"0", "1", "2"}
|
||||
assert expected_nodes.issubset(
|
||||
node_names
|
||||
), f"Expected nodes {expected_nodes} not found."
|
||||
|
||||
expected_edges = {("0", "1"), ("1", "2")}
|
||||
assert expected_edges.issubset(
|
||||
edge_pairs
|
||||
), f"Expected edges {expected_edges} not found."
|
||||
|
||||
cleanup_files("compiled_graph")
|
||||
|
||||
|
||||
def test_visualize_multi_return(ray_start_regular, cleanup_files):
|
||||
"""
|
||||
Expect output or dot_source:
|
||||
MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
|
||||
0 -> 1 [label=SharedMemoryType]
|
||||
1 -> 2 [label=SharedMemoryType]
|
||||
1 -> 3 [label=SharedMemoryType]
|
||||
2 -> 4 [label=SharedMemoryType]
|
||||
3 -> 4 [label=SharedMemoryType]
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
@ray.method(num_returns=2)
|
||||
def return_two(self, x):
|
||||
return x, x + 1
|
||||
|
||||
actor = Actor.remote()
|
||||
|
||||
with InputNode() as i:
|
||||
o1, o2 = actor.return_two.bind(i)
|
||||
dag = MultiOutputNode([o1, o2])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Get the DOT source
|
||||
dot_source = compiled_dag.visualize()
|
||||
|
||||
graphs = pydot.graph_from_dot_data(dot_source)
|
||||
graph = graphs[0]
|
||||
|
||||
node_names = {node.get_name() for node in graph.get_nodes()}
|
||||
edge_pairs = {
|
||||
(edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
|
||||
}
|
||||
|
||||
expected_nodes = {"0", "1", "2", "3", "4"}
|
||||
assert expected_nodes.issubset(
|
||||
node_names
|
||||
), f"Expected nodes {expected_nodes} not found."
|
||||
|
||||
expected_edges = {("0", "1"), ("1", "2"), ("1", "3"), ("2", "4"), ("3", "4")}
|
||||
assert expected_edges.issubset(
|
||||
edge_pairs
|
||||
), f"Expected edges {expected_edges} not found."
|
||||
cleanup_files("compiled_graph")
|
||||
|
||||
|
||||
def test_visualize_multi_return2(ray_start_regular, cleanup_files):
|
||||
"""
|
||||
Expect output or dot_source:
|
||||
MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
|
||||
0 -> 1 [label=SharedMemoryType]
|
||||
1 -> 2 [label=SharedMemoryType]
|
||||
1 -> 3 [label=SharedMemoryType]
|
||||
2 -> 4 [label=SharedMemoryType]
|
||||
3 -> 5 [label=SharedMemoryType]
|
||||
4 -> 6 [label=SharedMemoryType]
|
||||
5 -> 6 [label=SharedMemoryType]
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
@ray.method(num_returns=2)
|
||||
def return_two(self, x):
|
||||
return x, x + 1
|
||||
|
||||
def echo(self, x):
|
||||
return x
|
||||
|
||||
a = Actor.remote()
|
||||
b = Actor.remote()
|
||||
with InputNode() as i:
|
||||
o1, o2 = a.return_two.bind(i)
|
||||
o3 = b.echo.bind(o1)
|
||||
o4 = b.echo.bind(o2)
|
||||
dag = MultiOutputNode([o3, o4])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Get the DOT source
|
||||
dot_source = compiled_dag.visualize()
|
||||
|
||||
graphs = pydot.graph_from_dot_data(dot_source)
|
||||
graph = graphs[0]
|
||||
|
||||
node_names = {node.get_name() for node in graph.get_nodes()}
|
||||
edge_pairs = {
|
||||
(edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
|
||||
}
|
||||
|
||||
expected_nodes = {"0", "1", "2", "3", "4", "5", "6"}
|
||||
assert expected_nodes.issubset(
|
||||
node_names
|
||||
), f"Expected nodes {expected_nodes} not found."
|
||||
|
||||
expected_edges = {
|
||||
("0", "1"),
|
||||
("1", "2"),
|
||||
("1", "3"),
|
||||
("2", "4"),
|
||||
("3", "5"),
|
||||
("4", "6"),
|
||||
("5", "6"),
|
||||
}
|
||||
assert expected_edges.issubset(
|
||||
edge_pairs
|
||||
), f"Expected edges {expected_edges} not found."
|
||||
cleanup_files("compiled_graph")
|
||||
|
||||
|
||||
def test_visualize_multi_input_nodes(ray_start_regular, cleanup_files):
|
||||
"""
|
||||
Expect output or dot_source:
|
||||
MultiOutputNode" fillcolor=yellow shape=rectangle style=filled]
|
||||
0 -> 1
|
||||
0 -> 2
|
||||
0 -> 3
|
||||
1 -> 4
|
||||
2 -> 5
|
||||
3 -> 6
|
||||
4 -> 7
|
||||
5 -> 7
|
||||
6 -> 7
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def echo(self, x):
|
||||
return x
|
||||
|
||||
actor = Actor.remote()
|
||||
|
||||
with InputNode() as inp:
|
||||
o1 = actor.echo.bind(inp.x)
|
||||
o2 = actor.echo.bind(inp.y)
|
||||
o3 = actor.echo.bind(inp.z)
|
||||
dag = MultiOutputNode([o1, o2, o3])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Get the DOT source
|
||||
dot_source = compiled_dag.visualize()
|
||||
|
||||
graphs = pydot.graph_from_dot_data(dot_source)
|
||||
graph = graphs[0]
|
||||
|
||||
node_names = {node.get_name() for node in graph.get_nodes()}
|
||||
edge_pairs = {
|
||||
(edge.get_source(), edge.get_destination()) for edge in graph.get_edges()
|
||||
}
|
||||
|
||||
expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7"}
|
||||
assert expected_nodes.issubset(
|
||||
node_names
|
||||
), f"Expected nodes {expected_nodes} not found."
|
||||
|
||||
expected_edges = {
|
||||
("0", "1"),
|
||||
("0", "2"),
|
||||
("0", "3"),
|
||||
("1", "4"),
|
||||
("2", "5"),
|
||||
("3", "6"),
|
||||
("4", "7"),
|
||||
("5", "7"),
|
||||
("6", "7"),
|
||||
}
|
||||
assert expected_edges.issubset(
|
||||
edge_pairs
|
||||
), f"Expected edges {expected_edges} not found."
|
||||
cleanup_files("compiled_graph")
|
||||
|
||||
|
||||
class TestVisualizationAscii:
|
||||
|
||||
"""Tests for the visualize_ascii method of compiled DAGs."""
|
||||
|
||||
@staticmethod
|
||||
def parse_ascii_visualization(ascii_visualization: str):
|
||||
"""
|
||||
Parses the ASCII visualization output to extract node names and edge pairs.
|
||||
|
||||
Args:
|
||||
ascii_visualization: The ASCII visualization
|
||||
output generated by the `visualize` function.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing:
|
||||
- node_names: A set of strings representing node names.
|
||||
- edge_pairs: A set of tuples representing edge
|
||||
pairs with type hints.
|
||||
"""
|
||||
import re
|
||||
|
||||
# Sets to store unique nodes and edges
|
||||
node_names = set()
|
||||
edge_pairs = set()
|
||||
|
||||
# Extract nodes from "Nodes Information" section
|
||||
node_pattern = re.compile(r'^(\d+) \[label="Task \d+')
|
||||
edge_pattern = re.compile(r"^(\d+) (--->|\+\+\+>) (\d+)")
|
||||
|
||||
lines = ascii_visualization.splitlines()
|
||||
in_nodes_section = False
|
||||
in_edges_section = False
|
||||
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
|
||||
# Check for nodes section
|
||||
if line.startswith("Nodes Information:"):
|
||||
in_nodes_section = True
|
||||
in_edges_section = False
|
||||
continue
|
||||
|
||||
# Check for edges section
|
||||
if line.startswith("Edges Information:"):
|
||||
in_edges_section = True
|
||||
in_nodes_section = False
|
||||
continue
|
||||
|
||||
# Collect nodes
|
||||
if in_nodes_section:
|
||||
node_match = node_pattern.match(line)
|
||||
if node_match:
|
||||
node_id = node_match.group(1)
|
||||
node_names.add(node_id)
|
||||
|
||||
# Collect edges
|
||||
if in_edges_section:
|
||||
edge_match = edge_pattern.match(line)
|
||||
if edge_match:
|
||||
from_node, _, to_node = edge_match.groups()
|
||||
edge_pairs.add((from_node, to_node))
|
||||
|
||||
return node_names, edge_pairs
|
||||
|
||||
def test_visualize_ascii_basic(self, ray_start_regular):
|
||||
"""
|
||||
Expect output:
|
||||
Nodes Information:
|
||||
0 [label="Task 0 InputNode"]
|
||||
1 [label="Task 1 Actor: d6c5c4... Method: echo"]
|
||||
2 [label="Task 2 MultiOutputNode"]
|
||||
|
||||
Edges Information:
|
||||
0 ---> 1
|
||||
1 ---> 2
|
||||
|
||||
Legend:
|
||||
+++> : Represents Nccl-type data channels
|
||||
---> : Represents Shared Memory data channels
|
||||
|
||||
Experimental Graph:
|
||||
0:InputNode
|
||||
|
|
||||
1:Actor_d6c5c4:echo
|
||||
|
|
||||
2:MultiOutputNode
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def echo(self, x):
|
||||
return x
|
||||
|
||||
actor = Actor.remote()
|
||||
|
||||
with InputNode() as i:
|
||||
dag = actor.echo.bind(i)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Call the visualize method
|
||||
ascii_visualization = compiled_dag.visualize(format="ascii")
|
||||
|
||||
node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
|
||||
ascii_visualization
|
||||
)
|
||||
print(node_names, edge_pairs)
|
||||
expected_nodes = {"0", "1", "2"}
|
||||
assert expected_nodes.issubset(
|
||||
node_names
|
||||
), f"Expected nodes {expected_nodes} not found."
|
||||
|
||||
expected_edges = {("0", "1"), ("1", "2")}
|
||||
assert expected_edges.issubset(
|
||||
edge_pairs
|
||||
), f"Expected edges {expected_edges} not found."
|
||||
|
||||
def test_visualize_ascii_multi_return(self, ray_start_regular):
|
||||
"""
|
||||
Expect output:
|
||||
Nodes Information:
|
||||
0 [label="Task 0 InputNode"]
|
||||
1 [label="Task 1 Actor: 885f1d... Method: return_two"]
|
||||
2 [label="Task 2 ClassMethodOutputNode[0]"]
|
||||
3 [label="Task 3 ClassMethodOutputNode[1]"]
|
||||
4 [label="Task 4 MultiOutputNode"]
|
||||
|
||||
Edges Information:
|
||||
0 ---> 1
|
||||
1 ---> 2
|
||||
1 ---> 3
|
||||
2 ---> 4
|
||||
3 ---> 4
|
||||
|
||||
Legend:
|
||||
+++> : Represents Nccl-type data channels
|
||||
---> : Represents Shared Memory data channels
|
||||
|
||||
Graph Built:
|
||||
0:InputNode
|
||||
|
|
||||
1:Actor_885f1d:return_two
|
||||
|---------------------------->|
|
||||
2:Output[0] 3:Output[1]
|
||||
|<----------------------------|
|
||||
4:MultiOutputNode
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
@ray.method(num_returns=2)
|
||||
def return_two(self, x):
|
||||
return x, x + 1
|
||||
|
||||
actor = Actor.remote()
|
||||
|
||||
with InputNode() as i:
|
||||
o1, o2 = actor.return_two.bind(i)
|
||||
dag = MultiOutputNode([o1, o2])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
ascii_visualization = compiled_dag.visualize(format="ascii")
|
||||
|
||||
node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
|
||||
ascii_visualization
|
||||
)
|
||||
|
||||
expected_nodes = {"0", "1", "2", "3", "4"}
|
||||
assert expected_nodes.issubset(
|
||||
node_names
|
||||
), f"Expected nodes {expected_nodes} not found."
|
||||
|
||||
expected_edges = {("0", "1"), ("1", "2"), ("1", "3"), ("2", "4"), ("3", "4")}
|
||||
assert expected_edges.issubset(
|
||||
edge_pairs
|
||||
), f"Expected edges {expected_edges} not found."
|
||||
|
||||
def test_visualize_ascii_multi_return2(self, ray_start_regular):
|
||||
"""
|
||||
Expect output:
|
||||
Nodes Information:
|
||||
0 [label="Task 0 InputNode"]
|
||||
1 [label="Task 1 Actor: f3e919... Method: return_two"]
|
||||
2 [label="Task 2 ClassMethodOutputNode[0]"]
|
||||
3 [label="Task 3 ClassMethodOutputNode[1]"]
|
||||
4 [label="Task 4 Actor: 15ec69... Method: echo"]
|
||||
5 [label="Task 5 Actor: 15ec69... Method: echo"]
|
||||
6 [label="Task 6 MultiOutputNode"]
|
||||
|
||||
Edges Information:
|
||||
0 ---> 1
|
||||
1 ---> 2
|
||||
1 ---> 3
|
||||
2 ---> 4
|
||||
3 ---> 5
|
||||
4 ---> 6
|
||||
5 ---> 6
|
||||
|
||||
Legend:
|
||||
+++> : Represents Nccl-type data channels
|
||||
---> : Represents Shared Memory data channels
|
||||
|
||||
Graph Built:
|
||||
0:InputNode
|
||||
|
|
||||
1:Actor_f3e919:return_two
|
||||
|---------------------------->|
|
||||
2:Output[0] 3:Output[1]
|
||||
| |
|
||||
4:Actor_15ec69:echo 5:Actor_15ec69:echo
|
||||
|<----------------------------|
|
||||
6:MultiOutputNode
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
@ray.method(num_returns=2)
|
||||
def return_two(self, x):
|
||||
return x, x + 1
|
||||
|
||||
def echo(self, x):
|
||||
return x
|
||||
|
||||
a = Actor.remote()
|
||||
b = Actor.remote()
|
||||
with InputNode() as i:
|
||||
o1, o2 = a.return_two.bind(i)
|
||||
o3 = b.echo.bind(o1)
|
||||
o4 = b.echo.bind(o2)
|
||||
dag = MultiOutputNode([o3, o4])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
ascii_visualization = compiled_dag.visualize(format="ascii")
|
||||
|
||||
node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
|
||||
ascii_visualization
|
||||
)
|
||||
|
||||
expected_nodes = {"0", "1", "2", "3", "4", "5", "6"}
|
||||
assert expected_nodes.issubset(
|
||||
node_names
|
||||
), f"Expected nodes {expected_nodes} not found."
|
||||
|
||||
expected_edges = {
|
||||
("0", "1"),
|
||||
("1", "2"),
|
||||
("1", "3"),
|
||||
("2", "4"),
|
||||
("3", "5"),
|
||||
("4", "6"),
|
||||
("5", "6"),
|
||||
}
|
||||
assert expected_edges.issubset(
|
||||
edge_pairs
|
||||
), f"Expected edges {expected_edges} not found."
|
||||
|
||||
def test_visualize_ascii_complicate(self, ray_start_regular):
|
||||
"""
|
||||
Expect output:
|
||||
Nodes Information:
|
||||
0 [label="Task 0 InputNode"]
|
||||
1 [label="Task 1 Actor: 54777d... Method: return_three"]
|
||||
2 [label="Task 2 ClassMethodOutputNode[0]"]
|
||||
3 [label="Task 3 ClassMethodOutputNode[1]"]
|
||||
4 [label="Task 4 ClassMethodOutputNode[2]"]
|
||||
5 [label="Task 5 Actor: c927c9... Method: echo"]
|
||||
6 [label="Task 6 Actor: c927c9... Method: echo"]
|
||||
7 [label="Task 7 Actor: c927c9... Method: return_two"]
|
||||
8 [label="Task 8 MultiOutputNode"]
|
||||
9 [label="Task 9 ClassMethodOutputNode[0]"]
|
||||
10 [label="Task 10 ClassMethodOutputNode[1]"]
|
||||
|
||||
Edges Information:
|
||||
0 ---> 1
|
||||
1 ---> 2
|
||||
1 ---> 3
|
||||
1 ---> 4
|
||||
2 ---> 5
|
||||
3 ---> 6
|
||||
4 ---> 7
|
||||
5 ---> 8
|
||||
6 ---> 8
|
||||
9 ---> 8
|
||||
10 ---> 8
|
||||
7 ---> 9
|
||||
7 ---> 10
|
||||
|
||||
Legend:
|
||||
+++> : Represents Nccl-type data channels
|
||||
---> : Represents Shared Memory data channels
|
||||
|
||||
Graph Built:
|
||||
0:InputNode
|
||||
|
|
||||
1:Actor_54777d:return_three
|
||||
|---------------------------->|---------------------------->| # noqa
|
||||
2:Output[0] 3:Output[1] 4:Output[2] # noqa
|
||||
| | | # noqa
|
||||
5:Actor_c927c9:echo 6:Actor_c927c9:echo 7:Actor_c927c9:return_two # noqa
|
||||
| | |---------------------------->| # noqa
|
||||
| | 9:Output[0] 10:Output[1] # noqa
|
||||
|<----------------------------|-----------------------------|-----------------------------| # noqa
|
||||
8:MultiOutputNode
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
@ray.method(num_returns=3)
|
||||
def return_three(self, x):
|
||||
return x, x + 1, x + 2
|
||||
|
||||
def echo(self, x):
|
||||
return x
|
||||
|
||||
@ray.method(num_returns=2)
|
||||
def return_two(self, x):
|
||||
return x, x + 1
|
||||
|
||||
a = Actor.remote()
|
||||
b = Actor.remote()
|
||||
with InputNode() as i:
|
||||
o1, o2, o3 = a.return_three.bind(i)
|
||||
o4 = b.echo.bind(o1)
|
||||
o5 = b.echo.bind(o2)
|
||||
o6, o7 = b.return_two.bind(o3)
|
||||
dag = MultiOutputNode([o4, o5, o6, o7])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
ascii_visualization = compiled_dag.visualize(format="ascii")
|
||||
|
||||
node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
|
||||
ascii_visualization
|
||||
)
|
||||
|
||||
expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"}
|
||||
assert expected_nodes.issubset(
|
||||
node_names
|
||||
), f"Expected nodes {expected_nodes} not found."
|
||||
|
||||
expected_edges = {
|
||||
("0", "1"),
|
||||
("1", "2"),
|
||||
("1", "3"),
|
||||
("1", "4"),
|
||||
("2", "5"),
|
||||
("3", "6"),
|
||||
("4", "7"),
|
||||
("5", "8"),
|
||||
("6", "8"),
|
||||
("9", "8"),
|
||||
("10", "8"),
|
||||
("7", "9"),
|
||||
("7", "10"),
|
||||
}
|
||||
assert expected_edges.issubset(
|
||||
edge_pairs
|
||||
), f"Expected edges {expected_edges} not found."
|
||||
|
||||
def test_visualize_ascii_cross_line(self, ray_start_regular):
|
||||
"""
|
||||
Expect output:
|
||||
Nodes Information:
|
||||
0 [label="Task 0 InputNode"]
|
||||
1 [label="Task 1 Actor: 84835a... Method: return_three"]
|
||||
2 [label="Task 2 ClassMethodOutputNode[0]"]
|
||||
3 [label="Task 3 ClassMethodOutputNode[1]"]
|
||||
4 [label="Task 4 ClassMethodOutputNode[2]"]
|
||||
5 [label="Task 5 Actor: 02a6a1... Method: echo"]
|
||||
6 [label="Task 6 Actor: 02a6a1... Method: return_two"]
|
||||
7 [label="Task 7 Actor: 02a6a1... Method: echo"]
|
||||
8 [label="Task 8 MultiOutputNode"]
|
||||
9 [label="Task 9 ClassMethodOutputNode[0]"]
|
||||
10 [label="Task 10 ClassMethodOutputNode[1]"]
|
||||
|
||||
Edges Information:
|
||||
0 ---> 1
|
||||
1 ---> 2
|
||||
1 ---> 3
|
||||
1 ---> 4
|
||||
2 ---> 5
|
||||
3 ---> 6
|
||||
4 ---> 7
|
||||
5 ---> 8
|
||||
7 ---> 8
|
||||
9 ---> 8
|
||||
10 ---> 8
|
||||
6 ---> 9
|
||||
6 ---> 10
|
||||
|
||||
Legend:
|
||||
+++> : Represents Nccl-type data channels
|
||||
---> : Represents Shared Memory data channels
|
||||
|
||||
Graph Built:
|
||||
0:InputNode
|
||||
|
|
||||
1:Actor_84835a:return_three
|
||||
|---------------------------->|---------------------------->| # noqa
|
||||
2:Output[0] 3:Output[1] 4:Output[2] # noqa
|
||||
| | | # noqa
|
||||
5:Actor_02a6a1:echo 6:Actor_02a6a1:return_two 7:Actor_02a6a1:echo # noqa
|
||||
| |---------------------------->| # noqa
|
||||
| 9:Output[0] 10:Output[1] # noqa
|
||||
|<----------------------------------------------------------| # noqa
|
||||
8:MultiOutputNod
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
@ray.method(num_returns=3)
|
||||
def return_three(self, x):
|
||||
return x, x + 1, x + 2
|
||||
|
||||
def echo(self, x):
|
||||
return x
|
||||
|
||||
@ray.method(num_returns=2)
|
||||
def return_two(self, x):
|
||||
return x, x + 1
|
||||
|
||||
a = Actor.remote()
|
||||
b = Actor.remote()
|
||||
with InputNode() as i:
|
||||
o1, o2, o3 = a.return_three.bind(i)
|
||||
o4 = b.echo.bind(o1)
|
||||
o5 = b.echo.bind(o3)
|
||||
o6, o7 = b.return_two.bind(o2)
|
||||
dag = MultiOutputNode([o4, o5, o6, o7])
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
ascii_visualization = compiled_dag.visualize(format="ascii")
|
||||
|
||||
node_names, edge_pairs = TestVisualizationAscii.parse_ascii_visualization(
|
||||
ascii_visualization
|
||||
)
|
||||
|
||||
expected_nodes = {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10"}
|
||||
assert expected_nodes.issubset(
|
||||
node_names
|
||||
), f"Expected nodes {expected_nodes} not found."
|
||||
|
||||
expected_edges = {
|
||||
("0", "1"),
|
||||
("1", "2"),
|
||||
("1", "3"),
|
||||
("1", "4"),
|
||||
("2", "5"),
|
||||
("3", "6"),
|
||||
("4", "7"),
|
||||
("5", "8"),
|
||||
("7", "8"),
|
||||
("9", "8"),
|
||||
("10", "8"),
|
||||
("6", "9"),
|
||||
("6", "10"),
|
||||
}
|
||||
assert expected_edges.issubset(
|
||||
edge_pairs
|
||||
), f"Expected edges {expected_edges} not found."
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if os.environ.get("PARALLEL_CI"):
|
||||
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
||||
else:
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,469 @@
|
||||
# coding: utf-8
|
||||
import os
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
import ray.cluster_utils
|
||||
from ray.dag import InputNode, MultiOutputNode
|
||||
from ray.dag.compiled_dag_node import CompiledDAG
|
||||
from ray.dag.dag_node_operation import _DAGNodeOperationType
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
if sys.platform != "linux" and sys.platform != "darwin":
|
||||
pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
|
||||
|
||||
USE_GPU = os.environ.get("RAY_PYTEST_USE_GPU") == "1"
|
||||
|
||||
if not USE_GPU:
|
||||
pytest.skip("Skipping, these tests require GPUs.", allow_module_level=True)
|
||||
|
||||
|
||||
@ray.remote(num_cpus=0, num_gpus=1)
|
||||
class Worker:
|
||||
def __init__(self, rank: Optional[int] = None):
|
||||
self.rank = rank
|
||||
self.trace = []
|
||||
|
||||
def fwd(self, value):
|
||||
self.trace.append(("FWD", self.rank))
|
||||
return value
|
||||
|
||||
def bwd(self, value):
|
||||
self.trace.append(("BWD", self.rank))
|
||||
return value
|
||||
|
||||
def pop_trace(self):
|
||||
trace = self.trace
|
||||
self.trace = []
|
||||
return trace
|
||||
|
||||
def read_input(self, input):
|
||||
return input
|
||||
|
||||
def send(self, shape, dtype, value: int, send_tensor=True):
|
||||
if not send_tensor:
|
||||
return 1
|
||||
return torch.ones(shape, dtype=dtype, device=self.device) * value
|
||||
|
||||
def recv(self, tensor):
|
||||
# Check that tensor got loaded to the correct device.
|
||||
assert tensor.device == self.device
|
||||
return (tensor[0].item(), tensor.shape, tensor.dtype)
|
||||
|
||||
def no_op(self, value):
|
||||
return value
|
||||
|
||||
def no_op_two(self, value1, value2):
|
||||
return value1, value2
|
||||
|
||||
|
||||
def generate_1f1b_dag(
|
||||
num_workers: int, num_microbatches: int, num_lead_microbatches: int
|
||||
) -> CompiledDAG:
|
||||
workers = [Worker.remote(rank) for rank in range(num_workers)]
|
||||
|
||||
with ray.dag.InputNode() as inp:
|
||||
fwd_queues = [[] for _ in range(num_workers)]
|
||||
bwd_queues = [[] for _ in range(num_workers)]
|
||||
# Once a worker's counter reaches 0, it cannot execute another fwd until it
|
||||
# executes a bwd first.
|
||||
fwd_counter = [num_lead_microbatches - i for i in range(num_workers)]
|
||||
# All of the done batches.
|
||||
done = []
|
||||
|
||||
# FWD on worker 0.
|
||||
input_data = workers[0].read_input.bind(inp)
|
||||
for i in range(num_microbatches):
|
||||
fwd_queues[0].append(input_data)
|
||||
|
||||
while len(done) < num_microbatches:
|
||||
for i, worker in enumerate(workers):
|
||||
if fwd_counter[i] > 0 and fwd_queues[i]:
|
||||
b = fwd_queues[i].pop(0)
|
||||
b = worker.fwd.bind(b)
|
||||
if i < num_workers - 1:
|
||||
fwd_queues[i + 1].append(b)
|
||||
# Use NCCL channel for communication between workers.
|
||||
b.with_tensor_transport(transport="nccl")
|
||||
else:
|
||||
bwd_queues[i].append(b)
|
||||
fwd_counter[i] -= 1
|
||||
elif bwd_queues[i]:
|
||||
b = bwd_queues[i].pop(0)
|
||||
b = worker.bwd.bind(b)
|
||||
if i > 0:
|
||||
bwd_queues[i - 1].append(b)
|
||||
# Use NCCL channel for communication between workers.
|
||||
b.with_tensor_transport(transport="nccl")
|
||||
else:
|
||||
done.append(b)
|
||||
fwd_counter[i] += 1
|
||||
dag = ray.dag.MultiOutputNode(done)
|
||||
compiled_dag = dag.experimental_compile()
|
||||
return compiled_dag
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True)
|
||||
@pytest.mark.parametrize("single_fetch", [True, False])
|
||||
def test_simulate_pp_2workers_2batches_1f1b(
|
||||
ray_start_regular, single_fetch, monkeypatch
|
||||
):
|
||||
"""
|
||||
This test simulates a simple 1F1B pipeline parallelism for training with
|
||||
2 workers and 2 batches.
|
||||
|
||||
w1: fwd_b1 fwd_b2 bwd_b1 bwd_b2
|
||||
w2: fwd_b1 bwd_b1 fwd_b2 bwd_b2
|
||||
|
||||
The communication between workers is done using NCCL. The communication
|
||||
within the worker actor is done using IntraProcessChannel.
|
||||
"""
|
||||
if not USE_GPU:
|
||||
pytest.skip("NCCL tests require GPUs")
|
||||
|
||||
w1 = Worker.remote()
|
||||
w2 = Worker.remote()
|
||||
|
||||
with InputNode() as inp:
|
||||
w1_input = w1.read_input.bind(inp)
|
||||
batch_1 = w1.fwd.bind(w1_input)
|
||||
batch_1.with_tensor_transport(transport="nccl")
|
||||
batch_2 = w1.fwd.bind(w1_input)
|
||||
batch_2.with_tensor_transport(transport="nccl")
|
||||
batch_1 = w2.fwd.bind(batch_1)
|
||||
batch_1 = w2.bwd.bind(batch_1)
|
||||
batch_1.with_tensor_transport(transport="nccl")
|
||||
batch_2 = w2.fwd.bind(batch_2)
|
||||
batch_1 = w1.bwd.bind(batch_1)
|
||||
batch_2 = w2.bwd.bind(batch_2)
|
||||
batch_2.with_tensor_transport(transport="nccl")
|
||||
batch_2 = w1.bwd.bind(batch_2)
|
||||
dag = MultiOutputNode([batch_1, batch_2])
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
w1_expected_schedule = [
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
(1, _DAGNodeOperationType.READ),
|
||||
(1, _DAGNodeOperationType.COMPUTE),
|
||||
(1, _DAGNodeOperationType.WRITE),
|
||||
# `w1 (3, READ)` (P2P recv) is scheduled together with
|
||||
# `w2 (1, WRITE)` (P2P send).
|
||||
(3, _DAGNodeOperationType.READ),
|
||||
(2, _DAGNodeOperationType.READ),
|
||||
(2, _DAGNodeOperationType.COMPUTE),
|
||||
(2, _DAGNodeOperationType.WRITE),
|
||||
# `w1 (4, READ)` (P2P recv) is scheduled together with
|
||||
# `w2 (3, WRITE)` (P2P send).
|
||||
(4, _DAGNodeOperationType.READ),
|
||||
(3, _DAGNodeOperationType.COMPUTE),
|
||||
(3, _DAGNodeOperationType.WRITE),
|
||||
(4, _DAGNodeOperationType.COMPUTE),
|
||||
(4, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
w2_expected_schedule = [
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
(1, _DAGNodeOperationType.READ),
|
||||
(1, _DAGNodeOperationType.COMPUTE),
|
||||
(1, _DAGNodeOperationType.WRITE),
|
||||
(2, _DAGNodeOperationType.READ),
|
||||
(2, _DAGNodeOperationType.COMPUTE),
|
||||
(2, _DAGNodeOperationType.WRITE),
|
||||
(3, _DAGNodeOperationType.READ),
|
||||
(3, _DAGNodeOperationType.COMPUTE),
|
||||
(3, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
w1_schedule = compiled_dag.actor_to_execution_schedule[w1]
|
||||
w2_schedule = compiled_dag.actor_to_execution_schedule[w2]
|
||||
|
||||
for schedule, expected_schedule in zip(
|
||||
[w1_schedule, w2_schedule], [w1_expected_schedule, w2_expected_schedule]
|
||||
):
|
||||
assert len(schedule) == len(expected_schedule)
|
||||
for i, operation in enumerate(schedule):
|
||||
assert operation.exec_task_idx == expected_schedule[i][0]
|
||||
assert operation.type == expected_schedule[i][1]
|
||||
|
||||
tensor_cpu = torch.zeros(10, 10)
|
||||
tensor_cuda = tensor_cpu.to("cuda:0")
|
||||
refs = compiled_dag.execute(tensor_cuda)
|
||||
|
||||
if single_fetch:
|
||||
assert len(refs) == 2
|
||||
for ref in refs:
|
||||
tensor = ray.get(ref)
|
||||
assert torch.equal(tensor.cpu(), tensor_cpu)
|
||||
else:
|
||||
tensors = ray.get(refs)
|
||||
assert len(tensors) == 2
|
||||
for tensor in tensors:
|
||||
assert torch.equal(tensor.cpu(), tensor_cpu)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 4}], indirect=True)
|
||||
def test_simulate_pp_4workers_8batches_1f1b(ray_start_regular, monkeypatch):
|
||||
"""
|
||||
This test simulates a 1F1B pipeline parallelism for training with
|
||||
4 workers and 8 batches.
|
||||
"""
|
||||
if not USE_GPU:
|
||||
pytest.skip("NCCL tests require GPUs")
|
||||
|
||||
num_workers, num_microbatches, num_lead_microbatches = 4, 8, 4
|
||||
compiled_dag = generate_1f1b_dag(
|
||||
num_workers, num_microbatches, num_lead_microbatches
|
||||
)
|
||||
|
||||
tensor_cpu = torch.zeros(10, 10)
|
||||
tensor_cuda = tensor_cpu.to("cuda:0")
|
||||
tensors = ray.get(compiled_dag.execute(tensor_cuda))
|
||||
|
||||
assert len(tensors) == num_microbatches
|
||||
for t in tensors:
|
||||
assert torch.equal(t.cpu(), tensor_cpu)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 3}], indirect=True)
|
||||
def test_three_actors_with_nccl_1(ray_start_regular):
|
||||
"""
|
||||
Driver -> a.no_op -> b.no_op -> a.no_op_two -> Driver
|
||||
| |
|
||||
-> c.no_op -
|
||||
"""
|
||||
if not USE_GPU:
|
||||
pytest.skip("NCCL tests require GPUs")
|
||||
|
||||
a = Worker.remote()
|
||||
b = Worker.remote()
|
||||
c = Worker.remote()
|
||||
|
||||
with InputNode() as inp:
|
||||
dag = a.no_op.bind(inp)
|
||||
dag.with_tensor_transport(transport="nccl")
|
||||
branch1 = b.no_op.bind(dag)
|
||||
branch1.with_tensor_transport(transport="nccl")
|
||||
branch2 = c.no_op.bind(dag)
|
||||
branch2.with_tensor_transport(transport="nccl")
|
||||
dag = a.no_op_two.bind(branch1, branch2)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
a_expected_schedule = [
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
(1, _DAGNodeOperationType.READ),
|
||||
(1, _DAGNodeOperationType.COMPUTE),
|
||||
(1, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
b_expected_schedule = [
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
c_expected_schedule = [
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
a_schedule = compiled_dag.actor_to_execution_schedule[a]
|
||||
b_schedule = compiled_dag.actor_to_execution_schedule[b]
|
||||
c_schedule = compiled_dag.actor_to_execution_schedule[c]
|
||||
|
||||
for schedule, expected_schedule in zip(
|
||||
[a_schedule, b_schedule, c_schedule],
|
||||
[a_expected_schedule, b_expected_schedule, c_expected_schedule],
|
||||
):
|
||||
assert len(schedule) == len(expected_schedule)
|
||||
for i, operation in enumerate(schedule):
|
||||
assert operation.exec_task_idx == expected_schedule[i][0]
|
||||
assert operation.type == expected_schedule[i][1]
|
||||
|
||||
tensor_cpu = torch.zeros(10, 10)
|
||||
tensor_cuda = tensor_cpu.to("cuda:0")
|
||||
ref = compiled_dag.execute(tensor_cuda)
|
||||
tensors = ray.get(ref)
|
||||
|
||||
assert len(tensors) == 2
|
||||
for t in tensors:
|
||||
assert torch.equal(t.cpu(), tensor_cpu)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 3}], indirect=True)
|
||||
@pytest.mark.parametrize("single_fetch", [True, False])
|
||||
def test_three_actors_with_nccl_2(ray_start_regular, single_fetch, monkeypatch):
|
||||
"""
|
||||
Driver --> a.no_op -> b.no_op --> Driver
|
||||
| |
|
||||
-> b.no_op -> c.no_op -
|
||||
| |
|
||||
-> c.no_op -> a.no_op -
|
||||
"""
|
||||
if not USE_GPU:
|
||||
pytest.skip("NCCL tests require GPUs")
|
||||
|
||||
a = Worker.remote()
|
||||
b = Worker.remote()
|
||||
c = Worker.remote()
|
||||
|
||||
with InputNode() as inp:
|
||||
branch1 = a.no_op.bind(inp)
|
||||
branch1.with_tensor_transport(transport="nccl")
|
||||
branch2 = b.no_op.bind(inp)
|
||||
branch2.with_tensor_transport(transport="nccl")
|
||||
branch3 = c.no_op.bind(inp)
|
||||
branch3.with_tensor_transport(transport="nccl")
|
||||
dag = MultiOutputNode(
|
||||
[
|
||||
a.no_op.bind(branch3),
|
||||
b.no_op.bind(branch1),
|
||||
c.no_op.bind(branch2),
|
||||
]
|
||||
)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
a_expected_schedule = [
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
(1, _DAGNodeOperationType.READ),
|
||||
(1, _DAGNodeOperationType.COMPUTE),
|
||||
(1, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
b_expected_schedule = [
|
||||
# `b (1, READ)` (P2P recv) is scheduled together with
|
||||
# `a (0, WRITE)` (P2P send).
|
||||
(1, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
(1, _DAGNodeOperationType.COMPUTE),
|
||||
(1, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
c_expected_schedule = [
|
||||
# `c (1, READ)` (P2P recv) is scheduled together with
|
||||
# `a (0, WRITE)` (P2P send).
|
||||
(1, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
(1, _DAGNodeOperationType.COMPUTE),
|
||||
(1, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
|
||||
a_schedule = compiled_dag.actor_to_execution_schedule[a]
|
||||
b_schedule = compiled_dag.actor_to_execution_schedule[b]
|
||||
c_schedule = compiled_dag.actor_to_execution_schedule[c]
|
||||
|
||||
for schedule, expected_schedule in zip(
|
||||
[a_schedule, b_schedule, c_schedule],
|
||||
[a_expected_schedule, b_expected_schedule, c_expected_schedule],
|
||||
):
|
||||
assert len(schedule) == len(expected_schedule)
|
||||
for i, operation in enumerate(schedule):
|
||||
assert operation.exec_task_idx == expected_schedule[i][0]
|
||||
assert operation.type == expected_schedule[i][1]
|
||||
|
||||
tensor_cpu = torch.zeros(10, 10)
|
||||
tensor_cuda = tensor_cpu.to("cuda:0")
|
||||
refs = compiled_dag.execute(tensor_cuda)
|
||||
|
||||
if single_fetch:
|
||||
assert len(refs) == 3
|
||||
for ref in refs:
|
||||
tensor = ray.get(ref)
|
||||
assert torch.equal(tensor.cpu(), tensor_cpu)
|
||||
else:
|
||||
tensors = ray.get(refs)
|
||||
assert len(tensors) == 3
|
||||
for tensor in tensors:
|
||||
assert torch.equal(tensor.cpu(), tensor_cpu)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 3}], indirect=True)
|
||||
@pytest.mark.parametrize("overlap_gpu_communication", [True, False])
|
||||
def test_overlap_gpu_communication(ray_start_regular, overlap_gpu_communication):
|
||||
"""
|
||||
Driver --> sender1.send -> receiver.recv --> Driver
|
||||
| |
|
||||
-> sender2.send -> receiver.recv -
|
||||
"""
|
||||
if not USE_GPU:
|
||||
pytest.skip("NCCL tests require GPUs")
|
||||
|
||||
sender1 = Worker.remote()
|
||||
sender2 = Worker.remote()
|
||||
receiver = Worker.remote()
|
||||
|
||||
shape = (10000,)
|
||||
dtype = torch.float16
|
||||
|
||||
with InputNode() as inp:
|
||||
branch1 = sender1.send.bind(shape, dtype, inp)
|
||||
|
||||
branch1 = branch1.with_tensor_transport(
|
||||
transport="nccl", _static_shape=True, _direct_return=True
|
||||
)
|
||||
branch1 = receiver.recv.bind(branch1)
|
||||
|
||||
branch2 = sender2.send.bind(shape, dtype, inp)
|
||||
branch2 = branch2.with_tensor_transport(
|
||||
transport="nccl", _static_shape=True, _direct_return=True
|
||||
)
|
||||
branch2 = receiver.recv.bind(branch2)
|
||||
dag = MultiOutputNode([branch1, branch2])
|
||||
|
||||
# Test normal execution.
|
||||
compiled_dag = dag.experimental_compile(
|
||||
_overlap_gpu_communication=overlap_gpu_communication
|
||||
)
|
||||
|
||||
# Check receiver schedule
|
||||
expected_no_overlap_schedule = [
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
# `receiver (1, READ)` (P2P recv) is scheduled together with
|
||||
# `sender2 (0, WRITE)` (P2P send).
|
||||
(1, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
(1, _DAGNodeOperationType.COMPUTE),
|
||||
(1, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
expected_overlap_schedule = [
|
||||
(0, _DAGNodeOperationType.READ),
|
||||
# `receiver (1, READ)` (P2P recv) is scheduled together with
|
||||
# `sender2 (0, WRITE)` (P2P send).
|
||||
(1, _DAGNodeOperationType.READ),
|
||||
(0, _DAGNodeOperationType.COMPUTE),
|
||||
(0, _DAGNodeOperationType.WRITE),
|
||||
(1, _DAGNodeOperationType.COMPUTE),
|
||||
(1, _DAGNodeOperationType.WRITE),
|
||||
]
|
||||
if overlap_gpu_communication:
|
||||
expected_receiver_schedule = expected_overlap_schedule
|
||||
else:
|
||||
expected_receiver_schedule = expected_no_overlap_schedule
|
||||
|
||||
receiver_schedule = compiled_dag.actor_to_execution_schedule[receiver]
|
||||
|
||||
assert len(receiver_schedule) == len(expected_receiver_schedule)
|
||||
for i, operation in enumerate(receiver_schedule):
|
||||
assert operation.exec_task_idx == expected_receiver_schedule[i][0]
|
||||
assert operation.type == expected_receiver_schedule[i][1]
|
||||
|
||||
compiled_dag.teardown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if os.environ.get("PARALLEL_CI"):
|
||||
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
||||
else:
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,408 @@
|
||||
# coding: utf-8
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
import ray.cluster_utils
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
from ray.dag import InputNode
|
||||
from ray.exceptions import RayChannelError, RayTaskError
|
||||
from ray.experimental.channel.conftest import (
|
||||
Barrier,
|
||||
start_nccl_mock,
|
||||
)
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
def error_logged(capsys, msg):
|
||||
out, err = capsys.readouterr()
|
||||
# Write captured back to stdout, stderr for easier test debugging.
|
||||
sys.stdout.write(out)
|
||||
sys.stderr.write(err)
|
||||
return msg in err
|
||||
|
||||
|
||||
@ray.remote(num_cpus=0, num_gpus=1)
|
||||
class MockedWorker:
|
||||
def __init__(self):
|
||||
self.chan = None
|
||||
|
||||
def start_mock(self):
|
||||
"""
|
||||
Patch methods that require CUDA.
|
||||
"""
|
||||
start_nccl_mock()
|
||||
|
||||
def send(self, shape, dtype, value: int, send_as_dict=False):
|
||||
if send_as_dict:
|
||||
return self.send_dict([(value, value, shape, dtype)])
|
||||
|
||||
return torch.ones(shape, dtype=dtype) * value
|
||||
|
||||
def recv(self, tensor):
|
||||
if isinstance(tensor, dict):
|
||||
assert len(tensor) == 1
|
||||
tensor = list(tensor.values())[0]
|
||||
|
||||
return (tensor[0].item(), tensor.shape, tensor.dtype)
|
||||
|
||||
def send_dict(self, entries):
|
||||
results = {}
|
||||
for key, value, shape, dtype in entries:
|
||||
results[key] = torch.ones(shape, dtype=dtype) * value
|
||||
return results
|
||||
|
||||
def recv_dict(self, tensor_dict):
|
||||
results = []
|
||||
for key in sorted(tensor_dict.keys()):
|
||||
tensor = tensor_dict[key]
|
||||
results.append((key, tensor[0].item(), tensor.shape, tensor.dtype))
|
||||
return results
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 2,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_p2p(ray_start_cluster):
|
||||
"""
|
||||
Test simple sender -> receiver pattern. Check that receiver receives
|
||||
correct results.
|
||||
"""
|
||||
# Barrier name should be barrier-{lower rank}-{higher rank}.
|
||||
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
||||
|
||||
sender = MockedWorker.remote()
|
||||
receiver = MockedWorker.remote()
|
||||
|
||||
ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
|
||||
|
||||
shape = (10,)
|
||||
dtype = torch.float16
|
||||
|
||||
# Test torch.Tensor sent between actors.
|
||||
with InputNode() as inp:
|
||||
dag = sender.send.bind(inp.shape, inp.dtype, inp[0], inp.send_as_dict)
|
||||
dag = dag.with_tensor_transport(transport="nccl")
|
||||
dag = receiver.recv.bind(dag)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
for i in range(3):
|
||||
ref = compiled_dag.execute(i, shape=shape, dtype=dtype, send_as_dict=False)
|
||||
assert ray.get(ref) == (i, shape, dtype)
|
||||
|
||||
# Sending tensors of different shape also works.
|
||||
for i in range(3):
|
||||
ref = compiled_dag.execute(i, shape=(20,), dtype=dtype, send_as_dict=False)
|
||||
assert ray.get(ref) == (i, (20,), dtype)
|
||||
|
||||
# Sending tensors inside a dictionary also works.
|
||||
for i in range(3):
|
||||
ref = compiled_dag.execute(i, shape=shape, dtype=dtype, send_as_dict=True)
|
||||
assert ray.get(ref) == (i, shape, dtype)
|
||||
|
||||
compiled_dag.teardown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 2,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize("send_as_dict", [True, False])
|
||||
def test_p2p_static_shape(ray_start_cluster, send_as_dict):
|
||||
"""
|
||||
Test simple send -> recv pattern with
|
||||
_static_shape=True. If sender always sends tensors of
|
||||
the same shape, then it works.
|
||||
"""
|
||||
# Barrier name should be barrier-{lower rank}-{higher rank}.
|
||||
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
||||
|
||||
sender = MockedWorker.remote()
|
||||
receiver = MockedWorker.remote()
|
||||
|
||||
ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
|
||||
|
||||
shape = (10,)
|
||||
dtype = torch.float16
|
||||
|
||||
# Test torch.Tensor sent between actors.
|
||||
with InputNode() as inp:
|
||||
dag = sender.send.bind(inp.shape, inp.dtype, inp[0], send_as_dict=send_as_dict)
|
||||
dag = dag.with_tensor_transport(transport="nccl", _static_shape=True)
|
||||
dag = receiver.recv.bind(dag)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
for i in range(3):
|
||||
ref = compiled_dag.execute(i, shape=shape, dtype=dtype)
|
||||
assert ray.get(ref) == (i, shape, dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 2,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize("send_as_dict", [True, False])
|
||||
def test_p2p_static_shape_error(capsys, ray_start_cluster, send_as_dict):
|
||||
"""
|
||||
Test that when static_shape=True, an error is thrown when a tensor with a
|
||||
different shape or dtype is found.
|
||||
"""
|
||||
# Barrier name should be barrier-{lower rank}-{higher rank}.
|
||||
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
||||
|
||||
sender = MockedWorker.remote()
|
||||
receiver = MockedWorker.remote()
|
||||
|
||||
ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
|
||||
|
||||
shape = (10,)
|
||||
dtype = torch.float16
|
||||
|
||||
# Test torch.Tensor sent between actors.
|
||||
with InputNode() as inp:
|
||||
dag = sender.send.bind(inp.shape, inp.dtype, inp[0], send_as_dict=send_as_dict)
|
||||
dag = dag.with_tensor_transport(transport="nccl", _static_shape=True)
|
||||
dag = receiver.recv.bind(dag)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
for i in range(3):
|
||||
ref = compiled_dag.execute(i, shape=shape, dtype=dtype)
|
||||
assert ray.get(ref) == (i, shape, dtype)
|
||||
|
||||
# Sending wrong shape errors.
|
||||
ref = compiled_dag.execute(i, shape=(20,), dtype=dtype)
|
||||
with pytest.raises(RayTaskError):
|
||||
ray.get(ref)
|
||||
|
||||
# Sending correct shape still errors because the DAG has already been torn
|
||||
# down after the previous error.
|
||||
with pytest.raises(RayChannelError):
|
||||
ref = compiled_dag.execute(i, shape=shape, dtype=dtype)
|
||||
|
||||
wait_for_condition(
|
||||
lambda: error_logged(
|
||||
capsys,
|
||||
"ValueError: Expected torch.Tensors with shapes and dtypes: "
|
||||
"[(shape=torch.Size([10]), dtype=torch.float16)], found: "
|
||||
"[(shape=torch.Size([20]), dtype=torch.float16)]",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 2,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_p2p_direct_return(ray_start_cluster):
|
||||
"""
|
||||
Test simple sender -> receiver pattern with _direct_return=True
|
||||
"""
|
||||
# Barrier name should be barrier-{lower rank}-{higher rank}.
|
||||
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
||||
|
||||
sender = MockedWorker.remote()
|
||||
receiver = MockedWorker.remote()
|
||||
|
||||
ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
|
||||
|
||||
# Test torch.Tensor sent between actors.
|
||||
with InputNode() as inp:
|
||||
dag = sender.send.bind(inp.shape, inp.dtype, inp.value, inp.send_as_dict)
|
||||
dag = dag.with_tensor_transport(
|
||||
transport="nccl",
|
||||
_direct_return=True,
|
||||
)
|
||||
dag = receiver.recv.bind(dag)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
dtype = torch.float16
|
||||
for i in range(3):
|
||||
shape = (10 * (i + 1),)
|
||||
ref = compiled_dag.execute(
|
||||
shape=shape, dtype=dtype, value=i, send_as_dict=False
|
||||
)
|
||||
assert ray.get(ref) == (i, shape, dtype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 2,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
def test_p2p_direct_return_error(capsys, ray_start_cluster):
|
||||
"""
|
||||
Test simple sender -> receiver pattern with
|
||||
_direct_return=True. Test that error is thrown when
|
||||
actor task does not return a tensor directly.
|
||||
"""
|
||||
# Barrier name should be barrier-{lower rank}-{higher rank}.
|
||||
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
||||
|
||||
sender = MockedWorker.remote()
|
||||
receiver = MockedWorker.remote()
|
||||
|
||||
ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
|
||||
|
||||
# Test torch.Tensor sent between actors.
|
||||
with InputNode() as inp:
|
||||
dag = sender.send.bind(inp.shape, inp.dtype, inp.value, inp.send_as_dict)
|
||||
dag = dag.with_tensor_transport(
|
||||
transport="nccl",
|
||||
_direct_return=True,
|
||||
)
|
||||
dag = receiver.recv.bind(dag)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
dtype = torch.float16
|
||||
for i in range(3):
|
||||
shape = (10 * (i + 1),)
|
||||
ref = compiled_dag.execute(
|
||||
shape=shape, dtype=dtype, value=i, send_as_dict=False
|
||||
)
|
||||
assert ray.get(ref) == (i, shape, dtype)
|
||||
|
||||
# Error is thrown if we do not send a tensor.
|
||||
ref = compiled_dag.execute(shape=shape, dtype=dtype, value=1, send_as_dict=True)
|
||||
with pytest.raises(RayTaskError):
|
||||
ray.get(ref)
|
||||
|
||||
# Currently the receiver cannot catch the exception so the DAG cannot be
|
||||
# used again.
|
||||
with pytest.raises(RayChannelError):
|
||||
ref = compiled_dag.execute(
|
||||
shape=shape, dtype=dtype, value=1, send_as_dict=False
|
||||
)
|
||||
|
||||
wait_for_condition(
|
||||
lambda: error_logged(
|
||||
capsys,
|
||||
"Task annotated with _direct_return=True must "
|
||||
"return a CUDA torch.Tensor",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ray_start_cluster",
|
||||
[
|
||||
{
|
||||
"num_cpus": 2,
|
||||
"num_gpus": 2,
|
||||
"num_nodes": 1,
|
||||
}
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize("check_static_shape", [True, False])
|
||||
def test_p2p_static_shape_and_direct_return(
|
||||
capsys, ray_start_cluster, check_static_shape
|
||||
):
|
||||
"""
|
||||
Test simple sender -> receiver pattern with both _static_shape=True and
|
||||
_direct_return=True. Check errors are thrown if tensors with wrong shape
|
||||
are passed (check_static_shape=True) OR if non-tensor value is returned
|
||||
(check_static_shape=False).
|
||||
"""
|
||||
# Barrier name should be barrier-{lower rank}-{higher rank}.
|
||||
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
||||
|
||||
sender = MockedWorker.remote()
|
||||
receiver = MockedWorker.remote()
|
||||
|
||||
ray.get([sender.start_mock.remote(), receiver.start_mock.remote()])
|
||||
|
||||
# Test torch.Tensor sent between actors.
|
||||
with InputNode() as inp:
|
||||
dag = sender.send.bind(inp.shape, inp.dtype, inp.value, inp.send_as_dict)
|
||||
dag = dag.with_tensor_transport(
|
||||
transport="nccl",
|
||||
_static_shape=True,
|
||||
_direct_return=True,
|
||||
)
|
||||
dag = receiver.recv.bind(dag)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
shape = (10,)
|
||||
dtype = torch.float16
|
||||
for i in range(3):
|
||||
ref = compiled_dag.execute(
|
||||
shape=shape, dtype=dtype, value=i, send_as_dict=False
|
||||
)
|
||||
assert ray.get(ref) == (i, shape, dtype)
|
||||
|
||||
if check_static_shape:
|
||||
# Error is thrown if we send the wrong shape.
|
||||
ref = compiled_dag.execute(
|
||||
shape=(20,), dtype=dtype, value=1, send_as_dict=False
|
||||
)
|
||||
else:
|
||||
# Error is thrown if we do not send a tensor.
|
||||
ref = compiled_dag.execute(shape=shape, dtype=dtype, value=1, send_as_dict=True)
|
||||
|
||||
with pytest.raises(RayTaskError):
|
||||
ray.get(ref)
|
||||
|
||||
# Currently the receiver cannot catch either kind of
|
||||
# exception so the DAG cannot be used again.
|
||||
with pytest.raises(RayChannelError):
|
||||
ref = compiled_dag.execute(
|
||||
shape=shape, dtype=dtype, value=1, send_as_dict=False
|
||||
)
|
||||
|
||||
if check_static_shape:
|
||||
msg = (
|
||||
"ValueError: Expected torch.Tensors with shapes and dtypes: "
|
||||
"[(shape=torch.Size([10]), dtype=torch.float16)], found: "
|
||||
"[(shape=torch.Size([20]), dtype=torch.float16)]"
|
||||
)
|
||||
else:
|
||||
msg = "Task annotated with _direct_return=True must return a CUDA torch.Tensor"
|
||||
wait_for_condition(lambda: error_logged(capsys, msg))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if os.environ.get("PARALLEL_CI"):
|
||||
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
||||
else:
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,75 @@
|
||||
# coding: utf-8
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
import ray.cluster_utils
|
||||
from ray.dag import InputNode, MultiOutputNode
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
if sys.platform != "linux" and sys.platform != "darwin":
|
||||
pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
|
||||
|
||||
USE_GPU = os.environ.get("RAY_PYTEST_USE_GPU") == "1"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_gpus": 2}], indirect=True)
|
||||
def test_multi_args_simulate_pp(ray_start_regular):
|
||||
if not USE_GPU:
|
||||
pytest.skip("NCCL tests require GPUs")
|
||||
|
||||
@ray.remote(num_cpus=0, num_gpus=1)
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def forward(self, data):
|
||||
return data
|
||||
|
||||
def backward(self, data):
|
||||
return data
|
||||
|
||||
NUM_MICROBATCHES = 2
|
||||
w0 = Worker.remote()
|
||||
w1 = Worker.remote()
|
||||
with InputNode() as dag_input:
|
||||
dag_outs = []
|
||||
for microbatch_idx in range(NUM_MICROBATCHES):
|
||||
microbatch = dag_input[microbatch_idx]
|
||||
stage_fwd_out = w0.forward.bind(microbatch)
|
||||
stage_fwd_out.with_tensor_transport(transport="nccl")
|
||||
stage_fwd_out = w1.forward.bind(stage_fwd_out)
|
||||
dag_outs.append(stage_fwd_out)
|
||||
|
||||
grad_out = dag_input[NUM_MICROBATCHES]
|
||||
for _ in range(NUM_MICROBATCHES):
|
||||
stage_bwd_out = w1.backward.bind(grad_out)
|
||||
stage_bwd_out.with_tensor_transport(transport="nccl")
|
||||
stage_bwd_out = w0.backward.bind(stage_bwd_out)
|
||||
dag_outs.append(stage_bwd_out)
|
||||
|
||||
dag = MultiOutputNode(dag_outs)
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
tensor_cpu_list = [torch.zeros(1, i + 1) for i in range(3)]
|
||||
tensor_cuda_list = [t.to("cuda:0") for t in tensor_cpu_list]
|
||||
ref = compiled_dag.execute(
|
||||
tensor_cuda_list[0], tensor_cuda_list[1], tensor_cuda_list[2]
|
||||
)
|
||||
tensors = ray.get(ref)
|
||||
|
||||
assert len(tensors) == 4
|
||||
assert torch.equal(tensors[0], tensor_cpu_list[0])
|
||||
assert torch.equal(tensors[1], tensor_cpu_list[1])
|
||||
assert torch.equal(tensors[2], tensor_cpu_list[2])
|
||||
assert torch.equal(tensors[3], tensor_cpu_list[2])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if os.environ.get("PARALLEL_CI"):
|
||||
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
||||
else:
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,414 @@
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
import ray.remote_function
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
from ray.dag import InputNode, MultiOutputNode
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
if sys.platform != "linux" and sys.platform != "darwin":
|
||||
pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self, init_value, fail_after=None, sys_exit=False):
|
||||
self.i = init_value
|
||||
self.fail_after = fail_after
|
||||
self.sys_exit = sys_exit
|
||||
|
||||
self.count = 0
|
||||
|
||||
def _fail_if_needed(self):
|
||||
if self.fail_after and self.count > self.fail_after:
|
||||
# Randomize the failures to better cover multi actor scenarios.
|
||||
if random.random() > 0.5:
|
||||
if self.sys_exit:
|
||||
os._exit(1)
|
||||
else:
|
||||
raise RuntimeError("injected fault")
|
||||
|
||||
def inc(self, x):
|
||||
self.i += x
|
||||
self.count += 1
|
||||
self._fail_if_needed()
|
||||
return self.i
|
||||
|
||||
def double_and_inc(self, x):
|
||||
self.i *= 2
|
||||
self.i += x
|
||||
return self.i
|
||||
|
||||
def echo(self, x):
|
||||
print("ECHO!")
|
||||
self.count += 1
|
||||
self._fail_if_needed()
|
||||
return x
|
||||
|
||||
def append_to(self, lst):
|
||||
lst.append(self.i)
|
||||
return lst
|
||||
|
||||
def inc_two(self, x, y):
|
||||
self.i += x
|
||||
self.i += y
|
||||
return self.i
|
||||
|
||||
def sleep(self, x):
|
||||
time.sleep(x)
|
||||
return x
|
||||
|
||||
@ray.method(num_returns=2)
|
||||
def return_two(self, x):
|
||||
return x, x + 1
|
||||
|
||||
|
||||
def test_readers_on_different_nodes(ray_start_cluster):
|
||||
cluster = ray_start_cluster
|
||||
# This node is for the driver (including the CompiledDAG.DAGDriverProxyActor) and
|
||||
# one of the readers.
|
||||
cluster.add_node(num_cpus=1)
|
||||
ray.init(address=cluster.address)
|
||||
# 2 more nodes for other readers.
|
||||
cluster.add_node(num_cpus=1)
|
||||
cluster.add_node(num_cpus=1)
|
||||
cluster.wait_for_nodes()
|
||||
# Wait until nodes actually start, otherwise the code below will fail.
|
||||
wait_for_condition(lambda: len(ray.nodes()) == 3)
|
||||
|
||||
a = Actor.options(num_cpus=1).remote(0)
|
||||
b = Actor.options(num_cpus=1).remote(0)
|
||||
c = Actor.options(num_cpus=1).remote(0)
|
||||
actors = [a, b, c]
|
||||
|
||||
def _get_node_id(self) -> "ray.NodeID":
|
||||
return ray.get_runtime_context().get_node_id()
|
||||
|
||||
node_ids = ray.get([act.__ray_call__.remote(_get_node_id) for act in actors])
|
||||
assert len(set(node_ids)) == 3
|
||||
|
||||
with InputNode() as inp:
|
||||
x = a.inc.bind(inp)
|
||||
y = b.inc.bind(inp)
|
||||
z = c.inc.bind(inp)
|
||||
dag = MultiOutputNode([x, y, z])
|
||||
|
||||
cdag = dag.experimental_compile()
|
||||
|
||||
for i in range(1, 10):
|
||||
assert ray.get(cdag.execute(1)) == [i, i, i]
|
||||
|
||||
|
||||
def test_bunch_readers_on_different_nodes(ray_start_cluster):
|
||||
cluster = ray_start_cluster
|
||||
ACTORS_PER_NODE = 2
|
||||
NUM_REMOTE_NODES = 2
|
||||
# driver node
|
||||
cluster.add_node(num_cpus=ACTORS_PER_NODE)
|
||||
ray.init(address=cluster.address)
|
||||
# additional nodes for multi readers in multi nodes
|
||||
for _ in range(NUM_REMOTE_NODES):
|
||||
cluster.add_node(num_cpus=ACTORS_PER_NODE)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
wait_for_condition(lambda: len(ray.nodes()) == NUM_REMOTE_NODES + 1)
|
||||
|
||||
actors = [
|
||||
Actor.options(num_cpus=1).remote(0)
|
||||
for _ in range(ACTORS_PER_NODE * (NUM_REMOTE_NODES + 1))
|
||||
]
|
||||
|
||||
def _get_node_id(self) -> "ray.NodeID":
|
||||
return ray.get_runtime_context().get_node_id()
|
||||
|
||||
node_ids = ray.get([act.__ray_call__.remote(_get_node_id) for act in actors])
|
||||
assert len(set(node_ids)) == NUM_REMOTE_NODES + 1
|
||||
|
||||
with InputNode() as inp:
|
||||
outputs = []
|
||||
for actor in actors:
|
||||
outputs.append(actor.inc.bind(inp))
|
||||
dag = MultiOutputNode(outputs)
|
||||
|
||||
cdag = dag.experimental_compile()
|
||||
|
||||
for i in range(1, 10):
|
||||
assert ray.get(cdag.execute(1)) == [
|
||||
i for _ in range(ACTORS_PER_NODE * (NUM_REMOTE_NODES + 1))
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("single_fetch", [True, False])
|
||||
def test_pp(ray_start_cluster, single_fetch):
|
||||
cluster = ray_start_cluster
|
||||
# This node is for the driver.
|
||||
cluster.add_node(num_cpus=0)
|
||||
ray.init(address=cluster.address)
|
||||
TP = 2
|
||||
# This node is for the PP stage 1.
|
||||
cluster.add_node(resources={"pp1": TP})
|
||||
# This node is for the PP stage 2.
|
||||
cluster.add_node(resources={"pp2": TP})
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def execute_model(self, val):
|
||||
return val
|
||||
|
||||
pp1_workers = [
|
||||
Worker.options(num_cpus=0, resources={"pp1": 1}).remote() for _ in range(TP)
|
||||
]
|
||||
pp2_workers = [
|
||||
Worker.options(num_cpus=0, resources={"pp2": 1}).remote() for _ in range(TP)
|
||||
]
|
||||
|
||||
with InputNode() as inp:
|
||||
outputs = [inp for _ in range(TP)]
|
||||
outputs = [pp1_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
|
||||
outputs = [pp2_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
|
||||
dag = MultiOutputNode(outputs)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
refs = compiled_dag.execute(1)
|
||||
if single_fetch:
|
||||
for i in range(TP):
|
||||
assert ray.get(refs[i]) == 1
|
||||
else:
|
||||
assert ray.get(refs) == [1] * TP
|
||||
|
||||
# So that raylets' error messages are printed to the driver
|
||||
time.sleep(2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("single_fetch", [True, False])
|
||||
def test_pp_exception(ray_start_cluster, single_fetch):
|
||||
"""
|
||||
This test is to verify that the exception can be passed properly
|
||||
through pipeline parallel workers on different nodes.
|
||||
"""
|
||||
cluster = ray_start_cluster
|
||||
# This node is for the driver.
|
||||
cluster.add_node(num_cpus=0)
|
||||
ray.init(address=cluster.address)
|
||||
TP = 2
|
||||
# This node is for the PP stage 1.
|
||||
cluster.add_node(resources={"pp1": TP})
|
||||
# This node is for the PP stage 2.
|
||||
cluster.add_node(resources={"pp2": TP})
|
||||
# This node is for the PP stage 3.
|
||||
cluster.add_node(resources={"pp3": TP})
|
||||
|
||||
# Simulate a large error message (e.g., those with a long stack trace)
|
||||
large_error_message = "Model execution failed" * 10000
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def execute_model(self, val):
|
||||
if val == "exception_trigger":
|
||||
# Simulate an exception happened during model execution
|
||||
raise RuntimeError(large_error_message)
|
||||
return val
|
||||
|
||||
pp1_workers = [
|
||||
Worker.options(num_cpus=0, resources={"pp1": 1}).remote() for _ in range(TP)
|
||||
]
|
||||
pp2_workers = [
|
||||
Worker.options(num_cpus=0, resources={"pp2": 1}).remote() for _ in range(TP)
|
||||
]
|
||||
pp3_workers = [
|
||||
Worker.options(num_cpus=0, resources={"pp3": 1}).remote() for _ in range(TP)
|
||||
]
|
||||
|
||||
with InputNode() as inp:
|
||||
outputs = [inp for _ in range(TP)]
|
||||
outputs = [pp1_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
|
||||
outputs = [pp2_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
|
||||
outputs = [pp3_workers[i].execute_model.bind(outputs[i]) for i in range(TP)]
|
||||
dag = MultiOutputNode(outputs)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
refs = compiled_dag.execute("exception_trigger")
|
||||
|
||||
# Without the fix in this PR, we will encounter the following exception:
|
||||
# File "/Users/ruiqiao/repos2/ray/python/ray/_private/serialization.py",
|
||||
# line 460, in deserialize_objects
|
||||
# obj = self._deserialize_object(data, metadata, object_ref)
|
||||
# raise Exception(
|
||||
# Exception: Can't deserialize object:
|
||||
# ObjectRef(00a33d534c5b0ce51bdf175790467da3114801680100000002e1f505), metadata: b'\x00'
|
||||
# With this fix, the original exception will be propagated.
|
||||
if single_fetch:
|
||||
for i in range(TP):
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
ray.get(refs[i])
|
||||
assert "Can't deserialize object" not in str(exc_info.value)
|
||||
assert large_error_message in str(exc_info.value)
|
||||
else:
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
ray.get(refs)
|
||||
assert "Can't deserialize object" not in str(exc_info.value)
|
||||
assert large_error_message in str(exc_info.value)
|
||||
|
||||
|
||||
def test_payload_large(ray_start_cluster, monkeypatch):
|
||||
GRPC_MAX_SIZE = 1024 * 1024 * 5
|
||||
monkeypatch.setenv("RAY_max_grpc_message_size", str(GRPC_MAX_SIZE))
|
||||
cluster = ray_start_cluster
|
||||
# This node is for the driver (including the CompiledDAG.DAGDriverProxyActor).
|
||||
first_node_handle = cluster.add_node(num_cpus=1)
|
||||
# This node is for the reader.
|
||||
second_node_handle = cluster.add_node(num_cpus=1)
|
||||
ray.init(address=cluster.address)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
nodes = [first_node_handle.node_id, second_node_handle.node_id]
|
||||
# We want to check that there are two nodes. Thus, we convert `nodes` to a set and
|
||||
# then back to a list to remove duplicates. Then we check that the length of `nodes`
|
||||
# is 2.
|
||||
nodes = list(set(nodes))
|
||||
assert len(nodes) == 2
|
||||
|
||||
def create_actor(node):
|
||||
return Actor.options(label_selector={ray._raylet.RAY_NODE_ID_KEY: node}).remote(
|
||||
0
|
||||
)
|
||||
|
||||
def get_node_id(self):
|
||||
return ray.get_runtime_context().get_node_id()
|
||||
|
||||
driver_node = get_node_id(None)
|
||||
nodes.remove(driver_node)
|
||||
|
||||
a = create_actor(nodes[0])
|
||||
a_node = ray.get(a.__ray_call__.remote(get_node_id))
|
||||
assert a_node == nodes[0]
|
||||
# Check that the driver and actor are on different nodes.
|
||||
assert driver_node != a_node
|
||||
|
||||
with InputNode() as i:
|
||||
dag = a.echo.bind(i)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
size = GRPC_MAX_SIZE + (1024 * 1024 * 2)
|
||||
val = b"x" * size
|
||||
|
||||
for i in range(3):
|
||||
ref = compiled_dag.execute(val)
|
||||
result = ray.get(ref)
|
||||
assert result == val
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_actors", [1, 4])
|
||||
@pytest.mark.parametrize("num_nodes", [1, 4])
|
||||
def test_multi_node_multi_reader_large_payload(
|
||||
ray_start_cluster, num_actors, num_nodes, monkeypatch
|
||||
):
|
||||
# Set max grpc size to 5mb.
|
||||
GRPC_MAX_SIZE = 1024 * 1024 * 5
|
||||
monkeypatch.setenv("RAY_max_grpc_message_size", str(GRPC_MAX_SIZE))
|
||||
cluster = ray_start_cluster
|
||||
ACTORS_PER_NODE = num_actors
|
||||
NUM_REMOTE_NODES = num_nodes
|
||||
cluster.add_node(num_cpus=ACTORS_PER_NODE)
|
||||
ray.init(address=cluster.address)
|
||||
# This node is for the other two readers.
|
||||
for _ in range(NUM_REMOTE_NODES):
|
||||
cluster.add_node(num_cpus=ACTORS_PER_NODE)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
wait_for_condition(lambda: len(ray.nodes()) == NUM_REMOTE_NODES + 1)
|
||||
|
||||
actors = [
|
||||
Actor.options(num_cpus=1).remote(0)
|
||||
for _ in range(ACTORS_PER_NODE * (NUM_REMOTE_NODES + 1))
|
||||
]
|
||||
|
||||
def _get_node_id(self) -> "ray.NodeID":
|
||||
return ray.get_runtime_context().get_node_id()
|
||||
|
||||
node_ids = ray.get([act.__ray_call__.remote(_get_node_id) for act in actors])
|
||||
assert len(set(node_ids)) == NUM_REMOTE_NODES + 1
|
||||
|
||||
with InputNode() as inp:
|
||||
outputs = []
|
||||
for actor in actors:
|
||||
outputs.append(actor.echo.bind(inp))
|
||||
dag = MultiOutputNode(outputs)
|
||||
|
||||
compiled_dag = dag.experimental_compile()
|
||||
|
||||
# Set the object size a little bigger than the gRPC size so that
|
||||
# it triggers chunking and resizing.
|
||||
size = GRPC_MAX_SIZE + (1024 * 1024 * 2)
|
||||
val = b"x" * size
|
||||
|
||||
for _ in range(3):
|
||||
ref = compiled_dag.execute(val)
|
||||
# In the CI environment, the object store may use /tmp instead of /dev/shm
|
||||
# due to limited size of /tmp/shm and therefore has degraded performance.
|
||||
# Therefore, we use a longer timeout to avoid flakiness.
|
||||
result = ray.get(ref, timeout=50)
|
||||
assert result == [val for _ in range(ACTORS_PER_NODE * (NUM_REMOTE_NODES + 1))]
|
||||
|
||||
|
||||
def test_multi_node_dag_from_actor(ray_start_cluster):
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(num_cpus=1)
|
||||
ray.init()
|
||||
cluster.add_node(num_cpus=1)
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
class SameNodeActor:
|
||||
def predict(self, x: str):
|
||||
return x
|
||||
|
||||
@ray.remote(num_cpus=1)
|
||||
class RemoteNodeActor:
|
||||
def predict(self, x: str, y: str):
|
||||
return y
|
||||
|
||||
@ray.remote(num_cpus=1)
|
||||
class DriverActor:
|
||||
def __init__(self):
|
||||
self._base_actor = SameNodeActor.options(
|
||||
label_selector={
|
||||
ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
|
||||
}
|
||||
).remote()
|
||||
self._refiner_actor = RemoteNodeActor.remote()
|
||||
|
||||
with InputNode() as inp:
|
||||
x = self._base_actor.predict.bind(inp)
|
||||
dag = self._refiner_actor.predict.bind(
|
||||
inp,
|
||||
x,
|
||||
)
|
||||
|
||||
self._cdag = dag.experimental_compile(
|
||||
_submit_timeout=120,
|
||||
)
|
||||
|
||||
def call(self, prompt: str) -> bytes:
|
||||
return ray.get(self._cdag.execute(prompt))
|
||||
|
||||
parallel = DriverActor.remote()
|
||||
assert ray.get(parallel.call.remote("abc")) == "abc"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if os.environ.get("PARALLEL_CI"):
|
||||
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
||||
else:
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,627 @@
|
||||
import os
|
||||
import sys
|
||||
from typing import Any, Dict
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.dag import InputNode
|
||||
from ray.exceptions import RaySystemError, RayTaskError
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
if sys.platform != "linux" and sys.platform != "darwin":
|
||||
pytest.skip("Skipping, requires Linux or Mac.", allow_module_level=True)
|
||||
|
||||
USE_GPU = os.environ.get("RAY_PYTEST_USE_GPU") == "1"
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def echo_device(self, tensor: torch.Tensor) -> str:
|
||||
if isinstance(tensor, RaySystemError):
|
||||
raise tensor
|
||||
return str(tensor.device)
|
||||
|
||||
def echo_dict_device(
|
||||
self, dict_of_tensors: Dict[str, torch.Tensor]
|
||||
) -> Dict[str, str]:
|
||||
if isinstance(dict_of_tensors, RaySystemError):
|
||||
raise dict_of_tensors
|
||||
return {k: str(v.device) for k, v in dict_of_tensors.items()}
|
||||
|
||||
def send(self, device: str) -> torch.Tensor:
|
||||
return torch.ones((100,), device=device)
|
||||
|
||||
def send_dict(self, name_device_pairs: Dict[str, str]) -> Dict[str, torch.Tensor]:
|
||||
tensor_dict = {}
|
||||
for name, device in name_device_pairs.items():
|
||||
tensor_dict[name] = torch.ones((100,), device=device)
|
||||
return tensor_dict
|
||||
|
||||
|
||||
def run_driver_to_worker_dag(
|
||||
actor: "ray.actor.ActorHandle",
|
||||
device: str,
|
||||
tensor_input: Any,
|
||||
is_dict: bool = False,
|
||||
):
|
||||
"""Create and execute a DAG with tensor transport for driver to worker tests.
|
||||
|
||||
Args:
|
||||
actor: Ray actor to use
|
||||
device: Target device ("cpu", "cuda", or "default")
|
||||
tensor_input: Input tensor(s) to execute with
|
||||
is_dict: Whether to use dict version of the method
|
||||
|
||||
Returns:
|
||||
ray.ObjectRef: Result reference
|
||||
"""
|
||||
with InputNode() as inp:
|
||||
method = actor.echo_dict_device if is_dict else actor.echo_device
|
||||
dag = method.bind(inp.with_tensor_transport(device=device))
|
||||
compiled_dag = dag.experimental_compile()
|
||||
return compiled_dag.execute(tensor_input)
|
||||
|
||||
|
||||
def run_worker_to_worker_dag(
|
||||
sender: "ray.actor.ActorHandle",
|
||||
receiver: "ray.actor.ActorHandle",
|
||||
device: str,
|
||||
input_device: str,
|
||||
is_dict: bool = False,
|
||||
):
|
||||
"""Create and execute a DAG with tensor transport for worker to worker tests.
|
||||
|
||||
Args:
|
||||
sender: Sender Ray actor
|
||||
receiver: Receiver Ray actor
|
||||
device: Target device for tensor transport
|
||||
input_device: Device string to pass to sender
|
||||
is_dict: Whether to use dict version of the methods
|
||||
|
||||
Returns:
|
||||
ray.ObjectRef: Result reference or ValueError for compilation errors
|
||||
"""
|
||||
with InputNode() as inp:
|
||||
if is_dict:
|
||||
tensor = sender.send_dict.bind(inp)
|
||||
dag = receiver.echo_dict_device.bind(
|
||||
tensor.with_tensor_transport(device=device)
|
||||
)
|
||||
else:
|
||||
tensor = sender.send.bind(inp)
|
||||
dag = receiver.echo_device.bind(tensor.with_tensor_transport(device=device))
|
||||
compiled_dag = dag.experimental_compile()
|
||||
return compiled_dag.execute(input_device)
|
||||
|
||||
|
||||
def run_worker_to_driver_dag(
|
||||
actor: "ray.actor.ActorHandle",
|
||||
device: str,
|
||||
input_device: str,
|
||||
is_dict: bool = False,
|
||||
):
|
||||
"""Create and execute a DAG with tensor transport for worker to driver tests.
|
||||
|
||||
Args:
|
||||
actor: Ray actor to use
|
||||
device: Target device for tensor transport
|
||||
input_device: Device string to pass to actor
|
||||
is_dict: Whether to use dict version of the method
|
||||
|
||||
Returns:
|
||||
ray.ObjectRef: Result reference
|
||||
"""
|
||||
with InputNode() as inp:
|
||||
if is_dict:
|
||||
dag = actor.send_dict.bind(inp).with_tensor_transport(device=device)
|
||||
else:
|
||||
dag = actor.send.bind(inp).with_tensor_transport(device=device)
|
||||
compiled_dag = dag.experimental_compile()
|
||||
return compiled_dag.execute(input_device)
|
||||
|
||||
|
||||
class TestDriverToWorkerDeviceCPU:
|
||||
"""Tests driver to worker tensor transport with CPU device."""
|
||||
|
||||
def create_and_execute_dag(self, actor, device, tensor_input, is_dict=False):
|
||||
"""Create a DAG with tensor transport and execute it."""
|
||||
with InputNode() as inp:
|
||||
method = actor.echo_dict_device if is_dict else actor.echo_device
|
||||
dag = method.bind(inp.with_tensor_transport(device=device))
|
||||
compiled_dag = dag.experimental_compile()
|
||||
return compiled_dag.execute(tensor_input)
|
||||
|
||||
def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
ref = run_driver_to_worker_dag(actor, "cpu", torch.tensor([1]))
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
ref = run_driver_to_worker_dag(actor, "cpu", torch.tensor([1], device="cuda"))
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_driver_to_worker_dag(actor, "cpu", torch.tensor([1]))
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_driver_to_worker_dag(actor, "cpu", torch.tensor([1], device="cuda"))
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
tensor_dict = {
|
||||
"cpu_tensor": torch.tensor([1]),
|
||||
"gpu_tensor": torch.tensor([1], device="cuda"),
|
||||
}
|
||||
ref = run_driver_to_worker_dag(actor, "cpu", tensor_dict, is_dict=True)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cpu"}
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
tensor_dict = {
|
||||
"cpu_tensor": torch.tensor([1]),
|
||||
"gpu_tensor": torch.tensor([1], device="cuda"),
|
||||
}
|
||||
ref = run_driver_to_worker_dag(actor, "cpu", tensor_dict, is_dict=True)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cpu"}
|
||||
|
||||
|
||||
class TestDriverToWorkerDeviceGPU:
|
||||
"""Tests driver to worker tensor transport with GPU device."""
|
||||
|
||||
def create_and_execute_dag(self, actor, device, tensor_input, is_dict=False):
|
||||
"""Create a DAG with tensor transport and execute it."""
|
||||
with InputNode() as inp:
|
||||
method = actor.echo_dict_device if is_dict else actor.echo_device
|
||||
dag = method.bind(inp.with_tensor_transport(device=device))
|
||||
compiled_dag = dag.experimental_compile()
|
||||
return compiled_dag.execute(tensor_input)
|
||||
|
||||
def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
ref = run_driver_to_worker_dag(actor, "cuda", torch.tensor([1]))
|
||||
if torch.cuda.is_available():
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
else:
|
||||
with pytest.raises(
|
||||
RayTaskError, match="RuntimeError: No CUDA GPUs are available"
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
ref = run_driver_to_worker_dag(actor, "cuda", torch.tensor([1], device="cuda"))
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_driver_to_worker_dag(actor, "cuda", torch.tensor([1]))
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_driver_to_worker_dag(actor, "cuda", torch.tensor([1], device="cuda"))
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
tensor_dict = {
|
||||
"cpu_tensor": torch.tensor([1]),
|
||||
"gpu_tensor": torch.tensor([1], device="cuda"),
|
||||
}
|
||||
ref = run_driver_to_worker_dag(actor, "cuda", tensor_dict, is_dict=True)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cuda:0", "gpu_tensor": "cuda:0"}
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
tensor_dict = {
|
||||
"cpu_tensor": torch.tensor([1]),
|
||||
"gpu_tensor": torch.tensor([1], device="cuda"),
|
||||
}
|
||||
ref = run_driver_to_worker_dag(actor, "cuda", tensor_dict, is_dict=True)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cuda:0", "gpu_tensor": "cuda:0"}
|
||||
|
||||
|
||||
class TestDriverToWorkerDeviceDefault:
|
||||
"""Tests driver to worker tensor transport with default device."""
|
||||
|
||||
def create_and_execute_dag(self, actor, device, tensor_input, is_dict=False):
|
||||
"""Create a DAG with tensor transport and execute it."""
|
||||
with InputNode() as inp:
|
||||
method = actor.echo_dict_device if is_dict else actor.echo_device
|
||||
dag = method.bind(inp.with_tensor_transport(device=device))
|
||||
compiled_dag = dag.experimental_compile()
|
||||
return compiled_dag.execute(tensor_input)
|
||||
|
||||
def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
ref = run_driver_to_worker_dag(actor, "default", torch.tensor([1]))
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
ref = run_driver_to_worker_dag(
|
||||
actor, "default", torch.tensor([1], device="cuda")
|
||||
)
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_driver_to_worker_dag(actor, "default", torch.tensor([1]))
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_driver_to_worker_dag(
|
||||
actor, "default", torch.tensor([1], device="cuda")
|
||||
)
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
tensor_dict = {
|
||||
"cpu_tensor": torch.tensor([1]),
|
||||
"gpu_tensor": torch.tensor([1], device="cuda"),
|
||||
}
|
||||
ref = run_driver_to_worker_dag(actor, "default", tensor_dict, is_dict=True)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cuda:0"}
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
tensor_dict = {
|
||||
"cpu_tensor": torch.tensor([1]),
|
||||
"gpu_tensor": torch.tensor([1], device="cuda"),
|
||||
}
|
||||
ref = run_driver_to_worker_dag(actor, "default", tensor_dict, is_dict=True)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cuda:0"}
|
||||
|
||||
|
||||
class TestWorkerToWorkerDeviceCPU:
|
||||
"""Tests worker to worker tensor transport with CPU device."""
|
||||
|
||||
def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
sender = Actor.remote()
|
||||
receiver = Actor.remote()
|
||||
ref = run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
sender = Actor.remote()
|
||||
receiver = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.remote()
|
||||
ref = run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
|
||||
def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.options(num_gpus=1).remote()
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="accelerator transport is not supported with CPU target device.",
|
||||
):
|
||||
run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.options().remote()
|
||||
ref = run_worker_to_worker_dag(
|
||||
sender,
|
||||
receiver,
|
||||
"cpu",
|
||||
{"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
|
||||
is_dict=True,
|
||||
)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cpu"}
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
|
||||
def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.options(num_gpus=1).remote()
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="accelerator transport is not supported with CPU target device.",
|
||||
):
|
||||
run_worker_to_worker_dag(
|
||||
sender,
|
||||
receiver,
|
||||
"cpu",
|
||||
{"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
|
||||
is_dict=True,
|
||||
)
|
||||
|
||||
|
||||
class TestWorkerToWorkerDeviceGPU:
|
||||
"""Tests worker to worker tensor transport with GPU device."""
|
||||
|
||||
@pytest.mark.parametrize("gpu_device", ["gpu", "cuda"])
|
||||
def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular, gpu_device):
|
||||
sender = Actor.remote()
|
||||
receiver = Actor.remote()
|
||||
ref = run_worker_to_worker_dag(sender, receiver, gpu_device, "cpu")
|
||||
if torch.cuda.is_available():
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
else:
|
||||
with pytest.raises(
|
||||
RayTaskError, match="RuntimeError: No CUDA GPUs are available"
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
sender = Actor.remote()
|
||||
receiver = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_worker_dag(sender, receiver, "cuda", "cpu")
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.remote()
|
||||
ref = run_worker_to_worker_dag(sender, receiver, "cuda", "cuda")
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
|
||||
def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.options(num_gpus=1).remote()
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="accelerator transport is not supported with CPU target device.",
|
||||
):
|
||||
run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.options().remote()
|
||||
ref = run_worker_to_worker_dag(
|
||||
sender,
|
||||
receiver,
|
||||
"cuda",
|
||||
{"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
|
||||
is_dict=True,
|
||||
)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cuda:0", "gpu_tensor": "cuda:0"}
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
|
||||
@pytest.mark.parametrize("gpu_device", ["gpu", "cuda"])
|
||||
def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular, gpu_device):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.options(num_gpus=1).remote()
|
||||
|
||||
ref = run_worker_to_worker_dag(
|
||||
sender,
|
||||
receiver,
|
||||
gpu_device,
|
||||
{"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
|
||||
is_dict=True,
|
||||
)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cuda:0", "gpu_tensor": "cuda:0"}
|
||||
|
||||
|
||||
class TestWorkerToWorkerDeviceDefault:
|
||||
"""Tests worker to worker tensor transport with default device."""
|
||||
|
||||
def test_src_cpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
sender = Actor.remote()
|
||||
receiver = Actor.remote()
|
||||
ref = run_worker_to_worker_dag(sender, receiver, "default", "cpu")
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_cpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
sender = Actor.remote()
|
||||
receiver = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_worker_dag(sender, receiver, "default", "cpu")
|
||||
assert ray.get(ref) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor_dst_cpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.remote()
|
||||
ref = run_worker_to_worker_dag(sender, receiver, "default", "cuda")
|
||||
assert ray.get(ref) == "cuda:0"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
|
||||
def test_src_gpu_tensor_dst_gpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.options(num_gpus=1).remote()
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="accelerator transport is not supported with CPU target device.",
|
||||
):
|
||||
run_worker_to_worker_dag(sender, receiver, "cpu", "cpu")
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors_dst_cpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.options().remote()
|
||||
ref = run_worker_to_worker_dag(
|
||||
sender,
|
||||
receiver,
|
||||
"default",
|
||||
{"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
|
||||
is_dict=True,
|
||||
)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cuda:0"}
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
@pytest.mark.parametrize("ray_start_regular", [{"num_cpus": 2}], indirect=True)
|
||||
def test_src_mix_tensors_dst_gpu_node(self, ray_start_regular):
|
||||
sender = Actor.options(num_gpus=1).remote()
|
||||
receiver = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_worker_dag(
|
||||
sender,
|
||||
receiver,
|
||||
"default",
|
||||
{"cpu_tensor": "cpu", "gpu_tensor": "cuda"},
|
||||
is_dict=True,
|
||||
)
|
||||
assert ray.get(ref) == {"cpu_tensor": "cpu", "gpu_tensor": "cuda:0"}
|
||||
|
||||
|
||||
class TestWorkerToDriverDeviceCPU:
|
||||
"""Tests worker to driver tensor transport with CPU device."""
|
||||
|
||||
def test_src_cpu_tensor(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
ref = run_worker_to_driver_dag(actor, "cpu", "cpu")
|
||||
tensor = ray.get(ref)
|
||||
assert str(tensor.device) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_driver_dag(actor, "cpu", "cuda")
|
||||
tensor = ray.get(ref)
|
||||
assert str(tensor.device) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_driver_dag(
|
||||
actor, "cpu", {"cpu_tensor": "cpu", "gpu_tensor": "cuda"}, is_dict=True
|
||||
)
|
||||
tensor = ray.get(ref)
|
||||
assert str(tensor["cpu_tensor"].device) == "cpu"
|
||||
assert str(tensor["gpu_tensor"].device) == "cpu"
|
||||
|
||||
|
||||
class TestWorkerToDriverDeviceGPU:
|
||||
"""Tests worker to driver tensor transport with GPU device."""
|
||||
|
||||
def test_src_cpu_tensor(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
ref = run_worker_to_driver_dag(actor, "cuda", "cpu")
|
||||
|
||||
# different behavior between a driver node with GPU and without GPU
|
||||
if torch.cuda.is_available():
|
||||
tensor = ray.get(ref)
|
||||
assert str(tensor.device) == "cuda:0"
|
||||
else:
|
||||
with pytest.raises(
|
||||
RayTaskError, match="RuntimeError: No CUDA GPUs are available"
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_driver_dag(actor, "cuda", "cuda")
|
||||
|
||||
# different behavior between a driver node with GPU and without GPU
|
||||
if torch.cuda.is_available():
|
||||
tensor = ray.get(ref)
|
||||
assert str(tensor.device) == "cuda:0"
|
||||
else:
|
||||
with pytest.raises(
|
||||
RayTaskError, match="RuntimeError: No CUDA GPUs are available"
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_driver_dag(
|
||||
actor, "cuda", {"cpu_tensor": "cpu", "gpu_tensor": "cuda"}, is_dict=True
|
||||
)
|
||||
|
||||
# different behavior between a driver node with GPU and without GPU
|
||||
if torch.cuda.is_available():
|
||||
tensor = ray.get(ref)
|
||||
assert str(tensor["cpu_tensor"].device) == "cuda:0"
|
||||
assert str(tensor["gpu_tensor"].device) == "cuda:0"
|
||||
else:
|
||||
with pytest.raises(
|
||||
RayTaskError, match="RuntimeError: No CUDA GPUs are available"
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
class TestWorkerToDriverDeviceDefault:
|
||||
"""Tests worker to driver tensor transport with default device."""
|
||||
|
||||
def test_src_cpu_tensor(self, ray_start_regular):
|
||||
actor = Actor.remote()
|
||||
ref = run_worker_to_driver_dag(actor, "default", "cpu")
|
||||
tensor = ray.get(ref)
|
||||
assert str(tensor.device) == "cpu"
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_gpu_tensor(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_driver_dag(actor, "default", "cuda")
|
||||
|
||||
# different behavior between a driver node with GPU and without GPU
|
||||
if torch.cuda.is_available():
|
||||
tensor = ray.get(ref)
|
||||
assert str(tensor.device) == "cuda:0"
|
||||
else:
|
||||
with pytest.raises(
|
||||
RayTaskError, match="RuntimeError: No CUDA GPUs are available"
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
@pytest.mark.skipif(not USE_GPU, reason="Test requires GPU")
|
||||
def test_src_mix_tensors(self, ray_start_regular):
|
||||
actor = Actor.options(num_gpus=1).remote()
|
||||
ref = run_worker_to_driver_dag(
|
||||
actor, "default", {"cpu_tensor": "cpu", "gpu_tensor": "cuda"}, is_dict=True
|
||||
)
|
||||
|
||||
# different behavior between a driver node with GPU and without GPU
|
||||
if torch.cuda.is_available():
|
||||
tensor = ray.get(ref)
|
||||
assert str(tensor["cpu_tensor"].device) == "cpu"
|
||||
assert str(tensor["gpu_tensor"].device) == "cuda:0"
|
||||
else:
|
||||
with pytest.raises(
|
||||
RayTaskError, match="RuntimeError: No CUDA GPUs are available"
|
||||
):
|
||||
ray.get(ref)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if os.environ.get("PARALLEL_CI"):
|
||||
sys.exit(pytest.main(["-n", "auto", "--boxed", "-vs", __file__]))
|
||||
else:
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
Reference in New Issue
Block a user