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

Code before:
            timeout: The maximum time in seconds to wait for the object to be
                present in the GPU object store. If not specified, wait
                indefinitely.

        Returns:
            The GPU 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 GPU object to be present in the GPU object 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 GPU object 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["torch.Tensor"]:
        with self._lock:
            queue = self._gpu_object_store.get(obj_id)
            assert queue is not None, f"obj_id={obj_id} not found in GPU object store"
            gpu_object = queue.popleft()
            if len(queue) == 0:
                del self._gpu_object_store[obj_id]
            if gpu_object.error:
                raise gpu_object.error
            for tensor in gpu_object.data:
                self._tensor_to_object_ids[tensor.data_ptr()].remove(obj_id)
                if len(self._tensor_to_object_ids[tensor.data_ptr()]) == 0:
                    self._tensor_to_object_ids.pop(tensor.data_ptr())
            self._object_freed_cv.notify_all()
            return gpu_object.data

    def wait_tensor_freed(
        self, tensor: "torch.Tensor", timeout: Optional[float] = None
    ) -> None:
        """

Review comment (bug):
Oh I guess there must be a potential bug here if you combine this with #60509. @Sparks0219, can you add a test to #60509 that checks that if two views of the same tensor are added, calling `wait_tensor_free` only waits until that specific view is freed?

Code after:

    def pop_object(self, obj_id: str) -> List["torch.Tensor"]:
        with self._lock:
            queue = self._gpu_object_store.get(obj_id)
            assert queue is not None, f"obj_id={obj_id} not found in GPU object store"
            gpu_object = queue.popleft()
            if len(queue) == 0:
                del self._gpu_object_store[obj_id]
            if gpu_object.error:
                raise gpu_object.error
            for tensor in gpu_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 gpu_object.data

    def wait_tensor_freed(
        self, tensor: "torch.Tensor", timeout: Optional[float] = None
    ) -> None:
        """
        Wait for the object to be freed from the GPU object 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 GPU object store.
        """
        with self._lock:
            # Count total objects across all queues
            return sum(len(queue) for queue in self._gpu_object_store.values())