gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

docs: KVC v1-v4 debug journey + raise session soft_cap to 16

Browse Source 
Document the iterative debugging from v1 (broken KVC) through v4
(routing fixed + session cap raised), with code-level analysis of
the two main bugs encountered:

1. v2 root cause (mis-diagnosed previously as `allow_local_prefill`):
   `--policy default` for KVC mechanism caused replay's round-robin
   policy and the PD router's round-robin to diverge, sending requests
   with `session_params` to a D worker that did not have the session
   open. Resulted in 56-61% truncation with finish_reason
   "session id X does not exist".
   Fix: use `--policy kv-aware` (sweep_tp1_v3_kvaware.sh) so replay
   emits `x-smg-target-worker` and PD router uses consistent_hashing.

2. v3 new bottleneck: `pd-router-fallback-large-append-session-cap`
   dominated 52-65% of requests. Root cause was hardcoded
   `min(4, ...)` in `_decode_session_soft_cap`. With 7 D workers x 4
   sessions = 28 slots for 52 trace sessions, ~24 sessions starved
   permanently (bimodal direct-to-D rate of 0% or 99%).
   Fix: raise the cap to 16 (replay.py).

Also includes the v3 finding that direct-to-d-session path P50=0.495s
and TTFT P50=0.043s already beats the 8-way DP baseline (0.65s/0.093s)
- the KVC core mechanism works when fallback paths are avoided.

Files:
- docs/KVC_DEBUG_JOURNEY_V1_TO_V4.md: full journey + code location index
- docs/SWEBENCH_EXPERIMENT_{PROGRESS,RESULTS}.md: prior session notes
- scripts/sweep_tp1_v{2,3,4}*.sh: experiment driver scripts
- src/agentic_pd_hybrid/replay.py: cap 4 -> 16, audit fields
- src/agentic_pd_hybrid/pd_router.py: strip session_params from prefill
- src/agentic_pd_hybrid/metrics.py: truncated_request_count

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
kzlin
2026-04-28 21:10:41 +08:00
parent e9062b1d6e
commit c9d350b372
24 changed files with 1672 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

193
src/agentic_pd_hybrid/pd_router.py
View File

@@ -74,8 +74,58 @@ class RouterState:
return idx
@dataclass
class DpRouterConfig:
host: str
port: int
backend_urls: list[str]
backend_policy: str = "round_robin"
request_timeout_s: float = 1800.0
class DpRouterState:
"""DP (data-parallel) router: forward each request to exactly one backend."""
def __init__(self, config: DpRouterConfig):
if not config.backend_urls:
raise ValueError("At least one backend worker is required")
self.config = config
self.cursor = 0
self.sticky_map: dict[str, int] = {}
def select_backend(self, headers: dict[str, str]) -> str:
idx = self._select_index(headers)
return self.config.backend_urls[idx]
def _select_index(self, headers: dict[str, str]) -> int:
target_worker = headers.get("x-smg-target-worker")
routing_key = headers.get("x-smg-routing-key")
if (
self.config.backend_policy == "consistent_hashing"
and target_worker is not None
):
idx = int(target_worker)
if 0 <= idx < len(self.config.backend_urls):
return idx
if self.config.backend_policy == "manual" and routing_key:
cached = self.sticky_map.get(routing_key)
if cached is not None:
return cached
idx = self.cursor % len(self.config.backend_urls)
self.cursor += 1
self.sticky_map[routing_key] = idx
return idx
idx = self.cursor % len(self.config.backend_urls)
self.cursor += 1
return idx
app = FastAPI()
router_state: RouterState | None = None
dp_state: DpRouterState | None = None
@app.get("/health")
@@ -85,6 +135,16 @@ async def health() -> Response:
@app.get("/health_generate")
async def health_generate() -> Response:
if dp_state is not None:
async with aiohttp.ClientSession() as session:
tasks = [
session.get(f"{url}/health_generate")
for url in dp_state.config.backend_urls
]
for response in asyncio.as_completed(tasks):
async with await response:
pass
return Response(status_code=200)
state = _require_state()
async with aiohttp.ClientSession() as session:
tasks = []
@@ -101,6 +161,11 @@ async def health_generate() -> Response:
@app.get("/v1/models")
async def models() -> ORJSONResponse:
if dp_state is not None:
async with aiohttp.ClientSession() as session:
async with session.get(f"{dp_state.config.backend_urls[0]}/v1/models") as resp:
payload = await resp.json()
return ORJSONResponse(payload, status_code=resp.status)
state = _require_state()
async with aiohttp.ClientSession() as session:
async with session.get(f"{state.config.prefill_urls[0][0]}/v1/models") as response:
@@ -147,6 +212,15 @@ async def _forward_to_backend(
headers: dict[str, str],
endpoint_name: str,
) -> Response:
# DP mode: forward to a single backend
if dp_state is not None:
return await _forward_to_dp_backend(
request_data=request_data,
headers=headers,
endpoint_name=endpoint_name,
)
# PD mode: coordinate prefill + decode
state = _require_state()
prefill_server, bootstrap_port, decode_server = state.select_pair(headers)
prefill_request, decode_request = _build_backend_requests(
@@ -186,6 +260,63 @@ async def _forward_to_backend(
)
async def _forward_to_dp_backend(
*,
request_data: dict,
headers: dict[str, str],
endpoint_name: str,
) -> Response:
assert dp_state is not None
backend_server = dp_state.select_backend(headers)
cleaned = _strip_internal_fields(request_data)
timeout_s = dp_state.config.request_timeout_s
if request_data.get("stream", False):
return StreamingResponse(
_stream_dp_generate(
request_data=cleaned,
backend_server=backend_server,
endpoint_name=endpoint_name,
timeout_s=timeout_s,
),
media_type="text/event-stream",
)
async with aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=timeout_s)
) as session:
async with session.post(
f"{backend_server}/{endpoint_name}", json=cleaned
) as response:
body = await response.read()
return Response(
content=body,
status_code=response.status,
media_type=response.content_type,
)
async def _stream_dp_generate(
*,
request_data: dict,
backend_server: str,
endpoint_name: str,
timeout_s: float,
) -> AsyncIterator[bytes]:
async with aiohttp.ClientSession(
timeout=aiohttp.ClientTimeout(total=timeout_s)
) as session:
async with session.post(
f"{backend_server}/{endpoint_name}", json=request_data
) as response:
if response.status != HTTPStatus.OK:
payload = await response.read()
yield payload
return
async for chunk in response.content.iter_chunked(_STREAM_CHUNK_SIZE):
yield chunk
async def _stream_generate(
*,
prefill_request: dict,
@@ -241,6 +372,12 @@ def _build_backend_requests(
prefill_request.update(bootstrap_payload)
decode_request.update(bootstrap_payload)
# session_params is only meaningful for the decode worker (streaming session
# KV reuse). Sending it to the prefill worker causes the D side to
# short-circuit with local-prefill on already-open sessions, returning
# truncated responses while P's KV transfer gets aborted.
prefill_request.pop("session_params", None)
if prefill_priority is not None:
prefill_request["priority"] = int(prefill_priority)
if decode_priority is not None:
@@ -262,7 +399,7 @@ def _require_state() -> RouterState:
def main() -> None:
parser = argparse.ArgumentParser(description="Minimal local PD router")
parser = argparse.ArgumentParser(description="Minimal local PD / DP router")
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument(
@@ -270,30 +407,58 @@ def main() -> None:
nargs=2,
metavar=("URL", "BOOTSTRAP_PORT"),
action="append",
required=True,
default=None,
)
parser.add_argument(
"--decode",
action="append",
required=True,
default=None,
)
parser.add_argument("--prefill-policy", default="round_robin")
parser.add_argument("--decode-policy", default="manual")
parser.add_argument(
"--backend",
action="append",
default=None,
help="Backend URL for DP (data-parallel) mode. Repeat for each worker.",
)
parser.add_argument(
"--backend-policy",
default="round_robin",
help="Routing policy for DP mode: round_robin, manual, consistent_hashing.",
)
parser.add_argument("--request-timeout-s", type=float, default=1800.0)
args = parser.parse_args()
global router_state
router_state = RouterState(
RouterConfig(
host=args.host,
port=args.port,
prefill_urls=[(url, int(port)) for url, port in args.prefill],
decode_urls=list(args.decode),
prefill_policy=args.prefill_policy,
decode_policy=args.decode_policy,
request_timeout_s=args.request_timeout_s,
global router_state, dp_state
if args.backend:
# DP mode: simple forward to one of N backends
dp_state = DpRouterState(
DpRouterConfig(
host=args.host,
port=args.port,
backend_urls=list(args.backend),
backend_policy=args.backend_policy,
request_timeout_s=args.request_timeout_s,
)
)
)
elif args.prefill and args.decode:
# PD mode: prefill/decode coordination
router_state = RouterState(
RouterConfig(
host=args.host,
port=args.port,
prefill_urls=[(url, int(port)) for url, port in args.prefill],
decode_urls=list(args.decode),
prefill_policy=args.prefill_policy,
decode_policy=args.decode_policy,
request_timeout_s=args.request_timeout_s,
)
)
else:
parser.error("Either --backend (DP mode) or both --prefill and --decode (PD mode) are required")
uvicorn.run(app, host=args.host, port=args.port, log_level="info")