Repository navigation
UCT/CUDA_IPC: uct_cuda_ipc_mkey_pack leaves struct padding in the packed rkey uninitialized #12099
Description
Describe the bug
On x86_64, uct_cuda_ipc_rkey_t has 4 bytes of padding between pid_t pid and CUdeviceptr d_bptr (cuda_ipc_md.h). uct_cuda_ipc_mkey_pack assigns the fields one by one into mkey_buffer and never clears it (cuda_ipc_md.c), so those 4 bytes carry whatever the buffer held before.
As a result, packing the same memory twice can produce different bytes. We registered one cudaMalloc'd buffer, packed it, deregistered it, and registered and packed it again. The two packed keys differ only in those 4 bytes:
first: ... 8a 76 08 00 | 00 00 00 00 | 00 00 20 99 9d 7f 00 00 ...
second: ... 8a 76 08 00 | 2e 74 61 72 | 00 00 20 99 9d 7f 00 00 ...
pid (554634) padding d_bptr
This has two effects:
- Code that compares packed rkeys sees a change where there is none. NIXL rejects reloading a remote agent's metadata when a known region's packed key differs, and a rejection currently breaks that peer for the rest of the process (see the linked NIXL issue).
- A few bytes of process memory are sent to the peer. Here they were ASCII (
.tar), so likely heap leftovers.
A possible fix is to zero mkey_buffer (the size of uct_cuda_ipc_extended_rkey_t when it is used) before filling in the fields. ph is copied from key->ph, so any padding inside uct_cuda_ipc_md_handle_t follows whatever key holds. We did not see differences there.
Steps to Reproduce
- Seen through NIXL, which packs this rkey into its agent metadata. The two-process script below registers the same CuPy buffer twice and prints the differing bytes of the region's packed metadata. It differed in 3 of 3 runs, with no UCX environment variables set. With
UCX_TLS=^cuda_ipc, the metadata is identical across rounds. - UCX 1.23.0 as bundled in the
nixl-cu131.5.0 wheel (revision65e342a). The code in question is unchanged onmaster(07ac81e94).
Reproducer (NIXL 1.5.0, CuPy)
"""A rejected loadRemoteMD leaves the peer impossible to load again.
The producer registers the same GPU buffer twice (deregistering in between) and
sends its agent metadata each time. The consumer loads the first, reads, then
loads the second.
"""
import multiprocessing as mp
import time
def producer(q_out, q_in):
import cupy as cp
from nixl._api import nixl_agent, nixl_agent_config
agent = nixl_agent("producer", nixl_agent_config(backends=["UCX"]))
buf = cp.empty(1 << 18, dtype=cp.float32)
for i in range(2):
buf.fill(i)
cp.cuda.Device().synchronize()
reg = agent.register_memory([(buf.data.ptr, buf.nbytes, 0, "")], "VRAM")
region_md = agent.get_partial_agent_metadata(reg, False)
q_out.put((agent.get_agent_metadata(), region_md, buf.data.ptr, buf.nbytes))
q_in.get()
agent.deregister_memory(reg)
q_out.put(None)
def consumer(q_in, q_out):
import cupy as cp
from nixl._api import nixl_agent, nixl_agent_config
agent = nixl_agent("consumer", nixl_agent_config(backends=["UCX"]))
out = cp.empty(1 << 18, dtype=cp.float32)
regions = []
def read(addr, nbytes):
local = agent.register_memory([(out.data.ptr, nbytes, 0, "")], "VRAM")
remote = agent.get_xfer_descs([(addr, nbytes, 0)], "VRAM")
handle = agent.initialize_xfer("READ", local.trim(), remote, "producer", b"r")
agent.transfer(handle)
while (state := agent.check_xfer_state(handle)) == "PROC":
time.sleep(0.0001)
agent.release_xfer_handle(handle)
agent.deregister_memory(local)
return state, float(out[0])
i = 0
while (item := q_in.get()) is not None:
md, region_md, addr, nbytes = item
regions.append(region_md)
try:
agent.add_remote_agent(md)
print(f"round {i}: load ok, read {read(addr, nbytes)}", flush=True)
except Exception as e:
print(f"round {i}: load FAILED: {e}", flush=True)
agent.remove_remote_agent("producer")
try:
agent.add_remote_agent(md)
print(f"round {i}: reload after remove ok, read {read(addr, nbytes)}", flush=True)
except Exception as e2:
print(f"round {i}: reload after remove FAILED: {e2}", flush=True)
q_out.put("ack")
i += 1
print("region metadata identical across rounds:", regions[0] == regions[1], flush=True)
if regions[0] != regions[1]:
diff = [k for k in range(len(regions[0])) if regions[0][k] != regions[1][k]]
print(f" {len(diff)} bytes differ at offsets {diff}:",
regions[0][diff[0]:diff[-1] + 1].hex(), "->", regions[1][diff[0]:diff[-1] + 1].hex(), flush=True)
if __name__ == "__main__":
ctx = mp.get_context("spawn")
q1, q2 = ctx.Queue(), ctx.Queue()
procs = [ctx.Process(target=producer, args=(q1, q2)), ctx.Process(target=consumer, args=(q1, q2))]
for p in procs:
p.start()
for p in procs:
p.join(120)Setup and versions
- Ubuntu 22.04.5 LTS, x86_64, kernel 6.8.0-1059-azure
- NVIDIA A100 80GB PCIe, driver 580.159.03, CUDA runtime 13.2 (from the CuPy wheel)
- No
nvidia_peermemorgdrdrvloaded
Activity