chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,2 @@
|
||||
# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
|
||||
from ray.tests.conftest import pytest_runtest_makereport # noqa
|
||||
@@ -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__]))
|
||||
@@ -0,0 +1,362 @@
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.dag import (
|
||||
PARENT_CLASS_NODE_KEY,
|
||||
PREV_CLASS_METHOD_CALL_KEY,
|
||||
InputNode,
|
||||
MultiOutputNode,
|
||||
)
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Counter:
|
||||
def __init__(self, init_value=0):
|
||||
self.i = init_value
|
||||
|
||||
def inc(self):
|
||||
self.i += 1
|
||||
|
||||
def get(self):
|
||||
return self.i
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self, init_value):
|
||||
self.i = init_value
|
||||
|
||||
def inc(self, x):
|
||||
self.i += x
|
||||
|
||||
def get(self):
|
||||
return self.i
|
||||
|
||||
def echo(self, x):
|
||||
return x
|
||||
|
||||
@ray.method(num_returns=2)
|
||||
def return_two(self, x):
|
||||
return x, x + 1
|
||||
|
||||
@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
|
||||
|
||||
|
||||
def test_basic_actor_dag(shared_ray_instance):
|
||||
@ray.remote
|
||||
def combine(x, y):
|
||||
return x + y
|
||||
|
||||
a1 = Actor.bind(10)
|
||||
res = a1.get.bind()
|
||||
print(res)
|
||||
assert ray.get(res.execute()) == 10
|
||||
|
||||
a2 = Actor.bind(10)
|
||||
a1.inc.bind(2)
|
||||
a1.inc.bind(4)
|
||||
a2.inc.bind(6)
|
||||
dag = combine.bind(a1.get.bind(), a2.get.bind())
|
||||
|
||||
print(dag)
|
||||
assert ray.get(dag.execute()) == 32
|
||||
|
||||
|
||||
def test_class_as_class_constructor_arg(shared_ray_instance):
|
||||
@ray.remote
|
||||
class OuterActor:
|
||||
def __init__(self, inner_actor):
|
||||
self.inner_actor = inner_actor
|
||||
|
||||
def inc(self, x):
|
||||
self.inner_actor.inc.remote(x)
|
||||
|
||||
def get(self):
|
||||
return ray.get(self.inner_actor.get.remote())
|
||||
|
||||
outer = OuterActor.bind(Actor.bind(10))
|
||||
outer.inc.bind(2)
|
||||
dag = outer.get.bind()
|
||||
print(dag)
|
||||
assert ray.get(dag.execute()) == 12
|
||||
|
||||
|
||||
def test_class_as_function_constructor_arg(shared_ray_instance):
|
||||
@ray.remote
|
||||
def f(actor_handle):
|
||||
return ray.get(actor_handle.get.remote())
|
||||
|
||||
dag = f.bind(Actor.bind(10))
|
||||
print(dag)
|
||||
assert ray.get(dag.execute()) == 10
|
||||
|
||||
|
||||
def test_basic_actor_dag_constructor_options(shared_ray_instance):
|
||||
a1 = Actor.bind(10)
|
||||
dag = a1.get.bind()
|
||||
print(dag)
|
||||
assert ray.get(dag.execute()) == 10
|
||||
|
||||
a1 = Actor.options(name="Actor", namespace="test", max_pending_calls=10).bind(10)
|
||||
dag = a1.get.bind()
|
||||
print(dag)
|
||||
# Ensure execution result is identical with .options() in init()
|
||||
assert ray.get(dag.execute()) == 10
|
||||
# Ensure options are passed in
|
||||
assert a1.get_options().get("name") == "Actor"
|
||||
assert a1.get_options().get("namespace") == "test"
|
||||
assert a1.get_options().get("max_pending_calls") == 10
|
||||
|
||||
|
||||
def test_actor_method_options(shared_ray_instance):
|
||||
a1 = Actor.bind(10)
|
||||
dag = a1.get.options(name="actor_method_options").bind()
|
||||
print(dag)
|
||||
assert ray.get(dag.execute()) == 10
|
||||
assert dag.get_options().get("name") == "actor_method_options"
|
||||
|
||||
|
||||
def test_basic_actor_dag_constructor_invalid_options(shared_ray_instance):
|
||||
with pytest.raises(
|
||||
ValueError, match=r".*quantity of resource num_cpus cannot be negative.*"
|
||||
):
|
||||
a1 = Actor.options(num_cpus=-1).bind(10)
|
||||
invalid_dag = a1.get.bind()
|
||||
ray.get(invalid_dag.execute())
|
||||
|
||||
|
||||
def test_actor_options_complicated(shared_ray_instance):
|
||||
"""Test a more complicated setup where we apply .options() in both
|
||||
constructor and method call with overlapping keys, and ensure end to end
|
||||
options correctness.
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
def combine(x, y):
|
||||
return x + y
|
||||
|
||||
a1 = Actor.options(name="a1_v0").bind(10)
|
||||
res = a1.get.options(name="v1").bind()
|
||||
print(res)
|
||||
assert ray.get(res.execute()) == 10
|
||||
assert a1.get_options().get("name") == "a1_v0"
|
||||
assert res.get_options().get("name") == "v1"
|
||||
|
||||
a1 = Actor.options(name="a1_v1").bind(10) # Cannot
|
||||
a2 = Actor.options(name="a2_v0").bind(10)
|
||||
a1.inc.options(name="v1").bind(2)
|
||||
a1.inc.options(name="v2").bind(4)
|
||||
a2.inc.options(name="v3").bind(6)
|
||||
dag = combine.options(name="v4").bind(a1.get.bind(), a2.get.bind())
|
||||
|
||||
print(dag)
|
||||
assert ray.get(dag.execute()) == 32
|
||||
test_a1 = dag.get_args()[0] # call graph for a1.get.bind()
|
||||
test_a2 = dag.get_args()[1] # call graph for a2.get.bind()
|
||||
assert test_a2.get_options() == {} # No .options() at outer call
|
||||
# refer to a2 constructor .options() call
|
||||
assert (
|
||||
test_a2.get_other_args_to_resolve()[PARENT_CLASS_NODE_KEY]
|
||||
.get_options()
|
||||
.get("name")
|
||||
== "a2_v0"
|
||||
)
|
||||
# refer to actor method a2.inc.options() call
|
||||
assert (
|
||||
test_a2.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
|
||||
.get_options()
|
||||
.get("name")
|
||||
== "v3"
|
||||
)
|
||||
# refer to a1 constructor .options() call
|
||||
assert (
|
||||
test_a1.get_other_args_to_resolve()[PARENT_CLASS_NODE_KEY]
|
||||
.get_options()
|
||||
.get("name")
|
||||
== "a1_v1"
|
||||
)
|
||||
# refer to latest actor method a1.inc.options() call
|
||||
assert (
|
||||
test_a1.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
|
||||
.get_options()
|
||||
.get("name")
|
||||
== "v2"
|
||||
)
|
||||
# refer to first bound actor method a1.inc.options() call
|
||||
assert (
|
||||
test_a1.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
|
||||
.get_other_args_to_resolve()[PREV_CLASS_METHOD_CALL_KEY]
|
||||
.get_options()
|
||||
.get("name")
|
||||
== "v1"
|
||||
)
|
||||
|
||||
|
||||
def test_pass_actor_handle(shared_ray_instance):
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def ping(self):
|
||||
return "hello"
|
||||
|
||||
@ray.remote
|
||||
def caller(handle):
|
||||
assert isinstance(handle, ray.actor.ActorHandle), handle
|
||||
return ray.get(handle.ping.remote())
|
||||
|
||||
a1 = Actor.bind()
|
||||
dag = caller.bind(a1)
|
||||
print(dag)
|
||||
assert ray.get(dag.execute()) == "hello"
|
||||
|
||||
|
||||
def test_dynamic_pipeline(shared_ray_instance):
|
||||
@ray.remote
|
||||
class Model:
|
||||
def __init__(self, arg):
|
||||
self.arg = arg
|
||||
|
||||
def forward(self, x):
|
||||
return self.arg + str(x)
|
||||
|
||||
@ray.remote
|
||||
class ModelSelection:
|
||||
def is_even(self, x):
|
||||
return x % 2 == 0
|
||||
|
||||
@ray.remote
|
||||
def pipeline(x, m1, m2, selection):
|
||||
sel = selection.is_even.remote(x)
|
||||
if ray.get(sel):
|
||||
result = m1.forward.remote(x)
|
||||
else:
|
||||
result = m2.forward.remote(x)
|
||||
return ray.get(result)
|
||||
|
||||
m1 = Model.bind("Even: ")
|
||||
m2 = Model.bind("Odd: ")
|
||||
selection = ModelSelection.bind()
|
||||
|
||||
even_input = pipeline.bind(20, m1, m2, selection)
|
||||
print(even_input)
|
||||
assert ray.get(even_input.execute()) == "Even: 20"
|
||||
|
||||
odd_input = pipeline.bind(21, m1, m2, selection)
|
||||
print(odd_input)
|
||||
assert ray.get(odd_input.execute()) == "Odd: 21"
|
||||
|
||||
|
||||
def test_unsupported_bind():
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def ping(self):
|
||||
return "hello"
|
||||
|
||||
with pytest.raises(
|
||||
AttributeError,
|
||||
match=r"\.bind\(\) cannot be used again on",
|
||||
):
|
||||
actor = Actor.bind()
|
||||
_ = actor.bind()
|
||||
|
||||
with pytest.raises(
|
||||
AttributeError,
|
||||
match=r"\.remote\(\) cannot be used on ClassMethodNodes",
|
||||
):
|
||||
actor = Actor.bind()
|
||||
_ = actor.ping.remote()
|
||||
|
||||
|
||||
def test_unsupported_remote():
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def ping(self):
|
||||
return "hello"
|
||||
|
||||
with pytest.raises(AttributeError, match="'Actor' has no attribute 'remote'"):
|
||||
_ = Actor.bind().remote()
|
||||
|
||||
@ray.remote
|
||||
def func():
|
||||
return 1
|
||||
|
||||
with pytest.raises(AttributeError, match=r"\.remote\(\) cannot be used on"):
|
||||
_ = func.bind().remote()
|
||||
|
||||
|
||||
def test_two_returns_first():
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
o1, o2 = a.return_two.bind(i)
|
||||
dag = o1
|
||||
|
||||
for _ in range(3):
|
||||
res = ray.get(dag.execute(1))
|
||||
assert res == 1
|
||||
|
||||
|
||||
def test_two_returns_second():
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
o1, o2 = a.return_two.bind(i)
|
||||
dag = o2
|
||||
|
||||
for _ in range(3):
|
||||
res = ray.get(dag.execute(1))
|
||||
assert res == 2
|
||||
|
||||
|
||||
def test_two_returns_one_reader_multi_times():
|
||||
a = Actor.remote(0)
|
||||
b = Actor.remote(0)
|
||||
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])
|
||||
|
||||
for _ in range(3):
|
||||
res = ray.get(dag.execute(1))
|
||||
assert res == [1, 2]
|
||||
|
||||
|
||||
def test_two_returns_two_readers():
|
||||
a = Actor.remote(0)
|
||||
b = Actor.remote(0)
|
||||
c = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
o1, o2 = a.return_two.bind(i)
|
||||
o3 = b.echo.bind(o1)
|
||||
o4 = c.echo.bind(o2)
|
||||
dag = MultiOutputNode([o3, o4])
|
||||
|
||||
for _ in range(3):
|
||||
res = ray.get(dag.execute(1))
|
||||
assert res == [1, 2]
|
||||
|
||||
|
||||
def test_inc_two_returns():
|
||||
a = Actor.remote(0)
|
||||
with InputNode() as i:
|
||||
o1, o2 = a.inc_and_return_two.bind(i)
|
||||
dag = MultiOutputNode([o1, o2])
|
||||
|
||||
for i in range(3):
|
||||
res = ray.get(dag.execute(1))
|
||||
assert res == [i + 1, i + 2]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,215 @@
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Counter:
|
||||
def __init__(self, init_value=0):
|
||||
self.i = init_value
|
||||
|
||||
def inc(self):
|
||||
self.i += 1
|
||||
|
||||
def get(self):
|
||||
return self.i
|
||||
|
||||
|
||||
def test_basic_task_dag(shared_ray_instance):
|
||||
ct = Counter.remote()
|
||||
|
||||
@ray.remote
|
||||
def a():
|
||||
ray.get(ct.inc.remote())
|
||||
return 2
|
||||
|
||||
@ray.remote
|
||||
def b(x):
|
||||
ray.get(ct.inc.remote())
|
||||
return x * 2
|
||||
|
||||
@ray.remote
|
||||
def c(x):
|
||||
ray.get(ct.inc.remote())
|
||||
return x + 1
|
||||
|
||||
@ray.remote
|
||||
def d(x, y):
|
||||
ray.get(ct.inc.remote())
|
||||
return x + y
|
||||
|
||||
a_ref = a.bind()
|
||||
b_ref = b.bind(a_ref)
|
||||
c_ref = c.bind(a_ref)
|
||||
d_ref = d.bind(b_ref, c_ref)
|
||||
d1_ref = d.bind(d_ref, d_ref)
|
||||
d2_ref = d.bind(d1_ref, d_ref)
|
||||
dag = d.bind(d2_ref, d_ref)
|
||||
print(dag)
|
||||
|
||||
assert ray.get(dag.execute()) == 28
|
||||
assert ray.get(ct.get.remote()) == 7
|
||||
|
||||
|
||||
def test_basic_task_dag_with_options(shared_ray_instance):
|
||||
ct = Counter.remote()
|
||||
|
||||
@ray.remote
|
||||
def a():
|
||||
ray.get(ct.inc.remote())
|
||||
return 2
|
||||
|
||||
@ray.remote
|
||||
def b(x):
|
||||
ray.get(ct.inc.remote())
|
||||
return x * 2
|
||||
|
||||
@ray.remote
|
||||
def c(x):
|
||||
ray.get(ct.inc.remote())
|
||||
return x + 1
|
||||
|
||||
@ray.remote
|
||||
def d(x, y):
|
||||
ray.get(ct.inc.remote())
|
||||
return x + y
|
||||
|
||||
a_ref = a.bind()
|
||||
b_ref = b.options(name="b", num_returns=1).bind(a_ref)
|
||||
c_ref = c.options(name="c", max_retries=3).bind(a_ref)
|
||||
dag = d.options(name="d", num_cpus=2).bind(b_ref, c_ref)
|
||||
|
||||
print(dag)
|
||||
|
||||
assert ray.get(dag.execute()) == 7
|
||||
assert ray.get(ct.get.remote()) == 4
|
||||
|
||||
assert b_ref.get_options().get("name") == "b"
|
||||
assert b_ref.get_options().get("num_returns") == 1
|
||||
assert c_ref.get_options().get("name") == "c"
|
||||
assert c_ref.get_options().get("max_retries") == 3
|
||||
assert dag.get_options().get("name") == "d"
|
||||
assert dag.get_options().get("num_cpus") == 2
|
||||
|
||||
|
||||
def test_invalid_task_options(shared_ray_instance):
|
||||
"""
|
||||
Test to ensure options used in DAG binding are applied, and will throw
|
||||
as expected even given invalid values.
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
def a():
|
||||
return 2
|
||||
|
||||
@ray.remote
|
||||
def b(x):
|
||||
return x * 2
|
||||
|
||||
a_ref = a.bind()
|
||||
dag = b.bind(a_ref)
|
||||
|
||||
# Ensure current DAG is executable
|
||||
assert ray.get(dag.execute()) == 4
|
||||
with pytest.raises(
|
||||
ValueError, match=r".*quantity of resource num_cpus cannot be negative.*"
|
||||
):
|
||||
invalid_dag = b.options(num_cpus=-1).bind(a_ref)
|
||||
ray.get(invalid_dag.execute())
|
||||
|
||||
|
||||
def test_node_accessors(shared_ray_instance):
|
||||
@ray.remote
|
||||
def a(*a, **kw):
|
||||
pass
|
||||
|
||||
tmp1 = a.bind()
|
||||
tmp2 = a.bind()
|
||||
tmp3 = a.bind()
|
||||
node = a.bind(1, tmp1, x=tmp2, y={"foo": tmp3})
|
||||
assert node.get_args() == (1, tmp1)
|
||||
assert node.get_kwargs() == {"x": tmp2, "y": {"foo": tmp3}}
|
||||
assert node._get_toplevel_child_nodes() == [tmp1, tmp2]
|
||||
assert node._get_all_child_nodes() == [tmp1, tmp2, tmp3]
|
||||
|
||||
tmp4 = a.bind()
|
||||
tmp5 = a.bind()
|
||||
replace = {tmp1: tmp4, tmp2: tmp4, tmp3: tmp5}
|
||||
n2 = node._apply_and_replace_all_child_nodes(lambda x: replace[x])
|
||||
assert n2._get_all_child_nodes() == [tmp4, tmp5]
|
||||
|
||||
|
||||
def test_nested_args(shared_ray_instance):
|
||||
ct = Counter.remote()
|
||||
|
||||
@ray.remote
|
||||
def a():
|
||||
ray.get(ct.inc.remote())
|
||||
return 2
|
||||
|
||||
@ray.remote
|
||||
def b(**kwargs):
|
||||
ray.get(ct.inc.remote())
|
||||
return kwargs["x"] * 2
|
||||
|
||||
@ray.remote
|
||||
def c(**kwargs):
|
||||
ray.get(ct.inc.remote())
|
||||
return kwargs["x"] + 1
|
||||
|
||||
@ray.remote
|
||||
def d(nested):
|
||||
ray.get(ct.inc.remote())
|
||||
return ray.get(nested["x"]) + ray.get(nested["y"])
|
||||
|
||||
a_ref = a.bind()
|
||||
b_ref = b.bind(x=a_ref)
|
||||
c_ref = c.bind(x=a_ref)
|
||||
dag = d.bind({"x": b_ref, "y": c_ref})
|
||||
print(dag)
|
||||
|
||||
assert ray.get(dag.execute()) == 7
|
||||
assert ray.get(ct.get.remote()) == 4
|
||||
|
||||
|
||||
def test_dag_options(shared_ray_instance):
|
||||
@ray.remote(num_gpus=100)
|
||||
def foo():
|
||||
pass
|
||||
|
||||
assert foo.bind().get_options() == {"max_calls": 1, "num_gpus": 100}
|
||||
assert foo.options(num_gpus=300).bind().get_options() == {"num_gpus": 300}
|
||||
assert foo.options(num_cpus=500).bind().get_options() == {
|
||||
"num_gpus": 100,
|
||||
"num_cpus": 500,
|
||||
}
|
||||
|
||||
@ray.remote
|
||||
def bar():
|
||||
pass
|
||||
|
||||
assert bar.bind().get_options() == {}
|
||||
assert bar.options(num_gpus=100).bind().get_options() == {"num_gpus": 100}
|
||||
|
||||
@ray.remote(num_gpus=100)
|
||||
class Foo:
|
||||
pass
|
||||
|
||||
assert Foo.bind().get_options() == {"num_gpus": 100}
|
||||
assert Foo.options(num_gpus=300).bind().get_options() == {"num_gpus": 300}
|
||||
assert Foo.options(num_cpus=500).bind().get_options() == {
|
||||
"num_gpus": 100,
|
||||
"num_cpus": 500,
|
||||
}
|
||||
|
||||
@ray.remote
|
||||
class Bar:
|
||||
pass
|
||||
|
||||
assert Bar.bind().get_options() == {}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,386 @@
|
||||
"""
|
||||
Tests to ensure ray DAG can correctly mark its input(s) to take user
|
||||
request, for all DAGNode types.
|
||||
"""
|
||||
|
||||
from typing import Any, TypeVar
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.dag.dag_node import DAGNode
|
||||
from ray.dag.input_node import InputNode
|
||||
|
||||
RayHandleLike = TypeVar("RayHandleLike")
|
||||
|
||||
|
||||
def test_no_args_to_input_node(shared_ray_instance):
|
||||
@ray.remote
|
||||
def f(input):
|
||||
return input
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="InputNode should not take any args or kwargs"
|
||||
):
|
||||
with InputNode(0) as dag_input:
|
||||
f.bind(dag_input)
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="InputNode should not take any args or kwargs",
|
||||
):
|
||||
with InputNode(key=1) as dag_input:
|
||||
f.bind(dag_input)
|
||||
|
||||
|
||||
def test_simple_func(shared_ray_instance):
|
||||
@ray.remote
|
||||
def a(input: str):
|
||||
return f"{input} -> a"
|
||||
|
||||
@ray.remote
|
||||
def b(a: "RayHandleLike"):
|
||||
# At runtime, a is replaced with execution result of a.
|
||||
return f"{a} -> b"
|
||||
|
||||
# input -> a - > b -> ouput
|
||||
with InputNode() as dag_input:
|
||||
a_node = a.bind(dag_input)
|
||||
dag = b.bind(a_node)
|
||||
|
||||
assert ray.get(dag.execute("input")) == "input -> a -> b"
|
||||
assert ray.get(dag.execute("test")) == "test -> a -> b"
|
||||
|
||||
|
||||
def test_func_dag(shared_ray_instance):
|
||||
@ray.remote
|
||||
def a(user_input):
|
||||
return user_input
|
||||
|
||||
@ray.remote
|
||||
def b(x):
|
||||
return x * 2
|
||||
|
||||
@ray.remote
|
||||
def c(x):
|
||||
return x + 1
|
||||
|
||||
@ray.remote
|
||||
def d(x, y):
|
||||
return x + y
|
||||
|
||||
with InputNode() as dag_input:
|
||||
a_ref = a.bind(dag_input)
|
||||
b_ref = b.bind(a_ref)
|
||||
c_ref = c.bind(a_ref)
|
||||
d_ref = d.bind(b_ref, c_ref)
|
||||
d1_ref = d.bind(d_ref, d_ref)
|
||||
d2_ref = d.bind(d1_ref, d_ref)
|
||||
dag = d.bind(d2_ref, d_ref)
|
||||
|
||||
# [(2*2 + 2+1) + (2*2 + 2+1)] + [(2*2 + 2+1) + (2*2 + 2+1)]
|
||||
assert ray.get(dag.execute(2)) == 28
|
||||
# [(3*2 + 3+1) + (3*2 + 3+1)] + [(3*2 + 3+1) + (3*2 + 3+1)]
|
||||
assert ray.get(dag.execute(3)) == 40
|
||||
|
||||
|
||||
def test_multi_input_func_dag(shared_ray_instance):
|
||||
@ray.remote
|
||||
def a(user_input):
|
||||
return user_input * 2
|
||||
|
||||
@ray.remote
|
||||
def b(user_input):
|
||||
return user_input + 1
|
||||
|
||||
@ray.remote
|
||||
def c(x, y):
|
||||
return x + y
|
||||
|
||||
with InputNode() as dag_input:
|
||||
a_ref = a.bind(dag_input)
|
||||
b_ref = b.bind(dag_input)
|
||||
dag = c.bind(a_ref, b_ref)
|
||||
|
||||
# (2*2) + (2*1)
|
||||
assert ray.get(dag.execute(2)) == 7
|
||||
# (3*2) + (3*1)
|
||||
assert ray.get(dag.execute(3)) == 10
|
||||
|
||||
|
||||
def test_invalid_input_node_as_class_constructor(shared_ray_instance):
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self, val):
|
||||
self.val = val
|
||||
|
||||
def get(self):
|
||||
return self.val
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"InputNode handles user dynamic input the DAG, and "
|
||||
"cannot be used as args, kwargs, or other_args_to_resolve "
|
||||
"in ClassNode constructor because it is not available at "
|
||||
"class construction or binding time."
|
||||
),
|
||||
):
|
||||
with InputNode() as dag_input:
|
||||
Actor.bind(dag_input)
|
||||
|
||||
|
||||
def test_class_method_input(shared_ray_instance):
|
||||
@ray.remote
|
||||
class Model:
|
||||
def __init__(self, weight: int):
|
||||
self.weight = weight
|
||||
|
||||
def forward(self, input: "RayHandleLike"):
|
||||
return self.weight * input
|
||||
|
||||
@ray.remote
|
||||
class FeatureProcessor:
|
||||
def __init__(self, scale):
|
||||
self.scale = scale
|
||||
|
||||
def process(self, input: int):
|
||||
return input * self.scale
|
||||
|
||||
with InputNode() as dag_input:
|
||||
preprocess = FeatureProcessor.bind(0.5)
|
||||
feature = preprocess.process.bind(dag_input)
|
||||
model = Model.bind(4)
|
||||
dag = model.forward.bind(feature)
|
||||
|
||||
# 2 * 0.5 * 4
|
||||
assert ray.get(dag.execute(2)) == 4
|
||||
# 6 * 0.5 * 4
|
||||
assert ray.get(dag.execute(6)) == 12
|
||||
|
||||
|
||||
def test_multi_class_method_input(shared_ray_instance):
|
||||
"""
|
||||
Test a multiple class methods can all be used as inputs in a dag.
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Model:
|
||||
def __init__(self, weight: int):
|
||||
self.weight = weight
|
||||
|
||||
def forward(self, input: int):
|
||||
return self.weight * input
|
||||
|
||||
@ray.remote
|
||||
def combine(m1: "RayHandleLike", m2: "RayHandleLike"):
|
||||
return m1 + m2
|
||||
|
||||
with InputNode() as dag_input:
|
||||
m1 = Model.bind(2)
|
||||
m2 = Model.bind(3)
|
||||
|
||||
m1_output = m1.forward.bind(dag_input)
|
||||
m2_output = m2.forward.bind(dag_input)
|
||||
|
||||
dag = combine.bind(m1_output, m2_output)
|
||||
|
||||
# 1*2 + 1*3
|
||||
assert ray.get(dag.execute(1)) == 5
|
||||
# 2*2 + 2*3
|
||||
assert ray.get(dag.execute(2)) == 10
|
||||
|
||||
|
||||
def test_func_class_mixed_input(shared_ray_instance):
|
||||
"""
|
||||
Test both class method and function are used as input in the
|
||||
same dag.
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Model:
|
||||
def __init__(self, weight: int):
|
||||
self.weight = weight
|
||||
|
||||
def forward(self, input: int):
|
||||
return self.weight * input
|
||||
|
||||
@ray.remote
|
||||
def model_func(input: int):
|
||||
return input * 2
|
||||
|
||||
@ray.remote
|
||||
def combine(m1: "RayHandleLike", m2: "RayHandleLike"):
|
||||
return m1 + m2
|
||||
|
||||
with InputNode() as dag_input:
|
||||
m1 = Model.bind(3)
|
||||
m1_output = m1.forward.bind(dag_input)
|
||||
m2_output = model_func.bind(dag_input)
|
||||
|
||||
dag = combine.bind(m1_output, m2_output)
|
||||
# 2*3 + 2*2
|
||||
assert ray.get(dag.execute(2)) == 10
|
||||
# 3*3 + 3*2
|
||||
assert ray.get(dag.execute(3)) == 15
|
||||
|
||||
|
||||
def test_input_attr_partial_access(shared_ray_instance):
|
||||
@ray.remote
|
||||
class Model:
|
||||
def __init__(self, weight: int):
|
||||
self.weight = weight
|
||||
|
||||
def forward(self, input: int):
|
||||
return self.weight * input
|
||||
|
||||
@ray.remote
|
||||
def combine(a, b, c, d=None):
|
||||
if not d:
|
||||
return a + b + c
|
||||
else:
|
||||
return a + b + c + d["deep"]["nested"]
|
||||
|
||||
# 1) Test default wrapping of args and kwargs into internal python object
|
||||
with InputNode() as dag_input:
|
||||
m1 = Model.bind(1)
|
||||
m2 = Model.bind(2)
|
||||
m1_output = m1.forward.bind(dag_input[0])
|
||||
m2_output = m2.forward.bind(dag_input[1])
|
||||
dag = combine.bind(m1_output, m2_output, dag_input.m3, dag_input.m4)
|
||||
# 1*1 + 2*2 + 3 + 4 = 12
|
||||
assert ray.get(dag.execute(1, 2, m3=3, m4={"deep": {"nested": 4}})) == 12
|
||||
|
||||
# 2) Test user passed data object as only input to the dag.execute()
|
||||
class UserDataObj:
|
||||
user_object_field_0: Any
|
||||
user_object_field_1: Any
|
||||
field_3: Any
|
||||
|
||||
def __init__(
|
||||
self, user_object_field_0: Any, user_object_field_1: Any, field_3: Any
|
||||
) -> None:
|
||||
self.user_object_field_0 = user_object_field_0
|
||||
self.user_object_field_1 = user_object_field_1
|
||||
self.field_3 = field_3
|
||||
|
||||
with InputNode() as dag_input:
|
||||
m1 = Model.bind(1)
|
||||
m2 = Model.bind(2)
|
||||
m1_output = m1.forward.bind(dag_input.user_object_field_0)
|
||||
m2_output = m2.forward.bind(dag_input.user_object_field_1)
|
||||
dag = combine.bind(m1_output, m2_output, dag_input.field_3)
|
||||
|
||||
# 1*1 + 2*2 + 3
|
||||
assert ray.get(dag.execute(UserDataObj(1, 2, 3))) == 8
|
||||
|
||||
# 3) Test user passed only one list object with regular list index accessor
|
||||
with InputNode() as dag_input:
|
||||
m1 = Model.bind(1)
|
||||
m2 = Model.bind(2)
|
||||
m1_output = m1.forward.bind(dag_input[0])
|
||||
m2_output = m2.forward.bind(dag_input[1])
|
||||
dag = combine.bind(m1_output, m2_output, dag_input[2])
|
||||
# 1*1 + 2*2 + 3 + 4 = 12
|
||||
assert ray.get(dag.execute([1, 2, 3])) == 8
|
||||
|
||||
# 4) Test user passed only one dict object with key str accessor
|
||||
with InputNode() as dag_input:
|
||||
m1 = Model.bind(1)
|
||||
m2 = Model.bind(2)
|
||||
m1_output = m1.forward.bind(dag_input["m1"])
|
||||
m2_output = m2.forward.bind(dag_input["m2"])
|
||||
dag = combine.bind(m1_output, m2_output, dag_input["m3"])
|
||||
# 1*1 + 2*2 + 3 + 4 = 12
|
||||
assert ray.get(dag.execute({"m1": 1, "m2": 2, "m3": 3})) == 8
|
||||
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match="Please only use int index or str as first-level key",
|
||||
):
|
||||
with InputNode() as dag_input:
|
||||
m1 = Model.bind(1)
|
||||
dag = m1.forward.bind(dag_input[(1, 2)])
|
||||
|
||||
|
||||
def test_ensure_in_context_manager(shared_ray_instance):
|
||||
# No enforcement on creation given __enter__ executes after __init__
|
||||
input = InputNode()
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=(
|
||||
"InputNode is a singleton instance that should be only used "
|
||||
"in context manager"
|
||||
),
|
||||
):
|
||||
input.execute()
|
||||
|
||||
@ray.remote
|
||||
def f(input):
|
||||
return input
|
||||
|
||||
# No enforcement on creation given __enter__ executes after __init__
|
||||
dag = f.bind(InputNode())
|
||||
with pytest.raises(
|
||||
AssertionError,
|
||||
match=(
|
||||
"InputNode is a singleton instance that should be only used "
|
||||
"in context manager"
|
||||
),
|
||||
):
|
||||
dag.execute()
|
||||
|
||||
|
||||
def test_ensure_input_node_singleton(shared_ray_instance):
|
||||
@ray.remote
|
||||
def f(input):
|
||||
return input
|
||||
|
||||
@ray.remote
|
||||
def combine(a, b):
|
||||
return a + b
|
||||
|
||||
with InputNode() as input_1:
|
||||
a = f.bind(input_1)
|
||||
with InputNode() as input_2:
|
||||
b = f.bind(input_2)
|
||||
dag = combine.bind(a, b)
|
||||
|
||||
with pytest.raises(
|
||||
AssertionError, match="Each DAG should only have one unique InputNode"
|
||||
):
|
||||
_ = ray.get(dag.execute(2))
|
||||
|
||||
|
||||
def test_apply_recursive_caching(shared_ray_instance):
|
||||
@ray.remote
|
||||
def f(input):
|
||||
return input
|
||||
|
||||
input = InputNode()
|
||||
f_node = f.bind(input)
|
||||
|
||||
a, b = input, f_node
|
||||
for _ in range(10):
|
||||
a, b = f.bind(a, b), f.bind(a, b)
|
||||
|
||||
counter = 0
|
||||
original_apply_recursive = DAGNode.apply_recursive
|
||||
|
||||
def _apply_recursive_with_counter(self, fn):
|
||||
nonlocal counter
|
||||
counter += 1
|
||||
return original_apply_recursive(self, fn)
|
||||
|
||||
DAGNode.apply_recursive = _apply_recursive_with_counter
|
||||
|
||||
a.apply_recursive(lambda node: node)
|
||||
|
||||
# Prior to #40337; count was 2559
|
||||
assert counter == 40
|
||||
DAGNode.apply_recursive = original_apply_recursive
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,203 @@
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
from ray.dag.input_node import InputNode
|
||||
from ray.dag.output_node import MultiOutputNode
|
||||
from ray.util.state import list_tasks
|
||||
|
||||
|
||||
def test_output_node(shared_ray_instance):
|
||||
@ray.remote
|
||||
def f(input):
|
||||
return input
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
with InputNode() as input_data:
|
||||
dag = MultiOutputNode(f.bind(input_data))
|
||||
|
||||
with InputNode() as input_data:
|
||||
dag = MultiOutputNode([f.bind(input_data)])
|
||||
|
||||
assert ray.get(dag.execute(1)) == [1]
|
||||
assert ray.get(dag.execute(2)) == [2]
|
||||
|
||||
with InputNode() as input_data:
|
||||
dag = MultiOutputNode([f.bind(input_data["x"]), f.bind(input_data["y"])])
|
||||
|
||||
refs = dag.execute({"x": 1, "y": 2})
|
||||
assert len(refs) == 2
|
||||
assert ray.get(refs) == [1, 2]
|
||||
|
||||
with InputNode() as input_data:
|
||||
dag = MultiOutputNode(
|
||||
[f.bind(input_data["x"]), f.bind(input_data["y"]), f.bind(input_data["x"])]
|
||||
)
|
||||
|
||||
refs = dag.execute({"x": 1, "y": 2})
|
||||
assert len(refs) == 3
|
||||
assert ray.get(refs) == [1, 2, 1]
|
||||
|
||||
|
||||
def test_dag_with_actor_handle(shared_ray_instance):
|
||||
"""Verify DAG API works with actor created by .remote"""
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self):
|
||||
self.forward_called = 0
|
||||
self.init_called = 0
|
||||
|
||||
def forward(self, input):
|
||||
print("forward")
|
||||
self.forward_called += 1
|
||||
return input
|
||||
|
||||
def initialize(self, input):
|
||||
print("initialize")
|
||||
self.init_called += 1
|
||||
return input
|
||||
|
||||
def get(self):
|
||||
return (self.forward_called, self.init_called)
|
||||
|
||||
worker = Worker.remote()
|
||||
with InputNode() as input_node:
|
||||
init_dag = worker.initialize.bind(input_node)
|
||||
with InputNode() as input_node:
|
||||
forward_dag = worker.forward.bind(input_node)
|
||||
|
||||
assert ray.get(init_dag.execute(1)) == 1
|
||||
assert ray.get(forward_dag.execute(2)) == 2
|
||||
|
||||
# Make sure both forward/initialize called only once
|
||||
assert ray.get(worker.get.remote()) == (1, 1)
|
||||
|
||||
# Double check the actor is resued.
|
||||
assert ray.get(init_dag.execute(1)) == 1
|
||||
assert ray.get(worker.get.remote()) == (1, 2)
|
||||
|
||||
|
||||
def test_dag_with_alive_actors_chained(shared_ray_instance):
|
||||
"""Verify we can have multiple DAGs to the
|
||||
same actor that are chained.
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Actor:
|
||||
def __init__(self, init_value):
|
||||
self.i = init_value
|
||||
|
||||
def add(self, x):
|
||||
return self.i + x
|
||||
|
||||
@ray.remote
|
||||
def combine(x, y):
|
||||
return x + y
|
||||
|
||||
a1 = Actor.remote(10)
|
||||
a1_dag = a1.add.bind(a1.add.bind(2)) # 22
|
||||
a1_dag_2 = a1.add.bind(a1.add.bind(6)) # 26
|
||||
dag = combine.bind(a1_dag, a1_dag_2)
|
||||
|
||||
assert ray.get(dag.execute()) == 48
|
||||
|
||||
|
||||
def test_tensor_parallel_dag(shared_ray_instance):
|
||||
"""Simulate the TP DAG with N workers.
|
||||
Input -> forward -> MultiOutput
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
class Worker:
|
||||
def __init__(self, rank):
|
||||
self.rank = rank
|
||||
self.forwarded = 0
|
||||
|
||||
def forward(self, input_data: int):
|
||||
print(input_data)
|
||||
self.forwarded += 1
|
||||
return self.rank + input_data
|
||||
|
||||
def initialize(self):
|
||||
pass
|
||||
|
||||
def get_forwarded(self):
|
||||
return self.forwarded
|
||||
|
||||
NUM_WORKERS = 4
|
||||
workers = [Worker.remote(i) for i in range(NUM_WORKERS)]
|
||||
# Init multiple times.
|
||||
for _ in range(4):
|
||||
ray.get([worker.initialize.remote() for worker in workers])
|
||||
|
||||
with InputNode() as input_data:
|
||||
dag = MultiOutputNode([worker.forward.bind(input_data) for worker in workers])
|
||||
|
||||
# Run DAG repetitively.
|
||||
ITER = 4
|
||||
assert ITER > 1
|
||||
for i in range(ITER):
|
||||
ref = dag.execute(i)
|
||||
all_outputs = ray.get(ref)
|
||||
assert len(all_outputs) == NUM_WORKERS
|
||||
assert all_outputs == [i + j for j in range(NUM_WORKERS)]
|
||||
|
||||
forwarded = ray.get([worker.get_forwarded.remote() for worker in workers])
|
||||
assert forwarded == [ITER for _ in range(NUM_WORKERS)]
|
||||
|
||||
|
||||
def test_shared_output(shared_ray_instance):
|
||||
"""Verify when an upstream task output is shared by
|
||||
multi output, the upstream task runs only once.
|
||||
"""
|
||||
|
||||
@ray.remote
|
||||
def shared_f():
|
||||
return 1
|
||||
|
||||
@ray.remote
|
||||
def g(input):
|
||||
return input + 1
|
||||
|
||||
@ray.remote
|
||||
def h(input):
|
||||
return input + 2
|
||||
|
||||
x = shared_f.bind()
|
||||
dag = MultiOutputNode([g.bind(x), h.bind(x)])
|
||||
|
||||
assert ray.get(dag.execute()) == [2, 3]
|
||||
|
||||
# Verify f ran only once.
|
||||
def verify():
|
||||
tasks = list_tasks(filters=[("name", "=", "shared_f")])
|
||||
return len(tasks) == 1
|
||||
|
||||
wait_for_condition(verify)
|
||||
|
||||
|
||||
def test_bind_survives_handle_deletion(shared_ray_instance):
|
||||
"""Verify that .bind().execute() still works even if the original handle was dropped."""
|
||||
|
||||
@ray.remote
|
||||
class A:
|
||||
def f(self):
|
||||
return 1
|
||||
|
||||
# Grab the handle and the bound method node
|
||||
actor = A.remote()
|
||||
method_node = actor.f.bind()
|
||||
|
||||
# Destroy the only Python variable reference and force collection
|
||||
del actor
|
||||
|
||||
# Executing should now succeed because the node holds the ref
|
||||
result = ray.get(method_node.execute())
|
||||
assert result == 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,67 @@
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def test_basic_dag_with_names_plot():
|
||||
@ray.remote
|
||||
def a(*args, **kwargs):
|
||||
pass
|
||||
|
||||
tmp1 = a.options(name="tmp1").bind()
|
||||
tmp2 = a.options(name="tmp2").bind()
|
||||
tmp3 = a.options(name="tmp3").bind(tmp1, tmp2)
|
||||
tmp4 = a.options(name="tmp4").bind()
|
||||
tmp5 = a.options(name="tmp5").bind(tmp4)
|
||||
tmp6 = a.options(name="tmp6").bind()
|
||||
dag = a.bind(tmp3, tmp5, tmp6)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
to_file = os.path.join(tmpdir, "tmp.png")
|
||||
ray.dag.plot(dag, to_file)
|
||||
assert os.path.isfile(to_file)
|
||||
|
||||
graph = ray.dag.vis_utils._dag_to_dot(dag)
|
||||
to_string = graph.to_string()
|
||||
assert "tmp1 -> tmp3" in to_string
|
||||
assert "tmp2 -> tmp3" in to_string
|
||||
assert "tmp4 -> tmp5" in to_string
|
||||
assert "tmp3 -> a" in to_string
|
||||
assert "tmp5 -> a" in to_string
|
||||
assert "tmp6 -> a" in to_string
|
||||
|
||||
|
||||
def test_basic_dag_without_names_plot():
|
||||
@ray.remote
|
||||
def a(*args, **kwargs):
|
||||
pass
|
||||
|
||||
tmp1 = a.bind()
|
||||
tmp2 = a.bind()
|
||||
tmp3 = a.bind(tmp1, tmp2)
|
||||
tmp4 = a.bind()
|
||||
tmp5 = a.bind(tmp4)
|
||||
tmp6 = a.bind()
|
||||
dag = a.bind(tmp3, tmp5, tmp6)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
to_file = os.path.join(tmpdir, "tmp.png")
|
||||
ray.dag.plot(dag, to_file)
|
||||
assert os.path.isfile(to_file)
|
||||
|
||||
graph = ray.dag.vis_utils._dag_to_dot(dag)
|
||||
to_string = graph.to_string()
|
||||
assert "a -> a_2" in to_string
|
||||
assert "a_1 -> a_2" in to_string
|
||||
assert "a_3 -> a_4" in to_string
|
||||
assert "a_2 -> a_6" in to_string
|
||||
assert "a_4 -> a_6" in to_string
|
||||
assert "a_5 -> a_6" in to_string
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", "-s", __file__]))
|
||||
@@ -0,0 +1,88 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.dag.py_obj_scanner import _instances, _PyObjScanner
|
||||
|
||||
|
||||
class Source:
|
||||
pass
|
||||
|
||||
|
||||
def test_simple_replace():
|
||||
scanner = _PyObjScanner(source_type=Source)
|
||||
my_objs = [Source(), [Source(), {"key": Source()}]]
|
||||
|
||||
found = scanner.find_nodes(my_objs)
|
||||
assert len(found) == 3
|
||||
|
||||
replaced = scanner.replace_nodes({obj: 1 for obj in found})
|
||||
assert replaced == [1, [1, {"key": 1}]]
|
||||
|
||||
|
||||
def test_replace_multiple_types():
|
||||
class OtherSource:
|
||||
pass
|
||||
|
||||
scanner = _PyObjScanner(source_type=(Source, OtherSource))
|
||||
my_objs = [Source(), [Source(), {"key": Source(), "key2": OtherSource()}]]
|
||||
|
||||
found = scanner.find_nodes(my_objs)
|
||||
assert len(found) == 4
|
||||
|
||||
replaced = scanner.replace_nodes(
|
||||
{obj: 1 if isinstance(obj, Source) else 2 for obj in found}
|
||||
)
|
||||
assert replaced == [1, [1, {"key": 1, "key2": 2}]]
|
||||
|
||||
|
||||
def test_replace_nested_in_obj():
|
||||
"""Test that the source can be nested in arbitrary objects."""
|
||||
scanner = _PyObjScanner(source_type=Source)
|
||||
|
||||
class Outer:
|
||||
def __init__(self, inner: Any):
|
||||
self._inner = inner
|
||||
|
||||
def __eq__(self, other):
|
||||
return self._inner == other._inner
|
||||
|
||||
my_objs = [Outer(Source()), Outer(Outer(Source())), Outer((Source(),))]
|
||||
|
||||
found = scanner.find_nodes(my_objs)
|
||||
assert len(found) == 3
|
||||
|
||||
replaced = scanner.replace_nodes({obj: 1 for obj in found})
|
||||
assert replaced == [Outer(1), Outer(Outer(1)), Outer((1,))]
|
||||
|
||||
|
||||
def test_scanner_clear():
|
||||
"""Test scanner clear to make the scanner GCable"""
|
||||
prev_len = len(_instances)
|
||||
|
||||
def call_find_nodes():
|
||||
scanner = _PyObjScanner(source_type=Source)
|
||||
my_objs = [Source(), [Source(), {"key": Source()}]]
|
||||
scanner.find_nodes(my_objs)
|
||||
scanner.clear()
|
||||
assert id(scanner) not in _instances
|
||||
|
||||
call_find_nodes()
|
||||
assert prev_len == len(_instances)
|
||||
|
||||
def call_find_and_replace_nodes():
|
||||
scanner = _PyObjScanner(source_type=Source)
|
||||
my_objs = [Source(), [Source(), {"key": Source()}]]
|
||||
found = scanner.find_nodes(my_objs)
|
||||
scanner.replace_nodes({obj: 1 for obj in found})
|
||||
scanner.clear()
|
||||
assert id(scanner) not in _instances
|
||||
|
||||
call_find_and_replace_nodes()
|
||||
assert prev_len == len(_instances)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user