891 lines
34 KiB
Python
891 lines
34 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import shlex
|
|
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.profile import ProfileConfig, print_profile_summary, write_profile
|
|
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."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--kvcache-seed-max-resident-tokens",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For kvcache-centric routing, do not seed/reseed a decode session "
|
|
"when input+output tokens exceed this value."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--kvcache-seed-max-output-tokens",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For kvcache-centric routing, do not seed/reseed a decode session "
|
|
"when output tokens exceed this value."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--kvcache-seed-min-turn-id",
|
|
type=int,
|
|
default=1,
|
|
help=(
|
|
"For kvcache-centric routing, do not seed/reseed a decode session "
|
|
"before this turn id."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--kvcache-seed-only-multiturn-sessions",
|
|
action="store_true",
|
|
help=(
|
|
"Oracle ablation for kvcache-centric routing: only seed sessions "
|
|
"that have more than one turn in the replay trace."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--kvcache-prefill-backup-policy",
|
|
choices=["release-after-transfer", "capacity-backup"],
|
|
default="release-after-transfer",
|
|
help=(
|
|
"For kvcache-centric seed/reseed, release the P-side session after "
|
|
"P->D transfer or keep a capacity-limited P-side backup."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--kvcache-seed-max-inflight-decode",
|
|
type=int,
|
|
default=3,
|
|
help=(
|
|
"For kvcache-centric routing, skip seed/reseed when the router "
|
|
"shadow inflight decode load assigned to the target D exceeds this value. "
|
|
"Use a negative value to disable this filter."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--kvcache-seed-max-decode-transfer-queue-reqs",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For kvcache-centric worker admission, skip seed/reseed when the "
|
|
"target D reports more transfer-queue requests than this value. "
|
|
"Use a negative value to disable this filter."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--kvcache-direct-max-decode-transfer-queue-reqs",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For kvcache-centric worker admission, skip direct-to-D append when "
|
|
"the target D reports more transfer-queue requests than this value. "
|
|
"Use a negative value to disable this filter."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--kvcache-prefill-priority-eviction",
|
|
action="store_true",
|
|
help=(
|
|
"For kvcache-centric routing, mark P-side prefixes predicted to move "
|
|
"direct-to-D as lower priority. Requires P workers to use "
|
|
"--radix-eviction-policy priority."
|
|
),
|
|
)
|
|
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)."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--enable-backpressure",
|
|
action="store_true",
|
|
help=(
|
|
"Honor recommended_pause_ms hints from D's admission endpoint. "
|
|
"When set, replay sleeps before issuing requests to a saturated D. "
|
|
"Default off — preserves baseline behavior."
|
|
),
|
|
)
|
|
replay.add_argument(
|
|
"--backpressure-max-pause-s",
|
|
type=float,
|
|
default=2.0,
|
|
help="Cap on per-request backpressure sleep, regardless of D hint.",
|
|
)
|
|
replay.add_argument(
|
|
"--progress-interval-s",
|
|
type=float,
|
|
default=30.0,
|
|
help=(
|
|
"Write client-side replay progress to <output_dir>/replay-progress.jsonl "
|
|
"every N seconds. 0 disables the heartbeat."
|
|
),
|
|
)
|
|
|
|
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)
|
|
|
|
profile = subparsers.add_parser(
|
|
"profile",
|
|
help="Profile a trace and optional request metrics for routing analysis",
|
|
)
|
|
profile.add_argument("--trace", type=Path, required=True)
|
|
profile.add_argument("--output", type=Path, default=None)
|
|
profile.add_argument("--metrics", type=Path, default=None)
|
|
profile.add_argument("--baseline-metrics", type=Path, default=None)
|
|
profile.add_argument("--candidate-metrics", type=Path, default=None)
|
|
profile.add_argument("--direct-max-uncached-tokens", type=int, default=2048)
|
|
|
|
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(
|
|
"--request-timeout-s",
|
|
type=float,
|
|
default=None,
|
|
help=(
|
|
"Per-request replay/router timeout. If unset, --timeout-s is used. "
|
|
"--timeout-s still controls stack startup."
|
|
),
|
|
)
|
|
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(
|
|
"--kvcache-seed-max-resident-tokens",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For kvcache-centric routing, do not seed/reseed a decode session "
|
|
"when input+output tokens exceed this value."
|
|
),
|
|
)
|
|
benchmark.add_argument(
|
|
"--kvcache-seed-max-output-tokens",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For kvcache-centric routing, do not seed/reseed a decode session "
|
|
"when output tokens exceed this value."
|
|
),
|
|
)
|
|
benchmark.add_argument(
|
|
"--kvcache-seed-min-turn-id",
|
|
type=int,
|
|
default=1,
|
|
help=(
|
|
"For kvcache-centric routing, do not seed/reseed a decode session "
|
|
"before this turn id."
|
|
),
|
|
)
|
|
benchmark.add_argument(
|
|
"--kvcache-seed-only-multiturn-sessions",
|
|
action="store_true",
|
|
help=(
|
|
"Oracle ablation for kvcache-centric routing: only seed sessions "
|
|
"that have more than one turn in the replay trace."
|
|
),
|
|
)
|
|
benchmark.add_argument(
|
|
"--kvcache-prefill-backup-policy",
|
|
choices=["release-after-transfer", "capacity-backup"],
|
|
default="release-after-transfer",
|
|
help=(
|
|
"For kvcache-centric seed/reseed, release the P-side session after "
|
|
"P->D transfer or keep a capacity-limited P-side backup."
|
|
),
|
|
)
|
|
benchmark.add_argument(
|
|
"--kvcache-seed-max-inflight-decode",
|
|
type=int,
|
|
default=3,
|
|
help=(
|
|
"For kvcache-centric routing, skip seed/reseed when the router "
|
|
"shadow inflight decode load assigned to the target D exceeds this value. "
|
|
"Use a negative value to disable this filter."
|
|
),
|
|
)
|
|
benchmark.add_argument(
|
|
"--kvcache-seed-max-decode-transfer-queue-reqs",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For kvcache-centric worker admission, skip seed/reseed when the "
|
|
"target D reports more transfer-queue requests than this value. "
|
|
"Use a negative value to disable this filter."
|
|
),
|
|
)
|
|
benchmark.add_argument(
|
|
"--kvcache-direct-max-decode-transfer-queue-reqs",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"For kvcache-centric worker admission, skip direct-to-D append when "
|
|
"the target D reports more transfer-queue requests than this value. "
|
|
"Use a negative value to disable this filter."
|
|
),
|
|
)
|
|
benchmark.add_argument(
|
|
"--kvcache-prefill-priority-eviction",
|
|
action="store_true",
|
|
help=(
|
|
"For kvcache-centric benchmark-live, launch P workers with priority "
|
|
"radix eviction and mark direct-to-D predicted prefixes lower priority."
|
|
),
|
|
)
|
|
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(
|
|
"--enable-backpressure",
|
|
action="store_true",
|
|
help=(
|
|
"Honor recommended_pause_ms hints from D's admission endpoint."
|
|
),
|
|
)
|
|
benchmark.add_argument(
|
|
"--backpressure-max-pause-s",
|
|
type=float,
|
|
default=2.0,
|
|
help="Cap on per-request backpressure sleep, regardless of D hint.",
|
|
)
|
|
benchmark.add_argument(
|
|
"--progress-interval-s",
|
|
type=float,
|
|
default=30.0,
|
|
help=(
|
|
"Write client-side replay progress to <run_dir>/replay-progress.jsonl "
|
|
"every N seconds. 0 disables the heartbeat."
|
|
),
|
|
)
|
|
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)
|
|
benchmark.add_argument(
|
|
"--use-trace-as-sample",
|
|
action="store_true",
|
|
help=(
|
|
"Replay the provided --trace exactly instead of sampling sessions into "
|
|
"a new trace. Use this for prebuilt real-workload samples."
|
|
),
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.command == "print-launch":
|
|
topology = _topology_from_args(args)
|
|
has_pd = bool(topology.prefill_workers and topology.decode_workers)
|
|
has_direct_only = bool(
|
|
topology.direct_workers
|
|
and not topology.prefill_workers
|
|
and not topology.decode_workers
|
|
)
|
|
plan = build_launch_plan(
|
|
topology,
|
|
prefill_policy=args.prefill_policy,
|
|
decode_policy=args.decode_policy,
|
|
include_router=has_pd or has_direct_only,
|
|
naive_dp=has_direct_only,
|
|
)
|
|
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,
|
|
kvcache_seed_max_resident_tokens=args.kvcache_seed_max_resident_tokens,
|
|
kvcache_seed_max_output_tokens=args.kvcache_seed_max_output_tokens,
|
|
kvcache_seed_min_turn_id=args.kvcache_seed_min_turn_id,
|
|
kvcache_seed_only_multiturn_sessions=(
|
|
args.kvcache_seed_only_multiturn_sessions
|
|
),
|
|
kvcache_prefill_backup_policy=args.kvcache_prefill_backup_policy,
|
|
kvcache_seed_max_inflight_decode=(
|
|
None
|
|
if args.kvcache_seed_max_inflight_decode < 0
|
|
else args.kvcache_seed_max_inflight_decode
|
|
),
|
|
kvcache_seed_max_decode_transfer_queue_reqs=(
|
|
None
|
|
if args.kvcache_seed_max_decode_transfer_queue_reqs is None
|
|
or args.kvcache_seed_max_decode_transfer_queue_reqs < 0
|
|
else args.kvcache_seed_max_decode_transfer_queue_reqs
|
|
),
|
|
kvcache_direct_max_decode_transfer_queue_reqs=(
|
|
None
|
|
if args.kvcache_direct_max_decode_transfer_queue_reqs is None
|
|
or args.kvcache_direct_max_decode_transfer_queue_reqs < 0
|
|
else args.kvcache_direct_max_decode_transfer_queue_reqs
|
|
),
|
|
kvcache_prefill_priority_eviction=(
|
|
args.kvcache_prefill_priority_eviction
|
|
),
|
|
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,
|
|
enable_backpressure=args.enable_backpressure,
|
|
backpressure_max_pause_s=args.backpressure_max_pause_s,
|
|
progress_interval_s=args.progress_interval_s,
|
|
)
|
|
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 == "profile":
|
|
if (args.baseline_metrics is None) != (args.candidate_metrics is None):
|
|
raise ValueError(
|
|
"--baseline-metrics and --candidate-metrics must be provided together"
|
|
)
|
|
report = write_profile(
|
|
ProfileConfig(
|
|
trace_path=args.trace,
|
|
output_path=args.output,
|
|
metrics_path=args.metrics,
|
|
baseline_metrics_path=args.baseline_metrics,
|
|
candidate_metrics_path=args.candidate_metrics,
|
|
direct_max_uncached_tokens=args.direct_max_uncached_tokens,
|
|
)
|
|
)
|
|
print_profile_summary(report)
|
|
if args.output is not None:
|
|
print(f"wrote profile to {args.output}")
|
|
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,
|
|
request_timeout_s=args.request_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,
|
|
kvcache_seed_max_resident_tokens=args.kvcache_seed_max_resident_tokens,
|
|
kvcache_seed_max_output_tokens=args.kvcache_seed_max_output_tokens,
|
|
kvcache_seed_min_turn_id=args.kvcache_seed_min_turn_id,
|
|
kvcache_seed_only_multiturn_sessions=(
|
|
args.kvcache_seed_only_multiturn_sessions
|
|
),
|
|
kvcache_prefill_backup_policy=args.kvcache_prefill_backup_policy,
|
|
kvcache_seed_max_inflight_decode=(
|
|
None
|
|
if args.kvcache_seed_max_inflight_decode < 0
|
|
else args.kvcache_seed_max_inflight_decode
|
|
),
|
|
kvcache_seed_max_decode_transfer_queue_reqs=(
|
|
None
|
|
if args.kvcache_seed_max_decode_transfer_queue_reqs is None
|
|
or args.kvcache_seed_max_decode_transfer_queue_reqs < 0
|
|
else args.kvcache_seed_max_decode_transfer_queue_reqs
|
|
),
|
|
kvcache_direct_max_decode_transfer_queue_reqs=(
|
|
None
|
|
if args.kvcache_direct_max_decode_transfer_queue_reqs is None
|
|
or args.kvcache_direct_max_decode_transfer_queue_reqs < 0
|
|
else args.kvcache_direct_max_decode_transfer_queue_reqs
|
|
),
|
|
kvcache_prefill_priority_eviction=(
|
|
args.kvcache_prefill_priority_eviction
|
|
),
|
|
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,
|
|
enable_backpressure=args.enable_backpressure,
|
|
backpressure_max_pause_s=args.backpressure_max_pause_s,
|
|
progress_interval_s=args.progress_interval_s,
|
|
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,
|
|
use_trace_as_sample=args.use_trace_as_sample,
|
|
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",
|
|
)
|
|
parser.add_argument(
|
|
"--extra-server-args",
|
|
default="",
|
|
help="Extra arguments appended to every sglang.launch_server command.",
|
|
)
|
|
parser.add_argument(
|
|
"--prefill-extra-server-args",
|
|
default="",
|
|
help="Extra arguments appended only to prefill launch_server commands.",
|
|
)
|
|
parser.add_argument(
|
|
"--decode-extra-server-args",
|
|
default="",
|
|
help="Extra arguments appended only to decode launch_server commands.",
|
|
)
|
|
parser.add_argument(
|
|
"--direct-extra-server-args",
|
|
default="",
|
|
help="Extra arguments appended only to direct launch_server commands.",
|
|
)
|
|
|
|
|
|
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,
|
|
extra_server_args=tuple(shlex.split(args.extra_server_args)),
|
|
prefill_extra_server_args=tuple(shlex.split(args.prefill_extra_server_args)),
|
|
decode_extra_server_args=tuple(shlex.split(args.decode_extra_server_args)),
|
|
direct_extra_server_args=(
|
|
"--enable-streaming-session",
|
|
*tuple(shlex.split(args.direct_extra_server_args)),
|
|
),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|