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

Code before:
        if gpu_object:
            # We assume all tensors in one GPU object have the same device type, but we
            # don't assume they're all on the same device.
            devices = set()
            device = gpu_object[0].device
            for t in gpu_object:
                if t.device.type != device.type:
                    raise ValueError(
                        "All tensors in an RDT object must have the same device type."
                    )
                tensor_meta.append((t.shape, t.dtype))
                devices.add(t.device)
            if device.type == "cuda":
                # We have to synchronize before memory registration to assure the object
                # has been created because nixl doesn't guarantee it will.
                for dev in devices:
                    torch.cuda.synchronize(dev)
            with self._nixl_memory_lock:
                nixl_agent = self.get_nixl_agent()
                reg_descs = nixl_agent.register_memory(gpu_object)
                serialized_descs = nixl_agent.get_serialized_descs(reg_descs.trim())
                agent_meta = nixl_agent.get_agent_metadata()
                agent_name = nixl_agent.name
                # If there is any memory deregistered before the registration of this gpu object, we need to tell the receiver to reset the remote agent.
                # Note, we can remove this limit once nixl supports metadata updates.
                reset_remote_agent = (
                    self._memory_deregistered and NIXL_REMOTE_AGENT_CACHE_MAXSIZE > 0
                )
                self._memory_deregistered = False

        else:
            reg_descs, serialized_descs, agent_meta, agent_name, reset_remote_agent = (
                None,
                None,
                None,
                None,
                None,
            )

        ret = NixlTransportMetadata(
            tensor_meta=tensor_meta,
            tensor_device=device,
            nixl_reg_descs=reg_descs,
            nixl_serialized_descs=serialized_descs,
            nixl_agent_meta=agent_meta,
            nixl_agent_name=agent_name,
            reset_remote_agent=reset_remote_agent,
        )
        gpu_object_store.record_managed_meta_nixl(obj_id, ret)
        return ret


Review comment (performance):
Since we just registered a new memory region at line 150, shouldn't we need to reset in that case too, not just if a deregistration happened? I'm also a bit confused about the and condition to check NIXL_REMOTE_AGENT_CACHE_MAXSIZE > 0; why do we need that check?

Code after:
    ) -> NixlTransportMetadata:
        import torch

        with self._cache_lock:
            device = None
            tensor_meta = []
            duplicate_meta = self._record_and_get_meta_if_duplicate(obj_id, gpu_object)
            if duplicate_meta is not None:
                return duplicate_meta

            if gpu_object:
                # We assume all tensors in one GPU object have the same device type,
                # but we don't assume they're all on the same device.
                devices = set()
                device = gpu_object[0].device
                for t in gpu_object:
                    if t.device.type != device.type:
                        raise ValueError(
                            "All tensors in an RDT object must have the same device type."
                        )
                    if not t.is_contiguous():
                        raise ValueError(
                            "All tensors in an RDT object must be contiguous."
                        )
                    if (
                        t.data_ptr() in self._tensor_desc_cache
                        and self._tensor_desc_cache[t.data_ptr()].nbytes != t.nbytes
                    ):
                        raise ValueError(
                            "Tensors in an RDT object can not be partially overlapping with already registered tensors"
                        )
                    tensor_meta.append((t.shape, t.dtype))
                    devices.add(t.device)
                if device.type == "cuda":
                    # We have to synchronize before memory registration to assure the
                    # object has been created because nixl doesn't guarantee it will.
                    for dev in devices:
                        torch.cuda.synchronize(dev)

                nixl_agent = self.get_nixl_agent()
                self._add_tensor_descs(gpu_object)
                xfer_descs = nixl_agent.get_xfer_descs(gpu_object)
                serialized_descs = nixl_agent.get_serialized_descs(xfer_descs)
                agent_meta = nixl_agent.get_agent_metadata()
                agent_name = nixl_agent.name
                agent_meta_version = self._nixl_agent_meta_version
            else:
                serialized_descs, agent_meta = None, None
                agent_name, agent_meta_version = None, None

            ret = NixlTransportMetadata(