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

Code before:
    ):
        from ray._private.worker import global_worker

        with self._cache_lock:
            assert isinstance(tensor_transport_meta, NixlTransportMetadata)
            gpu_object_store = global_worker.gpu_object_manager.gpu_object_store
            self._managed_meta_nixl.pop(obj_id, None)
            tensors = gpu_object_store.get_object(obj_id)
            for tensor in tensors:
                key = tensor.data_ptr()
                if key in self._tensor_desc_cache:
                    tensor_desc = self._tensor_desc_cache[key]
                    tensor_desc.ref_count -= 1
                    if tensor_desc.ref_count == 0:
                        self._tensor_desc_cache.pop(key)
                        self.get_nixl_agent().deregister_memory(tensor_desc.reg_desc)

    def abort_transport(
        self,
        obj_id: str,
        communicator_metadata: CommunicatorMetadata,
    ):
        with self._aborted_transfer_obj_ids_lock:
            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 (refactor):
```suggestion
    def _record_and_get_meta_if_duplicate(
```
usually internal methods are prefixed with _ bc python doesn't have private

Code after:
            self._managed_meta_nixl.pop(obj_id, None)
            tensors = gpu_object_store.get_object(obj_id)
            for tensor in tensors:
                key = tensor.data_ptr()
                if key in self._tensor_desc_cache:
                    tensor_desc = self._tensor_desc_cache[key]
                    tensor_desc.metadata_count -= 1
                    if tensor_desc.metadata_count == 0:
                        self._tensor_desc_cache.pop(key)
                        self.get_nixl_agent().deregister_memory(tensor_desc.reg_desc)

    def abort_transport(
        self,
        obj_id: str,
        communicator_metadata: CommunicatorMetadata,
    ):
        with self._aborted_transfer_obj_ids_lock:
            self._aborted_transfer_obj_ids.add(obj_id)

    # NOTE: The below methods are intended to be used internally hence they assume the caller is already holding the cache lock.
    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
        )
        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