feat(kvc): add real Ali replay workflow

This commit is contained in:
2026-05-12 05:28:00 +00:00
parent 1d51704dad
commit 4e8f943875
5 changed files with 901 additions and 21 deletions

View File

@@ -2,6 +2,7 @@ from __future__ import annotations
import argparse
import asyncio
import shlex
from pathlib import Path
from agentic_pd_hybrid.benchmark import BenchmarkConfig, run_live_benchmark
@@ -260,6 +261,15 @@ def main() -> None:
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",
@@ -501,6 +511,15 @@ def main() -> None:
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"],
@@ -512,6 +531,14 @@ def main() -> 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()
@@ -586,6 +613,7 @@ def main() -> None:
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(
@@ -732,12 +760,14 @@ def main() -> None:
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,
)
)
@@ -797,6 +827,26 @@ def _add_topology_arguments(parser: argparse.ArgumentParser) -> None:
"--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):
@@ -826,7 +876,13 @@ def _topology_from_args(args: argparse.Namespace):
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",),
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)),
),
)