dc4867c270
A thin wrapper around mooncake.engine.TransferEngine that does one-sided RDMA writes between two SnapshotPeer endpoints. Bypasses SGLang's MooncakeKVManager (which is hard-gated to PREFILL/DECODE roles via add_transfer_request assertion at conn.py:1563) so the D→P direction doesn't require invasive role-axis changes upstream. Smoke test (two subprocess.Popen processes, mlx5_60, 127.0.0.1): 1 KB 9.0 ms (one-time openSegment handshake) 16 KB 0.04 ms 3.5 Gbps 1 MB 0.10 ms 82 Gbps 16 MB 0.58 ms 232 Gbps 64 MB 1.70 ms 316 Gbps (~80% of NDR 400G line rate) All 5 sizes pass SHA256 verification end-to-end. Files: src/agentic_pd_hybrid/snapshot_link.py — SnapshotPeer, SnapshotEndpoint scripts/snapshot_link_receiver.py — child-process receiver scripts/smoke_snapshot_link.py — sender + verifier docs/D_TO_P_PHASE1_LINK_ZH.md — phase 1 acceptance doc Next: Phase 2 (D-side scheduler commit hook), Phase 3 (P-side prefill bypass with snapshot KV). See docs/D_TO_P_SYNC_DESIGN_ZH.md §5.
245 lines
8.2 KiB
Python
Executable File
245 lines
8.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Two-process smoke test for snapshot_link D→P RDMA byte transfer.
|
|
|
|
Spawns scripts/snapshot_link_receiver.py via subprocess.Popen with stderr
|
|
piped to ``<tmpdir>/recv.stderr.log`` for post-mortem if something dies.
|
|
|
|
Sender (this process):
|
|
1. Spawns receiver child, waits for endpoint.json
|
|
2. Brings up own SnapshotPeer (no recv buffer), registers a send buffer
|
|
3. For each size: fill pattern, batch_transfer_sync_write, signal child,
|
|
wait for child's ack
|
|
4. Reads child's stdout (one JSON event per line) for verification
|
|
|
|
Pass = every size yields a child "verify" event with ok=true.
|
|
|
|
Usage:
|
|
bash scripts/setup_env.sh && uv run --no-sync python scripts/smoke_snapshot_link.py
|
|
|
|
Env (optional):
|
|
SNAPSHOT_LINK_HOST default 127.0.0.1
|
|
SNAPSHOT_LINK_IB default mlx5_60
|
|
SNAPSHOT_LINK_RECV_PORT default 17777
|
|
SNAPSHOT_LINK_SEND_PORT default 17778
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import ctypes
|
|
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 << 10, # 1 KB
|
|
1 << 14, # 16 KB
|
|
1 << 18, # 256 KB
|
|
1 << 20, # 1 MB
|
|
1 << 22, # 4 MB
|
|
1 << 24, # 16 MB
|
|
1 << 26, # 64 MB
|
|
]
|
|
|
|
|
|
def _pattern_byte(i: int, seed: int) -> int:
|
|
return (i * 2654435761 + seed) & 0xFF
|
|
|
|
|
|
def _fill_pattern(buf, length: int, seed: int) -> None:
|
|
tile_size = 4096
|
|
tile = bytes(_pattern_byte(i, seed) for i in range(tile_size))
|
|
tile_arr = (ctypes.c_ubyte * tile_size).from_buffer_copy(tile)
|
|
n_full = length // tile_size
|
|
rem = length - n_full * tile_size
|
|
base = ctypes.addressof(buf)
|
|
src_addr = ctypes.addressof(tile_arr)
|
|
for k in range(n_full):
|
|
ctypes.memmove(base + k * tile_size, src_addr, tile_size)
|
|
if rem:
|
|
ctypes.memmove(base + n_full * tile_size, src_addr, rem)
|
|
|
|
|
|
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", "17777")))
|
|
ap.add_argument("--send-port", type=int,
|
|
default=int(os.environ.get("SNAPSHOT_LINK_SEND_PORT", "17778")))
|
|
ap.add_argument("--max-bytes", type=int, default=128 * 1024 * 1024)
|
|
ap.add_argument("--sizes", default=",".join(str(s) for s in SIZES_BYTES_DEFAULT))
|
|
args = ap.parse_args()
|
|
|
|
sizes = [int(s) for s in args.sizes.split(",")]
|
|
tmpdir = Path(tempfile.mkdtemp(prefix="snapshot_link_smoke_"))
|
|
control_path = tmpdir / "endpoint.json"
|
|
recv_stderr_log = tmpdir / "recv.stderr.log"
|
|
|
|
recv_cmd = [
|
|
sys.executable,
|
|
str(_HERE / "snapshot_link_receiver.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,
|
|
]
|
|
recv_stderr = open(recv_stderr_log, "w")
|
|
print(f"[sender] launching receiver: {' '.join(recv_cmd)}", flush=True)
|
|
print(f"[sender] receiver stderr → {recv_stderr_log}", flush=True)
|
|
recv_proc = subprocess.Popen(
|
|
recv_cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=recv_stderr,
|
|
bufsize=1,
|
|
universal_newlines=True,
|
|
)
|
|
|
|
try:
|
|
# Wait for endpoint metadata
|
|
deadline = time.time() + 60.0
|
|
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 early (rc={recv_proc.returncode})")
|
|
return 1
|
|
time.sleep(0.1)
|
|
else:
|
|
print("[sender] FAIL: timed out waiting for receiver endpoint", flush=True)
|
|
return 1
|
|
|
|
print(f"[sender] receiver endpoint: {meta}", 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,
|
|
)
|
|
send_buf = (ctypes.c_byte * args.max_bytes)()
|
|
send_addr = ctypes.addressof(send_buf)
|
|
peer.register_send_buffer(send_addr, args.max_bytes)
|
|
print(f"[sender] own session_id={peer.session_id}, send_buf @ {hex(send_addr)} ({args.max_bytes} B)", flush=True)
|
|
|
|
transfers = []
|
|
for size in sizes:
|
|
if size > args.max_bytes:
|
|
continue
|
|
seed = int(time.time() * 1e6) & 0xFFFFFFFF
|
|
_fill_pattern(send_buf, size, seed)
|
|
t0 = time.perf_counter()
|
|
ret = peer.push(endpoint, send_addr, 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_path = control_path.with_suffix(f".do{size}")
|
|
ack_path = control_path.with_suffix(f".ack{size}")
|
|
signal_path.write_text(str(seed))
|
|
ack_deadline = time.time() + 60.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}", flush=True)
|
|
_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(),
|
|
})
|
|
|
|
peer.close()
|
|
|
|
# Drain child stdout — each line is a JSON event
|
|
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)}")
|
|
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 40) ---") -> None:
|
|
try:
|
|
text = path.read_text()
|
|
except FileNotFoundError:
|
|
return
|
|
print(header, flush=True)
|
|
for line in text.splitlines()[-40:]:
|
|
print(f" {line}", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|