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

profile(kvc): add D KV pool timeseries poller + analyzer for v6 root-cause

Browse Source 
v5 dropped errors but pushed session-cap fallback to 46-51%. Before adding
v6 mitigations we need to attribute that capacity loss to one of:
  (a) active sessions — real footprint
  (b) idle-evictable sessions — LRU not aggressive enough
  (c) prefill backup blocks / in-flight / fragmentation — release timing

Without this it's all guessing. Plumb a 1Hz poller into replay that hits
each P/D worker's /server_info, captures session_cache + memory_usage, and
writes a per-worker time-series JSONL to <run_dir>/d-pool-timeseries.jsonl.
Off by default (--pool-poll-interval-s 0); v5+profile sweep enables it at
1.0s. Per-tick HTTP cost is ~8 parallel /server_info calls — negligible
relative to the 50min run.

Analyzer (scripts/analysis/analyze_pool_timeseries.py) decomposes each D's
capacity into active_held / idle_evictable / other (= cap-held-avail, the
backup-blocks bucket) / free, and reports session residency churn across
workers as a starvation/thrashing signal.

Mock-tested poller end-to-end (cancellation clean, file flushed, sessions
captured); analyzer validated against synthetic timeseries.

Next: run scripts/sweep_tp1_v5_optD_profile.sh on hardware (~90min), then
analyze results to pick a v6 direction.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
kzlin
2026-04-29 20:04:21 +08:00
parent 6572d7f3f4
commit 51f5386691
5 changed files with 587 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/agentic_pd_hybrid/benchmark.py
View File

@@ -43,6 +43,8 @@ class BenchmarkConfig:
kvcache_prefill_priority_eviction: bool = False
kvcache_prefill_direct_priority: int = -100
kvcache_prefill_normal_priority: int = 100
pool_poll_interval_s: float = 0.0
pool_poll_include_sessions: bool = True
sample_profile: str = "default"
min_initial_input_tokens: int | None = None
max_initial_input_tokens: int | None = None
@@ -190,6 +192,8 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
),
kvcache_prefill_direct_priority=config.kvcache_prefill_direct_priority,
kvcache_prefill_normal_priority=config.kvcache_prefill_normal_priority,
pool_poll_interval_s=config.pool_poll_interval_s,
pool_poll_include_sessions=config.pool_poll_include_sessions,
)
if config.request_timeout_s is not None:
replay_config = replace(
@@ -246,6 +250,8 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts:
"kvcache_prefill_normal_priority": (
config.kvcache_prefill_normal_priority
),
"pool_poll_interval_s": config.pool_poll_interval_s,
"pool_poll_include_sessions": config.pool_poll_include_sessions,
"sample_profile": config.sample_profile,
"min_initial_input_tokens": config.min_initial_input_tokens,
"max_initial_input_tokens": config.max_initial_input_tokens,

38
src/agentic_pd_hybrid/cli.py
View File

@@ -228,6 +228,23 @@ def main() -> None:
)
replay.add_argument("--kvcache-prefill-direct-priority", type=int, default=-100)
replay.add_argument("--kvcache-prefill-normal-priority", type=int, default=100)
replay.add_argument(
"--pool-poll-interval-s",
type=float,
default=0.0,
help=(
"Poll each P/D worker's /server_info every N seconds and write a "
"time-series snapshot to <run_dir>/d-pool-timeseries.jsonl. "
"0 disables polling."
),
)
replay.add_argument(
"--pool-poll-no-sessions",
action="store_true",
help=(
"Disable per-session detail in the pool timeseries (smaller files)."
),
)
sample = subparsers.add_parser(
"sample-sessions",
@@ -439,6 +456,23 @@ def main() -> None:
)
benchmark.add_argument("--kvcache-prefill-direct-priority", type=int, default=-100)
benchmark.add_argument("--kvcache-prefill-normal-priority", type=int, default=100)
benchmark.add_argument(
"--pool-poll-interval-s",
type=float,
default=0.0,
help=(
"Poll each P/D worker's /server_info every N seconds and write a "
"time-series snapshot to <run_dir>/d-pool-timeseries.jsonl. "
"0 disables polling."
),
)
benchmark.add_argument(
"--pool-poll-no-sessions",
action="store_true",
help=(
"Disable per-session detail in the pool timeseries (smaller files)."
),
)
benchmark.add_argument(
"--sample-profile",
choices=["default", "small-append"],
@@ -520,6 +554,8 @@ def main() -> None:
),
kvcache_prefill_direct_priority=args.kvcache_prefill_direct_priority,
kvcache_prefill_normal_priority=args.kvcache_prefill_normal_priority,
pool_poll_interval_s=args.pool_poll_interval_s,
pool_poll_include_sessions=not args.pool_poll_no_sessions,
)
results = asyncio.run(replay_trace(config))
print(
@@ -662,6 +698,8 @@ def main() -> None:
kvcache_prefill_normal_priority=(
args.kvcache_prefill_normal_priority
),
pool_poll_interval_s=args.pool_poll_interval_s,
pool_poll_include_sessions=not args.pool_poll_no_sessions,
sample_profile=args.sample_profile,
min_initial_input_tokens=args.min_initial_input_tokens,
max_initial_input_tokens=args.max_initial_input_tokens,

143
src/agentic_pd_hybrid/replay.py
View File

@@ -64,6 +64,8 @@ class ReplayConfig:
kvcache_prefill_priority_eviction: bool = False
kvcache_prefill_direct_priority: int = -100
kvcache_prefill_normal_priority: int = 100
pool_poll_interval_s: float = 0.0
pool_poll_include_sessions: bool = True
@dataclass
@@ -155,6 +157,25 @@ async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]:
client=client,
config=config,
)
poll_task: asyncio.Task[None] | None = None
if config.pool_poll_interval_s > 0:
poll_workers: list[tuple[str, str, str]] = []
for worker in config.topology.decode_workers:
poll_workers.append((worker.worker_id, "decode", worker.url))
for worker in config.topology.prefill_workers:
poll_workers.append((worker.worker_id, "prefill", worker.url))
if poll_workers:
poll_output = config.output_path.parent / "d-pool-timeseries.jsonl"
poll_task = asyncio.create_task(
_poll_pool_timeseries(
client=client,
workers=poll_workers,
interval_s=config.pool_poll_interval_s,
output_path=poll_output,
start_time=start_time,
include_sessions=config.pool_poll_include_sessions,
)
)
tasks = []
for request in requests:
if config.pace:
@@ -182,6 +203,12 @@ async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]:
session_tail_tasks[request.session_id] = tasks[-1]
results = await asyncio.gather(*tasks)
if poll_task is not None:
poll_task.cancel()
try:
await poll_task
except asyncio.CancelledError:
pass
for session in direct_sessions.values():
if session.opened:
try:
@@ -644,6 +671,122 @@ async def _fetch_decode_server_state(
)
async def _query_pool_snapshot(
*,
client: httpx.AsyncClient,
server_url: str,
include_sessions: bool,
) -> dict[str, Any]:
try:
response = await client.get(
f"{server_url.rstrip('/')}/server_info",
timeout=_ADMISSION_PROBE_TIMEOUT_S,
)
response.raise_for_status()
payload = response.json()
except Exception as exc:
return {"error": type(exc).__name__}
internal = _extract_internal_state(payload)
session_cache = _extract_session_cache(payload)
sessions: list[dict[str, Any]] = []
if include_sessions and isinstance(session_cache.get("sessions"), list):
for entry in session_cache["sessions"]:
if not isinstance(entry, dict):
continue
sessions.append(
{
"session_id": entry.get("session_id"),
"resident": bool(entry.get("resident")),
"resident_tokens": int(entry.get("resident_tokens") or 0),
"idle_evictable": bool(entry.get("idle_evictable")),
"timed_out": bool(entry.get("timed_out")),
}
)
memory_usage = internal.get("memory_usage") if isinstance(internal, dict) else None
if not isinstance(memory_usage, dict):
memory_usage = {}
return {
"session_cache_enabled": bool(session_cache.get("enabled")),
"session_count": int(session_cache.get("session_count") or 0),
"resident_session_count": int(session_cache.get("resident_session_count") or 0),
"held_tokens": int(session_cache.get("held_tokens") or 0),
"available_tokens": int(session_cache.get("available_tokens") or 0),
"capacity_tokens": int(session_cache.get("capacity_tokens") or 0),
"idle_evictable_session_count": int(
session_cache.get("idle_evictable_session_count") or 0
),
"idle_evictable_tokens": int(session_cache.get("idle_evictable_tokens") or 0),
"kvcache_mem_gb": float(memory_usage.get("kvcache") or 0.0),
"token_capacity": int(memory_usage.get("token_capacity") or 0),
"max_total_num_tokens": int(internal.get("max_total_num_tokens") or 0)
if isinstance(internal, dict)
else 0,
"last_gen_throughput": float(internal.get("last_gen_throughput") or 0.0)
if isinstance(internal, dict)
else 0.0,
"sessions": sessions,
}
async def _poll_pool_timeseries(
*,
client: httpx.AsyncClient,
workers: list[tuple[str, str, str]],
interval_s: float,
output_path: Path,
start_time: float,
include_sessions: bool,
) -> None:
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w", encoding="utf-8") as handle:
try:
while True:
tick_started = time.perf_counter()
ts = time.time()
wall_s = tick_started - start_time
snapshots = await asyncio.gather(
*(
_query_pool_snapshot(
client=client,
server_url=url,
include_sessions=include_sessions,
)
for _, _, url in workers
),
return_exceptions=True,
)
for (worker_id, role, url), snap in zip(workers, snapshots):
if isinstance(snap, BaseException):
row: dict[str, Any] = {
"ts": ts,
"wall_s": wall_s,
"worker_id": worker_id,
"worker_role": role,
"worker_url": url,
"error": type(snap).__name__,
}
else:
row = {
"ts": ts,
"wall_s": wall_s,
"worker_id": worker_id,
"worker_role": role,
"worker_url": url,
**snap,
}
handle.write(json.dumps(row, sort_keys=True) + "\n")
handle.flush()
elapsed = time.perf_counter() - tick_started
sleep_s = interval_s - elapsed
if sleep_s > 0:
await asyncio.sleep(sleep_s)
except asyncio.CancelledError:
return
async def _query_decode_direct_admission(
*,
client: httpx.AsyncClient,