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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user