feat(agentic): D→P snapshot orchestration in reseed path + CLI flag

Phase 3 — wires the SGLang-side snapshot RPCs (committed in 86412bb)
into the agentic reseed slow-path. On _invoke_kvcache_seeded_router:

  1. POST {prefill_url}/_snapshot/prepare_receive   alloc P-side slots
  2. POST {old_decode_url}/_snapshot/dump           RDMA push session KV
  3. POST {prefill_url}/_snapshot/finalize_ingest   insert into P radix

After step 3 P's radix tree has the session prefix cached; the subsequent
SGLang router-driven prefill on P hits cache instead of re-computing.

Any RPC failure short-circuits to the existing seeded_router fallback
(re-prefill from scratch). All steps are best-effort and structurally
logged for post-hoc analysis.

Flag plumbing:
  cli.py             --enable-d-to-p-sync          (replay + benchmark)
  topology.py        SingleNodeTopology.enable_d_to_p_sync
  stack.py           SGLANG_SNAPSHOT_LINK_ENABLE=1 injection per worker
  replay.py          ReplayConfig.enable_d_to_p_sync +
                     _attempt_d_to_p_sync helper

Snapshot port per worker derives from disaggregation_bootstrap_port +
1000 (set in third_party/.../snapshot/controller.py), so different
workers get distinct mooncake snapshot engines on the same node.

Smoke (next): scripts/smoke_snapshot_sglang_integration.py spawns one
D + one P, exercises the 3 RPCs end-to-end, checks cache_tokens on a
follow-up generate request.

See docs/D_TO_P_SYNC_DESIGN_ZH.md for the full design.
This commit is contained in:
Claude Code Agent
2026-05-13 08:16:46 +08:00
parent 86412bb174
commit b9b0cf0fac
5 changed files with 450 additions and 0 deletions

View File

@@ -116,6 +116,11 @@ class ReplayConfig:
# with shared cross-session prefix. 0 disables. See
# docs/E1_E2_FIX_DESIGN_ZH.md §Q2.
kvcache_load_floor_bonus: int = 0
# D→P snapshot push: when True and reseed fires, agentic will RDMA-dump
# the session's KV from the D-side worker that last held it onto the P
# worker and insert into P's radix tree, so the subsequent P prefill
# hits cache. See docs/D_TO_P_SYNC_DESIGN_ZH.md.
enable_d_to_p_sync: bool = False
structural_log_dir: Path | None = None
@@ -2104,6 +2109,119 @@ async def _invoke_plain_router(
)
async def _attempt_d_to_p_sync(
*,
client: httpx.AsyncClient,
request: TraceRequest,
config: ReplayConfig,
prefill_url: str,
decode_session: DirectSessionState,
) -> dict | None:
"""Try to RDMA-dump session KV from the D that last held it to ``prefill_url``.
Returns a dict with status info on success/skip, or ``None`` on a
non-recoverable error. The caller falls back to normal re-prefill on
any failure.
"""
if not config.enable_d_to_p_sync:
return None
source_d_url = decode_session.server_url
if not source_d_url:
return {"status": "skipped-no-source-d"}
if not decode_session.opened:
return {"status": "skipped-d-closed"}
# Compose token list for radix insert: we don't have the actual token_ids
# on the agentic side in a stable form; use the request's prompt_token_ids
# via the residency bookkeeping. For now we use a length proxy.
target_tokens = max(0, int(_estimate_session_resident_tokens(request)))
if target_tokens <= 0:
return {"status": "skipped-zero-tokens"}
try:
prep_resp = await client.post(
f"{prefill_url}/_snapshot/prepare_receive",
json={
"session_id": request.session_id,
"num_tokens": target_tokens,
},
timeout=30.0,
)
prep_resp.raise_for_status()
prep = prep_resp.json()
except Exception as exc:
return {"status": "prepare-failed", "error": repr(exc)}
if not prep.get("ok"):
return {"status": "prepare-not-ok", "reason": prep.get("reason")}
try:
dump_resp = await client.post(
f"{source_d_url}/_snapshot/dump",
json={
"session_id": request.session_id,
"target_snapshot_session_id": prep["snapshot_session_id"],
"target_k_base_ptrs": prep["k_base_ptrs"],
"target_v_base_ptrs": prep["v_base_ptrs"],
"target_slot_indices": prep["slot_indices"],
"target_stride_k_bytes": prep["stride_k_bytes"],
"target_stride_v_bytes": prep["stride_v_bytes"],
},
timeout=60.0,
)
dump_resp.raise_for_status()
dump = dump_resp.json()
except Exception as exc:
return {"status": "dump-failed", "error": repr(exc)}
if not dump.get("ok"):
return {"status": "dump-not-ok", "reason": dump.get("reason"),
"bytes_pushed": dump.get("bytes_pushed", 0)}
# We need token_ids for radix insert. The caller has request.input_token_ids
# for the first N — use that as best-available approximation.
tokens = list(getattr(request, "input_token_ids", []) or [])
if not tokens:
# No token_ids available — can't insert into radix. P will fall back
# to normal prefill but will have wasted slots. Discard.
try:
await client.post(
f"{prefill_url}/_snapshot/finalize_ingest",
json={
"session_id": request.session_id,
"token_ids": [],
"slot_indices": prep["slot_indices"],
},
timeout=15.0,
)
except Exception:
pass
return {"status": "no-tokens-discard", "bytes_pushed": dump.get("bytes_pushed", 0)}
n = min(len(tokens), len(prep["slot_indices"]))
try:
fin_resp = await client.post(
f"{prefill_url}/_snapshot/finalize_ingest",
json={
"session_id": request.session_id,
"token_ids": tokens[:n],
"slot_indices": prep["slot_indices"][:n],
},
timeout=30.0,
)
fin_resp.raise_for_status()
fin = fin_resp.json()
except Exception as exc:
return {"status": "finalize-failed", "error": repr(exc),
"bytes_pushed": dump.get("bytes_pushed", 0)}
if not fin.get("ok"):
return {"status": "finalize-not-ok", "reason": fin.get("reason"),
"bytes_pushed": dump.get("bytes_pushed", 0)}
return {
"status": "ok",
"bytes_pushed": int(dump.get("bytes_pushed", 0)),
"inserted_prefix_len": int(fin.get("inserted_prefix_len", 0)),
"snapshot_session_id": prep.get("snapshot_session_id"),
}
async def _invoke_kvcache_seeded_router(
*,
client: httpx.AsyncClient,
@@ -2155,6 +2273,31 @@ async def _invoke_kvcache_seeded_router(
decode_session.prefill_server_url = prefill_url
prefill_session_newly_opened = True
# D→P snapshot push (Phase 3) — best-effort; on any failure we silently
# fall back to the existing re-prefill path. The result is logged for
# post-hoc analysis but does not affect correctness.
if config.enable_d_to_p_sync:
sync_result = await _attempt_d_to_p_sync(
client=client,
request=request,
config=config,
prefill_url=prefill_url,
decode_session=decode_session,
)
if sync_result is not None and sync_result.get("status") != "ok":
logger.info(
"d_to_p_sync sid=%s rid=%s skipped: %s",
request.session_id, request.request_id, sync_result,
)
elif sync_result and sync_result.get("status") == "ok":
logger.info(
"d_to_p_sync sid=%s rid=%s pushed=%d ingested_prefix=%d",
request.session_id,
request.request_id,
sync_result.get("bytes_pushed", 0),
sync_result.get("inserted_prefix_len", 0),
)
decode_session_newly_opened = False
try:
prefill_priority = _prefill_priority_for_router_request(