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.
124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Receiver-side child process for the snapshot_link smoke test.
|
|
|
|
Reads CLI args, brings up a SnapshotPeer with a registered recv buffer,
|
|
writes endpoint metadata to a control file, then loops: wait for size
|
|
signal, verify recv buffer, write ack.
|
|
|
|
Status events are printed as single-line JSON to stdout for parent to
|
|
parse.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import ctypes
|
|
import hashlib
|
|
import json
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
|
|
|
|
|
|
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 _emit(d: dict) -> None:
|
|
print(json.dumps(d), flush=True)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--host", required=True)
|
|
ap.add_argument("--port", type=int, required=True)
|
|
ap.add_argument("--ib", required=True)
|
|
ap.add_argument("--max-bytes", type=int, required=True)
|
|
ap.add_argument("--control-path", required=True)
|
|
ap.add_argument("--sizes", required=True, help="comma-separated bytes")
|
|
args = ap.parse_args()
|
|
|
|
sizes = [int(s) for s in args.sizes.split(",")]
|
|
|
|
from agentic_pd_hybrid.snapshot_link import SnapshotPeer
|
|
|
|
try:
|
|
peer = SnapshotPeer(
|
|
host=args.host,
|
|
port=args.port,
|
|
ib_device=args.ib,
|
|
receive_capacity_bytes=args.max_bytes,
|
|
)
|
|
except Exception as e:
|
|
import traceback
|
|
_emit({"event": "init-failed", "error": repr(e), "tb": traceback.format_exc()})
|
|
sys.exit(2)
|
|
|
|
endpoint = peer.endpoint
|
|
Path(args.control_path).write_text(json.dumps({
|
|
"session_id": endpoint.session_id,
|
|
"base_ptr": endpoint.base_ptr,
|
|
"capacity_bytes": endpoint.capacity_bytes,
|
|
"ready": True,
|
|
}))
|
|
_emit({"event": "endpoint-ready", "session_id": endpoint.session_id,
|
|
"base_ptr": endpoint.base_ptr, "capacity": endpoint.capacity_bytes})
|
|
|
|
cp = Path(args.control_path)
|
|
for size in sizes:
|
|
if size > args.max_bytes:
|
|
continue
|
|
signal_path = cp.with_suffix(f".do{size}")
|
|
ack_path = cp.with_suffix(f".ack{size}")
|
|
deadline = time.time() + 120.0
|
|
while time.time() < deadline:
|
|
if signal_path.exists():
|
|
break
|
|
time.sleep(0.05)
|
|
else:
|
|
_emit({"event": "no-signal-timeout", "size": size})
|
|
continue
|
|
try:
|
|
seed = int(signal_path.read_text().strip())
|
|
except Exception as e:
|
|
_emit({"event": "signal-parse-error", "size": size, "err": repr(e)})
|
|
continue
|
|
expected_arr = (ctypes.c_ubyte * size)()
|
|
_fill_pattern(expected_arr, size, seed)
|
|
expected_hash = hashlib.sha256(bytes(expected_arr)).hexdigest()
|
|
recv_bytes = peer.read_bytes(0, size)
|
|
recv_hash = hashlib.sha256(recv_bytes).hexdigest()
|
|
ok = recv_hash == expected_hash
|
|
_emit({
|
|
"event": "verify",
|
|
"size": size,
|
|
"ok": ok,
|
|
"expected_sha": expected_hash[:16],
|
|
"got_sha": recv_hash[:16],
|
|
"first8_recv": recv_bytes[:8].hex(),
|
|
"last8_recv": recv_bytes[-8:].hex(),
|
|
})
|
|
ack_path.write_text("done")
|
|
|
|
peer.close()
|
|
_emit({"event": "receiver-done"})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|