feat: add agentic pd hybrid benchmark prototype
This commit is contained in:
484
src/agentic_pd_hybrid/cli.py
Normal file
484
src/agentic_pd_hybrid/cli.py
Normal file
@@ -0,0 +1,484 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
from agentic_pd_hybrid.benchmark import BenchmarkConfig, run_live_benchmark
|
||||
from agentic_pd_hybrid.launcher import build_launch_plan
|
||||
from agentic_pd_hybrid.microbench import SmallAppendTraceConfig, write_small_append_trace
|
||||
from agentic_pd_hybrid.replay import ReplayConfig, replay_trace
|
||||
from agentic_pd_hybrid.sampling import SessionSampleConfig, sample_trace_sessions
|
||||
from agentic_pd_hybrid.trace_profiles import (
|
||||
NormalizeTraceLengthsConfig,
|
||||
normalize_trace_lengths,
|
||||
)
|
||||
from agentic_pd_hybrid.topology import build_single_node_topology
|
||||
|
||||
|
||||
def _normalize_mechanism_name(name: str) -> str:
|
||||
normalized = name.strip().lower()
|
||||
aliases = {
|
||||
"pd-disagg": "pd-disaggregation",
|
||||
"pd-disaggregation": "pd-disaggregation",
|
||||
"pd-hybrid": "pd-disaggregation",
|
||||
"baseline-pd-disagg": "pd-disaggregation",
|
||||
"pd-colo": "pd-colo",
|
||||
"direct-d": "pd-colo",
|
||||
"colocation": "pd-colo",
|
||||
"kvcache-centric": "kvcache-centric",
|
||||
"turn2+-direct-to-d": "kvcache-centric",
|
||||
"pd-with-d-append": "kvcache-centric",
|
||||
}
|
||||
if normalized not in aliases:
|
||||
raise ValueError(f"Unsupported mechanism: {name}")
|
||||
return aliases[normalized]
|
||||
|
||||
|
||||
def _parse_gpu_id_list(value: str | None) -> tuple[int, ...] | None:
|
||||
if value is None:
|
||||
return None
|
||||
items = [item.strip() for item in value.split(",") if item.strip()]
|
||||
if not items:
|
||||
return tuple()
|
||||
return tuple(int(item) for item in items)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Agentic PD hybrid prototype")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
print_launch = subparsers.add_parser(
|
||||
"print-launch",
|
||||
help="Print one-node SGLang PD launch commands",
|
||||
)
|
||||
_add_topology_arguments(print_launch)
|
||||
print_launch.add_argument("--prefill-policy", default="round_robin")
|
||||
print_launch.add_argument("--decode-policy", default="manual")
|
||||
|
||||
replay = subparsers.add_parser(
|
||||
"replay",
|
||||
help="Replay trace and log request-level metrics",
|
||||
)
|
||||
_add_topology_arguments(replay)
|
||||
replay.add_argument("--trace", type=Path, required=True)
|
||||
replay.add_argument("--output", type=Path, required=True)
|
||||
replay.add_argument(
|
||||
"--policy",
|
||||
choices=["default", "sticky", "kv-aware"],
|
||||
default="sticky",
|
||||
)
|
||||
replay.add_argument(
|
||||
"--mechanism",
|
||||
choices=[
|
||||
"pd-disaggregation",
|
||||
"pd-hybrid",
|
||||
"pd-disagg",
|
||||
"pd-colo",
|
||||
"direct-d",
|
||||
"kvcache-centric",
|
||||
"turn2+-direct-to-d",
|
||||
"pd-with-d-append",
|
||||
],
|
||||
default="pd-disaggregation",
|
||||
)
|
||||
replay.add_argument("--router-url")
|
||||
replay.add_argument("--model")
|
||||
replay.add_argument(
|
||||
"--header-mode",
|
||||
choices=["auto", "none", "routing-key", "target-worker"],
|
||||
default="auto",
|
||||
)
|
||||
replay.add_argument(
|
||||
"--request-limit",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Replay at most this many requests",
|
||||
)
|
||||
replay.add_argument(
|
||||
"--no-pace",
|
||||
action="store_true",
|
||||
help="Disable wall-clock pacing from trace timestamps",
|
||||
)
|
||||
replay.add_argument(
|
||||
"--time-scale",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help="Scale trace timing by this factor when pacing is enabled",
|
||||
)
|
||||
replay.add_argument(
|
||||
"--concurrency-limit",
|
||||
type=int,
|
||||
default=32,
|
||||
)
|
||||
replay.add_argument(
|
||||
"--timeout-s",
|
||||
type=float,
|
||||
default=600.0,
|
||||
)
|
||||
replay.add_argument(
|
||||
"--no-stream",
|
||||
action="store_true",
|
||||
help="Use non-streaming OpenAI responses for more robust E2E-only replay.",
|
||||
)
|
||||
replay.add_argument(
|
||||
"--stream-idle-timeout-s",
|
||||
type=float,
|
||||
default=900.0,
|
||||
help="Abort a streaming request if no SSE line arrives within this many seconds.",
|
||||
)
|
||||
replay.add_argument(
|
||||
"--kvcache-direct-max-uncached-tokens",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="For kvcache-centric routing, bypass P when the uncached suffix is at most this many tokens.",
|
||||
)
|
||||
replay.add_argument(
|
||||
"--kvcache-admission-mode",
|
||||
choices=["router", "worker"],
|
||||
default="router",
|
||||
help=(
|
||||
"For kvcache-centric routing, use router shadow-state admission "
|
||||
"or query the decode worker on the critical path."
|
||||
),
|
||||
)
|
||||
|
||||
sample = subparsers.add_parser(
|
||||
"sample-sessions",
|
||||
help="Sample a session-granularity trace shard for live benchmarking",
|
||||
)
|
||||
sample.add_argument("--trace", type=Path, required=True)
|
||||
sample.add_argument("--output", type=Path, required=True)
|
||||
sample.add_argument("--target-duration-s", type=float, default=600.0)
|
||||
sample.add_argument("--start-time-s", type=float, default=0.0)
|
||||
sample.add_argument("--session-sample-rate", type=float, default=0.01)
|
||||
sample.add_argument("--min-turns", type=int, default=1)
|
||||
sample.add_argument("--max-requests", type=int, default=None)
|
||||
sample.add_argument(
|
||||
"--profile",
|
||||
choices=["default", "small-append"],
|
||||
default="default",
|
||||
help="Optional workload-shape filter for live benchmarks.",
|
||||
)
|
||||
sample.add_argument("--min-initial-input-tokens", type=int, default=None)
|
||||
sample.add_argument("--max-initial-input-tokens", type=int, default=None)
|
||||
sample.add_argument("--max-append-input-tokens", type=int, default=None)
|
||||
sample.add_argument("--max-output-tokens", type=int, default=None)
|
||||
sample.add_argument("--min-overlap-ratio", type=float, default=None)
|
||||
|
||||
normalize = subparsers.add_parser(
|
||||
"normalize-trace-lengths",
|
||||
help="Rewrite a trace to a fixed turn1/append/output length profile",
|
||||
)
|
||||
normalize.add_argument("--trace", type=Path, required=True)
|
||||
normalize.add_argument("--output", type=Path, required=True)
|
||||
normalize.add_argument("--initial-input-length", type=int, default=10_000)
|
||||
normalize.add_argument("--append-input-length", type=int, default=1_000)
|
||||
normalize.add_argument("--output-length", type=int, default=1_000)
|
||||
normalize.add_argument("--max-requests", type=int, default=None)
|
||||
|
||||
micro = subparsers.add_parser(
|
||||
"make-small-append-trace",
|
||||
help="Generate a synthetic multi-turn trace with small turn2+ appends",
|
||||
)
|
||||
micro.add_argument("--output", type=Path, required=True)
|
||||
micro.add_argument("--session-count", type=int, default=8)
|
||||
micro.add_argument("--turns-per-session", type=int, default=3)
|
||||
micro.add_argument("--initial-input-length", type=int, default=10_000)
|
||||
micro.add_argument("--append-input-length", type=int, default=1_000)
|
||||
micro.add_argument("--output-length", type=int, default=1_000)
|
||||
micro.add_argument("--inter-turn-gap-s", type=float, default=1.0)
|
||||
micro.add_argument("--session-stagger-s", type=float, default=0.1)
|
||||
|
||||
benchmark = subparsers.add_parser(
|
||||
"benchmark-live",
|
||||
help="Launch a real PD stack, sample sessions, and collect live E2E numbers",
|
||||
)
|
||||
_add_topology_arguments(benchmark)
|
||||
benchmark.add_argument("--trace", type=Path, required=True)
|
||||
benchmark.add_argument(
|
||||
"--policy",
|
||||
choices=["default", "sticky", "kv-aware"],
|
||||
default="sticky",
|
||||
)
|
||||
benchmark.add_argument(
|
||||
"--mechanism",
|
||||
choices=[
|
||||
"pd-disaggregation",
|
||||
"pd-hybrid",
|
||||
"pd-disagg",
|
||||
"pd-colo",
|
||||
"direct-d",
|
||||
"kvcache-centric",
|
||||
"turn2+-direct-to-d",
|
||||
"pd-with-d-append",
|
||||
],
|
||||
default="pd-disaggregation",
|
||||
)
|
||||
benchmark.add_argument("--output-root", type=Path, default=Path("outputs/live"))
|
||||
benchmark.add_argument("--target-duration-s", type=float, default=600.0)
|
||||
benchmark.add_argument("--start-time-s", type=float, default=0.0)
|
||||
benchmark.add_argument("--session-sample-rate", type=float, default=0.01)
|
||||
benchmark.add_argument("--min-turns", type=int, default=1)
|
||||
benchmark.add_argument("--time-scale", type=float, default=1.0)
|
||||
benchmark.add_argument("--concurrency-limit", type=int, default=32)
|
||||
benchmark.add_argument("--timeout-s", type=float, default=1200.0)
|
||||
benchmark.add_argument(
|
||||
"--no-stream",
|
||||
action="store_true",
|
||||
help="Use non-streaming OpenAI responses for E2E-only live benchmarking.",
|
||||
)
|
||||
benchmark.add_argument(
|
||||
"--stream-idle-timeout-s",
|
||||
type=float,
|
||||
default=900.0,
|
||||
help="Abort a streaming request if no SSE line arrives within this many seconds.",
|
||||
)
|
||||
benchmark.add_argument(
|
||||
"--kvcache-direct-max-uncached-tokens",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="For kvcache-centric routing, bypass P when the uncached suffix is at most this many tokens.",
|
||||
)
|
||||
benchmark.add_argument(
|
||||
"--kvcache-admission-mode",
|
||||
choices=["router", "worker"],
|
||||
default="router",
|
||||
help=(
|
||||
"For kvcache-centric routing, use router shadow-state admission "
|
||||
"or query the decode worker on the critical path."
|
||||
),
|
||||
)
|
||||
benchmark.add_argument(
|
||||
"--sample-profile",
|
||||
choices=["default", "small-append"],
|
||||
default="default",
|
||||
help="Optional session-shape filter applied before live replay.",
|
||||
)
|
||||
benchmark.add_argument("--min-initial-input-tokens", type=int, default=None)
|
||||
benchmark.add_argument("--max-initial-input-tokens", type=int, default=None)
|
||||
benchmark.add_argument("--max-append-input-tokens", type=int, default=None)
|
||||
benchmark.add_argument("--max-output-tokens", type=int, default=None)
|
||||
benchmark.add_argument("--min-overlap-ratio", type=float, default=None)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "print-launch":
|
||||
topology = _topology_from_args(args)
|
||||
plan = build_launch_plan(
|
||||
topology,
|
||||
prefill_policy=args.prefill_policy,
|
||||
decode_policy=args.decode_policy,
|
||||
include_router=bool(topology.prefill_workers and topology.decode_workers),
|
||||
)
|
||||
print(plan.render())
|
||||
return
|
||||
|
||||
if args.command == "replay":
|
||||
topology = _topology_from_args(args)
|
||||
config = ReplayConfig(
|
||||
trace_path=args.trace,
|
||||
output_path=args.output,
|
||||
policy_name=args.policy,
|
||||
mechanism_name=_normalize_mechanism_name(args.mechanism),
|
||||
topology=topology,
|
||||
router_url=args.router_url,
|
||||
model_name=args.model,
|
||||
pace=not args.no_pace,
|
||||
time_scale=args.time_scale,
|
||||
request_limit=args.request_limit,
|
||||
concurrency_limit=args.concurrency_limit,
|
||||
header_mode=args.header_mode,
|
||||
timeout_s=args.timeout_s,
|
||||
stream=not args.no_stream,
|
||||
stream_idle_timeout_s=args.stream_idle_timeout_s,
|
||||
kvcache_direct_max_uncached_tokens=args.kvcache_direct_max_uncached_tokens,
|
||||
kvcache_admission_mode=args.kvcache_admission_mode,
|
||||
)
|
||||
results = asyncio.run(replay_trace(config))
|
||||
print(
|
||||
f"wrote {len(results)} request records to {args.output} and "
|
||||
f"{args.output}{'.summary.json'}"
|
||||
)
|
||||
return
|
||||
|
||||
if args.command == "sample-sessions":
|
||||
summary = sample_trace_sessions(
|
||||
SessionSampleConfig(
|
||||
trace_path=args.trace,
|
||||
output_path=args.output,
|
||||
target_duration_s=args.target_duration_s,
|
||||
start_time_s=args.start_time_s,
|
||||
session_sample_rate=args.session_sample_rate,
|
||||
min_turns=args.min_turns,
|
||||
max_requests=args.max_requests,
|
||||
profile=args.profile,
|
||||
min_initial_input_tokens=args.min_initial_input_tokens,
|
||||
max_initial_input_tokens=args.max_initial_input_tokens,
|
||||
max_append_input_tokens=args.max_append_input_tokens,
|
||||
max_output_tokens=args.max_output_tokens,
|
||||
min_overlap_ratio=args.min_overlap_ratio,
|
||||
)
|
||||
)
|
||||
print(
|
||||
f"wrote {summary.request_count} requests from {summary.session_count} sessions "
|
||||
f"covering {summary.sampled_duration_s:.3f}s to {args.output}"
|
||||
)
|
||||
return
|
||||
|
||||
if args.command == "normalize-trace-lengths":
|
||||
summary = normalize_trace_lengths(
|
||||
NormalizeTraceLengthsConfig(
|
||||
trace_path=args.trace,
|
||||
output_path=args.output,
|
||||
initial_input_length=args.initial_input_length,
|
||||
append_input_length=args.append_input_length,
|
||||
output_length=args.output_length,
|
||||
max_requests=args.max_requests,
|
||||
)
|
||||
)
|
||||
print(
|
||||
f"wrote {summary.request_count} normalized requests from "
|
||||
f"{summary.session_count} sessions to {args.output}"
|
||||
)
|
||||
return
|
||||
|
||||
if args.command == "make-small-append-trace":
|
||||
summary = write_small_append_trace(
|
||||
SmallAppendTraceConfig(
|
||||
output_path=args.output,
|
||||
session_count=args.session_count,
|
||||
turns_per_session=args.turns_per_session,
|
||||
initial_input_length=args.initial_input_length,
|
||||
append_input_length=args.append_input_length,
|
||||
output_length=args.output_length,
|
||||
inter_turn_gap_s=args.inter_turn_gap_s,
|
||||
session_stagger_s=args.session_stagger_s,
|
||||
)
|
||||
)
|
||||
print(
|
||||
f"wrote {summary.request_count} requests across {summary.session_count} sessions "
|
||||
f"to {args.output}"
|
||||
)
|
||||
return
|
||||
|
||||
if args.command == "benchmark-live":
|
||||
topology = _topology_from_args(args)
|
||||
artifacts = run_live_benchmark(
|
||||
BenchmarkConfig(
|
||||
trace_path=args.trace,
|
||||
output_root=args.output_root,
|
||||
topology=topology,
|
||||
policy_name=args.policy,
|
||||
mechanism_name=_normalize_mechanism_name(args.mechanism),
|
||||
target_duration_s=args.target_duration_s,
|
||||
start_time_s=args.start_time_s,
|
||||
session_sample_rate=args.session_sample_rate,
|
||||
min_turns=args.min_turns,
|
||||
time_scale=args.time_scale,
|
||||
concurrency_limit=args.concurrency_limit,
|
||||
timeout_s=args.timeout_s,
|
||||
stream=not args.no_stream,
|
||||
stream_idle_timeout_s=args.stream_idle_timeout_s,
|
||||
kvcache_direct_max_uncached_tokens=args.kvcache_direct_max_uncached_tokens,
|
||||
kvcache_admission_mode=args.kvcache_admission_mode,
|
||||
sample_profile=args.sample_profile,
|
||||
min_initial_input_tokens=args.min_initial_input_tokens,
|
||||
max_initial_input_tokens=args.max_initial_input_tokens,
|
||||
max_append_input_tokens=args.max_append_input_tokens,
|
||||
max_output_tokens=args.max_output_tokens,
|
||||
min_overlap_ratio=args.min_overlap_ratio,
|
||||
launch_stack=True,
|
||||
)
|
||||
)
|
||||
print(
|
||||
f"benchmark artifacts written under {artifacts.run_dir}; "
|
||||
f"metrics={artifacts.metrics_path} summary={artifacts.summary_path}"
|
||||
)
|
||||
return
|
||||
|
||||
raise AssertionError(f"Unhandled command: {args.command}")
|
||||
|
||||
|
||||
def _add_topology_arguments(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
default="~/models/Qwen/Qwen3-Coder-30B-A3B-Instruct",
|
||||
)
|
||||
parser.add_argument("--prefill-workers", type=int, default=1)
|
||||
parser.add_argument("--decode-workers", type=int, default=1)
|
||||
parser.add_argument("--direct-workers", type=int, default=0)
|
||||
parser.add_argument("--prefill-tp-size", type=int, default=1)
|
||||
parser.add_argument("--decode-tp-size", type=int, default=1)
|
||||
parser.add_argument("--direct-tp-size", type=int, default=1)
|
||||
parser.add_argument("--gpu-budget", type=int, default=8)
|
||||
parser.add_argument(
|
||||
"--prefill-gpu-ids",
|
||||
default=None,
|
||||
help="Comma-separated GPU IDs for prefill workers, e.g. 3,4",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--decode-gpu-ids",
|
||||
default=None,
|
||||
help="Comma-separated GPU IDs for decode workers, e.g. 5",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--direct-gpu-ids",
|
||||
default=None,
|
||||
help="Comma-separated GPU IDs for direct workers, e.g. 6",
|
||||
)
|
||||
parser.add_argument("--host", default="127.0.0.1")
|
||||
parser.add_argument("--router-port", type=int, default=8000)
|
||||
parser.add_argument("--prefill-port-base", type=int, default=30000)
|
||||
parser.add_argument("--decode-port-base", type=int, default=31000)
|
||||
parser.add_argument("--direct-port-base", type=int, default=32000)
|
||||
parser.add_argument("--bootstrap-port-base", type=int, default=8998)
|
||||
parser.add_argument("--transfer-backend", default="nixl")
|
||||
parser.add_argument(
|
||||
"--force-rdma",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Force real RDMA transport for PD KV transfer. "
|
||||
"Currently this requires Mooncake plus --ib-device."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--ib-device", default=None)
|
||||
parser.add_argument(
|
||||
"--no-trust-remote-code",
|
||||
action="store_true",
|
||||
)
|
||||
|
||||
|
||||
def _topology_from_args(args: argparse.Namespace):
|
||||
transfer_backend = args.transfer_backend
|
||||
if args.force_rdma:
|
||||
transfer_backend = "mooncake"
|
||||
|
||||
return build_single_node_topology(
|
||||
model_path=str(Path(args.model_path).expanduser()),
|
||||
prefill_worker_count=args.prefill_workers,
|
||||
decode_worker_count=args.decode_workers,
|
||||
direct_worker_count=args.direct_workers,
|
||||
prefill_tp_size=args.prefill_tp_size,
|
||||
decode_tp_size=args.decode_tp_size,
|
||||
direct_tp_size=args.direct_tp_size,
|
||||
prefill_gpu_ids=_parse_gpu_id_list(args.prefill_gpu_ids),
|
||||
decode_gpu_ids=_parse_gpu_id_list(args.decode_gpu_ids),
|
||||
direct_gpu_ids=_parse_gpu_id_list(args.direct_gpu_ids),
|
||||
total_gpu_budget=args.gpu_budget,
|
||||
host=args.host,
|
||||
router_port=args.router_port,
|
||||
prefill_port_base=args.prefill_port_base,
|
||||
decode_port_base=args.decode_port_base,
|
||||
direct_port_base=args.direct_port_base,
|
||||
bootstrap_port_base=args.bootstrap_port_base,
|
||||
transfer_backend=transfer_backend,
|
||||
force_rdma=args.force_rdma,
|
||||
trust_remote_code=not args.no_trust_remote_code,
|
||||
ib_device=args.ib_device,
|
||||
direct_extra_server_args=("--enable-streaming-session",),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user