wehub-resource-sync
964 lines
40 KiB
Python
964 lines
40 KiB
Python
import logging
|
|
import threading
|
|
import time
|
|
import warnings
|
|
import weakref
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from queue import Queue
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Dict,
|
|
List,
|
|
NamedTuple,
|
|
Optional,
|
|
Set,
|
|
Tuple,
|
|
Union,
|
|
)
|
|
|
|
import ray
|
|
from ray._private import ray_constants
|
|
from ray._raylet import ObjectRef
|
|
from ray.experimental.rdt.tensor_transport_manager import FetchRequest
|
|
from ray.util.annotations import PublicAPI
|
|
|
|
|
|
@dataclass
|
|
class ObjectStoreFetchRequest(FetchRequest):
|
|
"""Pending fetch via the object store. Holds the remote ObjectRef to ray.get on.
|
|
|
|
Args:
|
|
obj_id: The RDT object ID being fetched.
|
|
object_ref: The ObjectRef returned by the __ray_fetch_rdt_object__ remote call.
|
|
tensors: Unused. Tensors are returned directly by ray.get.
|
|
"""
|
|
|
|
object_ref: Optional[ObjectRef] = None
|
|
tensors: Optional[List[Any]] = None
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
from ray.experimental.rdt.rdt_store import (
|
|
RDTStore,
|
|
)
|
|
from ray.experimental.rdt.tensor_transport_manager import (
|
|
CommunicatorMetadata,
|
|
FetchRequest,
|
|
TensorTransportMetadata,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# RDTMeta is a named tuple containing the source actor, tensor transport
|
|
# backend, tensor metadata, and other information that needs to be recorded.
|
|
# - The tensor transport backend is the backend used to transport the tensors.
|
|
# - The tensor metadata is a list of tuples, each containing the shape and dtype
|
|
# of a tensor in the RDT store.
|
|
class RDTMeta(NamedTuple):
|
|
src_actor: "ray.actor.ActorHandle"
|
|
tensor_transport_backend: str
|
|
# This is set when the actual object is created and the metadata makes it back to the owner.
|
|
# For ray.put the owner is the creator so it's immediately set.
|
|
tensor_transport_meta: Optional["TensorTransportMetadata"]
|
|
# sent_dest_actors tracks the set of actor IDs that this object has been sent to.
|
|
# Note that since the set is mutable, it shouldn't be accessed without a lock.
|
|
sent_dest_actors: Set[str]
|
|
# sent_to_src_actor_and_others_warned indicates whether the object has already triggered a warning about being sent back to the source actor and other actors simultaneously.
|
|
sent_to_src_actor_and_others_warned: bool
|
|
# If the user set buffers for the object, the object will be fetched directly into the buffers on a ray.get
|
|
target_buffers: Optional[List[weakref.ReferenceType[Any]]]
|
|
|
|
|
|
# This is used to periodically check in on the RDT transfer through the refs from
|
|
# __ray_send__ and __ray_recv__ and abort operations in case of failures / timeouts.
|
|
class TransferMetadata(NamedTuple):
|
|
src_actor: "ray.actor.ActorHandle"
|
|
dst_actor: "ray.actor.ActorHandle"
|
|
send_ref: Optional[ObjectRef]
|
|
recv_ref: ObjectRef
|
|
communicator_meta: "CommunicatorMetadata"
|
|
backend: str
|
|
obj_id: str
|
|
timeout: float
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
def wait_tensor_freed(tensor: Any, timeout: Optional[float] = None):
|
|
"""
|
|
Wait for the tensor to be freed.
|
|
|
|
This function is useful for cases where an actor keeps a reference to a
|
|
tensor after returning the tensor from a task annotated with
|
|
`@ray.method(tensor_transport=...)`. In this case, Ray will store a
|
|
*reference* to the tensor, so any in-place modifications made by the actor
|
|
that returned the tensor could be seen by other actors. See
|
|
:ref:`Ray Direct Transport (RDT) <direct-transport>` for more details.
|
|
|
|
Call this function for RDT objects to ensure that all corresponding
|
|
`ray.ObjectRefs` have gone out of scope and therefore the tensor is safe to
|
|
write to again.
|
|
|
|
Args:
|
|
tensor: The tensor to wait to be freed. This should be a tensor that was
|
|
previously returned by a task annotated with
|
|
`@ray.method(tensor_transport=...)` or stored via
|
|
`ray.put(_tensor_transport="...")`.
|
|
timeout: The timeout in seconds to wait for all references to the tensor
|
|
to go out of scope. Set to None to wait indefinitely. Note that if
|
|
None is used, this function could hang if the `ray.ObjectRefs` that
|
|
refer to this tensor never go out of scope.
|
|
"""
|
|
rdt_manager = ray.worker.global_worker.rdt_manager
|
|
rdt_manager.rdt_store.wait_tensor_freed(tensor, timeout)
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
def set_target_for_ref(ref: ObjectRef, target: List[Any]):
|
|
"""
|
|
Set target buffers for an RDT ObjectRef to fetch tensors into when `ray.get` is called.
|
|
|
|
This is only supported by some transports (e.g., NIXL). If the transport
|
|
does not support this feature, an exception will be raised during ray.get.
|
|
|
|
Before receiving, Ray validates that the provided target buffers match the metadata
|
|
of the tensors in the object (e.g., shape, dtype, device). If validation fails,
|
|
a `ValueError` is raised. We recommend sending over lists of tensors and passing a list
|
|
of the same length here because the serialization order from the sender-side must match
|
|
the order of the target tensors here.
|
|
|
|
Args:
|
|
ref: The ObjectRef to set the target buffers for. The ref must be for an RDT object.
|
|
target: A list of tensors to be used as the target buffers to receive into.
|
|
"""
|
|
rdt_manager = ray.worker.global_worker.rdt_manager
|
|
rdt_manager.set_target_buffers_for_ref(ref, target)
|
|
|
|
|
|
class RDTManager:
|
|
def __init__(self):
|
|
# This lock protects _managed_rdt_metadata, _queued_transfers, and _queued_frees since
|
|
# they can be accessed from the user's python thread or the CoreWorker's main io service thread.
|
|
self._lock = threading.Lock()
|
|
|
|
# A dictionary that maps from owned object's ID to RDTMeta.
|
|
# This dictionary is hosted on the "driver" process of the actors that
|
|
# store and send/receive RDT objects.
|
|
self._managed_rdt_metadata: Dict[str, RDTMeta] = {}
|
|
# Condition variable to wait for the tensor transport meta to be set.
|
|
self._tensor_transport_meta_cv = threading.Condition(self._lock)
|
|
|
|
# A dictionary that maps from an object id to a list of actors
|
|
# that are queued to receive the object.
|
|
self._queued_transfers: Dict[str, List["ray.actor.ActorHandle"]] = defaultdict(
|
|
list
|
|
)
|
|
# A set of object ids that are queued to be freed. This is used when the object is freed
|
|
# before the owner knows it's created (the tensor transport metadata is not available yet).
|
|
self._queued_frees: Set[str] = set()
|
|
|
|
# This lock makes sure the _rdt_store and _monitor_failures_thread are only created once.
|
|
self._init_lock = threading.Lock()
|
|
|
|
# Per-actor local storage for RDT objects. We create the RDT store
|
|
# lazily, if a user specifies a non-default tensor_transport, to avoid
|
|
# circular import and because it imports third-party dependencies like
|
|
# PyTorch.
|
|
self._rdt_store: Optional["RDTStore"] = None
|
|
|
|
# Thread safe queue of transport refs that the monitor thread needs to start monitoring
|
|
self._unmonitored_transfers: Queue[TransferMetadata] = Queue()
|
|
# Background thread to poll on the transfer operation.
|
|
self._monitor_failures_thread = None
|
|
# Event to signal the monitor_failures thread to shutdown
|
|
self._monitor_failures_shutdown_event = threading.Event()
|
|
|
|
# If the actor isn't in the dict, the task to launch the custom transport registration task hasn't been submitted yet.
|
|
# If the value is an object ref, we have to wait for the registration task to complete.
|
|
# If the value is True, the actor has registered any custom transports.
|
|
# The value should never be False.
|
|
# TODO: This is a short-term solution. In the future, we'll do registration with actor initialization
|
|
# to make actor restarts and submitting from another worker work.
|
|
self.actor_id_to_transports_registered: Dict[str, Union[ObjectRef, bool]] = {}
|
|
|
|
def register_custom_transports_on_actor(self, actor: "ray.actor.ActorHandle"):
|
|
from ray.experimental.rdt.util import (
|
|
register_custom_tensor_transports_on_actor,
|
|
)
|
|
|
|
ref = register_custom_tensor_transports_on_actor(actor)
|
|
# ref is None if there are no custom transports registered.
|
|
self.actor_id_to_transports_registered[actor._actor_id] = (
|
|
True if ref is None else ref
|
|
)
|
|
|
|
def wait_until_custom_transports_registered(self, actor: "ray.actor.ActorHandle"):
|
|
actor_id = actor._actor_id
|
|
if actor_id not in self.actor_id_to_transports_registered:
|
|
self.register_custom_transports_on_actor(actor)
|
|
|
|
if self.actor_id_to_transports_registered[actor_id] is not True:
|
|
ray.get(self.actor_id_to_transports_registered[actor_id])
|
|
self.actor_id_to_transports_registered[actor_id] = True
|
|
|
|
@property
|
|
def rdt_store(self) -> "ray.experimental.RDTStore":
|
|
with self._init_lock:
|
|
if self._rdt_store is None:
|
|
from ray.experimental.rdt.rdt_store import (
|
|
RDTStore,
|
|
)
|
|
|
|
self._rdt_store = RDTStore()
|
|
return self._rdt_store
|
|
|
|
def shutdown(self):
|
|
"""
|
|
Interrupt and join the _monitor_failures_thread
|
|
"""
|
|
with self._init_lock:
|
|
if self._monitor_failures_thread:
|
|
self._monitor_failures_shutdown_event.set()
|
|
self._monitor_failures_thread.join()
|
|
self._monitor_failures_shutdown_event.clear()
|
|
self._monitor_failures_thread = None
|
|
|
|
def start_monitor_thread_if_needed(self):
|
|
with self._init_lock:
|
|
# To make sure _monitor_failures_thread is started only once
|
|
if self._monitor_failures_thread is None:
|
|
self._monitor_failures_thread = threading.Thread(
|
|
target=self._monitor_failures, daemon=True
|
|
)
|
|
self._monitor_failures_thread.start()
|
|
|
|
def is_managed_object(self, obj_id: str) -> bool:
|
|
"""
|
|
Check if the RDT object is owned or borrowed by this process.
|
|
"""
|
|
with self._lock:
|
|
return obj_id in self._managed_rdt_metadata
|
|
|
|
def set_rdt_metadata(self, obj_id: str, rdt_meta: RDTMeta):
|
|
with self._lock:
|
|
self._managed_rdt_metadata[obj_id] = rdt_meta
|
|
|
|
def get_rdt_metadata(self, obj_id: str) -> Optional[RDTMeta]:
|
|
with self._lock:
|
|
return self._managed_rdt_metadata.get(obj_id, None)
|
|
|
|
def wait_for_tensor_transport_metadata(
|
|
self, obj_id: str, timeout: float
|
|
) -> Optional["TensorTransportMetadata"]:
|
|
with self._tensor_transport_meta_cv:
|
|
if self._tensor_transport_meta_cv.wait_for(
|
|
lambda: self._managed_rdt_metadata[obj_id].tensor_transport_meta
|
|
is not None,
|
|
timeout=timeout,
|
|
):
|
|
return self._managed_rdt_metadata[obj_id].tensor_transport_meta
|
|
else:
|
|
return None
|
|
|
|
def _monitor_failures(self):
|
|
"""
|
|
Monitor the refs from send and recv tasks and abort the transfers
|
|
if they error out or timeout to prevent hanging.
|
|
"""
|
|
not_done = []
|
|
done = []
|
|
ref_info_map = {}
|
|
while not self._monitor_failures_shutdown_event.is_set():
|
|
while not self._unmonitored_transfers.empty():
|
|
ref_info = self._unmonitored_transfers.get()
|
|
if ref_info.send_ref:
|
|
not_done.append(ref_info.send_ref)
|
|
ref_info_map[ref_info.send_ref.hex()] = ref_info
|
|
not_done.append(ref_info.recv_ref)
|
|
ref_info_map[ref_info.recv_ref.hex()] = ref_info
|
|
if len(not_done) > 0:
|
|
done, not_done = ray.wait(not_done, num_returns=1, timeout=1)
|
|
if len(done) > 0:
|
|
try:
|
|
ray.get(done[0])
|
|
ref_info_map.pop(done[0].hex(), None)
|
|
except Exception as e:
|
|
self._abort_transport(done[0], ref_info_map, e)
|
|
|
|
while len(not_done) > 0:
|
|
if not_done[0].hex() not in ref_info_map:
|
|
# The associated transfer was already aborted.
|
|
not_done.pop(0)
|
|
elif ref_info_map[not_done[0].hex()].timeout < time.time():
|
|
self._abort_transport(
|
|
not_done[0],
|
|
ref_info_map,
|
|
TimeoutError(
|
|
f"RDT transfer failed after {ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS}s. "
|
|
"You can increase the timeout by setting RAY_rdt_fetch_fail_timeout_milliseconds"
|
|
),
|
|
)
|
|
else:
|
|
# wait returns lists in the same order they were passed in, so if
|
|
# the timeout of first hasn't been reached, neither have the others.
|
|
break
|
|
if len(not_done) == 0:
|
|
# If we emptied out _unmonitored_transfers on this iteration, wait for a bit.
|
|
self._monitor_failures_shutdown_event.wait(1)
|
|
|
|
def _abort_transport(
|
|
self,
|
|
failed_ref: ObjectRef,
|
|
ref_info_map: Dict[str, TransferMetadata],
|
|
exception: Exception,
|
|
):
|
|
"""
|
|
Cleans up the ref_info_map, kill the src and dst actors, and destroy the
|
|
collective group if necessary.
|
|
"""
|
|
from ray.experimental.collective import destroy_collective_group
|
|
from ray.experimental.rdt.collective_tensor_transport import (
|
|
CollectiveCommunicatorMetadata,
|
|
)
|
|
from ray.experimental.rdt.rdt_store import (
|
|
__ray_abort_transport__,
|
|
)
|
|
from ray.experimental.rdt.util import (
|
|
get_tensor_transport_manager,
|
|
)
|
|
|
|
ref_info = ref_info_map.pop(failed_ref.hex(), None)
|
|
if ref_info is None:
|
|
return
|
|
|
|
if ref_info.send_ref:
|
|
ref_info_map.pop(ref_info.send_ref.hex(), None)
|
|
ref_info_map.pop(ref_info.recv_ref.hex(), None)
|
|
|
|
tensor_transport_manager = get_tensor_transport_manager(ref_info.backend)
|
|
if tensor_transport_manager.can_abort_transport():
|
|
if not tensor_transport_manager.__class__.is_one_sided():
|
|
# This is dead code until we implement a NCCL abort since NIXL
|
|
# is the only abortable transport for now and is one-sided.
|
|
ref_info.src_actor.__ray_call__.options(
|
|
concurrency_group="_ray_system_error"
|
|
).remote(
|
|
__ray_abort_transport__,
|
|
ref_info.obj_id,
|
|
ref_info.communicator_meta,
|
|
ref_info.backend,
|
|
)
|
|
ref_info.dst_actor.__ray_call__.options(
|
|
concurrency_group="_ray_system_error"
|
|
).remote(
|
|
__ray_abort_transport__,
|
|
ref_info.obj_id,
|
|
ref_info.communicator_meta,
|
|
ref_info.backend,
|
|
)
|
|
logger.info(
|
|
"RDT transfer with src actor %s and dst actor %s failed due to %s.",
|
|
ref_info.src_actor,
|
|
ref_info.dst_actor,
|
|
exception,
|
|
)
|
|
else:
|
|
# TODO(#51276): Kill all actors in the collective group when we support more collective operations
|
|
ray.kill(ref_info.src_actor)
|
|
ray.kill(ref_info.dst_actor)
|
|
logger.error(
|
|
"RDT transfer with src actor %s and dst actor %s failed. Killing the actors. "
|
|
"Transfer failed with exception: %s",
|
|
ref_info.src_actor,
|
|
ref_info.dst_actor,
|
|
exception,
|
|
)
|
|
|
|
# isinstance does an implicit cast and makes communicator_name inaccessible
|
|
# so we have to get communicator_name before the cast.
|
|
if isinstance(ref_info.communicator_meta, CollectiveCommunicatorMetadata):
|
|
try:
|
|
collective_group_name = ref_info.communicator_meta.communicator_name
|
|
destroy_collective_group(collective_group_name)
|
|
logger.error(
|
|
"Destroyed collective group %s due to a hanging/failed RDT transfer",
|
|
collective_group_name,
|
|
)
|
|
except ValueError:
|
|
# Collective group was already destroyed
|
|
pass
|
|
|
|
def add_rdt_ref(
|
|
self,
|
|
obj_ref: ObjectRef,
|
|
src_actor: "ray.actor.ActorHandle",
|
|
tensor_transport: str,
|
|
tensor_transport_meta: Optional["TensorTransportMetadata"] = None,
|
|
):
|
|
"""Add an RDT object reference to the RDT manager. This should be
|
|
called whenever the current process calls a task that is annotated with
|
|
`@ray.method(tensor_transport=...)`.
|
|
|
|
Args:
|
|
obj_ref: The ObjectRef of the task output.
|
|
src_actor: The actor that executes the task and that creates the RDT object.
|
|
tensor_transport: The tensor transport protocol to use for the RDT object.
|
|
tensor_transport_meta: The tensor transport metadata that is pre-computed.
|
|
This is known at ref creation time if the object is created through ray.put.
|
|
"""
|
|
self.set_rdt_metadata(
|
|
obj_ref.hex(),
|
|
RDTMeta(
|
|
src_actor=src_actor,
|
|
tensor_transport_backend=tensor_transport,
|
|
tensor_transport_meta=tensor_transport_meta, # None if not from ray.put
|
|
sent_dest_actors=set(),
|
|
sent_to_src_actor_and_others_warned=False,
|
|
target_buffers=None,
|
|
),
|
|
)
|
|
|
|
def set_tensor_transport_metadata_and_trigger_queued_operations(
|
|
self, obj_id: str, tensor_transport_meta: "TensorTransportMetadata"
|
|
):
|
|
"""
|
|
Sets the tensor transport metadata for an object and triggers any queued
|
|
up transfers or frees for that object.
|
|
"""
|
|
dst_actors = None
|
|
free_object = False
|
|
with self._tensor_transport_meta_cv:
|
|
self._managed_rdt_metadata[obj_id] = self._managed_rdt_metadata[
|
|
obj_id
|
|
]._replace(tensor_transport_meta=tensor_transport_meta)
|
|
dst_actors = self._queued_transfers.pop(obj_id, None)
|
|
free_object = obj_id in self._queued_frees
|
|
if free_object:
|
|
self._queued_frees.remove(obj_id)
|
|
# There shouldn't be any transfers queued if the free was queued,
|
|
# since we clear the queued transfers when queueing the free.
|
|
assert dst_actors is None
|
|
self._tensor_transport_meta_cv.notify_all()
|
|
|
|
if free_object:
|
|
self.free_object_primary_copy(obj_id)
|
|
if dst_actors:
|
|
for dst_actor in dst_actors:
|
|
# Trigger the transfer now that the metadata is available.
|
|
self.trigger_out_of_band_tensor_transfer(dst_actor, obj_id)
|
|
|
|
def set_target_buffers_for_ref(self, ref: ObjectRef, target_buffers: List[Any]):
|
|
with self._lock:
|
|
if ref.hex() not in self._managed_rdt_metadata:
|
|
raise ValueError(f"Ref {ref} is not an RDT object.")
|
|
|
|
self._managed_rdt_metadata[ref.hex()] = self._managed_rdt_metadata[
|
|
ref.hex()
|
|
]._replace(
|
|
target_buffers=[
|
|
weakref.ref(target_buffer) for target_buffer in target_buffers
|
|
]
|
|
)
|
|
|
|
def _trigger_fetch(
|
|
self,
|
|
obj_id: str,
|
|
use_object_store: bool,
|
|
) -> FetchRequest:
|
|
"""
|
|
Start fetching an RDT object.
|
|
|
|
If the specified transport supports async fetches, this will trigger the
|
|
fetch without blocking. Note that this always triggers a fetch, even if
|
|
the object is already in the store.
|
|
|
|
Args:
|
|
obj_id: The object ID of the RDT object.
|
|
use_object_store: Whether to fetch through the object store or through
|
|
the designated one-sided tensor transport.
|
|
|
|
Returns:
|
|
A FetchRequest. Wait on the FetchRequest to get the tensors.
|
|
"""
|
|
from ray.experimental.rdt.rdt_store import (
|
|
__ray_fetch_rdt_object__,
|
|
)
|
|
from ray.experimental.rdt.util import (
|
|
get_tensor_transport_manager,
|
|
is_one_sided_transport,
|
|
)
|
|
|
|
rdt_meta = self.get_rdt_metadata(obj_id)
|
|
assert rdt_meta is not None
|
|
|
|
if use_object_store:
|
|
if rdt_meta.target_buffers:
|
|
logger.warning(
|
|
"Target buffers are not supported for use_object_store=True. Ignoring the target buffers."
|
|
)
|
|
|
|
src_actor = rdt_meta.src_actor
|
|
object_ref = src_actor.__ray_call__.options(
|
|
concurrency_group="_ray_system"
|
|
).remote(__ray_fetch_rdt_object__, obj_id)
|
|
return ObjectStoreFetchRequest(
|
|
obj_id=obj_id, object_ref=object_ref, tensors=[]
|
|
)
|
|
else:
|
|
tensor_transport = rdt_meta.tensor_transport_backend
|
|
if not is_one_sided_transport(tensor_transport):
|
|
raise ValueError(
|
|
f"ray.get is not allowed on RDT objects using the two-sided transport {tensor_transport}. "
|
|
"Either use a one-sided RDT transport or pass _use_object_store=True to ray.get to fetch the object through the object store instead."
|
|
)
|
|
tensor_transport_manager = get_tensor_transport_manager(tensor_transport)
|
|
communicator_meta = tensor_transport_manager.get_communicator_metadata(
|
|
None, None, tensor_transport
|
|
)
|
|
|
|
tensor_transport_meta = rdt_meta.tensor_transport_meta
|
|
if tensor_transport_meta is None:
|
|
# We can't fetch the object until we know the creator has actually created the object.
|
|
timeout = ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS
|
|
tensor_transport_meta = self.wait_for_tensor_transport_metadata(
|
|
obj_id, timeout
|
|
)
|
|
if tensor_transport_meta is None:
|
|
raise TimeoutError(
|
|
f"Timed out after {timeout}s waiting for object {obj_id} to be created while trying to get the object. "
|
|
"You can increase the timeout by setting RAY_rdt_fetch_fail_timeout_milliseconds."
|
|
)
|
|
|
|
target_buffers = None
|
|
if rdt_meta.target_buffers:
|
|
# Try to get the target buffers from the weak references. If any of the
|
|
# target buffers are not alive, we just won't use the target buffers.
|
|
target_buffers = []
|
|
for target_buffer in rdt_meta.target_buffers:
|
|
buffer = target_buffer()
|
|
if buffer is None:
|
|
target_buffers = None
|
|
break
|
|
else:
|
|
target_buffers.append(buffer)
|
|
|
|
if target_buffers is not None:
|
|
from ray.experimental.rdt.rdt_store import validate_tensor_buffers
|
|
|
|
device = tensor_transport_meta.tensor_device
|
|
tensor_meta = tensor_transport_meta.tensor_meta
|
|
validate_tensor_buffers(target_buffers, tensor_meta, device)
|
|
|
|
return tensor_transport_manager.fetch_multiple_tensors(
|
|
obj_id,
|
|
tensor_transport_meta,
|
|
communicator_meta,
|
|
target_buffers,
|
|
)
|
|
|
|
def _wait_fetch(
|
|
self, obj_id: str, fetch_request: FetchRequest, timeout: float = -1
|
|
) -> List[Any]:
|
|
"""
|
|
Waits for a previously triggered fetch to complete and returns the tensors.
|
|
|
|
Args:
|
|
obj_id: The object ID of the RDT object.
|
|
fetch_request: An ObjectStoreFetchRequest representing an object
|
|
transferred via Ray's object store or a FetchRequest
|
|
representing an object transferred via a tensor transport.
|
|
timeout: Maximum time in seconds to wait. -1 means wait indefinitely.
|
|
0 means return immediately if not ready.
|
|
|
|
Returns:
|
|
The list of tensors fetched.
|
|
"""
|
|
if isinstance(fetch_request, ObjectStoreFetchRequest):
|
|
return ray.get(fetch_request.object_ref, timeout=timeout)
|
|
else:
|
|
from ray.experimental.rdt.util import get_tensor_transport_manager
|
|
|
|
rdt_meta = self.get_rdt_metadata(obj_id)
|
|
tensor_transport_manager = get_tensor_transport_manager(
|
|
rdt_meta.tensor_transport_backend
|
|
)
|
|
return tensor_transport_manager.wait_fetch_complete(
|
|
fetch_request, timeout=timeout
|
|
)
|
|
|
|
def queue_or_trigger_out_of_band_tensor_transfer(
|
|
self, dst_actor: "ray.actor.ActorHandle", task_args: Tuple[Any, ...]
|
|
):
|
|
"""
|
|
Triggers the transfer if the tensor metadata is available for the object. If it's
|
|
not available, the transfer is queued up until the metadata is available.
|
|
"""
|
|
rdt_object_ids: Set[str] = set()
|
|
for arg in task_args:
|
|
# If an ObjectRef is managed, it means the actual value is a list of tensors stored
|
|
# on a remote actor. Therefore, this function will trigger a tensor communication
|
|
# operation between the sender and receiver actors.
|
|
if not isinstance(arg, ObjectRef):
|
|
continue
|
|
obj_id = arg.hex()
|
|
if self.is_managed_object(obj_id):
|
|
rdt_object_ids.add(obj_id)
|
|
if rdt_object_ids:
|
|
self.wait_until_custom_transports_registered(dst_actor)
|
|
for obj_id in rdt_object_ids:
|
|
# Atomically gets the tensor transport metadata for an object and queues up a transfer
|
|
# if the tensor transport metadata is not available.
|
|
with self._lock:
|
|
tensor_transport_meta = self._managed_rdt_metadata[
|
|
obj_id
|
|
].tensor_transport_meta
|
|
if tensor_transport_meta is None:
|
|
self._queued_transfers[obj_id].append(dst_actor)
|
|
if tensor_transport_meta is not None:
|
|
self.trigger_out_of_band_tensor_transfer(dst_actor, obj_id)
|
|
|
|
def trigger_out_of_band_tensor_transfer(
|
|
self, dst_actor: "ray.actor.ActorHandle", obj_id: str
|
|
):
|
|
"""
|
|
Triggers tensor communication operations between actors. When a managed ObjectRef is passed
|
|
to another actor task, CPU data will still be passed through the object store, but the in-actor
|
|
tensors will be passed out-of-band.
|
|
|
|
This function triggers the out-of-band tensor transfer by submitting Ray actor
|
|
tasks `__ray_send__` to the sender actor and `__ray_recv__` to the receiver actor to initiate
|
|
tensor communication using protocols like NCCL or GLOO.
|
|
|
|
Before the receiver actor executes the actor task, the deserializer combines the
|
|
CPU data with the tensors from the sender actor to reconstruct the original task output
|
|
generated by the sender actor.
|
|
|
|
Args:
|
|
dst_actor: The target actor to receive tensors
|
|
obj_id: ID of the object to send to the dst_actor.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
from ray.experimental.rdt.rdt_store import (
|
|
__ray_recv__,
|
|
__ray_send__,
|
|
)
|
|
from ray.experimental.rdt.util import (
|
|
get_tensor_transport_manager,
|
|
)
|
|
|
|
with self._lock:
|
|
# Since sent_dest_actors is mutable, this whole block needs to be protected.
|
|
rdt_meta = self._managed_rdt_metadata[obj_id]
|
|
src_actor = rdt_meta.src_actor
|
|
tensor_transport_meta = rdt_meta.tensor_transport_meta
|
|
|
|
# Update the set of destination actors for this object
|
|
# The set inside NamedTuple is mutable, so we can modify it directly
|
|
rdt_meta.sent_dest_actors.add(dst_actor._actor_id)
|
|
# Check if a warning should be triggered for this object:
|
|
# 1. object has not triggered a warning yet.
|
|
# 2. object is sent back to its source actor.
|
|
# 3. object is also sent to at least one other actor
|
|
if (
|
|
not rdt_meta.sent_to_src_actor_and_others_warned
|
|
and src_actor._actor_id in rdt_meta.sent_dest_actors
|
|
and len(rdt_meta.sent_dest_actors) > 1
|
|
):
|
|
warnings.warn(
|
|
f"RDT ObjectRef({obj_id}) is being passed back to the actor that created it {src_actor}. "
|
|
"Note that RDT objects are mutable. If the tensor is modified, Ray's internal copy will "
|
|
"also be updated, and subsequent passes to other actors will receive the updated version "
|
|
"instead of the original.",
|
|
UserWarning,
|
|
)
|
|
# Mark the object as warned so that we don't warn again for this object.
|
|
self._managed_rdt_metadata[obj_id] = rdt_meta._replace(
|
|
sent_to_src_actor_and_others_warned=True
|
|
)
|
|
|
|
if src_actor._actor_id == dst_actor._actor_id:
|
|
# If the source and destination actors are the same, the tensors can
|
|
# be transferred intra-process, so we skip the out-of-band tensor
|
|
# transfer.
|
|
return
|
|
|
|
tensor_transport_manager = get_tensor_transport_manager(
|
|
rdt_meta.tensor_transport_backend
|
|
)
|
|
communicator_meta = tensor_transport_manager.get_communicator_metadata(
|
|
src_actor,
|
|
dst_actor,
|
|
rdt_meta.tensor_transport_backend,
|
|
)
|
|
|
|
send_ref = None
|
|
if not tensor_transport_manager.__class__.is_one_sided():
|
|
# Send tensors stored in the `src_actor`'s GPU object store to the
|
|
# destination rank `dst_rank`.
|
|
# NOTE: We put this task on the background thread to avoid tasks
|
|
# executing on the main thread blocking the data transfer.
|
|
send_ref = src_actor.__ray_call__.options(
|
|
concurrency_group="_ray_system"
|
|
).remote(
|
|
__ray_send__,
|
|
obj_id,
|
|
tensor_transport_meta,
|
|
communicator_meta,
|
|
rdt_meta.tensor_transport_backend,
|
|
)
|
|
|
|
# Receive tensors from the source rank and store them in the
|
|
# `dst_actor`'s GPU object store.
|
|
# NOTE: Putting this task on the background thread is technically only
|
|
# needed for the sender task, but we put the receiver task on the same
|
|
# background thread to ensure that all communication operations are
|
|
# executed in a global order.
|
|
recv_ref = dst_actor.__ray_call__.options(
|
|
concurrency_group="_ray_system"
|
|
).remote(
|
|
__ray_recv__,
|
|
obj_id,
|
|
tensor_transport_meta,
|
|
communicator_meta,
|
|
rdt_meta.tensor_transport_backend,
|
|
)
|
|
|
|
self._unmonitored_transfers.put(
|
|
TransferMetadata(
|
|
src_actor=src_actor,
|
|
dst_actor=dst_actor,
|
|
send_ref=send_ref,
|
|
recv_ref=recv_ref,
|
|
communicator_meta=communicator_meta,
|
|
backend=rdt_meta.tensor_transport_backend,
|
|
obj_id=obj_id,
|
|
timeout=time.time() + ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS,
|
|
)
|
|
)
|
|
self.start_monitor_thread_if_needed()
|
|
|
|
def get_rdt_objects(
|
|
self,
|
|
object_ids: List[str],
|
|
) -> Dict[str, List[Any]]:
|
|
"""
|
|
Get RDT objects that have already been transferred (e.g. via __ray_recv__).
|
|
|
|
This is used in the task argument deserialization path where the
|
|
out-of-band tensor transfer has already been triggered by the caller.
|
|
It only waits on the local RDT store for the tensors to arrive.
|
|
|
|
Args:
|
|
object_ids: The object IDs of the RDT objects.
|
|
|
|
Returns:
|
|
A dict mapping object ID to the RDT object (list of tensors).
|
|
"""
|
|
rdt_store = self.rdt_store
|
|
result: Dict[str, List[Any]] = {}
|
|
for object_id in object_ids:
|
|
pop_object = not rdt_store.is_primary_copy(object_id)
|
|
if pop_object:
|
|
result[object_id] = rdt_store.wait_and_pop_object(
|
|
object_id, timeout=ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS
|
|
)
|
|
else:
|
|
result[object_id] = rdt_store.wait_and_get_object(
|
|
object_id, timeout=ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS
|
|
)
|
|
return result
|
|
|
|
def fetch_and_get_rdt_objects(
|
|
self,
|
|
object_ids: List[str],
|
|
timeout: Optional[float] = None,
|
|
use_object_store: bool = False,
|
|
) -> Dict[str, List[Any]]:
|
|
"""
|
|
Fetch and get RDT objects for a list of object IDs, pipelining async fetches.
|
|
|
|
This is used in the ray.get codepath where the caller initiates the
|
|
tensor fetch. For one-sided transports (e.g. NIXL), all transfers are
|
|
triggered first before waiting, eliminating serial transfer latency.
|
|
|
|
Args:
|
|
object_ids: The object IDs of the RDT objects.
|
|
timeout: The user-specified timeout from ray.get, or None for no
|
|
user timeout. The actual deadline is the minimum of this and
|
|
RDT_FETCH_FAIL_TIMEOUT_SECONDS.
|
|
use_object_store: Whether to fetch through the object store or through
|
|
the designated tensor transport.
|
|
|
|
Returns:
|
|
A dict mapping object ID to the RDT object (list of tensors).
|
|
|
|
Raises:
|
|
GetTimeoutError: If the user-specified timeout is exceeded.
|
|
ObjectFetchTimedOutError: If RDT_FETCH_FAIL_TIMEOUT_SECONDS is exceeded.
|
|
"""
|
|
from ray.exceptions import GetTimeoutError, ObjectFetchTimedOutError
|
|
|
|
rdt_timeout = ray_constants.RDT_FETCH_FAIL_TIMEOUT_SECONDS
|
|
now = time.time()
|
|
if timeout is not None and timeout >= 0:
|
|
rdt_deadline = now + rdt_timeout
|
|
user_deadline = now + timeout
|
|
if user_deadline < rdt_deadline:
|
|
deadline = user_deadline
|
|
user_timeout_is_smaller = True
|
|
else:
|
|
deadline = rdt_deadline
|
|
user_timeout_is_smaller = False
|
|
else:
|
|
deadline = now + rdt_timeout
|
|
user_timeout_is_smaller = False
|
|
|
|
rdt_store = self.rdt_store
|
|
result: Dict[str, List[Any]] = {}
|
|
|
|
# First, try to get objects that are already available in the store
|
|
# These are primary copies, or secondary copies created via
|
|
# __ray_recv__ that haven't been consumed yet.
|
|
if not use_object_store:
|
|
for object_id in object_ids:
|
|
try:
|
|
result[object_id] = rdt_store.wait_and_get_object(
|
|
object_id, timeout=0
|
|
)
|
|
except TimeoutError:
|
|
pass
|
|
|
|
# For remaining objects, trigger fetches.
|
|
fetch_requests: Dict[str, "FetchRequest"] = {}
|
|
for object_id in object_ids:
|
|
if object_id in result:
|
|
continue
|
|
assert self.is_managed_object(
|
|
object_id
|
|
), f"No metadata found for {object_id}"
|
|
|
|
fetch_requests[object_id] = self._trigger_fetch(object_id, use_object_store)
|
|
|
|
# Wait for all in-flight fetches to complete.
|
|
while fetch_requests:
|
|
object_id, fetch_request = fetch_requests.popitem()
|
|
remaining = deadline - time.time()
|
|
if remaining < 0:
|
|
if user_timeout_is_smaller:
|
|
# User passed a timeout to ray.get that expired.
|
|
raise GetTimeoutError(f"ray.get timed out after {timeout}s.")
|
|
else:
|
|
# Object fetch timeout expired. Throw an error in case we
|
|
# hung.
|
|
raise ObjectFetchTimedOutError(
|
|
object_ref_hex=object_id,
|
|
owner_address="",
|
|
call_site="",
|
|
)
|
|
try:
|
|
result[object_id] = self._wait_fetch(
|
|
object_id, fetch_request, timeout=remaining
|
|
)
|
|
except (TimeoutError, GetTimeoutError):
|
|
if user_timeout_is_smaller:
|
|
raise GetTimeoutError(f"ray.get timed out after {timeout}s.")
|
|
else:
|
|
raise ObjectFetchTimedOutError(
|
|
object_ref_hex=object_id,
|
|
owner_address="",
|
|
call_site="",
|
|
)
|
|
return result
|
|
|
|
def queue_or_free_object_primary_copy(self, object_id: str):
|
|
"""
|
|
Free the RDT object on the primary copy holder and free metadata
|
|
if the tensor metadata is available (the object has been created).
|
|
Otherwise, queue up the free operation until the tensor metadata is available.
|
|
"""
|
|
# NOTE: This may have to change if we support lineage reconstruction for RDT
|
|
# TODO(#57962): Metadata is currently not removed on borrowers that borrow through
|
|
# the NIXL ray.put / ray.get
|
|
with self._lock:
|
|
self._queued_transfers.pop(object_id, None)
|
|
rdt_meta = self._managed_rdt_metadata[object_id]
|
|
tensor_transport_meta = rdt_meta.tensor_transport_meta
|
|
if tensor_transport_meta is None:
|
|
# The object hasn't been created at the time of the free.
|
|
self._queued_frees.add(object_id)
|
|
|
|
if tensor_transport_meta is not None:
|
|
self.free_object_primary_copy(object_id)
|
|
|
|
def free_object_primary_copy(self, object_id: str):
|
|
from ray.experimental.rdt.rdt_store import (
|
|
__ray_free__,
|
|
)
|
|
|
|
with self._lock:
|
|
rdt_meta = self._managed_rdt_metadata.pop(object_id)
|
|
# TODO(#60434): Once the driver has some notion about where rdt args are stored, we can
|
|
# integrate __ray_free__ with the current free objects RPC and avoid having to call
|
|
# an actor task here.
|
|
if not ray.is_initialized():
|
|
return
|
|
src_actor = rdt_meta.src_actor
|
|
tensor_transport_backend = rdt_meta.tensor_transport_backend
|
|
tensor_transport_meta = rdt_meta.tensor_transport_meta
|
|
src_actor.__ray_call__.options(concurrency_group="_ray_system").remote(
|
|
__ray_free__,
|
|
object_id,
|
|
tensor_transport_backend,
|
|
tensor_transport_meta,
|
|
)
|
|
|
|
@staticmethod
|
|
def actor_has_tensor_transport(
|
|
actor: "ray.actor.ActorHandle", tensor_transport: str
|
|
):
|
|
"""
|
|
Check if the actor has a communicator for the given tensor transport backend.
|
|
|
|
Args:
|
|
actor: The actor to check.
|
|
tensor_transport: The tensor transport backend to check.
|
|
|
|
Returns:
|
|
True if the actor has a communicator for the given tensor transport backend, False otherwise.
|
|
"""
|
|
from ray.experimental.rdt.util import (
|
|
get_tensor_transport_manager,
|
|
)
|
|
|
|
tensor_transport_manager = get_tensor_transport_manager(tensor_transport)
|
|
return tensor_transport_manager.actor_has_tensor_transport(actor)
|
|
|
|
def put_object(
|
|
self,
|
|
obj_ref: ObjectRef,
|
|
tensor_transport: str,
|
|
tensors: List[Any],
|
|
):
|
|
"""
|
|
Put the RDT object into the RDT manager.
|
|
|
|
Args:
|
|
obj_ref: The object ref of the RDT object.
|
|
tensor_transport: The tensor transport backend to use.
|
|
tensors: The tensors to put into the RDT manager.
|
|
"""
|
|
src_actor = ray.get_runtime_context().current_actor
|
|
tensor_transport_meta = self.rdt_store.add_object_primary(
|
|
obj_ref.hex(), tensors, tensor_transport
|
|
)
|
|
self.add_rdt_ref(
|
|
obj_ref,
|
|
src_actor,
|
|
tensor_transport,
|
|
tensor_transport_meta=tensor_transport_meta,
|
|
)
|