Repository: ray-project/ray
File: python/ray/experimental/gpu_object_manager/gpu_object_manager.py (Python)

Code before:
    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.
    """
    gpu_object_manager = ray.worker.global_worker.gpu_object_manager
    gpu_object_manager.gpu_object_store.wait_tensor_freed(tensor, timeout)


@PublicAPI(stability="alpha")
def set_target_for_ref(ref: ObjectRef, target: List["torch.Tensor"]):
    """
    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, the target buffers are ignored, and a warning will be
    logged.

    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. Note that the order must match the order of the tensors
    from the sender side.

    Note that this only applies for the first `ray.get` call on the ref.
    Subsequent `ray.get` calls on the same ref will not use the target buffers unless
    they are set again after the first `ray.get` call.

    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.
    """
    gpu_object_manager = ray.worker.global_worker.gpu_object_manager
    gpu_object_manager.set_target_buffers_for_ref(ref, target)


class GPUObjectManager:
    def __init__(self):
        # This lock protects _managed_gpu_object_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 GPUObjectMeta.
        # This dictionary is hosted on the "driver" process of the actors that
        # store and send/receive GPU objects.
        self._managed_gpu_object_metadata: Dict[str, GPUObjectMeta] = {}
        # Condition variable to wait for the tensor transport meta to be set.

Review comment (suggestion):
"order" is a little bit hard to explain here because it comes from the sender's serialization order, right? We probably want to make a note about that and suggest putting lists of tensors as a safe way to ensure order of serialization.

Code after:
    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.
    """
    gpu_object_manager = ray.worker.global_worker.gpu_object_manager
    gpu_object_manager.gpu_object_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.
    """
    gpu_object_manager = ray.worker.global_worker.gpu_object_manager
    gpu_object_manager.set_target_buffers_for_ref(ref, target)


class GPUObjectManager:
    def __init__(self):
        # This lock protects _managed_gpu_object_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 GPUObjectMeta.
        # This dictionary is hosted on the "driver" process of the actors that
        # store and send/receive GPU objects.
        self._managed_gpu_object_metadata: Dict[str, GPUObjectMeta] = {}
        # 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.