#!/usr/bin/env python3 """GPU-aware smoke test for snapshot_link RDMA byte transfer. Sender on cuda:0, receiver subprocess on cuda:1. Tests whether mooncake's transfer_sync_write can move bytes between two GPUs via RDMA (which is what the real D→P flow will need for KV bytes). Usage: bash scripts/setup_env.sh && uv run --no-sync python scripts/smoke_snapshot_link_gpu.py The sender uses cuda:0 (--send-gpu); the receiver subprocess uses cuda:1 (--recv-gpu) by default. """ from __future__ import annotations import argparse import hashlib import json import os import subprocess import sys import tempfile import time from pathlib import Path _HERE = Path(__file__).resolve().parent sys.path.insert(0, str(_HERE.parent / "src")) SIZES_BYTES_DEFAULT = [ 1 << 14, # 16 KB 1 << 20, # 1 MB 1 << 24, # 16 MB 1 << 26, # 64 MB 1 << 28, # 256 MB ] def main(): ap = argparse.ArgumentParser() ap.add_argument("--host", default=os.environ.get("SNAPSHOT_LINK_HOST", "127.0.0.1")) ap.add_argument("--ib", default=os.environ.get("SNAPSHOT_LINK_IB", "mlx5_60")) ap.add_argument("--recv-port", type=int, default=int(os.environ.get("SNAPSHOT_LINK_RECV_PORT", "17787"))) ap.add_argument("--send-port", type=int, default=int(os.environ.get("SNAPSHOT_LINK_SEND_PORT", "17788"))) ap.add_argument("--max-bytes", type=int, default=256 * 1024 * 1024) ap.add_argument("--sizes", default=",".join(str(s) for s in SIZES_BYTES_DEFAULT)) ap.add_argument("--send-gpu", type=int, default=0) ap.add_argument("--recv-gpu", type=int, default=1) args = ap.parse_args() sizes = [int(s) for s in args.sizes.split(",")] tmpdir = Path(tempfile.mkdtemp(prefix="snapshot_link_gpu_smoke_")) control_path = tmpdir / "endpoint.json" recv_stderr_log = tmpdir / "recv.stderr.log" recv_cmd = [ sys.executable, str(_HERE / "snapshot_link_receiver_gpu.py"), "--host", args.host, "--port", str(args.recv_port), "--ib", args.ib, "--max-bytes", str(args.max_bytes), "--control-path", str(control_path), "--sizes", args.sizes, "--gpu-id", str(args.recv_gpu), ] recv_stderr = open(recv_stderr_log, "w") print(f"[sender] receiver cmd: {' '.join(recv_cmd)}", flush=True) recv_proc = subprocess.Popen( recv_cmd, stdout=subprocess.PIPE, stderr=recv_stderr, bufsize=1, universal_newlines=True, ) try: import torch if not torch.cuda.is_available(): print("[sender] FAIL: cuda not available") return 1 torch.cuda.set_device(args.send_gpu) deadline = time.time() + 90.0 meta = None while time.time() < deadline: if control_path.exists(): try: meta = json.loads(control_path.read_text()) if meta.get("ready"): break except Exception: pass if recv_proc.poll() is not None: _dump_recv_stderr(recv_stderr_log) print(f"[sender] FAIL: receiver exited (rc={recv_proc.returncode})") return 1 time.sleep(0.1) if meta is None: print("[sender] FAIL: receiver endpoint timeout") return 1 print(f"[sender] receiver endpoint: gpu={meta['gpu_id']}, " f"sid={meta['session_id']}, ptr={hex(int(meta['base_ptr']))}, " f"cap={meta['capacity_bytes']}", flush=True) from agentic_pd_hybrid.snapshot_link import SnapshotPeer, SnapshotEndpoint endpoint = SnapshotEndpoint( session_id=meta["session_id"], base_ptr=int(meta["base_ptr"]), capacity_bytes=int(meta["capacity_bytes"]), ) peer = SnapshotPeer( host=args.host, port=args.send_port, ib_device=args.ib, receive_capacity_bytes=0, ) # Allocate a sender buffer on cuda:0 send_tensor = torch.zeros(args.max_bytes, dtype=torch.uint8, device=f"cuda:{args.send_gpu}") send_ptr = send_tensor.data_ptr() ret = peer.engine.register_memory(send_ptr, args.max_bytes) if ret != 0: print(f"[sender] FAIL: register_memory ret={ret}") return 1 print(f"[sender] own gpu={args.send_gpu}, sid={peer.session_id}, " f"buf @ {hex(send_ptr)} ({args.max_bytes} B)", flush=True) transfers = [] for size in sizes: if size > args.max_bytes: continue # Fill with deterministic pattern on GPU seed = int(time.time() * 1e6) & 0xFFFFFFFF # Use a simple seeded pattern via torch ops gen = torch.Generator(device=f"cuda:{args.send_gpu}") gen.manual_seed(seed) send_tensor[:size] = torch.randint(0, 256, (size,), dtype=torch.uint8, device=f"cuda:{args.send_gpu}", generator=gen) torch.cuda.synchronize(args.send_gpu) # Compute expected hash (host-side) host_view = send_tensor[:size].cpu().numpy().tobytes() expected_sha = hashlib.sha256(host_view).hexdigest() # Push via RDMA t0 = time.perf_counter() ret = peer.push(endpoint, send_ptr, 0, size, remote_offset=0) t1 = time.perf_counter() dt_ms = (t1 - t0) * 1000.0 gbps = (size * 8.0 / 1e9) / max(t1 - t0, 1e-9) print(f"[sender] push size={size:>10d} ret={ret} " f"dur={dt_ms:>9.3f} ms thru={gbps:>6.3f} Gbps", flush=True) # Signal receiver to verify signal_path = control_path.with_suffix(f".do{size}") ack_path = control_path.with_suffix(f".ack{size}") signal_path.write_text(json.dumps({"sha": expected_sha})) ack_deadline = time.time() + 90.0 while time.time() < ack_deadline: if ack_path.exists(): break if recv_proc.poll() is not None: print(f"[sender] FAIL: receiver died after size={size}") _dump_recv_stderr(recv_stderr_log) return 1 time.sleep(0.05) transfers.append({ "size": size, "ret": ret, "dur_ms": round(dt_ms, 3), "thru_Gbps": round(gbps, 3), "ack": ack_path.exists(), }) try: recv_proc.wait(timeout=10) except subprocess.TimeoutExpired: recv_proc.terminate() recv_proc.wait(timeout=5) events = [] if recv_proc.stdout is not None: for raw in recv_proc.stdout: raw = raw.strip() if not raw: continue try: events.append(json.loads(raw)) except json.JSONDecodeError: events.append({"event": "non-json", "raw": raw}) print("=" * 78) print("[receiver] events:") verify_ok = 0 verify_fail = 0 for ev in events: print(f" {ev}") if ev.get("event") == "verify": if ev.get("ok"): verify_ok += 1 else: verify_fail += 1 recv_stderr.close() _dump_recv_stderr(recv_stderr_log, header="--- receiver stderr ---") overall = "PASS" if verify_fail == 0 and verify_ok == len(transfers) else "FAIL" print("=" * 78) print(f"OVERALL: {overall} verify_ok={verify_ok} verify_fail={verify_fail} " f"transfers={len(transfers)} send_gpu={args.send_gpu} recv_gpu={args.recv_gpu}") return 0 if overall == "PASS" else 1 finally: try: recv_proc.terminate() recv_proc.wait(timeout=5) except Exception: try: recv_proc.kill() except Exception: pass def _dump_recv_stderr(path: Path, header: str = "--- receiver stderr (last 60) ---") -> None: try: text = path.read_text() except FileNotFoundError: return print(header, flush=True) for line in text.splitlines()[-60:]: print(f" {line}", flush=True) if __name__ == "__main__": sys.exit(main())