wehub-resource-sync
820 lines
33 KiB
Python
820 lines
33 KiB
Python
import io
|
|
import logging
|
|
import time
|
|
from collections import defaultdict, namedtuple
|
|
from typing import Any, Dict, List, Optional, Set, Tuple, Union
|
|
|
|
import ray
|
|
import ray.exceptions
|
|
from ray._raylet import SerializedObject
|
|
from ray.experimental.channel import utils
|
|
from ray.experimental.channel.common import ChannelInterface, ChannelOutputType
|
|
from ray.experimental.channel.intra_process_channel import IntraProcessChannel
|
|
from ray.experimental.channel.utils import get_self_actor
|
|
from ray.util.annotations import DeveloperAPI, PublicAPI
|
|
|
|
# Logger for this module. It should be configured at the entry point
|
|
# into the program using Ray. Ray provides a default configuration at
|
|
# entry/init points.
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _create_channel_ref(
|
|
self: Any,
|
|
buffer_size_bytes: int,
|
|
) -> "ray.ObjectRef":
|
|
"""Create a channel that can be read and written through Ray's shared-memory
|
|
object store.
|
|
|
|
The channel has no buffer, so the writer will block until reader(s) have
|
|
read the previous value.
|
|
|
|
A writer and colocated readers can communicate via a shared memory buffer.
|
|
If the readers are remote, then RPC is used to synchronize the writer and
|
|
readers' buffers.
|
|
|
|
Args:
|
|
self: The actor on which to allocate the channel buffer. Passed via
|
|
``__ray_call__`` so this function executes on the actor.
|
|
buffer_size_bytes: The initial buffer size in bytes for messages
|
|
that can be passed between tasks in the DAG. The buffers will
|
|
be automatically resized if larger messages are written to the
|
|
channel.
|
|
|
|
Returns:
|
|
A wrapper around ``ray.ObjectRef`` backing the channel.
|
|
"""
|
|
worker = ray._private.worker.global_worker
|
|
worker.check_connected()
|
|
|
|
value = b"0" * buffer_size_bytes
|
|
|
|
try:
|
|
object_ref = worker.put_object(value, _is_experimental_channel=True)
|
|
except ray.exceptions.ObjectStoreFullError:
|
|
logger.info(
|
|
"Put failed since the value was either too large or the "
|
|
"store was full of pinned objects."
|
|
)
|
|
raise
|
|
return object_ref
|
|
|
|
|
|
# Compiled Graph maintains 1 reader object reference (also called buffer) per node.
|
|
# reader_ref: The object reference.
|
|
# ref_owner_actor_id: The actor who created the object reference.
|
|
# num_readers: The number of reader actors who reads this object reference.
|
|
ReaderRefInfo = namedtuple(
|
|
"ReaderRefInfo", ["reader_ref", "ref_owner_actor_id", "num_reader_actors"]
|
|
)
|
|
|
|
|
|
class _ResizeChannel:
|
|
"""Sentinel used to resize a channel's backing store on the readers.
|
|
|
|
When a channel must be resized, the channel backing store must be resized on both
|
|
the writer and the reader nodes. The writer first resizes its own backing store. The
|
|
writer then uses an instance of this class as a sentinel value to tell the reader to
|
|
resize its own backing store. The class instance is sent through the channel.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
_node_id_to_reader_ref_info: Dict[str, ReaderRefInfo],
|
|
):
|
|
"""Initialize the resize sentinel.
|
|
|
|
Args:
|
|
_node_id_to_reader_ref_info: Mapping from node id to ``ReaderRefInfo``
|
|
describing the new reader buffers per node.
|
|
"""
|
|
self._node_id_to_reader_ref_info = _node_id_to_reader_ref_info
|
|
|
|
|
|
class SharedMemoryType(ChannelOutputType):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
buffer_size_bytes: Optional[int] = None,
|
|
num_shm_buffers: Optional[int] = None,
|
|
):
|
|
"""Initialize a ``SharedMemoryType``.
|
|
|
|
Args:
|
|
buffer_size_bytes: The initial buffer size in bytes for messages
|
|
that can be passed between tasks in the DAG. The buffers will
|
|
be automatically resized if larger messages are written to the
|
|
channel.
|
|
num_shm_buffers: The number of shared memory buffers per channel.
|
|
Note: In the case of multiple nodes, we only support 1 shared
|
|
memory buffer.
|
|
"""
|
|
super().__init__()
|
|
|
|
from ray.dag import DAGContext
|
|
|
|
ctx = DAGContext.get_current()
|
|
|
|
if buffer_size_bytes is None:
|
|
buffer_size_bytes = ctx.buffer_size_bytes
|
|
self.buffer_size_bytes = buffer_size_bytes
|
|
if num_shm_buffers is None:
|
|
num_shm_buffers = 1
|
|
self._num_shm_buffers = num_shm_buffers
|
|
|
|
def create_channel(
|
|
self,
|
|
writer: Optional["ray.actor.ActorHandle"],
|
|
reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
|
|
driver_actor_id: Optional[str] = None,
|
|
) -> "Channel":
|
|
"""
|
|
Instantiate a ChannelInterface class that can be used
|
|
to pass data of this type.
|
|
|
|
Args:
|
|
writer: The actor that may write to the channel. None signifies the driver.
|
|
reader_and_node_list: A list of tuples, where each tuple contains a reader
|
|
actor handle and the node ID where the actor is located.
|
|
driver_actor_id: If this channel is read by a driver and that driver is an
|
|
actual actor, this will be the actor ID of that driver actor.
|
|
|
|
Returns:
|
|
A ChannelInterface that can be used to pass data
|
|
of this type.
|
|
"""
|
|
return CompositeChannel(
|
|
writer,
|
|
reader_and_node_list,
|
|
self._num_shm_buffers,
|
|
driver_actor_id,
|
|
)
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class Channel(ChannelInterface):
|
|
"""
|
|
A wrapper type for ray.ObjectRef. Currently supports ray.get but not
|
|
ray.wait.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
writer: Optional[ray.actor.ActorHandle],
|
|
reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
|
|
typ: Optional[Union[int, SharedMemoryType]] = None,
|
|
_writer_node_id: Optional["ray.NodeID"] = None,
|
|
_writer_ref: Optional["ray.ObjectRef"] = None,
|
|
_node_id_to_reader_ref_info: Optional[Dict[str, ReaderRefInfo]] = None,
|
|
_writer_registered: bool = False,
|
|
_reader_registered: bool = False,
|
|
):
|
|
"""Create a channel that can be read and written by co-located Ray processes.
|
|
|
|
Anyone may write to or read from the channel. The channel has no
|
|
buffer, so the writer will block until reader(s) have read the previous
|
|
value.
|
|
|
|
Args:
|
|
writer: The actor that may write to the channel. None signifies the driver.
|
|
reader_and_node_list: A list of tuples, where each tuple contains a reader
|
|
actor handle and the node ID where the actor is located.
|
|
typ: Type information about the values passed through the channel.
|
|
Either an integer representing the max buffer size in bytes
|
|
allowed, or a SharedMemoryType.
|
|
_writer_node_id: Internal. Node ID hosting the writer. Provided
|
|
when rehydrating a channel that was constructed on another
|
|
process.
|
|
_writer_ref: Internal. Pre-existing writer-side ``ObjectRef``.
|
|
When set, the constructor skips allocating a new writer buffer.
|
|
_node_id_to_reader_ref_info: Internal. Mapping from node id to
|
|
``ReaderRefInfo`` describing existing reader buffers per node.
|
|
_writer_registered: Internal. Whether the writer side has already
|
|
been registered with the core worker.
|
|
_reader_registered: Internal. Whether the reader side has already
|
|
been registered with the core worker.
|
|
"""
|
|
assert len(reader_and_node_list) > 0
|
|
for reader, _ in reader_and_node_list:
|
|
assert isinstance(reader, ray.actor.ActorHandle)
|
|
|
|
if typ is None:
|
|
typ = SharedMemoryType()
|
|
elif isinstance(typ, int):
|
|
typ = SharedMemoryType(buffer_size_bytes=typ)
|
|
|
|
# The min buffer size must be large enough to at least fit an instance of the
|
|
# _ResizeChannel class along with any metadata.
|
|
MIN_BUFFER_SIZE = int(1000) # 1000 bytes
|
|
if typ.buffer_size_bytes < MIN_BUFFER_SIZE:
|
|
raise ValueError(
|
|
"typ.buffer_size_bytes must be at least MIN_BUFFER_SIZE "
|
|
f"({MIN_BUFFER_SIZE} bytes)"
|
|
)
|
|
|
|
self._writer = writer
|
|
self._reader_and_node_list = reader_and_node_list
|
|
self._typ = typ
|
|
|
|
self._worker = ray._private.worker.global_worker
|
|
self._worker.check_connected()
|
|
|
|
self._writer_registered = _writer_registered
|
|
self._reader_registered = _reader_registered
|
|
# NodeID -> ReaderRefInfo on that node. Note that there's only 1
|
|
# reader ref per node.
|
|
self._node_id_to_reader_ref_info: Dict[str, ReaderRefInfo] = (
|
|
_node_id_to_reader_ref_info or {}
|
|
)
|
|
|
|
# Node ID -> a list of reader actors.
|
|
self._node_id_to_readers: Dict[str, "ray.actor.ActorHandle"] = defaultdict(list)
|
|
for reader, node_id in self._reader_and_node_list:
|
|
self._node_id_to_readers[node_id].append(reader)
|
|
|
|
# Number of readers in a local node.
|
|
self._num_local_readers = 0
|
|
|
|
if _writer_ref is None:
|
|
# We are the writer. Check that the passed handle matches the
|
|
# current actor (or it is the driver).
|
|
# TODO(swang): Channels must be initially constructed by the writer
|
|
# actor, so we shouldn't need to include `writer` in the
|
|
# constructor args. Either support Channels being constructed by
|
|
# someone other than the writer or remove it from the args.
|
|
self_actor = get_self_actor()
|
|
assert writer == self_actor
|
|
|
|
self._writer_node_id = (
|
|
ray.runtime_context.get_runtime_context().get_node_id()
|
|
)
|
|
self._writer_ref = _create_channel_ref(self, typ.buffer_size_bytes)
|
|
|
|
self._create_reader_refs(typ.buffer_size_bytes)
|
|
else:
|
|
assert (
|
|
_writer_node_id is not None
|
|
), "_writer_node_id must also be passed to the constructor when "
|
|
"_writer_ref is."
|
|
assert _node_id_to_reader_ref_info is not None, (
|
|
"_node_id_to_reader_ref_info must also be passed to the constructor "
|
|
"when _writer_ref is."
|
|
)
|
|
|
|
self._writer_ref = _writer_ref
|
|
self._writer_node_id = _writer_node_id
|
|
self._node_id_to_reader_ref_info = _node_id_to_reader_ref_info
|
|
|
|
assert self._num_local_readers == 0
|
|
remote_node_exists = False
|
|
for node_id, readers in self._node_id_to_readers.items():
|
|
if self.is_local_node(node_id):
|
|
self._num_local_readers += len(readers)
|
|
else:
|
|
remote_node_exists = True
|
|
# If remote node exists, we have 1 additional reader that listens
|
|
# to object changes and push them to remote nodes.
|
|
if remote_node_exists:
|
|
self._num_local_readers += 1
|
|
# There must be at least 1 local reader
|
|
assert self._num_local_readers > 0
|
|
|
|
self._local_reader_ref: Optional["ray.ObjectRef"] = self._get_local_reader_ref(
|
|
self._node_id_to_reader_ref_info
|
|
)
|
|
|
|
def _get_local_reader_ref(
|
|
self, _node_id_to_reader_ref_info: Dict[str, ReaderRefInfo]
|
|
) -> Optional["ray.ObjectRef"]:
|
|
for node_id, reader_ref_info in _node_id_to_reader_ref_info.items():
|
|
if self.is_local_node(node_id):
|
|
return reader_ref_info.reader_ref
|
|
return None
|
|
|
|
def _create_reader_refs(
|
|
self,
|
|
buffer_size_bytes: int,
|
|
):
|
|
# TODO(jhumphri): Free the current reader ref once the reference to it is
|
|
# destroyed below.
|
|
|
|
for node_id, readers in self._node_id_to_readers.items():
|
|
if not self.is_local_node(node_id):
|
|
# Find 1 reader in a remote node to create a reference that's
|
|
# shared by all readers. When a new value is written to a reference,
|
|
# it is sent to this reference.
|
|
reader = readers[0]
|
|
fn = reader.__ray_call__
|
|
self._node_id_to_reader_ref_info[node_id] = ReaderRefInfo(
|
|
reader_ref=ray.get(
|
|
fn.remote(_create_channel_ref, buffer_size_bytes)
|
|
),
|
|
ref_owner_actor_id=reader._actor_id,
|
|
num_reader_actors=len(readers),
|
|
)
|
|
else:
|
|
writer_id = ray.ActorID.nil()
|
|
if self._writer is not None:
|
|
writer_id = self._writer._actor_id
|
|
self._node_id_to_reader_ref_info[node_id] = ReaderRefInfo(
|
|
reader_ref=self._writer_ref,
|
|
ref_owner_actor_id=writer_id,
|
|
num_reader_actors=len(readers),
|
|
)
|
|
# There must be only 1 node reader reference per node.
|
|
assert len(self._node_id_to_reader_ref_info) == len(self._node_id_to_readers)
|
|
|
|
# We need to register the new writer_ref.
|
|
self._writer_registered = False
|
|
self.ensure_registered_as_writer()
|
|
|
|
@staticmethod
|
|
def is_local_node(node_id):
|
|
return ray.runtime_context.get_runtime_context().get_node_id() == node_id
|
|
|
|
def ensure_registered_as_writer(self) -> None:
|
|
if self._writer_registered:
|
|
return
|
|
|
|
if not self.is_local_node(self._writer_node_id):
|
|
raise ValueError(
|
|
"`ensure_registered_as_writer()` must only be called on the node that "
|
|
"the writer is on."
|
|
)
|
|
|
|
remote_reader_ref_info: Dict[str, ReaderRefInfo] = {}
|
|
for node_id, reader_ref_info in self._node_id_to_reader_ref_info.items():
|
|
if self.is_local_node(node_id):
|
|
continue
|
|
remote_reader_ref_info[node_id] = reader_ref_info
|
|
|
|
self._worker.core_worker.experimental_channel_register_writer(
|
|
self._writer_ref,
|
|
remote_reader_ref_info,
|
|
)
|
|
self._writer_registered = True
|
|
|
|
def ensure_registered_as_reader(self) -> None:
|
|
if self._reader_registered:
|
|
return
|
|
|
|
for node_id, reader_ref_info in self._node_id_to_reader_ref_info.items():
|
|
if self.is_local_node(node_id):
|
|
self._worker.core_worker.experimental_channel_register_reader(
|
|
reader_ref_info.reader_ref,
|
|
)
|
|
self._reader_registered = True
|
|
|
|
@staticmethod
|
|
def _deserialize_reader_channel(
|
|
writer: ray.actor.ActorHandle,
|
|
reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
|
|
typ: int,
|
|
writer_node_id,
|
|
writer_ref: "ray.ObjectRef",
|
|
node_id_to_reader_ref_info: Dict[str, ReaderRefInfo],
|
|
writer_registered: bool,
|
|
reader_registered: bool,
|
|
) -> "Channel":
|
|
chan = Channel(
|
|
writer,
|
|
reader_and_node_list,
|
|
typ,
|
|
_writer_node_id=writer_node_id,
|
|
_writer_ref=writer_ref,
|
|
_node_id_to_reader_ref_info=node_id_to_reader_ref_info,
|
|
_writer_registered=writer_registered,
|
|
_reader_registered=reader_registered,
|
|
)
|
|
return chan
|
|
|
|
def __reduce__(self):
|
|
assert self._node_id_to_reader_ref_info is not None
|
|
return self._deserialize_reader_channel, (
|
|
self._writer,
|
|
self._reader_and_node_list,
|
|
self._typ,
|
|
self._writer_node_id,
|
|
self._writer_ref,
|
|
self._node_id_to_reader_ref_info,
|
|
self._writer_registered,
|
|
self._reader_registered,
|
|
)
|
|
|
|
def __str__(self) -> str:
|
|
return (
|
|
f"Channel(_node_id_to_reader_ref_info={self._node_id_to_reader_ref_info}, "
|
|
f"_writer_ref={self._writer_ref})"
|
|
)
|
|
|
|
def _resize_channel_if_needed(self, serialized_value: str, timeout_ms: int):
|
|
# serialized_value.total_bytes *only* includes the size of the data. It does not
|
|
# include the size of the metadata, so we must account for the size of the
|
|
# metadata explicitly.
|
|
size = serialized_value.total_bytes + len(serialized_value.metadata)
|
|
if size > self._typ.buffer_size_bytes:
|
|
# Now make the channel backing store larger.
|
|
self._typ.buffer_size_bytes = size
|
|
# TODO(jhumphri): Free the current writer ref once the reference to it is
|
|
# destroyed below.
|
|
# TODO(sang): Support different policies such as 2X buffer size.
|
|
prev_writer_ref = self._writer_ref
|
|
self._writer_ref = _create_channel_ref(self, self._typ.buffer_size_bytes)
|
|
self._create_reader_refs(self._typ.buffer_size_bytes)
|
|
self._local_reader_ref = self._get_local_reader_ref(
|
|
self._node_id_to_reader_ref_info
|
|
)
|
|
|
|
# Write a special message to the channel so that the readers know to
|
|
# stop using the current reader_ref.
|
|
special_message = _ResizeChannel(self._node_id_to_reader_ref_info)
|
|
special_message_serialized = (
|
|
self._worker.get_serialization_context().serialize(special_message)
|
|
)
|
|
self._worker.core_worker.experimental_channel_put_serialized(
|
|
special_message_serialized,
|
|
prev_writer_ref,
|
|
self._num_local_readers,
|
|
timeout_ms,
|
|
)
|
|
# TODO(sang): Clean the previous ref that won't be used.
|
|
# Right now, if we just close it here, it will not work because
|
|
# of race conditions.
|
|
# self._worker.core_worker.experimental_channel_set_error(
|
|
# prev_writer_ref
|
|
# )
|
|
|
|
def write(self, value: Any, timeout: Optional[float] = None) -> None:
|
|
self.ensure_registered_as_writer()
|
|
assert (
|
|
timeout is None or timeout >= 0 or timeout == -1
|
|
), "Timeout must be non-negative or -1."
|
|
# -1 means no timeout (block indefinitely)
|
|
timeout_ms = int(timeout * 1000) if timeout is not None else -1
|
|
|
|
if not isinstance(value, SerializedObject):
|
|
try:
|
|
serialized_value = self._worker.get_serialization_context().serialize(
|
|
value
|
|
)
|
|
except TypeError as e:
|
|
sio = io.StringIO()
|
|
ray.util.inspect_serializability(value, print_file=sio)
|
|
msg = (
|
|
"Could not serialize the put value "
|
|
f"{repr(value)}:\n"
|
|
f"{sio.getvalue()}"
|
|
)
|
|
raise TypeError(msg) from e
|
|
else:
|
|
serialized_value = value
|
|
|
|
start_time = time.monotonic()
|
|
self._resize_channel_if_needed(serialized_value, timeout_ms)
|
|
if timeout is not None:
|
|
timeout_ms -= int((time.monotonic() - start_time) * 1000)
|
|
timeout_ms = max(timeout_ms, 0)
|
|
|
|
self._worker.core_worker.experimental_channel_put_serialized(
|
|
serialized_value,
|
|
self._writer_ref,
|
|
self._num_local_readers,
|
|
timeout_ms,
|
|
)
|
|
|
|
def read(self, timeout: Optional[float] = None) -> Any:
|
|
assert (
|
|
timeout is None or timeout >= 0 or timeout == -1
|
|
), "Timeout must be non-negative or -1."
|
|
self.ensure_registered_as_reader()
|
|
|
|
start_time = time.monotonic()
|
|
ret = self._worker.get_objects(
|
|
[self._local_reader_ref], timeout=timeout, return_exceptions=True
|
|
)[0][0]
|
|
|
|
if isinstance(ret, _ResizeChannel):
|
|
self._node_id_to_reader_ref_info = ret._node_id_to_reader_ref_info
|
|
self._local_reader_ref = self._get_local_reader_ref(
|
|
self._node_id_to_reader_ref_info
|
|
)
|
|
# We need to register the new reader_ref.
|
|
self._reader_registered = False
|
|
self.ensure_registered_as_reader()
|
|
if timeout is not None:
|
|
timeout -= time.monotonic() - start_time
|
|
timeout = max(timeout, 0)
|
|
ret = self._worker.get_objects(
|
|
[self._local_reader_ref], timeout=timeout, return_exceptions=True
|
|
)[0][0]
|
|
|
|
return ret
|
|
|
|
def release_buffer(self, timeout: Optional[float] = None) -> None:
|
|
assert (
|
|
timeout is None or timeout >= 0 or timeout == -1
|
|
), "Timeout must be non-negative or -1."
|
|
self.ensure_registered_as_reader()
|
|
self._worker.get_objects(
|
|
[self._local_reader_ref],
|
|
timeout=timeout,
|
|
return_exceptions=True,
|
|
skip_deserialization=True,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
"""
|
|
Close this channel by setting the error bit on both the writer_ref and the
|
|
reader_ref.
|
|
"""
|
|
self._worker.core_worker.experimental_channel_set_error(self._writer_ref)
|
|
is_local_node_reader = False
|
|
|
|
for node_id in self._node_id_to_readers.keys():
|
|
if self.is_local_node(node_id):
|
|
is_local_node_reader = True
|
|
if is_local_node_reader:
|
|
self.ensure_registered_as_reader()
|
|
|
|
for reader_ref_info in self._node_id_to_reader_ref_info.values():
|
|
self._worker.core_worker.experimental_channel_set_error(
|
|
reader_ref_info.reader_ref
|
|
)
|
|
|
|
|
|
@DeveloperAPI
|
|
class BufferedSharedMemoryChannel(ChannelInterface):
|
|
"""A channel that can be read and written by Ray processes.
|
|
|
|
It creates `num_shm_buffers` number of buffers and allows buffered read and
|
|
write APIs. I.e., read and write APIs are non-blocking as long as it can write to
|
|
next buffer or read from a next buffer. See `read` and `write` APIs for
|
|
more details.
|
|
|
|
Args:
|
|
writer: The actor that may write to the channel. None signifies the driver.
|
|
reader_and_node_list: A list of tuples, where each tuple contains a reader
|
|
actor handle and the node ID where the actor is located. Note that currently
|
|
we only support this for readers on the same node as the writer.
|
|
num_shm_buffers: Number of shared memory buffers to read/write.
|
|
typ: Type information about the values passed through the channel.
|
|
Either an integer representing the max buffer size in bytes
|
|
allowed, or a SharedMemoryType.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
writer: Optional[ray.actor.ActorHandle],
|
|
reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
|
|
num_shm_buffers: int,
|
|
typ: Optional[Union[int, SharedMemoryType]] = None,
|
|
):
|
|
self._num_shm_buffers = num_shm_buffers
|
|
self._buffers = [
|
|
# We use Channel directly as a buffer implementation as
|
|
# channel only allows to have 1 shared memory buffer.
|
|
Channel(writer, reader_and_node_list, typ)
|
|
for _ in range(num_shm_buffers)
|
|
]
|
|
# The next index to write from self._buffers.
|
|
self._next_write_index = 0
|
|
# The next index to read from self._buffers.
|
|
self._next_read_index = 0
|
|
|
|
def ensure_registered_as_writer(self):
|
|
"""
|
|
Check whether the process is a valid writer. This method must be idempotent.
|
|
"""
|
|
for buffer in self._buffers:
|
|
buffer.ensure_registered_as_writer()
|
|
|
|
def ensure_registered_as_reader(self):
|
|
"""
|
|
Check whether the process is a valid reader. This method must be idempotent.
|
|
"""
|
|
for buffer in self._buffers:
|
|
buffer.ensure_registered_as_reader()
|
|
|
|
def write(self, value: Any, timeout: Optional[float] = None) -> None:
|
|
"""Write a value to a channel.
|
|
|
|
If the next buffer is available, it returns immediately. If the next
|
|
buffer is not read by downstream consumers, it blocks until a buffer is
|
|
available to write. If a buffer is not available within timeout, it raises
|
|
RayChannelTimeoutError.
|
|
"""
|
|
self.ensure_registered_as_writer()
|
|
# A single channel is not supposed to read and write at the same time.
|
|
assert self._next_read_index == 0
|
|
self._buffers[self._next_write_index].write(value, timeout)
|
|
self._next_write_index += 1
|
|
self._next_write_index %= self._num_shm_buffers
|
|
|
|
def read(self, timeout: Optional[float] = None) -> Any:
|
|
"""Read a value from a channel.
|
|
|
|
If the next buffer is available, it returns immediately. If the next
|
|
buffer is not written by an upstream producer, it blocks until a buffer is
|
|
available to read. If a buffer is not available within timeout, it raises
|
|
RayChannelTimeoutError.
|
|
"""
|
|
self.ensure_registered_as_reader()
|
|
# A single channel is not supposed to read and write at the same time.
|
|
assert self._next_write_index == 0
|
|
output = self._buffers[self._next_read_index].read(timeout)
|
|
self._next_read_index += 1
|
|
self._next_read_index %= self._num_shm_buffers
|
|
return output
|
|
|
|
def release_buffer(self, timeout: Optional[float] = None):
|
|
"""Release the native buffer of the channel to allow the buffer to be reused for
|
|
future data.
|
|
|
|
If the next buffer is available, it returns immediately. If the next
|
|
buffer is not written by an upstream producer, it blocks until a buffer is
|
|
available to be released. If a buffer is not available within timeout, it raises
|
|
RayChannelTimeoutError.
|
|
"""
|
|
# A single channel is not supposed to read and write at the same time.
|
|
assert self._next_write_index == 0
|
|
self._buffers[self._next_read_index].release_buffer(timeout)
|
|
self._next_read_index += 1
|
|
self._next_read_index %= self._num_shm_buffers
|
|
|
|
def close(self) -> None:
|
|
for buffer in self._buffers:
|
|
buffer.close()
|
|
|
|
@property
|
|
def next_write_index(self):
|
|
# Testing only
|
|
return self._next_write_index
|
|
|
|
@property
|
|
def next_read_index(self):
|
|
# Testing only
|
|
return self._next_read_index
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class CompositeChannel(ChannelInterface):
|
|
"""Routes data to different readers via per-locality channels.
|
|
|
|
For example, if the reader is in the same worker process as the writer,
|
|
the data can be sent via IntraProcessChannel. If the reader is in a different
|
|
worker process, the data can be sent via shared memory channel.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
writer: Optional[ray.actor.ActorHandle],
|
|
reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
|
|
num_shm_buffers: int,
|
|
driver_actor_id: Optional[str] = None,
|
|
_channel_dict: Optional[Dict[ray.ActorID, ChannelInterface]] = None,
|
|
_channels: Optional[Set[ChannelInterface]] = None,
|
|
_writer_registered: bool = False,
|
|
_reader_registered: bool = False,
|
|
):
|
|
"""Initialize a ``CompositeChannel``.
|
|
|
|
Args:
|
|
writer: The actor that may write to the channel. None signifies the driver.
|
|
reader_and_node_list: A list of tuples, where each tuple contains a reader
|
|
actor handle and the node ID where the actor is located.
|
|
num_shm_buffers: The number of shared memory buffers per channel.
|
|
Note: In the case of multiple nodes, we only support 1 shared
|
|
memory buffer.
|
|
driver_actor_id: If this channel is read by a driver and that driver is an
|
|
actual actor, this will be the actor ID of that driver actor.
|
|
_channel_dict: Internal. Pre-populated mapping from actor id to
|
|
the underlying channel. When provided, channels are not
|
|
re-created (used during deserialization).
|
|
_channels: Internal. Deduplicated set of channels backing
|
|
``_channel_dict``. When provided, channels are not re-created.
|
|
_writer_registered: Internal. Whether the writer side has already
|
|
been registered with the core worker.
|
|
_reader_registered: Internal. Whether the reader side has already
|
|
been registered with the core worker.
|
|
"""
|
|
self._writer = writer
|
|
self._reader_and_node_list = reader_and_node_list
|
|
self._num_shm_buffers = num_shm_buffers
|
|
self._driver_actor_id = driver_actor_id
|
|
self._writer_registered = _writer_registered
|
|
self._reader_registered = _reader_registered
|
|
# A dictionary that maps the actor ID to the channel object.
|
|
self._channel_dict = _channel_dict or {}
|
|
# The set of channels is a deduplicated version of the _channel_dict values.
|
|
self._channels = _channels or set()
|
|
if self._channels:
|
|
# This CompositeChannel object is created by deserialization.
|
|
# We don't need to create channels again.
|
|
return
|
|
|
|
(
|
|
remote_reader_and_node_list,
|
|
local_reader_and_node_list,
|
|
) = utils.split_readers_by_locality(self._writer, self._reader_and_node_list)
|
|
# There are some local readers which are the same worker process as the writer.
|
|
# Create a local channel for the writer and the local readers.
|
|
num_local_readers = len(local_reader_and_node_list)
|
|
if num_local_readers > 0:
|
|
# Use num_readers = 1 when creating the local channel,
|
|
# because we have channel cache to support reading
|
|
# from the same channel multiple times.
|
|
local_channel = IntraProcessChannel(num_readers=1)
|
|
self._channels.add(local_channel)
|
|
actor_id = self._get_actor_id(self._writer)
|
|
self._channel_dict[actor_id] = local_channel
|
|
# There are some remote readers which are not the same Ray actor as the writer.
|
|
# We create a BufferedSharedMemoryChannel for readers on the same node, and
|
|
# a single Channel for readers on different nodes due to
|
|
# https://github.com/ray-project/ray/issues/49044
|
|
(
|
|
readers_same_node,
|
|
readers_different_node,
|
|
) = utils.split_actors_by_node_locality(
|
|
utils.get_actor_node(self._writer), remote_reader_and_node_list
|
|
)
|
|
|
|
if len(readers_same_node) != 0:
|
|
remote_channel = BufferedSharedMemoryChannel(
|
|
self._writer, readers_same_node, num_shm_buffers
|
|
)
|
|
self._channels.add(remote_channel)
|
|
for reader, _ in readers_same_node:
|
|
actor_id = self._get_actor_id(reader)
|
|
self._channel_dict[actor_id] = remote_channel
|
|
|
|
if len(readers_different_node) != 0:
|
|
remote_channel = Channel(self._writer, readers_different_node)
|
|
self._channels.add(remote_channel)
|
|
for reader, _ in readers_different_node:
|
|
actor_id = self._get_actor_id(reader)
|
|
self._channel_dict[actor_id] = remote_channel
|
|
|
|
def _get_actor_id(self, reader: ray.actor.ActorHandle) -> str:
|
|
return reader._actor_id.hex()
|
|
|
|
def ensure_registered_as_writer(self) -> None:
|
|
if self._writer_registered:
|
|
return
|
|
for channel in self._channels:
|
|
channel.ensure_registered_as_writer()
|
|
self._writer_registered = True
|
|
|
|
def ensure_registered_as_reader(self) -> None:
|
|
if self._reader_registered:
|
|
return
|
|
for channel in self._channels:
|
|
channel.ensure_registered_as_reader()
|
|
self._reader_registered = True
|
|
|
|
def __reduce__(self):
|
|
return CompositeChannel, (
|
|
self._writer,
|
|
self._reader_and_node_list,
|
|
self._num_shm_buffers,
|
|
self._driver_actor_id,
|
|
self._channel_dict,
|
|
self._channels,
|
|
self._writer_registered,
|
|
self._reader_registered,
|
|
)
|
|
|
|
def __str__(self) -> str:
|
|
return (
|
|
"CompositeChannel(_channels="
|
|
f"{[str(channel) for channel in self._channels]})"
|
|
)
|
|
|
|
def write(self, value: Any, timeout: Optional[float] = None) -> None:
|
|
self.ensure_registered_as_writer()
|
|
for channel in self._channels:
|
|
channel.write(value, timeout)
|
|
|
|
def read(self, timeout: Optional[float] = None) -> Any:
|
|
self.ensure_registered_as_reader()
|
|
return self._channel_dict[self._resolve_actor_id()].read(timeout)
|
|
|
|
def release_buffer(self, timeout: Optional[float] = None):
|
|
self.ensure_registered_as_reader()
|
|
self._channel_dict[self._resolve_actor_id()].release_buffer(timeout)
|
|
|
|
def _resolve_actor_id(self) -> str:
|
|
actor_id = ray.get_runtime_context().get_actor_id()
|
|
# If actor_id is None, read was called by the driver
|
|
# If the driver is an actor, driver_actor_id will be set to that actor id
|
|
if actor_id is None or actor_id == self._driver_actor_id:
|
|
# Use the actor ID of the DAGDriverProxyActor.
|
|
# The proxy actor is always the first actor in the reader_and_node_list.
|
|
assert len(self._reader_and_node_list) >= 1
|
|
driver_proxy_actor = self._reader_and_node_list[0][0]
|
|
actor_id = self._get_actor_id(driver_proxy_actor)
|
|
return actor_id
|
|
|
|
def close(self) -> None:
|
|
for channel in self._channels:
|
|
channel.close()
|