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

Code before:
    def wait_and_get_object(
        self, obj_id: str, timeout: Optional[float] = None
    ) -> List["torch.Tensor"]:
        """Atomically waits for the GPU object to be present in the GPU object
        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 GPU object store. If not specified, wait indefinitely.

        Returns:
            The tensors in the GPU object.
        """
        with self._lock:
            self._wait_object(obj_id, timeout)
            return self.get_object(obj_id)

    def get_duplicate_objects(
        self,
        src_obj_id: str,
        src_gpu_object: List["torch.Tensor"],
        tensor_transport: Optional[str] = None,
    ) -> Optional[str]:
        """Get another object ID of the GPU object that duplicates the given GPU object."""
        with self._lock:
            if len(src_gpu_object) == 0:
                return None
            obj_id_set = set()
            for tensor in src_gpu_object:
                for obj_id in self._tensor_to_object_ids[tensor.data_ptr()]:
                    obj_id_set.add(obj_id)

            for dst_obj_id in obj_id_set:
                if dst_obj_id != src_obj_id:
                    dst_gpu_object = self._gpu_object_store[dst_obj_id][0].data
                    is_same_tensors = len(src_gpu_object) == len(
                        dst_gpu_object
                    ) and all(
                        t1.data_ptr() == t2.data_ptr()
                        for t1, t2 in zip(src_gpu_object, dst_gpu_object)
                    )
                    if is_same_tensors:
                        return dst_obj_id
            return None

    def wait_and_pop_object(
        self, obj_id: str, timeout: Optional[float] = None
    ) -> List["torch.Tensor"]:
        """Atomically waits for the GPU object to be present in the GPU object

Review comment (refactor):
```suggestion
        """
        Get another object ID of the GPU object that duplicates the given GPU object.
        Returns the object id if there is a duplicate, None if there is no duplicate.
        """
```

Code after:
    def wait_and_get_object(
        self, obj_id: str, timeout: Optional[float] = None
    ) -> List["torch.Tensor"]:
        """Atomically waits for the GPU object to be present in the GPU object
        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 GPU object store. If not specified, wait indefinitely.

        Returns:
            The tensors in the GPU object.
        """
        with self._lock:
            self._wait_object(obj_id, timeout)
            return self.get_object(obj_id)

    def get_duplicate_objects(
        self,
        src_obj_id: str,
        src_gpu_object: List["torch.Tensor"],
    ) -> Optional[str]:
        """
        Get another object ID of the GPU object that duplicates the given GPU object.
        Returns the object id if there is a duplicate, None if there is no duplicate.
        """
        with self._lock:
            if len(src_gpu_object) == 0:
                return None
            obj_id_set = set()
            for tensor in src_gpu_object:
                for obj_id in self._tensor_to_object_ids[tensor.data_ptr()]:
                    obj_id_set.add(obj_id)

            for dst_obj_id in obj_id_set:
                if dst_obj_id != src_obj_id:
                    dst_gpu_object = self._gpu_object_store[dst_obj_id][0].data
                    is_same_tensors = len(src_gpu_object) == len(
                        dst_gpu_object
                    ) and all(
                        t1.data_ptr() == t2.data_ptr()
                        for t1, t2 in zip(src_gpu_object, dst_gpu_object)
                    )
                    if is_same_tensors:
                        return dst_obj_id
            return None

    def wait_and_pop_object(
        self, obj_id: str, timeout: Optional[float] = None