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"