c9d350b372
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>
301 lines
12 KiB
Python
301 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import signal
|
|
from dataclasses import asdict, dataclass, replace
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
|
|
from agentic_pd_hybrid.replay import ReplayConfig, replay_trace
|
|
from agentic_pd_hybrid.sampling import SessionSampleConfig, sample_trace_sessions
|
|
from agentic_pd_hybrid.stack import ManagedPdStack, launch_pd_stack
|
|
from agentic_pd_hybrid.topology import SingleNodeTopology
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BenchmarkConfig:
|
|
trace_path: Path
|
|
output_root: Path
|
|
topology: SingleNodeTopology
|
|
policy_name: str
|
|
mechanism_name: str = "pd-disaggregation"
|
|
target_duration_s: float = 600.0
|
|
start_time_s: float = 0.0
|
|
session_sample_rate: float = 0.01
|
|
min_turns: int = 1
|
|
time_scale: float = 1.0
|
|
concurrency_limit: int = 32
|
|
timeout_s: float = 1200.0
|
|
request_timeout_s: float | None = None
|
|
stream: bool = True
|
|
stream_idle_timeout_s: float | None = 900.0
|
|
kvcache_direct_max_uncached_tokens: int = 2048
|
|
kvcache_admission_mode: str = "router"
|
|
kvcache_seed_max_resident_tokens: int | None = None
|
|
kvcache_seed_max_output_tokens: int | None = None
|
|
kvcache_seed_min_turn_id: int = 1
|
|
kvcache_seed_only_multiturn_sessions: bool = False
|
|
kvcache_prefill_backup_policy: str = "release-after-transfer"
|
|
kvcache_seed_max_inflight_decode: int | None = 3
|
|
kvcache_seed_max_decode_transfer_queue_reqs: int | None = None
|
|
kvcache_direct_max_decode_transfer_queue_reqs: int | None = None
|
|
kvcache_prefill_priority_eviction: bool = False
|
|
kvcache_prefill_direct_priority: int = -100
|
|
kvcache_prefill_normal_priority: int = 100
|
|
sample_profile: str = "default"
|
|
min_initial_input_tokens: int | None = None
|
|
max_initial_input_tokens: int | None = None
|
|
max_append_input_tokens: int | None = None
|
|
max_output_tokens: int | None = None
|
|
min_overlap_ratio: float | None = None
|
|
launch_stack: bool = True
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BenchmarkArtifacts:
|
|
run_dir: Path
|
|
sampled_trace_path: Path
|
|
metrics_path: Path
|
|
summary_path: Path
|
|
benchmark_config_path: Path
|
|
|
|
|
|
def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
|
|
run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
|
run_label = f"{config.mechanism_name}-{config.policy_name}"
|
|
if config.mechanism_name == "kvcache-centric":
|
|
run_label = f"{run_label}-{config.kvcache_admission_mode}-admission"
|
|
run_dir = config.output_root / f"{run_label}-{run_id}"
|
|
run_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
topology = config.topology
|
|
if config.mechanism_name == "kvcache-centric":
|
|
prefill_extra_server_args = topology.prefill_extra_server_args + (
|
|
"--enable-streaming-session",
|
|
)
|
|
if config.kvcache_prefill_priority_eviction:
|
|
prefill_extra_server_args = prefill_extra_server_args + (
|
|
"--radix-eviction-policy",
|
|
"priority",
|
|
)
|
|
topology = replace(
|
|
topology,
|
|
prefill_extra_server_args=prefill_extra_server_args,
|
|
decode_extra_server_args=topology.decode_extra_server_args
|
|
+ (
|
|
"--enable-streaming-session",
|
|
"--disaggregation-decode-allow-local-prefill",
|
|
),
|
|
)
|
|
|
|
sampled_trace_path = run_dir / "sampled-trace.jsonl"
|
|
sample_summary = sample_trace_sessions(
|
|
SessionSampleConfig(
|
|
trace_path=config.trace_path,
|
|
output_path=sampled_trace_path,
|
|
target_duration_s=config.target_duration_s,
|
|
start_time_s=config.start_time_s,
|
|
session_sample_rate=config.session_sample_rate,
|
|
min_turns=config.min_turns,
|
|
profile=config.sample_profile, # type: ignore[arg-type]
|
|
min_initial_input_tokens=config.min_initial_input_tokens,
|
|
max_initial_input_tokens=config.max_initial_input_tokens,
|
|
max_append_input_tokens=config.max_append_input_tokens,
|
|
max_output_tokens=config.max_output_tokens,
|
|
min_overlap_ratio=config.min_overlap_ratio,
|
|
)
|
|
)
|
|
|
|
stack: ManagedPdStack | None = None
|
|
previous_sigint = signal.getsignal(signal.SIGINT)
|
|
previous_sigterm = signal.getsignal(signal.SIGTERM)
|
|
|
|
def _handle_termination(signum, _frame) -> None:
|
|
if stack is not None:
|
|
stack.stop()
|
|
raise SystemExit(128 + signum)
|
|
|
|
try:
|
|
signal.signal(signal.SIGINT, _handle_termination)
|
|
signal.signal(signal.SIGTERM, _handle_termination)
|
|
_mechanisms_with_router = {"pd-disaggregation", "kvcache-centric", "pd-colo"}
|
|
_naive_dp = config.mechanism_name == "pd-colo"
|
|
if config.launch_stack:
|
|
stack = launch_pd_stack(
|
|
topology=topology,
|
|
run_dir=run_dir,
|
|
prefill_policy="round_robin",
|
|
decode_policy=_decode_policy_for(config.policy_name),
|
|
timeout_s=config.timeout_s,
|
|
router_request_timeout_s=(
|
|
config.request_timeout_s
|
|
if config.request_timeout_s is not None
|
|
else config.timeout_s
|
|
),
|
|
include_router=(
|
|
config.mechanism_name in _mechanisms_with_router
|
|
),
|
|
naive_dp=_naive_dp,
|
|
)
|
|
router_url = (
|
|
stack.router_url
|
|
if config.mechanism_name in _mechanisms_with_router
|
|
else None
|
|
)
|
|
else:
|
|
router_url = (
|
|
topology.router_url
|
|
if config.mechanism_name in _mechanisms_with_router
|
|
else None
|
|
)
|
|
|
|
metrics_path = run_dir / "request-metrics.jsonl"
|
|
replay_config = ReplayConfig(
|
|
trace_path=sampled_trace_path,
|
|
output_path=metrics_path,
|
|
policy_name=config.policy_name,
|
|
mechanism_name=config.mechanism_name,
|
|
topology=topology,
|
|
router_url=router_url,
|
|
model_name=topology.model_name,
|
|
pace=True,
|
|
time_scale=config.time_scale,
|
|
request_limit=None,
|
|
concurrency_limit=config.concurrency_limit,
|
|
header_mode=_header_mode_for(config.policy_name),
|
|
timeout_s=config.timeout_s,
|
|
stream=config.stream,
|
|
stream_idle_timeout_s=config.stream_idle_timeout_s,
|
|
kvcache_direct_max_uncached_tokens=config.kvcache_direct_max_uncached_tokens,
|
|
kvcache_admission_mode=config.kvcache_admission_mode, # type: ignore[arg-type]
|
|
kvcache_seed_max_resident_tokens=config.kvcache_seed_max_resident_tokens,
|
|
kvcache_seed_max_output_tokens=config.kvcache_seed_max_output_tokens,
|
|
kvcache_seed_min_turn_id=config.kvcache_seed_min_turn_id,
|
|
kvcache_seed_only_multiturn_sessions=(
|
|
config.kvcache_seed_only_multiturn_sessions
|
|
),
|
|
kvcache_prefill_backup_policy=config.kvcache_prefill_backup_policy, # type: ignore[arg-type]
|
|
kvcache_seed_max_inflight_decode=(
|
|
config.kvcache_seed_max_inflight_decode
|
|
),
|
|
kvcache_seed_max_decode_transfer_queue_reqs=(
|
|
config.kvcache_seed_max_decode_transfer_queue_reqs
|
|
),
|
|
kvcache_direct_max_decode_transfer_queue_reqs=(
|
|
config.kvcache_direct_max_decode_transfer_queue_reqs
|
|
),
|
|
kvcache_prefill_priority_eviction=(
|
|
config.kvcache_prefill_priority_eviction
|
|
),
|
|
kvcache_prefill_direct_priority=config.kvcache_prefill_direct_priority,
|
|
kvcache_prefill_normal_priority=config.kvcache_prefill_normal_priority,
|
|
)
|
|
if config.request_timeout_s is not None:
|
|
replay_config = replace(
|
|
replay_config,
|
|
timeout_s=config.request_timeout_s,
|
|
)
|
|
asyncio.run(replay_trace(replay_config))
|
|
finally:
|
|
signal.signal(signal.SIGINT, previous_sigint)
|
|
signal.signal(signal.SIGTERM, previous_sigterm)
|
|
if stack is not None:
|
|
stack.stop()
|
|
|
|
benchmark_config_path = run_dir / "benchmark-config.json"
|
|
with benchmark_config_path.open("w", encoding="utf-8") as handle:
|
|
json.dump(
|
|
{
|
|
"policy_name": config.policy_name,
|
|
"mechanism_name": config.mechanism_name,
|
|
"target_duration_s": config.target_duration_s,
|
|
"start_time_s": config.start_time_s,
|
|
"session_sample_rate": config.session_sample_rate,
|
|
"min_turns": config.min_turns,
|
|
"time_scale": config.time_scale,
|
|
"concurrency_limit": config.concurrency_limit,
|
|
"timeout_s": config.timeout_s,
|
|
"request_timeout_s": config.request_timeout_s,
|
|
"stream": config.stream,
|
|
"stream_idle_timeout_s": config.stream_idle_timeout_s,
|
|
"kvcache_direct_max_uncached_tokens": config.kvcache_direct_max_uncached_tokens,
|
|
"kvcache_admission_mode": config.kvcache_admission_mode,
|
|
"kvcache_seed_max_resident_tokens": config.kvcache_seed_max_resident_tokens,
|
|
"kvcache_seed_max_output_tokens": config.kvcache_seed_max_output_tokens,
|
|
"kvcache_seed_min_turn_id": config.kvcache_seed_min_turn_id,
|
|
"kvcache_seed_only_multiturn_sessions": (
|
|
config.kvcache_seed_only_multiturn_sessions
|
|
),
|
|
"kvcache_prefill_backup_policy": config.kvcache_prefill_backup_policy,
|
|
"kvcache_seed_max_inflight_decode": (
|
|
config.kvcache_seed_max_inflight_decode
|
|
),
|
|
"kvcache_seed_max_decode_transfer_queue_reqs": (
|
|
config.kvcache_seed_max_decode_transfer_queue_reqs
|
|
),
|
|
"kvcache_direct_max_decode_transfer_queue_reqs": (
|
|
config.kvcache_direct_max_decode_transfer_queue_reqs
|
|
),
|
|
"kvcache_prefill_priority_eviction": (
|
|
config.kvcache_prefill_priority_eviction
|
|
),
|
|
"kvcache_prefill_direct_priority": (
|
|
config.kvcache_prefill_direct_priority
|
|
),
|
|
"kvcache_prefill_normal_priority": (
|
|
config.kvcache_prefill_normal_priority
|
|
),
|
|
"sample_profile": config.sample_profile,
|
|
"min_initial_input_tokens": config.min_initial_input_tokens,
|
|
"max_initial_input_tokens": config.max_initial_input_tokens,
|
|
"max_append_input_tokens": config.max_append_input_tokens,
|
|
"max_output_tokens": config.max_output_tokens,
|
|
"min_overlap_ratio": config.min_overlap_ratio,
|
|
"sample_summary": asdict(sample_summary),
|
|
"topology": {
|
|
"model_path": config.topology.model_path,
|
|
"router_url": topology.router_url,
|
|
"transfer_backend": topology.transfer_backend,
|
|
"force_rdma": topology.force_rdma,
|
|
"ib_device": topology.ib_device,
|
|
"prefill_workers": [
|
|
worker.worker_id for worker in topology.prefill_workers
|
|
],
|
|
"decode_workers": [
|
|
worker.worker_id for worker in topology.decode_workers
|
|
],
|
|
"direct_workers": [
|
|
worker.worker_id for worker in topology.direct_workers
|
|
],
|
|
},
|
|
},
|
|
handle,
|
|
indent=2,
|
|
sort_keys=True,
|
|
)
|
|
|
|
return BenchmarkArtifacts(
|
|
run_dir=run_dir,
|
|
sampled_trace_path=sampled_trace_path,
|
|
metrics_path=run_dir / "request-metrics.jsonl",
|
|
summary_path=run_dir / "request-metrics.jsonl.summary.json",
|
|
benchmark_config_path=benchmark_config_path,
|
|
)
|
|
|
|
|
|
def _decode_policy_for(policy_name: str) -> str:
|
|
if policy_name == "sticky":
|
|
return "manual"
|
|
if policy_name == "kv-aware":
|
|
return "consistent_hashing"
|
|
return "round_robin"
|
|
|
|
|
|
def _header_mode_for(policy_name: str) -> str:
|
|
if policy_name == "sticky":
|
|
return "routing-key"
|
|
if policy_name == "kv-aware":
|
|
return "target-worker"
|
|
return "none"
|