wehub-resource-sync
1426 lines
45 KiB
Python
1426 lines
45 KiB
Python
# coding: utf-8
|
|
import logging
|
|
import pickle
|
|
import sys
|
|
import time
|
|
import traceback
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
import ray
|
|
import ray.cluster_utils
|
|
import ray.exceptions
|
|
import ray.experimental.channel as ray_channel
|
|
from ray._private.test_utils import get_actor_node_id
|
|
from ray.dag.compiled_dag_node import CompiledDAG
|
|
from ray.exceptions import RayChannelError, RayChannelTimeoutError
|
|
from ray.experimental.channel.torch_tensor_type import TorchTensorType
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def create_driver_actor():
|
|
return CompiledDAG.DAGDriverProxyActor.options(
|
|
label_selector={
|
|
ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
|
|
}
|
|
).remote()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_put_local_get(ray_start_regular):
|
|
driver_actor = create_driver_actor()
|
|
chan = ray_channel.Channel(
|
|
None,
|
|
[
|
|
(driver_actor, get_actor_node_id(driver_actor)),
|
|
],
|
|
1000,
|
|
)
|
|
|
|
num_writes = 1000
|
|
for i in range(num_writes):
|
|
val = i.to_bytes(8, "little")
|
|
chan.write(val)
|
|
assert chan.read() == val
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_read_timeout(ray_start_regular):
|
|
driver_actor = create_driver_actor()
|
|
chan = ray_channel.Channel(
|
|
None,
|
|
[
|
|
(driver_actor, get_actor_node_id(driver_actor)),
|
|
],
|
|
1000,
|
|
)
|
|
|
|
with pytest.raises(RayChannelTimeoutError):
|
|
chan.read(timeout=1)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_write_timeout(ray_start_regular):
|
|
driver_actor = create_driver_actor()
|
|
chan = ray_channel.Channel(
|
|
None,
|
|
[
|
|
(driver_actor, get_actor_node_id(driver_actor)),
|
|
],
|
|
1000,
|
|
)
|
|
|
|
val = 1
|
|
bytes = val.to_bytes(8, "little")
|
|
chan.write(bytes, timeout=1)
|
|
with pytest.raises(RayChannelTimeoutError):
|
|
chan.write(bytes, timeout=1)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
@pytest.mark.parametrize("remote", [True, False])
|
|
def test_driver_as_reader(ray_start_cluster, remote):
|
|
cluster = ray_start_cluster
|
|
if remote:
|
|
# This node is for the driver. num_cpus is 1 because the
|
|
# CompiledDAG.DAGDriverProxyActor needs a place to run.
|
|
cluster.add_node(num_cpus=1)
|
|
ray.init(address=cluster.address)
|
|
# This node is for the writer actor.
|
|
cluster.add_node(num_cpus=1)
|
|
else:
|
|
# This node is for both the driver (including the
|
|
# CompiledDAG.DAGDriverProxyActor) and the writer actor.
|
|
cluster.add_node(num_cpus=2)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def setup(self, driver_actor):
|
|
self._channel = ray_channel.Channel(
|
|
ray.get_runtime_context().current_actor,
|
|
[(driver_actor, get_actor_node_id(driver_actor))],
|
|
1000,
|
|
)
|
|
|
|
def get_channel(self):
|
|
return self._channel
|
|
|
|
def write(self):
|
|
self._channel.write(b"x")
|
|
|
|
a = Actor.remote()
|
|
ray.get(a.setup.remote(create_driver_actor()))
|
|
chan = ray.get(a.get_channel.remote())
|
|
|
|
ray.get(a.write.remote())
|
|
assert chan.read() == b"x"
|
|
|
|
|
|
@pytest.mark.parametrize("remote", [True, False])
|
|
def test_driver_as_reader_with_resize(ray_start_cluster, remote):
|
|
cluster = ray_start_cluster
|
|
if remote:
|
|
# This node is for the driver. num_cpus is 1 because the
|
|
# CompiledDAG.DAGDriverProxyActor needs a place to run.
|
|
cluster.add_node(num_cpus=1)
|
|
ray.init(address=cluster.address)
|
|
# This node is for the writer actor.
|
|
cluster.add_node(num_cpus=1)
|
|
else:
|
|
# This node is for both the driver (including the
|
|
# CompiledDAG.DAGDriverProxyActor) and the writer actor.
|
|
cluster.add_node(num_cpus=2)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def setup(self, driver_actor):
|
|
self._channel = ray_channel.Channel(
|
|
ray.get_runtime_context().current_actor,
|
|
[(driver_actor, get_actor_node_id(driver_actor))],
|
|
1000,
|
|
)
|
|
|
|
def get_channel(self):
|
|
return self._channel
|
|
|
|
def write(self):
|
|
self._channel.write(b"x")
|
|
|
|
def write_large(self):
|
|
self._channel.write(b"x" * 2000)
|
|
|
|
a = Actor.remote()
|
|
ray.get(a.setup.remote(create_driver_actor()))
|
|
chan = ray.get(a.get_channel.remote())
|
|
|
|
ray.get(a.write.remote())
|
|
assert chan.read() == b"x"
|
|
|
|
ray.get(a.write_large.remote())
|
|
assert chan.read() == b"x" * 2000
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_set_error_before_read(ray_start_regular):
|
|
"""
|
|
Tests that if a channel is closed after a reader, a subsequent read does not block
|
|
forever.
|
|
"""
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def __init__(self):
|
|
self.arr = None
|
|
|
|
def create_channel(self, writer, reader_and_node_list):
|
|
self._channel = ray_channel.Channel(writer, reader_and_node_list, 1000)
|
|
return self._channel
|
|
|
|
def pass_channel(self, channel):
|
|
self._channel = channel
|
|
|
|
def close(self):
|
|
self._channel.close()
|
|
|
|
def write(self, arr):
|
|
self._channel.write(arr)
|
|
|
|
def read(self):
|
|
self.arr = self._channel.read()
|
|
# Keep self.arr in scope. While self.arr is in scope, its backing
|
|
# shared_ptr<MutableObjectBuffer> in C++ will also stay in scope.
|
|
# Under normal execution, this will block the next read() from
|
|
# returning, since we are still using the shared buffer.
|
|
|
|
# In this test we are checking that if the channel is closed, then
|
|
# the next read() will return an error immediately instead of
|
|
# blocking, even though we still have self.arr in scope.
|
|
return self.arr
|
|
|
|
for _ in range(10):
|
|
a = Actor.remote()
|
|
b = Actor.remote()
|
|
node_b = get_actor_node_id(b)
|
|
|
|
chan = ray.get(a.create_channel.remote(a, [(b, node_b)]))
|
|
ray.get(b.pass_channel.remote(chan))
|
|
|
|
# Use numpy to enable zero-copy deserialization.
|
|
arr = np.random.rand(100)
|
|
ray.get(a.write.remote(arr))
|
|
assert (arr == ray.get(b.read.remote())).all()
|
|
|
|
# Check that the thread does not block on the second call to read() below.
|
|
# read() acquires a lock, though if the lock is not released when
|
|
# read() fails (because the channel has been closed), then an additional
|
|
# call to read() *could* block.
|
|
|
|
# We wrap both calls to read() in pytest.raises() as both calls could
|
|
# trigger an RayChannelError exception if the channel has already been closed.
|
|
with pytest.raises(
|
|
ray.exceptions.RayTaskError, match=r"Channel closed"
|
|
) as exc_info:
|
|
ray.get([a.close.remote(), b.read.remote()])
|
|
assert isinstance(exc_info.value.as_instanceof_cause(), RayChannelError)
|
|
with pytest.raises(ray.exceptions.RayTaskError) as exc_info:
|
|
ray.get(b.read.remote())
|
|
assert isinstance(exc_info.value.as_instanceof_cause(), RayChannelError)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_errors(ray_start_regular):
|
|
"""
|
|
Tests that an exception is thrown when there are more readers than specificed in the
|
|
channel constructor.
|
|
"""
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def make_chan(self, readers, do_write=True):
|
|
self.chan = ray_channel.Channel(
|
|
ray.get_runtime_context().current_actor, readers, 1000
|
|
)
|
|
if do_write:
|
|
self.chan.write(b"hello")
|
|
return self.chan
|
|
|
|
a = Actor.remote()
|
|
# Multiple consecutive reads from the same process are fine.
|
|
driver_actor = create_driver_actor()
|
|
chan = ray.get(
|
|
a.make_chan.remote(
|
|
[(driver_actor, get_actor_node_id(driver_actor))], do_write=True
|
|
)
|
|
)
|
|
assert chan.read() == b"hello"
|
|
|
|
@ray.remote
|
|
class Reader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def read(self, chan):
|
|
return chan.read()
|
|
|
|
readers = [Reader.remote(), Reader.remote()]
|
|
# Check that an exception is thrown when there are more readers than specificed in
|
|
# the channel constructor.
|
|
chan = ray.get(
|
|
a.make_chan.remote([(readers[0], get_actor_node_id(readers[0]))], do_write=True)
|
|
)
|
|
# At least 1 reader.
|
|
with pytest.raises(ray.exceptions.RayTaskError) as exc_info:
|
|
ray.get([reader.read.remote(chan) for reader in readers])
|
|
assert "ray.exceptions.RaySystemError" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_put_different_meta(ray_start_regular):
|
|
driver_actor = create_driver_actor()
|
|
chan = ray_channel.Channel(
|
|
None, [(driver_actor, get_actor_node_id(driver_actor))], 1000
|
|
)
|
|
|
|
def _test(val):
|
|
chan.write(val)
|
|
|
|
read_val = chan.read()
|
|
if isinstance(val, np.ndarray):
|
|
assert np.array_equal(read_val, val)
|
|
else:
|
|
assert read_val == val
|
|
|
|
_test(b"hello")
|
|
_test("hello")
|
|
_test(1000)
|
|
_test(np.random.rand(10))
|
|
|
|
|
|
def test_multiple_channels_different_nodes(ray_start_cluster):
|
|
"""
|
|
Tests that multiple channels can be used at the same time between two nodes.
|
|
"""
|
|
cluster = ray_start_cluster
|
|
# This node is for the driver.
|
|
cluster.add_node(num_cpus=0)
|
|
ray.init(address=cluster.address)
|
|
# This node is for the Reader actors.
|
|
cluster.add_node(num_cpus=1)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def read(self, channel, val):
|
|
read_val = channel.read()
|
|
if isinstance(val, np.ndarray):
|
|
assert np.array_equal(read_val, val)
|
|
else:
|
|
assert read_val == val
|
|
|
|
a = Actor.remote()
|
|
node_a = get_actor_node_id(a)
|
|
chan_a = ray_channel.Channel(None, [(a, node_a)], 1000)
|
|
chan_b = ray_channel.Channel(None, [(a, node_a)], 1000)
|
|
channels = [chan_a, chan_b]
|
|
|
|
val = np.random.rand(5)
|
|
for i in range(10):
|
|
for channel in channels:
|
|
channel.write(val)
|
|
for channel in channels:
|
|
ray.get(a.read.remote(channel, val))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_resize_channel_on_same_node(ray_start_regular):
|
|
"""
|
|
Tests that the channel backing store is automatically increased when a large object
|
|
is written to it. The writer and reader are on the same node.
|
|
"""
|
|
driver_actor = create_driver_actor()
|
|
chan = ray_channel.Channel(
|
|
None, [(driver_actor, get_actor_node_id(driver_actor))], 1000
|
|
)
|
|
|
|
def _test(val):
|
|
chan.write(val)
|
|
|
|
read_val = chan.read()
|
|
if isinstance(val, np.ndarray):
|
|
assert np.array_equal(read_val, val)
|
|
else:
|
|
assert read_val == val
|
|
|
|
# `np.random.rand(100)` requires more than 1000 bytes of storage. The channel is
|
|
# allocated above with a backing store size of 1000 bytes.
|
|
_test(np.random.rand(100))
|
|
|
|
# Check that another write still works.
|
|
_test(np.random.rand(5))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_resize_channel_on_same_node_with_actor(ray_start_regular):
|
|
"""
|
|
Tests that the channel backing store is automatically increased when a large object
|
|
is written to it. The writer and reader are on the same node, and the reader is an
|
|
actor.
|
|
"""
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def read(self, channel, val):
|
|
read_val = channel.read()
|
|
if isinstance(val, np.ndarray):
|
|
assert np.array_equal(read_val, val)
|
|
else:
|
|
assert read_val == val
|
|
|
|
def _test(channel, actor, val):
|
|
channel.write(val)
|
|
ray.get(actor.read.remote(channel, val))
|
|
|
|
a = Actor.remote()
|
|
node_a = get_actor_node_id(a)
|
|
chan = ray_channel.Channel(None, [(a, node_a)], 1000)
|
|
|
|
# `np.random.rand(100)` requires more than 1000 bytes of storage. The channel is
|
|
# allocated above with a backing store size of 1000 bytes.
|
|
_test(chan, a, np.random.rand(100))
|
|
|
|
# Check that another write still works.
|
|
_test(chan, a, np.random.rand(5))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_resize_channel_on_different_nodes(ray_start_cluster):
|
|
"""
|
|
Tests that the channel backing store is automatically increased when a large object
|
|
is written to it. The writer and reader are on different nodes, and the reader is an
|
|
actor.
|
|
"""
|
|
cluster = ray_start_cluster
|
|
# This node is for the driver.
|
|
cluster.add_node(num_cpus=0)
|
|
ray.init(address=cluster.address)
|
|
# This node is for the Reader actors.
|
|
cluster.add_node(num_cpus=1)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def read(self, channel, val):
|
|
read_val = channel.read()
|
|
if isinstance(val, np.ndarray):
|
|
assert np.array_equal(read_val, val)
|
|
else:
|
|
assert read_val == val
|
|
|
|
def _test(channel, actor, val):
|
|
channel.write(val)
|
|
ray.get(actor.read.remote(channel, val))
|
|
|
|
a = Actor.remote()
|
|
node_a = get_actor_node_id(a)
|
|
chan = ray_channel.Channel(None, [(a, node_a)], 1000)
|
|
|
|
# `np.random.rand(100)` requires more than 1000 bytes of storage. The channel is
|
|
# allocated above with a backing store size of 1000 bytes.
|
|
_test(chan, a, np.random.rand(100))
|
|
|
|
# Check that another write still works.
|
|
_test(chan, a, np.random.rand(5))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
@pytest.mark.parametrize("num_readers", [1, 4])
|
|
def test_put_remote_get(ray_start_regular, num_readers):
|
|
"""
|
|
Tests that an actor can read objects/primitives of various types through a channel
|
|
when the reader is spawned with @ray.remote.
|
|
"""
|
|
|
|
@ray.remote(num_cpus=0)
|
|
class Reader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def read(self, chan, num_writes):
|
|
for i in range(num_writes):
|
|
val = i.to_bytes(8, "little")
|
|
assert chan.read() == val
|
|
|
|
for i in range(num_writes):
|
|
val = i.to_bytes(100, "little")
|
|
assert chan.read() == val
|
|
|
|
for val in [
|
|
b"hello world",
|
|
"hello again",
|
|
1000,
|
|
]:
|
|
assert chan.read() == val
|
|
|
|
num_writes = 1000
|
|
reader_and_node_list = []
|
|
for _ in range(num_readers):
|
|
handle = Reader.remote()
|
|
node = get_actor_node_id(handle)
|
|
reader_and_node_list.append((handle, node))
|
|
|
|
chan = ray_channel.Channel(None, reader_and_node_list, 1000)
|
|
chan.ensure_registered_as_writer()
|
|
|
|
done = [reader.read.remote(chan, num_writes) for reader, _ in reader_and_node_list]
|
|
for i in range(num_writes):
|
|
val = i.to_bytes(8, "little")
|
|
chan.write(val)
|
|
|
|
# Test different data size.
|
|
for i in range(num_writes):
|
|
val = i.to_bytes(100, "little")
|
|
chan.write(val)
|
|
|
|
# Test different metadata.
|
|
for val in [
|
|
b"hello world",
|
|
"hello again",
|
|
1000,
|
|
]:
|
|
chan.write(val)
|
|
|
|
ray.get(done)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
@pytest.mark.parametrize("remote", [True, False])
|
|
def test_remote_reader(ray_start_cluster, remote):
|
|
"""
|
|
Tests that an actor can read objects/primitives of various types through a channel
|
|
when the reader and writer are on the (1) same node (remote=False) along with (2)
|
|
different nodes (remote=True).
|
|
"""
|
|
num_readers = 10
|
|
num_writes = 1000
|
|
num_iterations = 3
|
|
|
|
cluster = ray_start_cluster
|
|
if remote:
|
|
# This node is for the driver.
|
|
cluster.add_node(num_cpus=0)
|
|
ray.init(address=cluster.address)
|
|
# This node is for the Reader actors.
|
|
cluster.add_node(num_cpus=num_readers)
|
|
else:
|
|
# This node is for both the driver and the Reader actors.
|
|
cluster.add_node(num_cpus=num_readers)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Reader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def get_node_id(self) -> str:
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
def pass_channel(self, channel):
|
|
self._reader_chan = channel
|
|
|
|
def read(self, num_reads):
|
|
for i in range(num_reads):
|
|
self._reader_chan.read()
|
|
|
|
reader_and_node_list = []
|
|
for _ in range(num_readers):
|
|
handle = Reader.remote()
|
|
node = get_actor_node_id(handle)
|
|
reader_and_node_list.append((handle, node))
|
|
|
|
channel = ray_channel.Channel(None, reader_and_node_list, 1000)
|
|
|
|
# All readers have received the channel.
|
|
ray.get([reader.pass_channel.remote(channel) for reader, _ in reader_and_node_list])
|
|
|
|
for _ in range(num_iterations):
|
|
work = [reader.read.remote(num_writes) for reader, _ in reader_and_node_list]
|
|
start = time.perf_counter()
|
|
for i in range(num_writes):
|
|
channel.write(b"x")
|
|
end = time.perf_counter()
|
|
ray.get(work)
|
|
print(end - start, 10_000 / (end - start))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
@pytest.mark.parametrize("remote", [True, False])
|
|
def test_remote_reader_close(ray_start_cluster, remote):
|
|
"""
|
|
Tests that readers do not block forever on read() when they close the channel.
|
|
Specifically, the following behavior should happen:
|
|
1. Each reader calls read() on one channel.
|
|
2. Each reader calls close() on the channel on a different thread.
|
|
3. Each reader should unblock and return from read().
|
|
|
|
Tests (1) the readers and writer on the same node (remote=False) along with
|
|
different nodes (remote=True).
|
|
"""
|
|
num_readers = 10
|
|
|
|
cluster = ray_start_cluster
|
|
if remote:
|
|
# This node is for the driver.
|
|
cluster.add_node(num_cpus=0)
|
|
ray.init(address=cluster.address)
|
|
# This node is for the Reader actors.
|
|
cluster.add_node(num_cpus=num_readers)
|
|
else:
|
|
# This node is for both the driver and the Reader actors.
|
|
cluster.add_node(num_cpus=num_readers)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Reader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def get_node_id(self) -> str:
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
def pass_channel(self, channel):
|
|
self._reader_chan = channel
|
|
|
|
def read(self):
|
|
try:
|
|
self._reader_chan.read()
|
|
except RayChannelError:
|
|
pass
|
|
|
|
def close(self):
|
|
self._reader_chan.close()
|
|
|
|
reader_and_node_list = []
|
|
for _ in range(num_readers):
|
|
handle = Reader.remote()
|
|
node = get_actor_node_id(handle)
|
|
reader_and_node_list.append((handle, node))
|
|
channel = ray_channel.Channel(None, reader_and_node_list, 1000)
|
|
|
|
# All readers have received the channel.
|
|
ray.get([reader.pass_channel.remote(channel) for reader, _ in reader_and_node_list])
|
|
|
|
reads = [
|
|
reader.read.options(concurrency_group="_ray_system").remote()
|
|
for reader, _ in reader_and_node_list
|
|
]
|
|
with pytest.raises(ray.exceptions.GetTimeoutError):
|
|
ray.get(reads, timeout=1.0)
|
|
|
|
ray.get([reader.close.remote() for reader, _ in reader_and_node_list])
|
|
ray.get(reads)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_intra_process_channel_single_reader(ray_start_cluster):
|
|
"""
|
|
(1) Test whether an actor can read/write from an IntraProcessChannel.
|
|
(2) Test whether the _SerializationContext cleans up the
|
|
data after all readers have read it.
|
|
(3) Test whether the actor can write again after reading 1 time.
|
|
"""
|
|
# This node is for both the driver and the Reader actors.
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=1)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def pass_channel(self, channel):
|
|
self._chan = channel
|
|
|
|
def read(self):
|
|
return self._chan.read()
|
|
|
|
def write(self, value):
|
|
self._chan.write(value)
|
|
|
|
def get_ctx_buffer_size(self):
|
|
ctx = ray_channel.ChannelContext.get_current().serialization_context
|
|
return len(ctx.intra_process_channel_buffers)
|
|
|
|
actor = Actor.remote()
|
|
channel = ray_channel.IntraProcessChannel(num_readers=1)
|
|
ray.get(actor.pass_channel.remote(channel))
|
|
|
|
ray.get(actor.write.remote("hello"))
|
|
assert ray.get(actor.read.remote()) == "hello"
|
|
|
|
# The _SerializationContext should clean up the data after a read.
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 0
|
|
|
|
# Write again after reading num_readers times.
|
|
ray.get(actor.write.remote("world"))
|
|
assert ray.get(actor.read.remote()) == "world"
|
|
|
|
# The _SerializationContext should clean up the data after a read.
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 0
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_intra_process_channel_multi_readers(ray_start_cluster):
|
|
"""
|
|
(1) Test whether an actor can read/write from an IntraProcessChannel.
|
|
(2) Test whether the _SerializationContext cleans up the
|
|
data after all readers have read it.
|
|
(3) Test whether the actor can write again after reading num_readers times.
|
|
(4) Test whether an exception is raised when calling write() before all readers
|
|
have read the data.
|
|
"""
|
|
# This node is for both the driver and the Reader actors.
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=1)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def pass_channel(self, channel):
|
|
self._chan = channel
|
|
|
|
def read(self):
|
|
return self._chan.read()
|
|
|
|
def write(self, value):
|
|
self._chan.write(value)
|
|
|
|
def get_ctx_buffer_size(self):
|
|
ctx = ray_channel.ChannelContext.get_current().serialization_context
|
|
return len(ctx.intra_process_channel_buffers)
|
|
|
|
actor = Actor.remote()
|
|
channel = ray_channel.IntraProcessChannel(num_readers=2)
|
|
ray.get(actor.pass_channel.remote(channel))
|
|
|
|
ray.get(actor.write.remote("hello"))
|
|
# first read
|
|
assert ray.get(actor.read.remote()) == "hello"
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 1
|
|
# second read
|
|
assert ray.get(actor.read.remote()) == "hello"
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 0
|
|
|
|
# Write again after reading num_readers times.
|
|
ray.get(actor.write.remote("world"))
|
|
# first read
|
|
assert ray.get(actor.read.remote()) == "world"
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 1
|
|
# second read
|
|
assert ray.get(actor.read.remote()) == "world"
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 0
|
|
|
|
# Write again
|
|
ray.get(actor.write.remote("hello world"))
|
|
# first read
|
|
assert ray.get(actor.read.remote()) == "hello world"
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 1
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
ray.get(actor.write.remote("world hello"))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_cached_channel_single_reader():
|
|
ray.init()
|
|
|
|
@ray.remote
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def pass_channel(self, channel):
|
|
self._chan = channel
|
|
|
|
def read(self):
|
|
return self._chan.read()
|
|
|
|
def get_ctx_buffer_size(self):
|
|
ctx = ray_channel.ChannelContext.get_current().serialization_context
|
|
return len(ctx.intra_process_channel_buffers)
|
|
|
|
actor = Actor.remote()
|
|
inner_channel = ray_channel.Channel(
|
|
None,
|
|
[
|
|
(actor, get_actor_node_id(actor)),
|
|
],
|
|
1000,
|
|
)
|
|
channel = ray_channel.CachedChannel(num_reads=1, inner_channel=inner_channel)
|
|
ray.get(actor.pass_channel.remote(channel))
|
|
|
|
channel.write("hello")
|
|
assert ray.get(actor.read.remote()) == "hello"
|
|
|
|
# The _SerializationContext should clean up the data after a read.
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 0
|
|
|
|
# Write again after reading num_readers times.
|
|
channel.write("world")
|
|
assert ray.get(actor.read.remote()) == "world"
|
|
|
|
# The _SerializationContext should clean up the data after a read.
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 0
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_cached_channel_multi_readers(ray_start_cluster):
|
|
@ray.remote
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def pass_channel(self, channel):
|
|
self._chan = channel
|
|
|
|
def read(self):
|
|
return self._chan.read()
|
|
|
|
def get_ctx_buffer_size(self):
|
|
ctx = ray_channel.ChannelContext.get_current().serialization_context
|
|
return len(ctx.intra_process_channel_buffers)
|
|
|
|
actor = Actor.remote()
|
|
inner_channel = ray_channel.Channel(
|
|
None,
|
|
[
|
|
(actor, get_actor_node_id(actor)),
|
|
],
|
|
1000,
|
|
)
|
|
channel = ray_channel.CachedChannel(num_reads=2, inner_channel=inner_channel)
|
|
ray.get(actor.pass_channel.remote(channel))
|
|
|
|
channel.write("hello")
|
|
# first read
|
|
assert ray.get(actor.read.remote()) == "hello"
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 1
|
|
# second read
|
|
assert ray.get(actor.read.remote()) == "hello"
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 0
|
|
|
|
# Write again after reading num_readers times.
|
|
channel.write("world")
|
|
# first read
|
|
assert ray.get(actor.read.remote()) == "world"
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 1
|
|
# second read
|
|
assert ray.get(actor.read.remote()) == "world"
|
|
assert ray.get(actor.get_ctx_buffer_size.remote()) == 0
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_composite_channel_single_reader(ray_start_cluster):
|
|
"""
|
|
(1) The driver can write data to CompositeChannel and an actor can read it.
|
|
(2) An actor can write data to CompositeChannel and the actor itself can read it.
|
|
(3) An actor can write data to CompositeChannel and another actor can read it.
|
|
(4) An actor can write data to CompositeChannel and the driver can read it.
|
|
"""
|
|
# This node is for both the driver and the Reader actors.
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=2)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def pass_channel(self, channel):
|
|
self._chan = channel
|
|
|
|
def create_composite_channel(
|
|
self, writer, reader_and_node_list, driver_actor_id
|
|
):
|
|
self._chan = ray_channel.CompositeChannel(
|
|
writer, reader_and_node_list, 10, driver_actor_id
|
|
)
|
|
return self._chan
|
|
|
|
def read(self):
|
|
return self._chan.read()
|
|
|
|
def write(self, value):
|
|
self._chan.write(value)
|
|
|
|
actor1 = Actor.remote()
|
|
actor2 = Actor.remote()
|
|
node1 = get_actor_node_id(actor1)
|
|
node2 = get_actor_node_id(actor2)
|
|
|
|
# Create a channel to communicate between driver process and actor1.
|
|
driver_to_actor1_channel = ray_channel.CompositeChannel(None, [(actor1, node1)], 10)
|
|
ray.get(actor1.pass_channel.remote(driver_to_actor1_channel))
|
|
driver_to_actor1_channel.write("hello")
|
|
assert ray.get(actor1.read.remote()) == "hello"
|
|
|
|
# Create a channel to communicate between two tasks in actor1.
|
|
ray.get(actor1.create_composite_channel.remote(actor1, [(actor1, node1)], None))
|
|
ray.get(actor1.write.remote("world"))
|
|
assert ray.get(actor1.read.remote()) == "world"
|
|
|
|
# Create a channel to communicate between actor1 and actor2.
|
|
actor1_to_actor2_channel = ray.get(
|
|
actor1.create_composite_channel.remote(actor1, [(actor2, node2)], None)
|
|
)
|
|
ray.get(actor2.pass_channel.remote(actor1_to_actor2_channel))
|
|
ray.get(actor1.write.remote("hello world"))
|
|
assert ray.get(actor2.read.remote()) == "hello world"
|
|
|
|
# Create a channel to communicate between actor2 and driver process.
|
|
driver_actor = create_driver_actor()
|
|
actor2_to_driver_channel = ray.get(
|
|
actor2.create_composite_channel.remote(
|
|
actor2,
|
|
[(driver_actor, get_actor_node_id(driver_actor))],
|
|
driver_actor._actor_id.hex(),
|
|
)
|
|
)
|
|
ray.get(actor2.write.remote("world hello"))
|
|
assert actor2_to_driver_channel.read() == "world hello"
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_composite_channel_multiple_readers(ray_start_cluster):
|
|
"""
|
|
Test the behavior of CompositeChannel when there are multiple readers.
|
|
|
|
(1) The driver can write data to CompositeChannel and two actors can read it.
|
|
(2) An actor can write data to CompositeChannel and another actor, as well as
|
|
itself, can read it.
|
|
(3) An actor writes data to CompositeChannel and two actor methods on the same
|
|
actor read it. This is not supported and should raise an exception.
|
|
"""
|
|
# This node is for both the driver and the Reader actors.
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=2)
|
|
ray.init(address=cluster.address)
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def pass_channel(self, channel):
|
|
self._chan = channel
|
|
|
|
def create_composite_channel(self, writer, reader_and_node_list):
|
|
self._chan = ray_channel.CompositeChannel(writer, reader_and_node_list, 10)
|
|
return self._chan
|
|
|
|
def read(self):
|
|
return self._chan.read()
|
|
|
|
def write(self, value):
|
|
self._chan.write(value)
|
|
|
|
actor1 = Actor.remote()
|
|
actor2 = Actor.remote()
|
|
node1 = get_actor_node_id(actor1)
|
|
node2 = get_actor_node_id(actor2)
|
|
|
|
# The driver writes data to CompositeChannel and actor1 and actor2 read it.
|
|
driver_output_channel = ray_channel.CompositeChannel(
|
|
None, [(actor1, node1), (actor2, node2)], 10
|
|
)
|
|
ray.get(actor1.pass_channel.remote(driver_output_channel))
|
|
ray.get(actor2.pass_channel.remote(driver_output_channel))
|
|
driver_output_channel.write("hello")
|
|
assert ray.get([actor1.read.remote(), actor2.read.remote()]) == ["hello"] * 2
|
|
|
|
# actor1 writes data to CompositeChannel and actor1 and actor2 read it.
|
|
actor1_output_channel = ray.get(
|
|
actor1.create_composite_channel.remote(
|
|
actor1, [(actor1, node1), (actor2, node2)]
|
|
)
|
|
)
|
|
ray.get(actor2.pass_channel.remote(actor1_output_channel))
|
|
ray.get(actor1.write.remote("world"))
|
|
assert ray.get([actor1.read.remote(), actor2.read.remote()]) == ["world"] * 2
|
|
|
|
actor1_output_channel = ray.get(
|
|
actor1.create_composite_channel.remote(
|
|
actor1, [(actor1, node1), (actor1, node1)]
|
|
)
|
|
)
|
|
ray.get(actor1.write.remote("hello world"))
|
|
assert ray.get(actor1.read.remote()) == "hello world"
|
|
|
|
with pytest.raises(ray.exceptions.RayTaskError):
|
|
# actor1_output_channel can be read only once if the readers
|
|
# are the same actor. Note that reading the channel multiple
|
|
# times is supported via channel cache mechanism.
|
|
ray.get(actor1.read.remote())
|
|
|
|
# actor1 writes data to CompositeChannel and actor2 and the driver reads it.
|
|
driver_actor = create_driver_actor()
|
|
actor1_output_channel = ray.get(
|
|
actor1.create_composite_channel.remote(
|
|
actor1,
|
|
[(driver_actor, get_actor_node_id(driver_actor)), (actor2, node2)],
|
|
)
|
|
)
|
|
ray.get(actor2.pass_channel.remote(actor1_output_channel))
|
|
ray.get(actor1.write.remote("world hello"))
|
|
assert ray.get(actor2.read.remote()) == "world hello"
|
|
assert actor1_output_channel.read() == "world hello"
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_put_error(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
# This node is for both the driver (including the CompiledDAG.DAGDriverProxyActor)
|
|
# and the writer actor.
|
|
cluster.add_node(num_cpus=2)
|
|
ray.init(address=cluster.address)
|
|
|
|
def _wrap_exception(exc):
|
|
backtrace = ray._private.utils.format_error_message(
|
|
"".join(traceback.format_exception(type(exc), exc, exc.__traceback__)),
|
|
task_exception=True,
|
|
)
|
|
wrapped = ray.exceptions.RayTaskError(
|
|
function_name="do_exec_tasks",
|
|
traceback_str=backtrace,
|
|
cause=exc,
|
|
)
|
|
return wrapped
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def setup(self, reader_and_node_list):
|
|
self._channel = ray_channel.Channel(
|
|
ray.get_runtime_context().current_actor,
|
|
reader_and_node_list,
|
|
1000,
|
|
)
|
|
|
|
def get_channel(self):
|
|
return self._channel
|
|
|
|
def write(self, write_error):
|
|
if write_error:
|
|
try:
|
|
raise ValueError("")
|
|
except Exception as exc:
|
|
self._channel.write(_wrap_exception(exc))
|
|
else:
|
|
self._channel.write(b"x")
|
|
|
|
a = Actor.remote()
|
|
driver_actor = create_driver_actor()
|
|
ray.get(a.setup.remote([(driver_actor, get_actor_node_id(driver_actor))]))
|
|
chan = ray.get(a.get_channel.remote())
|
|
|
|
# Putting a bytes object multiple times is okay.
|
|
for _ in range(3):
|
|
ray.get(a.write.remote(write_error=False))
|
|
assert chan.read() == b"x"
|
|
|
|
# Putting an exception multiple times is okay.
|
|
for _ in range(3):
|
|
ray.get(a.write.remote(write_error=True))
|
|
try:
|
|
assert chan.read()
|
|
except Exception as exc:
|
|
assert isinstance(exc, ValueError)
|
|
assert isinstance(exc, ray.exceptions.RayTaskError)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_payload_large(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
# This node is for the driver.
|
|
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
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def get_node_id(self):
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
def read(self, channel, val):
|
|
assert channel.read() == val
|
|
|
|
def create_actor(node):
|
|
return Actor.options(
|
|
label_selector={ray._raylet.RAY_NODE_ID_KEY: node}
|
|
).remote()
|
|
|
|
driver_node = ray.get_runtime_context().get_node_id()
|
|
actor_node = nodes[0] if nodes[0] != driver_node else nodes[1]
|
|
assert driver_node != actor_node
|
|
a = create_actor(actor_node)
|
|
node_a = ray.get(a.get_node_id.remote())
|
|
assert driver_node != ray.get(a.get_node_id.remote())
|
|
|
|
# Ray sets the gRPC payload max size to 512 MiB. We choose a size in this test that
|
|
# is a bit larger.
|
|
size = 1024 * 1024 * 600
|
|
ch = ray_channel.Channel(None, [(a, node_a)], size)
|
|
|
|
val = b"x" * size
|
|
ch.write(val)
|
|
ray.get(a.read.remote(ch, val))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_payload_resize_large(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
# This node is for the driver.
|
|
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
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def get_node_id(self):
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
def read(self, channel, val):
|
|
assert channel.read() == val
|
|
|
|
def create_actor(node):
|
|
return Actor.options(
|
|
label_selector={ray._raylet.RAY_NODE_ID_KEY: node}
|
|
).remote()
|
|
|
|
driver_node = ray.get_runtime_context().get_node_id()
|
|
actor_node = nodes[0] if nodes[0] != driver_node else nodes[1]
|
|
assert driver_node != actor_node
|
|
a = create_actor(actor_node)
|
|
assert driver_node != ray.get(a.get_node_id.remote())
|
|
|
|
ch = ray_channel.Channel(None, [(a, actor_node)], 1000)
|
|
|
|
# Ray sets the gRPC payload max size to 512 MiB. We choose a size in this test that
|
|
# is a bit larger.
|
|
size = 1024 * 1024 * 600
|
|
val = b"x" * size
|
|
ch.write(val)
|
|
ray.get(a.read.remote(ch, val))
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
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.
|
|
first_node_handle = cluster.add_node(num_cpus=2)
|
|
# This node is for the other 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
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def get_node_id(self):
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
def read(self, channel, val):
|
|
assert channel.read() == val
|
|
return val
|
|
|
|
def create_actor(node):
|
|
return Actor.options(
|
|
label_selector={ray._raylet.RAY_NODE_ID_KEY: node}
|
|
).remote()
|
|
|
|
a = create_actor(nodes[0])
|
|
b = create_actor(nodes[1])
|
|
actors = [a, b]
|
|
|
|
nodes_check = ray.get([act.get_node_id.remote() for act in actors])
|
|
a_node = nodes_check[0]
|
|
b_node = nodes_check[1]
|
|
assert a_node != b_node
|
|
|
|
driver_actor = create_driver_actor()
|
|
driver_node = get_actor_node_id(driver_actor)
|
|
|
|
ch = ray_channel.Channel(
|
|
None, [(driver_actor, driver_node), (a, a_node), (b, b_node)], 1000
|
|
)
|
|
val = 1
|
|
ch.write(val)
|
|
assert ray.get([a.read.remote(ch, val) for a in actors]) == [val, val]
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.platform != "linux" and sys.platform != "darwin",
|
|
reason="Requires Linux or Mac.",
|
|
)
|
|
def test_bunch_readers_on_different_nodes(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
# This node is for the driver (including the DriverHelperActor) and two of the
|
|
# readers.
|
|
first_node_handle = cluster.add_node(num_cpus=3)
|
|
# This node is for the other two readers.
|
|
second_node_handle = cluster.add_node(num_cpus=2)
|
|
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 the readers are on different 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
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def get_node_id(self):
|
|
return ray.get_runtime_context().get_node_id()
|
|
|
|
def read(self, channel, val):
|
|
assert channel.read() == val
|
|
return val
|
|
|
|
def create_actor(node):
|
|
return Actor.options(
|
|
label_selector={ray._raylet.RAY_NODE_ID_KEY: node}
|
|
).remote()
|
|
|
|
a = create_actor(nodes[0])
|
|
b = create_actor(nodes[0])
|
|
c = create_actor(nodes[1])
|
|
d = create_actor(nodes[1])
|
|
actors = [a, b, c, d]
|
|
|
|
nodes_check = ray.get([act.get_node_id.remote() for act in actors])
|
|
a_node = nodes_check[0]
|
|
b_node = nodes_check[1]
|
|
c_node = nodes_check[2]
|
|
d_node = nodes_check[3]
|
|
assert a_node == b_node
|
|
assert b_node != c_node
|
|
assert c_node == d_node
|
|
|
|
driver_actor = create_driver_actor()
|
|
driver_node = get_actor_node_id(driver_actor)
|
|
|
|
ch = ray_channel.Channel(
|
|
None,
|
|
[
|
|
(driver_actor, driver_node),
|
|
(a, a_node),
|
|
(b, b_node),
|
|
(c, c_node),
|
|
(d, d_node),
|
|
],
|
|
1000,
|
|
)
|
|
i = 1
|
|
ch.write(i)
|
|
assert ray.get([a.read.remote(ch, i) for a in actors]) == [
|
|
i for _ in range(len(actors))
|
|
]
|
|
|
|
|
|
def test_buffered_channel(shutdown_only):
|
|
"""Test buffered shared memory channel."""
|
|
BUFFER_SIZE = 5
|
|
|
|
@ray.remote(num_cpus=1)
|
|
class Actor:
|
|
def __init__(self):
|
|
self.write_index = 0
|
|
|
|
def setup(self, driver_actor):
|
|
self._channel = ray_channel.BufferedSharedMemoryChannel(
|
|
ray.get_runtime_context().current_actor,
|
|
[(driver_actor, get_actor_node_id(driver_actor))],
|
|
BUFFER_SIZE,
|
|
typ=1000,
|
|
)
|
|
|
|
def get_channel(self):
|
|
return self._channel
|
|
|
|
def write(self, i, timeout=None) -> bool:
|
|
"""Write to a channel Return False if channel times out.
|
|
Return true otherwise.
|
|
"""
|
|
self.write_index += 1
|
|
try:
|
|
self._channel.write(i, timeout)
|
|
except ray.exceptions.RayChannelTimeoutError:
|
|
return False
|
|
assert self._channel.next_write_index == self.write_index % BUFFER_SIZE
|
|
return True
|
|
|
|
a = Actor.remote()
|
|
ray.get(a.setup.remote(create_driver_actor()))
|
|
chan = ray.get(a.get_channel.remote())
|
|
|
|
print("Test basic.")
|
|
# Iterate more than buffer size to make sure it works over and over again.
|
|
read_idx = 0
|
|
for i in range(BUFFER_SIZE * 3):
|
|
read_idx += 1
|
|
assert ray.get(a.write.remote(i))
|
|
assert chan.read() == i
|
|
assert chan.next_read_index == read_idx % BUFFER_SIZE
|
|
|
|
print("Test Write timeout.")
|
|
# Test write timeout.
|
|
for i in range(BUFFER_SIZE):
|
|
# fill the buffer withtout read.
|
|
ray.get(a.write.remote(i))
|
|
# timeout because all the buffer is exhausted without being consumed.
|
|
assert ray.get(a.write.remote(1, timeout=1)) is False
|
|
|
|
print("Test Read timeout.")
|
|
# Test read timeout.
|
|
for i in range(BUFFER_SIZE):
|
|
# This reads all previous writes.
|
|
assert chan.read() == i
|
|
# This read times out because there's no new write, and the call blocks.
|
|
with pytest.raises(ray.exceptions.RayChannelTimeoutError):
|
|
chan.read(timeout=1)
|
|
|
|
print("Test serialization/deserialization works")
|
|
deserialized = pickle.loads(pickle.dumps(chan))
|
|
assert len(chan._buffers) == len(deserialized._buffers)
|
|
for i in range(len(chan._buffers)):
|
|
assert (
|
|
deserialized._buffers[i]._writer._actor_id
|
|
== chan._buffers[i]._writer._actor_id
|
|
)
|
|
|
|
|
|
def test_torch_dtype():
|
|
typ = TorchTensorType()
|
|
typ.register_custom_serializer()
|
|
|
|
t = torch.randn(5, 5, dtype=torch.bfloat16)
|
|
with pytest.raises(TypeError):
|
|
t.numpy()
|
|
|
|
ref = ray.put(t)
|
|
t_out = ray.get(ref)
|
|
assert (t == t_out).all()
|
|
assert t_out.dtype == t.dtype
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|