feat: add agentic pd hybrid benchmark prototype
This commit is contained in:
221
src/agentic_pd_hybrid/benchmark.py
Normal file
221
src/agentic_pd_hybrid/benchmark.py
Normal file
@@ -0,0 +1,221 @@
|
||||
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
|
||||
stream: bool = True
|
||||
stream_idle_timeout_s: float | None = 900.0
|
||||
kvcache_direct_max_uncached_tokens: int = 2048
|
||||
kvcache_admission_mode: str = "router"
|
||||
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":
|
||||
topology = replace(
|
||||
topology,
|
||||
prefill_extra_server_args=topology.prefill_extra_server_args
|
||||
+ ("--enable-streaming-session",),
|
||||
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)
|
||||
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,
|
||||
include_router=(
|
||||
config.mechanism_name in {"pd-disaggregation", "kvcache-centric"}
|
||||
),
|
||||
)
|
||||
router_url = (
|
||||
stack.router_url
|
||||
if config.mechanism_name in {"pd-disaggregation", "kvcache-centric"}
|
||||
else None
|
||||
)
|
||||
else:
|
||||
router_url = (
|
||||
topology.router_url
|
||||
if config.mechanism_name in {"pd-disaggregation", "kvcache-centric"}
|
||||
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]
|
||||
)
|
||||
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,
|
||||
"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,
|
||||
"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"
|
||||
Reference in New Issue
Block a user