import threading from collections import defaultdict, deque from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union from ray.experimental.rdt.tensor_transport_manager import ( CommunicatorMetadata, TensorTransportMetadata, ) from ray.experimental.rdt.util import ( device_match_transport, get_tensor_transport_manager, ) if TYPE_CHECKING: import torch def __ray_send__( self, obj_id: str, tensor_transport_meta: TensorTransportMetadata, communicator_meta: CommunicatorMetadata, backend: str, ): """Helper function that runs on the src actor to send tensors to the dst actor.""" from ray._private.worker import global_worker rdt_store = global_worker.rdt_manager._rdt_store assert rdt_store.has_object(obj_id), f"obj_id={obj_id} not found in RDT store" tensors = rdt_store.get_object(obj_id) tensor_transport_manager = get_tensor_transport_manager(backend) tensor_transport_manager.send_multiple_tensors( tensors, tensor_transport_meta, communicator_meta, ) def validate_tensor_buffers( tensor_buffers: List["torch.Tensor"], tensor_meta: List[Tuple["torch.Size", "torch.dtype"]], device: str, ): if len(tensor_buffers) != len(tensor_meta): raise ValueError( f"Length of tensor_buffers ({len(tensor_buffers)}) does not match length from object metadata ({len(tensor_meta)})." ) def tensor_buffer_mismatch_msg(prop, idx, actual, expected): return f"{prop} of tensor_buffer at index {idx} ({actual}) does not match {prop.lower()} from object metadata ({expected})." for idx, single_buffer in enumerate(tensor_buffers): shape, dtype = tensor_meta[idx] if single_buffer.shape != shape: raise ValueError( tensor_buffer_mismatch_msg("Shape", idx, single_buffer.shape, shape) ) if single_buffer.dtype != dtype: raise ValueError( tensor_buffer_mismatch_msg("Dtype", idx, single_buffer.dtype, dtype) ) if single_buffer.device.type != device: raise ValueError( tensor_buffer_mismatch_msg( "Device", idx, single_buffer.device.type, device ) ) if not single_buffer.is_contiguous(): raise ValueError(f"Tensor buffer at index {idx} is not contiguous.") def __ray_recv__( self, obj_id: str, tensor_transport_meta: TensorTransportMetadata, communicator_meta: CommunicatorMetadata, backend: str, target_buffers: Optional[List[Any]] = None, ): """Helper function that runs on the dst actor to receive tensors from the src actor.""" from ray._private.worker import global_worker rdt_store = global_worker.rdt_manager.rdt_store try: tensor_transport_manager = get_tensor_transport_manager(backend) if target_buffers: # Currently only torch tensors are supported as target buffers. We could make this # more generic in the future by adding a pluggable buffer validation function. validate_tensor_buffers( target_buffers, tensor_transport_meta.tensor_meta, tensor_transport_meta.tensor_device, ) tensors = tensor_transport_manager.recv_multiple_tensors( obj_id, tensor_transport_meta, communicator_meta, target_buffers, ) assert len(tensors) == len(tensor_transport_meta.tensor_meta) rdt_store.add_object(obj_id, tensors) except Exception as e: # Store the error as an RDT object if the recv fails, so waiters will raise the error. rdt_store.add_object(obj_id, e) def __ray_abort_transport__( self, obj_id: str, communicator_meta: CommunicatorMetadata, backend: str ): """Helper function that can run on an actor doing a send or recv to abort the transport.""" tensor_transport_manager = get_tensor_transport_manager(backend) tensor_transport_manager.abort_transport(obj_id, communicator_meta) def __ray_free__( self, obj_id: str, tensor_transport_backend: str, tensor_transport_meta: TensorTransportMetadata, ): try: from ray._private.worker import global_worker tensor_transport_manager = get_tensor_transport_manager( tensor_transport_backend ) rdt_manager = global_worker.rdt_manager rdt_store = rdt_manager.rdt_store if not rdt_store.has_object(obj_id): return tensors = rdt_store.get_object(obj_id) tensor_transport_manager.garbage_collect(obj_id, tensor_transport_meta, tensors) rdt_store.pop_object(obj_id) except AssertionError: # This could fail if this is a retry and it's already been freed. pass def __ray_fetch_rdt_object__(self, obj_id: str): """Helper function that runs on the src actor to fetch tensors from the RDT store via the object store.""" from ray._private.worker import global_worker rdt_store = global_worker.rdt_manager.rdt_store rdt_object = rdt_store.wait_and_get_object(obj_id) return rdt_object @dataclass class _RDTObject: # A list of tensors representing the RDT object. data: List[Any] # Whether the RDT object is the primary copy. is_primary: bool # If a recv failed, we store the error here. error: Optional[Exception] = None class RDTStore: """ This class is thread-safe. The GPU object store is meant to be read and written by the following threads: 1. The main thread, which is executing user code. This thread may get, put, and pop objects. 2. The background _ray_system thread, which executes data transfers. This thread may get and put objects. 3. The background CoreWorker server thread, which executes garbage collection callbacks that pop objects that are no longer in use. """ def __init__(self): # A dictionary that maps from an object ID to a queue of tensor lists. # # Note: Currently, `_rdt_store` is only supported for Ray Actors. self._rdt_store: Dict[str, deque[_RDTObject]] = defaultdict(deque) # Mapping from tensor data pointer to the IDs of objects that contain it. self._tensor_to_object_ids: Dict[int, Set[str]] = defaultdict[int, Set[str]]( set ) # Synchronization for the RDT store. self._lock = threading.RLock() # Signal when an object becomes present in the object store. self._object_present_cv = threading.Condition(self._lock) # Signal when an object is freed from the object store. self._object_freed_cv = threading.Condition(self._lock) def has_object(self, obj_id: str) -> bool: with self._lock: existed = obj_id in self._rdt_store if existed: return len(self._rdt_store[obj_id]) > 0 return existed def has_tensor(self, tensor: Any) -> bool: # Method only used for testing. with self._lock: return id(tensor) in self._tensor_to_object_ids def get_object(self, obj_id: str) -> Optional[List[Any]]: with self._lock: if self._rdt_store[obj_id][0].error: raise self._rdt_store[obj_id][0].error return self._rdt_store[obj_id][0].data def add_object( self, obj_id: str, rdt_object: Union[List[Any], Exception], is_primary: bool = False, ): """ Add an RDT object to the RDT store. Args: obj_id: The object ID of the RDT object. rdt_object: A list of tensors representing the RDT object. is_primary: Whether the RDT object is the primary copy. """ with self._object_present_cv: if isinstance(rdt_object, Exception): self._rdt_store[obj_id].append( _RDTObject([], is_primary, error=rdt_object) ) else: for tensor in rdt_object: self._tensor_to_object_ids[id(tensor)].add(obj_id) # Append to the queue instead of overwriting self._rdt_store[obj_id].append( _RDTObject( rdt_object, is_primary, ) ) self._object_present_cv.notify_all() def add_object_primary( self, obj_id: str, tensors: List[Any], tensor_transport: str ) -> TensorTransportMetadata: with self._object_present_cv: # A primary entry may already exist from a prior attempt of the # same task (e.g., a task that succeeded and populated the RDT # store but whose reply was lost, then got retried). Keep the # existing primary — do not re-store — and return metadata # derived from it so the metadata matches what `__ray_send__` # will actually transmit. queue = self._rdt_store.get(obj_id) if queue: tensors_to_describe = queue[0].data else: self.add_object(obj_id, tensors, is_primary=True) tensors_to_describe = tensors tensor_transport_manager = get_tensor_transport_manager(tensor_transport) tensor_transport_meta = ( tensor_transport_manager.extract_tensor_transport_metadata( obj_id, tensors_to_describe ) ) if tensor_transport_meta.tensor_meta and not device_match_transport( tensor_transport_meta.tensor_device, tensor_transport ): raise ValueError( f"Tensor transport backend {tensor_transport} does not support " f"tensor transfer on device {tensor_transport_meta.tensor_device}." ) return tensor_transport_meta def is_primary_copy(self, obj_id: str) -> bool: with self._lock: return self.has_object(obj_id) and self._rdt_store[obj_id][0].is_primary def wait_and_get_object( self, obj_id: str, timeout: Optional[float] = None ) -> List[Any]: """Atomically waits for the RDT object to be present in the RDT store, then gets it. If the object is not present after the optional timeout, raise a TimeoutError. Args: obj_id: The object ID to wait for. timeout: The maximum time in seconds to wait for the object to be present in the RDT store. If not specified, wait indefinitely. Returns: The tensors in the RDT object. """ with self._lock: self._wait_object(obj_id, timeout) return self.get_object(obj_id) def wait_and_pop_object( self, obj_id: str, timeout: Optional[float] = None ) -> List[Any]: """Atomically waits for the RDT object to be present in the RDT store, then pops it. If the object is not present after the optional timeout, raise a TimeoutError. Args: obj_id: The object ID to wait for. timeout: The maximum time in seconds to wait for the object to be present in the RDT store. If not specified, wait indefinitely. Returns: The RDT object. """ with self._lock: self._wait_object(obj_id, timeout) return self.pop_object(obj_id) def _wait_object(self, obj_id: str, timeout: Optional[float] = None) -> None: """Helper method to wait for the RDT object to be present in the RDT store. If the object is not present after the optional timeout, raise a TimeoutError. Args: obj_id: The object ID to wait for. timeout: The maximum time in seconds to wait for the object to be present in the RDT store. If not specified, wait indefinitely. """ with self._object_present_cv: if not self._object_present_cv.wait_for( lambda: self.has_object(obj_id), timeout=timeout, ): raise TimeoutError( f"ObjectRef({obj_id}) not found in RDT object store after {timeout}s, transfer may have failed. Please report this issue on GitHub: https://github.com/ray-project/ray/issues/new/choose" ) def pop_object(self, obj_id: str) -> List[Any]: with self._lock: queue = self._rdt_store.get(obj_id) assert queue is not None, f"obj_id={obj_id} not found in RDT store" rdt_object = queue.popleft() if len(queue) == 0: del self._rdt_store[obj_id] if rdt_object.error: raise rdt_object.error for tensor in rdt_object.data: self._tensor_to_object_ids[id(tensor)].remove(obj_id) if len(self._tensor_to_object_ids[id(tensor)]) == 0: self._tensor_to_object_ids.pop(id(tensor)) self._object_freed_cv.notify_all() return rdt_object.data def wait_tensor_freed(self, tensor: Any, timeout: Optional[float] = None) -> None: """ Wait for the object to be freed from the RDT store. """ with self._object_freed_cv: if not self._object_freed_cv.wait_for( lambda: id(tensor) not in self._tensor_to_object_ids, timeout=timeout, ): raise TimeoutError( f"Tensor {tensor} not freed from RDT object store after {timeout}s. The tensor will not be freed until all ObjectRefs containing the tensor have gone out of scope." ) def get_num_objects(self) -> int: """ Return the number of objects in the RDT store. """ with self._lock: # Count total objects across all queues return sum(len(queue) for queue in self._rdt_store.values())