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 self._memory_deregistered:
                    agent_partial_meta = None
                else:
                    agent_partial_meta = nixl_agent.get_partial_agent_metadata(
                        reg_descs, inc_conn_info=False, backends=["UCX"]
                    )
                self._memory_deregistered = False

        else:
            reg_descs, serialized_descs, agent_meta, agent_name, agent_partial_meta = (
                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,
            nixl_agent_partial_meta=agent_partial_meta,
        )
        gpu_object_store.record_managed_meta_nixl(obj_id, ret)

Review comment (suggestion):
If there was a `deregister_memory` called before, then we should not use the partial agent meta. Otherwise, when the receiver calls `add_remote_agent(partial_agent_metadata)`, there would be a NIXL_ERR_NOT_ALLOWED thrown.

This is actually a TODO in nixl (search `// TODO: Support metadata updates`). When the receiver uses partial agent meta to update an existing remote agent (the sender), it will check the address of the newly-registered memory (which may reuse the address of a (previously) de-registered memory at the sender). If the address conflicts with any other memory (which has been de-registered at the sender but the receiver has not been notified), it throws NIXL_ERR_NOT_ALLOWED.

Code after:
            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(
                tensor_meta=tensor_meta,
                tensor_device=device,
                nixl_serialized_descs=serialized_descs,
                nixl_agent_meta=agent_meta,
                nixl_agent_name=agent_name,
                nixl_agent_meta_version=agent_meta_version,
            )
            self._put_meta(obj_id, ret)
            return ret