wehub-resource-sync
521 lines
15 KiB
Python
521 lines
15 KiB
Python
# coding: utf-8
|
|
import logging
|
|
import sys
|
|
from typing import Dict, List, Tuple
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
import ray
|
|
import ray.cluster_utils
|
|
from ray._private.test_utils import get_actor_node_id
|
|
from ray.experimental.channel.conftest import (
|
|
Barrier,
|
|
TracedChannel,
|
|
start_nccl_mock,
|
|
)
|
|
from ray.experimental.channel.torch_tensor_accelerator_channel import (
|
|
_init_communicator,
|
|
)
|
|
from ray.experimental.channel.torch_tensor_type import TorchTensorType
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@ray.remote(num_cpus=0, num_gpus=1)
|
|
class Worker:
|
|
def __init__(self):
|
|
self.tensor_chan = None
|
|
# Key -> a patched channel that we can use to record write ops.
|
|
self.traced_channels: Dict[str, TracedChannel] = {}
|
|
|
|
def start_mock(self):
|
|
"""
|
|
Patch methods that require CUDA.
|
|
"""
|
|
start_nccl_mock()
|
|
|
|
def set_nccl_channel(self, typ, tensor_chan):
|
|
typ.register_custom_serializer()
|
|
self.tensor_chan = tensor_chan
|
|
|
|
def create_traced_channel(self, key, readers):
|
|
self.traced_channels[key] = TracedChannel(
|
|
ray.get_runtime_context().current_actor, readers
|
|
)
|
|
|
|
def get_num_channel_ops(self, key):
|
|
ops = self.traced_channels[key].ops
|
|
self.traced_channels[key].ops = []
|
|
return len(ops)
|
|
|
|
def create_nccl_channel(
|
|
self,
|
|
typ: TorchTensorType,
|
|
reader_and_node_list: List[Tuple[ray.actor.ActorHandle, str]],
|
|
tensor_metadata_channel_key=None,
|
|
cpu_data_channel_key=None,
|
|
):
|
|
typ.register_custom_serializer()
|
|
|
|
tensor_metadata_channel = None
|
|
if tensor_metadata_channel_key is not None:
|
|
tensor_metadata_channel = self.traced_channels[tensor_metadata_channel_key]
|
|
cpu_data_channel = None
|
|
if cpu_data_channel_key is not None:
|
|
cpu_data_channel = self.traced_channels[cpu_data_channel_key]
|
|
|
|
self.tensor_chan = typ.create_channel(
|
|
ray.get_runtime_context().current_actor,
|
|
reader_and_node_list,
|
|
driver_actor_id=None,
|
|
_cpu_data_channel=cpu_data_channel,
|
|
_tensor_metadata_channel=tensor_metadata_channel,
|
|
)
|
|
|
|
return self.tensor_chan
|
|
|
|
def send(self, val, shape, dtype):
|
|
t = torch.ones(shape, dtype=dtype) * val
|
|
self.tensor_chan.write(t)
|
|
|
|
def receive(self):
|
|
t = self.tensor_chan.read()
|
|
data = (t[0].item(), t.shape, t.dtype)
|
|
return data
|
|
|
|
def send_dict(self, tensor_dict):
|
|
self.tensor_chan.write(tensor_dict)
|
|
|
|
def receive_dict(self):
|
|
tensor_dict = self.tensor_chan.read()
|
|
vals = []
|
|
for key, t in tensor_dict.items():
|
|
vals.append((key, t[0].item(), t.shape, t.dtype))
|
|
return vals
|
|
|
|
|
|
@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 = Worker.remote()
|
|
receiver = Worker.remote()
|
|
receiver_node = get_actor_node_id(receiver)
|
|
|
|
ray.get(
|
|
[
|
|
sender.start_mock.remote(),
|
|
receiver.start_mock.remote(),
|
|
]
|
|
)
|
|
|
|
nccl_id = _init_communicator([sender, receiver])
|
|
|
|
chan_typ = TorchTensorType(transport="accelerator")
|
|
chan_typ.set_communicator_id(nccl_id)
|
|
chan_ref = sender.create_nccl_channel.remote(chan_typ, [(receiver, receiver_node)])
|
|
receiver_ready = receiver.set_nccl_channel.remote(chan_typ, chan_ref)
|
|
ray.get([chan_ref, receiver_ready])
|
|
|
|
shape = (10,)
|
|
dtype = torch.float16
|
|
|
|
refs = []
|
|
for i in range(3):
|
|
sender.send.remote(i, shape, dtype)
|
|
refs.append(receiver.receive.remote())
|
|
assert ray.get(refs) == [(i, shape, dtype) for i in range(3)]
|
|
|
|
shapes = [(10,), (20,), (30,)]
|
|
dtypes = [torch.float, torch.float16, torch.int]
|
|
|
|
refs = []
|
|
for i in range(3):
|
|
sender.send.remote(i, shapes[i], dtypes[i])
|
|
refs.append(receiver.receive.remote())
|
|
assert ray.get(refs) == [(i, shapes[i], dtypes[i]) for i in range(3)]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster",
|
|
[
|
|
{
|
|
"num_cpus": 4,
|
|
"num_gpus": 4,
|
|
"num_nodes": 1,
|
|
}
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_multiple_receivers(ray_start_cluster):
|
|
"""
|
|
Test sender with multiple receivers pattern. Check that all receivers
|
|
receive correct results.
|
|
"""
|
|
# Create one barrier per sender-receiver pair.
|
|
barriers = [ # noqa
|
|
Barrier.options(name=f"barrier-0-{i}").remote() for i in range(1, 4)
|
|
] # noqa
|
|
|
|
sender = Worker.remote()
|
|
receiver_to_node = []
|
|
for _ in range(3):
|
|
handle = Worker.remote()
|
|
node = get_actor_node_id(handle)
|
|
receiver_to_node.append((handle, node))
|
|
|
|
workers = [sender] + [receiver for receiver, _ in receiver_to_node]
|
|
|
|
ray.get([worker.start_mock.remote() for worker in workers])
|
|
|
|
nccl_id = _init_communicator(workers)
|
|
|
|
chan_typ = TorchTensorType(transport="accelerator")
|
|
chan_typ.set_communicator_id(nccl_id)
|
|
chan_ref = sender.create_nccl_channel.remote(chan_typ, receiver_to_node)
|
|
receiver_ready = [
|
|
receiver.set_nccl_channel.remote(chan_typ, chan_ref)
|
|
for receiver, _ in receiver_to_node
|
|
]
|
|
ray.get(receiver_ready)
|
|
|
|
shapes = [(10,), (20,), (30,)]
|
|
dtypes = [torch.float, torch.float16, torch.int]
|
|
|
|
all_refs = []
|
|
for i in range(3):
|
|
sender.send.remote(i, shapes[i], dtypes[i])
|
|
all_refs.append([receiver.receive.remote() for receiver, _ in receiver_to_node])
|
|
# Check that all receivers received the correct value.
|
|
for i, refs in enumerate(all_refs):
|
|
for val in ray.get(refs):
|
|
assert val == (i, shapes[i], dtypes[i])
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster",
|
|
[
|
|
{
|
|
"num_cpus": 2,
|
|
"num_gpus": 2,
|
|
"num_nodes": 1,
|
|
}
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_static_shape(ray_start_cluster):
|
|
"""
|
|
Test that when static_shape=True is passed, we only send metadata for the
|
|
first operation. Afterwards, we reuse the same shape and dtype for future
|
|
tensors. Sending a tensor of the wrong shape or dtype throws an error.
|
|
"""
|
|
# Barrier name should be barrier-{lower rank}-{higher rank}.
|
|
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
|
|
|
sender = Worker.remote()
|
|
receiver = Worker.remote()
|
|
|
|
ray.get(
|
|
[
|
|
sender.start_mock.remote(),
|
|
receiver.start_mock.remote(),
|
|
]
|
|
)
|
|
|
|
nccl_id = _init_communicator([sender, receiver])
|
|
|
|
chan_typ = TorchTensorType(
|
|
transport="accelerator",
|
|
_static_shape=True,
|
|
)
|
|
chan_typ.set_communicator_id(nccl_id)
|
|
receiver_to_node = [(receiver, get_actor_node_id(receiver))]
|
|
sender.create_traced_channel.remote("tensor_metadata", receiver_to_node)
|
|
sender.create_traced_channel.remote("cpu_data", receiver_to_node)
|
|
chan_ref = sender.create_nccl_channel.remote(
|
|
chan_typ,
|
|
receiver_to_node,
|
|
"tensor_metadata",
|
|
"cpu_data",
|
|
)
|
|
receiver_ready = receiver.set_nccl_channel.remote(chan_typ, chan_ref)
|
|
ray.get([chan_ref, receiver_ready])
|
|
|
|
shape = (10,)
|
|
dtype = torch.float16
|
|
|
|
refs = []
|
|
for i in range(3):
|
|
sender.send.remote(i, shape, dtype)
|
|
refs.append(receiver.receive.remote())
|
|
assert ray.get(refs) == [(i, shape, dtype) for i in range(3)]
|
|
|
|
# When static_shape=True, we should only send one metadata op. After the
|
|
# first metadata is sent, we reuse it for all future sends.
|
|
num_tensor_metadata_ops = ray.get(
|
|
sender.get_num_channel_ops.remote("tensor_metadata")
|
|
)
|
|
assert num_tensor_metadata_ops == 1
|
|
num_cpu_data_ops = ray.get(sender.get_num_channel_ops.remote("cpu_data"))
|
|
assert num_cpu_data_ops == 3
|
|
|
|
# Attempting to write tensors of the wrong shape or dtype will error.
|
|
with pytest.raises(ValueError):
|
|
ray.get(sender.send.remote(4, (20,), dtype))
|
|
with pytest.raises(ValueError):
|
|
ray.get(sender.send.remote(4, shape, torch.float32))
|
|
|
|
# Attempting to write a different number of tensors will error.
|
|
with pytest.raises(ValueError):
|
|
ray.get(sender.send_dict.remote({}))
|
|
with pytest.raises(ValueError):
|
|
ray.get(
|
|
sender.send_dict.remote(
|
|
{
|
|
"first": torch.zeros(10),
|
|
"second": torch.zeros(10),
|
|
}
|
|
)
|
|
)
|
|
|
|
# The channel is still usable for valid tensors after errors occur.
|
|
sender.send.remote(4, shape, dtype)
|
|
ray.get(receiver.receive.remote()) == (4, shape, dtype)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster",
|
|
[
|
|
{
|
|
"num_cpus": 2,
|
|
"num_gpus": 2,
|
|
"num_nodes": 1,
|
|
}
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_direct_return(ray_start_cluster):
|
|
"""
|
|
Test that when _direct_return=True is passed, we never send metadata.
|
|
"""
|
|
# Barrier name should be barrier-{lower rank}-{higher rank}.
|
|
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
|
|
|
sender = Worker.remote()
|
|
receiver = Worker.remote()
|
|
|
|
ray.get(
|
|
[
|
|
sender.start_mock.remote(),
|
|
receiver.start_mock.remote(),
|
|
]
|
|
)
|
|
|
|
nccl_id = _init_communicator([sender, receiver])
|
|
|
|
chan_typ = TorchTensorType(
|
|
transport="accelerator",
|
|
_direct_return=True,
|
|
)
|
|
chan_typ.set_communicator_id(nccl_id)
|
|
receiver_to_node = [(receiver, get_actor_node_id(receiver))]
|
|
sender.create_traced_channel.remote("tensor_metadata", receiver_to_node)
|
|
chan_ref = sender.create_nccl_channel.remote(
|
|
chan_typ,
|
|
receiver_to_node,
|
|
"tensor_metadata",
|
|
)
|
|
receiver_ready = receiver.set_nccl_channel.remote(chan_typ, chan_ref)
|
|
ray.get([chan_ref, receiver_ready])
|
|
|
|
shape = (10,)
|
|
dtype = torch.float16
|
|
|
|
# Sending tensors of different shapes is okay as long as the non-tensor
|
|
# data stays the same.
|
|
refs = []
|
|
values = []
|
|
for i in range(1, 4):
|
|
sender.send.remote(i, shape, dtype)
|
|
values.append((i, shape, dtype))
|
|
refs.append(receiver.receive.remote())
|
|
assert ray.get(refs) == values
|
|
|
|
# When _direct_return=True, we should never send non-tensor data.
|
|
num_tensor_metadata_ops = ray.get(
|
|
sender.get_num_channel_ops.remote("tensor_metadata")
|
|
)
|
|
assert num_tensor_metadata_ops == 3
|
|
|
|
# Attempting to write a different number of tensors will error.
|
|
with pytest.raises(ValueError):
|
|
ray.get(sender.send_dict.remote({}))
|
|
with pytest.raises(ValueError):
|
|
ray.get(
|
|
sender.send_dict.remote(
|
|
{
|
|
"first": torch.zeros(10),
|
|
"second": torch.zeros(10),
|
|
}
|
|
)
|
|
)
|
|
|
|
# The channel is still usable for valid tensors after errors occur.
|
|
sender.send.remote(1, shape, dtype)
|
|
assert ray.get(receiver.receive.remote()) == (1, shape, dtype)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster",
|
|
[
|
|
{
|
|
"num_cpus": 2,
|
|
"num_gpus": 2,
|
|
"num_nodes": 1,
|
|
}
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_static_shape_and_direct_return(ray_start_cluster):
|
|
"""
|
|
Test that when static_shape=True and direct_return=True are passed, we only
|
|
send metadata for the first operation.
|
|
"""
|
|
# Barrier name should be barrier-{lower rank}-{higher rank}.
|
|
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
|
|
|
sender = Worker.remote()
|
|
receiver = Worker.remote()
|
|
|
|
ray.get(
|
|
[
|
|
sender.start_mock.remote(),
|
|
receiver.start_mock.remote(),
|
|
]
|
|
)
|
|
|
|
nccl_id = _init_communicator([sender, receiver])
|
|
|
|
chan_typ = TorchTensorType(
|
|
transport="accelerator",
|
|
_static_shape=True,
|
|
_direct_return=True,
|
|
)
|
|
chan_typ.set_communicator_id(nccl_id)
|
|
receiver_to_node = [(receiver, get_actor_node_id(receiver))]
|
|
sender.create_traced_channel.remote("tensor_metadata", receiver_to_node)
|
|
chan_ref = sender.create_nccl_channel.remote(
|
|
chan_typ,
|
|
receiver_to_node,
|
|
"tensor_metadata",
|
|
)
|
|
receiver_ready = receiver.set_nccl_channel.remote(chan_typ, chan_ref)
|
|
ray.get([chan_ref, receiver_ready])
|
|
|
|
shape = (10,)
|
|
dtype = torch.float16
|
|
|
|
refs = []
|
|
for i in range(3):
|
|
sender.send.remote(i, shape, dtype)
|
|
refs.append(receiver.receive.remote())
|
|
assert ray.get(refs) == [(i, shape, dtype) for i in range(3)]
|
|
|
|
# When static_shape=True, we should only send one metadata op. After the
|
|
# first metadata is sent, we reuse it for all future sends.
|
|
num_tensor_metadata_ops = ray.get(
|
|
sender.get_num_channel_ops.remote("tensor_metadata")
|
|
)
|
|
assert num_tensor_metadata_ops == 1
|
|
|
|
# Attempting to write tensors of the wrong shape or dtype will error.
|
|
with pytest.raises(ValueError):
|
|
ray.get(sender.send.remote(4, (20,), dtype))
|
|
with pytest.raises(ValueError):
|
|
ray.get(sender.send.remote(4, shape, torch.float32))
|
|
|
|
# Attempting to write a different number of tensors will error.
|
|
with pytest.raises(ValueError):
|
|
ray.get(sender.send_dict.remote({}))
|
|
with pytest.raises(ValueError):
|
|
ray.get(
|
|
sender.send_dict.remote(
|
|
{
|
|
"first": torch.zeros(10),
|
|
"second": torch.zeros(10),
|
|
}
|
|
)
|
|
)
|
|
|
|
# The channel is still usable for valid tensors after errors occur.
|
|
sender.send.remote(4, shape, dtype)
|
|
ray.get(receiver.receive.remote()) == (4, shape, dtype)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"ray_start_cluster",
|
|
[
|
|
{
|
|
"num_cpus": 2,
|
|
"num_gpus": 2,
|
|
"num_nodes": 1,
|
|
}
|
|
],
|
|
indirect=True,
|
|
)
|
|
def test_direct_return_with_cpu_data_channel(ray_start_cluster):
|
|
"""
|
|
Test that when _direct_return=True is passed, cpu_data_channel must be None.
|
|
If it is not None, an exception should be raised.
|
|
"""
|
|
barrier = Barrier.options(name="barrier-0-1").remote() # noqa
|
|
|
|
sender = Worker.remote()
|
|
receiver = Worker.remote()
|
|
|
|
ray.get(
|
|
[
|
|
sender.start_mock.remote(),
|
|
receiver.start_mock.remote(),
|
|
]
|
|
)
|
|
|
|
nccl_id = _init_communicator([sender, receiver])
|
|
chan_typ = TorchTensorType(
|
|
transport="accelerator",
|
|
_direct_return=True,
|
|
)
|
|
chan_typ.set_communicator_id(nccl_id)
|
|
receiver_to_node = [(receiver, get_actor_node_id(receiver))]
|
|
sender.create_traced_channel.remote("cpu_data", receiver_to_node)
|
|
chan_ref = sender.create_nccl_channel.remote(
|
|
chan_typ,
|
|
receiver_to_node,
|
|
None,
|
|
"cpu_data",
|
|
)
|
|
with pytest.raises(
|
|
AssertionError, match="CPU channel should be None if direct return is enabled"
|
|
):
|
|
ray.get(chan_ref)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|