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

Code before:
            self._aborted_transfer_obj_ids.add(obj_id)

    def record_and_get_meta_if_duplicate(
        self, src_obj_id: str, src_gpu_object: List["torch.Tensor"]
    ) -> Optional[NixlTransportMetadata]:
        """
        Record the NIXL managed meta for the given object ID if it is a duplicate of another object, and return the meta if it is.
        Assumes that the caller is already holding the cache lock.
        """
        from ray._private.worker import global_worker

        gpu_object_store = global_worker.gpu_object_manager.gpu_object_store
        duplicate_obj_id = gpu_object_store.get_duplicate_objects(
            src_obj_id, src_gpu_object, tensor_transport="NIXL"
        )
        if duplicate_obj_id is not None:
            meta = self._managed_meta_nixl[duplicate_obj_id]
            self._managed_meta_nixl[src_obj_id] = meta
            for tensor in src_gpu_object:
                key = tensor.data_ptr()
                if key in self._tensor_desc_cache:
                    self._tensor_desc_cache[key].ref_count += 1
            return meta
        return None

    def get_num_managed_meta_nixl(self) -> int:
        with self._cache_lock:
            return len(self._managed_meta_nixl)


Review comment (suggestion):
```suggestion
    def _get_num_managed_meta_nixl(self) -> int:
```

Code after:
        """
        Record the NIXL managed meta for the given object ID if it is a duplicate of another object, and return the meta if it is.
        Assumes that the caller is already holding the cache lock.
        """
        from ray._private.worker import global_worker

        gpu_object_store = global_worker.gpu_object_manager.gpu_object_store
        duplicate_obj_id = gpu_object_store.get_duplicate_objects(
            src_obj_id, src_gpu_object
        )
        if duplicate_obj_id is not None:
            meta = self._get_meta(duplicate_obj_id)
            if meta is None:
                raise ValueError(
                    f"NIXL transport metadata for object id {duplicate_obj_id} not found"
                )
            self._put_meta(src_obj_id, meta)
            self._add_tensor_descs(src_gpu_object)
            return meta
        return None

    def _get_num_managed_meta_nixl(self) -> int:
        with self._cache_lock:
            return len(self._managed_meta_nixl)

    def _get_meta(self, object_id: str) -> Optional[NixlTransportMetadata]:
        """
        Get the NIXL transport metadata for the given object ID if it exists
        """
        if object_id in self._managed_meta_nixl:
            return self._managed_meta_nixl[object_id]
        return None

    def _put_meta(self, object_id: str, meta: NixlTransportMetadata):
        """
        Store the NIXL transport metadata for the given object ID
        """
        self._managed_meta_nixl[object_id] = meta

    def _add_tensor_descs(self, tensors: List["torch.Tensor"]):
        """
        If this is the first time the tensor is being added, we register the memory with NIXL.
        Otherwise, we increment the reference count.
        """
        for tensor in tensors:
            key = tensor.data_ptr()
            if key in self._tensor_desc_cache:
                if tensor.nbytes != self._tensor_desc_cache[key].nbytes:
                    raise ValueError(
                        "Tensors in an RDT object cannot partially overlap with each other."
                    )