From 4e8f9438758709b7aa0e16f97f6bd80aff6bd22b Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Tue, 12 May 2026 05:28:00 +0000 Subject: [PATCH] feat(kvc): add real Ali replay workflow --- scripts/prepare_real_ali_samples.py | 450 ++++++++++++++++++++++++++++ scripts/sweep_real_ali_kvc.sh | 170 +++++++++++ src/agentic_pd_hybrid/benchmark.py | 86 +++++- src/agentic_pd_hybrid/cli.py | 58 +++- src/agentic_pd_hybrid/replay.py | 158 +++++++++- 5 files changed, 901 insertions(+), 21 deletions(-) create mode 100755 scripts/prepare_real_ali_samples.py create mode 100755 scripts/sweep_real_ali_kvc.sh diff --git a/scripts/prepare_real_ali_samples.py b/scripts/prepare_real_ali_samples.py new file mode 100755 index 0000000..e1642f8 --- /dev/null +++ b/scripts/prepare_real_ali_samples.py @@ -0,0 +1,450 @@ +#!/usr/bin/env python3 +"""Prepare balanced real-Ali trace samples for KVC experiments. + +The generic sampler is duration-oriented and can be dominated by one long +session. This script keeps real request lengths/timestamps but caps turns per +session so live sweeps can compare policies on a repeatable multi-session +workload. +""" +from __future__ import annotations + +import argparse +import json +import statistics +from collections import defaultdict +from dataclasses import asdict, dataclass +from pathlib import Path + +from agentic_pd_hybrid.trace import TraceRequest, load_trace + + +@dataclass(frozen=True) +class SampleSummary: + input_trace_path: str + output_trace_path: str + profile: str + request_count: int + session_count: int + multi_turn_session_count: int + turn2plus_count: int + direct_eligible_turn2plus_count: int + direct_eligible_turn2plus_ratio: float + missing_parent_count: int + max_sessions: int + max_turns_per_session: int + start_time_s: float + end_time_s: float + sampled_duration_s: float + rebased_timestamps: bool + input_tokens: dict[str, float] | None + output_tokens: dict[str, float] | None + append_tokens: dict[str, float] | None + inter_turn_gap_s: dict[str, float] | None + overlap_ratio: dict[str, float] | None + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--trace", type=Path, required=True) + parser.add_argument("--output-root", type=Path, required=True) + parser.add_argument("--max-sessions", type=int, default=64) + parser.add_argument("--max-turns-per-session", type=int, default=12) + parser.add_argument("--start-time-s", type=float, default=0.0) + parser.add_argument( + "--window-duration-s", + type=float, + default=None, + help=( + "If set, also write continuous-window samples that keep only requests " + "inside [start-time, start-time + window-duration]." + ), + ) + parser.add_argument( + "--window-target-requests", + type=int, + default=None, + help=( + "For continuous-window samples, select whole sessions across time " + "buckets until at least this many requests are included. This keeps " + "the window span while making live runs tractable." + ), + ) + parser.add_argument( + "--window-buckets", + type=int, + default=15, + help="Number of time buckets used with --window-target-requests.", + ) + parser.add_argument( + "--window-min-turns", + type=int, + default=1, + help=( + "Minimum number of in-window turns per selected session for " + "continuous-window samples." + ), + ) + parser.add_argument( + "--window-output-name", + default="ali-window.jsonl", + help="Output filename for the continuous-window sample.", + ) + parser.add_argument( + "--max-sampled-duration-s", + type=float, + default=None, + help=( + "For balanced profile samples, drop requests after the first selected " + "timestamp plus this duration. Use only for quick smoke runs; headline " + "runs should preserve the full sampled span." + ), + ) + parser.add_argument( + "--profiles", + nargs="+", + default=["representative-mt", "kvc-fit-smallappend"], + choices=["representative-mt", "kvc-fit-smallappend"], + ) + parser.add_argument( + "--no-rebase-timestamps", + action="store_true", + help="Keep original timestamps instead of shifting the sample to start at 0.", + ) + args = parser.parse_args() + + requests = load_trace(args.trace) + sessions: dict[str, list[TraceRequest]] = defaultdict(list) + for request in requests: + sessions[request.session_id].append(request) + + args.output_root.mkdir(parents=True, exist_ok=True) + if args.window_duration_s is not None: + if args.window_target_requests is None: + selected = _select_window( + requests=requests, + start_time_s=args.start_time_s, + window_duration_s=args.window_duration_s, + ) + profile = "window" + else: + selected = _select_window_session_sample( + sessions=sessions, + start_time_s=args.start_time_s, + window_duration_s=args.window_duration_s, + target_requests=args.window_target_requests, + bucket_count=args.window_buckets, + min_turns=args.window_min_turns, + ) + profile = ( + "window-session-sample" + if args.window_min_turns <= 1 + else f"window-session-sample-min{args.window_min_turns}turns" + ) + output_path = args.output_root / args.window_output_name + summary = _write_sample( + selected=selected, + input_trace_path=args.trace, + output_path=output_path, + profile=profile, + max_sessions=args.max_sessions, + max_turns_per_session=args.max_turns_per_session, + rebase_timestamps=not args.no_rebase_timestamps, + ) + print( + f"window: wrote {summary.request_count} requests from " + f"{summary.session_count} sessions to {output_path}" + ) + + for profile in args.profiles: + selected = _select_profile( + sessions=sessions, + profile=profile, + start_time_s=args.start_time_s, + max_sessions=args.max_sessions, + max_turns_per_session=args.max_turns_per_session, + max_sampled_duration_s=args.max_sampled_duration_s, + ) + output_path = args.output_root / f"ali-{profile}.jsonl" + summary = _write_sample( + selected=selected, + input_trace_path=args.trace, + output_path=output_path, + profile=profile, + max_sessions=args.max_sessions, + max_turns_per_session=args.max_turns_per_session, + rebase_timestamps=not args.no_rebase_timestamps, + ) + print( + f"{profile}: wrote {summary.request_count} requests from " + f"{summary.session_count} sessions to {output_path}" + ) + + +def _select_profile( + *, + sessions: dict[str, list[TraceRequest]], + profile: str, + start_time_s: float, + max_sessions: int, + max_turns_per_session: int, + max_sampled_duration_s: float | None, +) -> list[TraceRequest]: + eligible: list[list[TraceRequest]] = [] + for session_requests in sessions.values(): + ordered = _ordered(session_requests) + if len(ordered) < 2: + continue + if ordered[0].timestamp_s < start_time_s: + continue + if profile == "kvc-fit-smallappend" and not _is_kvc_fit_smallappend(ordered): + continue + eligible.append(ordered[:max_turns_per_session]) + + eligible.sort(key=lambda items: (items[0].timestamp_s, items[0].session_id)) + selected_sessions = eligible[:max_sessions] + selected = [request for items in selected_sessions for request in items] + selected.sort(key=lambda request: (request.timestamp_s, request.chat_id)) + if selected and max_sampled_duration_s is not None: + first_ts = selected[0].timestamp_s + end_ts = first_ts + max_sampled_duration_s + selected = [ + request for request in selected if request.timestamp_s <= end_ts + ] + return selected + + +def _select_window( + *, + requests: list[TraceRequest], + start_time_s: float, + window_duration_s: float, +) -> list[TraceRequest]: + end_time_s = start_time_s + window_duration_s + selected = [ + request + for request in requests + if start_time_s <= request.timestamp_s <= end_time_s + ] + selected.sort(key=lambda request: (request.timestamp_s, request.chat_id)) + return selected + + +def _select_window_session_sample( + *, + sessions: dict[str, list[TraceRequest]], + start_time_s: float, + window_duration_s: float, + target_requests: int, + bucket_count: int, + min_turns: int, +) -> list[TraceRequest]: + if target_requests <= 0: + raise ValueError("--window-target-requests must be positive") + if bucket_count <= 0: + raise ValueError("--window-buckets must be positive") + if min_turns <= 0: + raise ValueError("--window-min-turns must be positive") + + end_time_s = start_time_s + window_duration_s + bucket_width_s = window_duration_s / bucket_count + buckets: list[list[list[TraceRequest]]] = [[] for _ in range(bucket_count)] + for session_requests in sessions.values(): + ordered = _ordered(session_requests) + if not ordered: + continue + first = ordered[0] + if first.timestamp_s < start_time_s or first.timestamp_s > end_time_s: + continue + in_window = [ + request + for request in ordered + if start_time_s <= request.timestamp_s <= end_time_s + ] + if len(in_window) < min_turns: + continue + bucket_index = min( + bucket_count - 1, + int((first.timestamp_s - start_time_s) / bucket_width_s), + ) + buckets[bucket_index].append(in_window) + + for bucket in buckets: + bucket.sort(key=lambda items: (items[0].timestamp_s, items[0].session_id)) + + selected_sessions: list[list[TraceRequest]] = [] + selected_count = 0 + positions = [0 for _ in range(bucket_count)] + while selected_count < target_requests: + progressed = False + for index, bucket in enumerate(buckets): + if positions[index] >= len(bucket): + continue + session_requests = bucket[positions[index]] + positions[index] += 1 + selected_sessions.append(session_requests) + selected_count += len(session_requests) + progressed = True + if selected_count >= target_requests: + break + if not progressed: + break + + selected = [request for items in selected_sessions for request in items] + selected.sort(key=lambda request: (request.timestamp_s, request.chat_id)) + if len(selected) < target_requests: + raise ValueError( + f"window session sample selected only {len(selected)} requests; " + f"target was {target_requests}" + ) + return selected + + +def _is_kvc_fit_smallappend(session_requests: list[TraceRequest]) -> bool: + initial = session_requests[0] + if initial.input_length < 2048 or initial.input_length > 16000: + return False + for request in session_requests: + if request.output_length > 2048: + return False + for previous, current in zip(session_requests, session_requests[1:], strict=False): + append_tokens = current.input_length - ( + previous.input_length + previous.output_length + ) + if append_tokens <= 0 or append_tokens > 2048: + return False + if _overlap_ratio(previous, current) < 0.75: + return False + return True + + +def _write_sample( + *, + selected: list[TraceRequest], + input_trace_path: Path, + output_path: Path, + profile: str, + max_sessions: int, + max_turns_per_session: int, + rebase_timestamps: bool, +) -> SampleSummary: + if not selected: + raise ValueError(f"profile {profile!r} selected no requests") + + first_ts = selected[0].timestamp_s + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w", encoding="utf-8") as handle: + for request in selected: + timestamp = request.timestamp_s - first_ts if rebase_timestamps else request.timestamp_s + payload = { + "chat_id": request.chat_id, + "parent_chat_id": request.parent_chat_id, + "timestamp": round(timestamp, 6), + "input_length": request.input_length, + "output_length": request.output_length, + "type": request.request_type, + "turn": request.turn_id, + "hash_ids": list(request.hash_ids), + } + handle.write(json.dumps(payload, sort_keys=True) + "\n") + + sessions = defaultdict(list) + for request in selected: + sessions[request.session_id].append(request) + + selected_chat_ids = {request.chat_id for request in selected} + missing_parent_count = sum( + 1 + for request in selected + if request.parent_chat_id >= 0 and request.parent_chat_id not in selected_chat_ids + ) + append_values: list[float] = [] + gap_values: list[float] = [] + overlap_values: list[float] = [] + direct_eligible_count = 0 + for session_requests in sessions.values(): + ordered = _ordered(session_requests) + for previous, current in zip(ordered, ordered[1:], strict=False): + append_tokens = current.input_length - ( + previous.input_length + previous.output_length + ) + overlap_ratio = _overlap_ratio(previous, current) + append_values.append(float(append_tokens)) + gap_values.append(float(current.timestamp_s - previous.timestamp_s)) + overlap_values.append(overlap_ratio) + if append_tokens > 0 and append_tokens <= 2048 and overlap_ratio > 0: + direct_eligible_count += 1 + + turn2plus_count = sum(max(0, len(items) - 1) for items in sessions.values()) + + start = min(request.timestamp_s for request in selected) + end = max(request.timestamp_s for request in selected) + summary = SampleSummary( + input_trace_path=str(input_trace_path), + output_trace_path=str(output_path), + profile=profile, + request_count=len(selected), + session_count=len(sessions), + multi_turn_session_count=sum(1 for items in sessions.values() if len(items) > 1), + turn2plus_count=turn2plus_count, + direct_eligible_turn2plus_count=direct_eligible_count, + direct_eligible_turn2plus_ratio=( + direct_eligible_count / turn2plus_count if turn2plus_count else 0.0 + ), + missing_parent_count=missing_parent_count, + max_sessions=max_sessions, + max_turns_per_session=max_turns_per_session, + start_time_s=0.0 if rebase_timestamps else start, + end_time_s=end - start if rebase_timestamps else end, + sampled_duration_s=end - start, + rebased_timestamps=rebase_timestamps, + input_tokens=_stats([float(request.input_length) for request in selected]), + output_tokens=_stats([float(request.output_length) for request in selected]), + append_tokens=_stats(append_values), + inter_turn_gap_s=_stats(gap_values), + overlap_ratio=_stats(overlap_values), + ) + with output_path.with_suffix(output_path.suffix + ".summary.json").open( + "w", encoding="utf-8" + ) as handle: + json.dump(asdict(summary), handle, indent=2, sort_keys=True) + return summary + + +def _ordered(session_requests: list[TraceRequest]) -> list[TraceRequest]: + return sorted( + session_requests, + key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id), + ) + + +def _overlap_ratio(previous: TraceRequest, current: TraceRequest) -> float: + if not current.hash_ids: + return 0.0 + previous_blocks = set(previous.hash_ids) + overlap = sum(1 for block in current.hash_ids if block in previous_blocks) + return overlap / len(current.hash_ids) + + +def _stats(values: list[float]) -> dict[str, float] | None: + if not values: + return None + ordered = sorted(values) + return { + "count": float(len(ordered)), + "mean": statistics.fmean(ordered), + "min": ordered[0], + "p50": _percentile(ordered, 0.50), + "p90": _percentile(ordered, 0.90), + "p99": _percentile(ordered, 0.99), + "max": ordered[-1], + } + + +def _percentile(sorted_values: list[float], percentile: float) -> float: + if len(sorted_values) == 1: + return sorted_values[0] + return sorted_values[round((len(sorted_values) - 1) * percentile)] + + +if __name__ == "__main__": + main() diff --git a/scripts/sweep_real_ali_kvc.sh b/scripts/sweep_real_ali_kvc.sh new file mode 100755 index 0000000..561139f --- /dev/null +++ b/scripts/sweep_real_ali_kvc.sh @@ -0,0 +1,170 @@ +#!/usr/bin/env bash +# Real Ali workload sweep for KVC pd-hybrid. +# +# This script expects a prebuilt sample trace and replays it exactly for every +# mechanism. It intentionally keeps pool polling disabled for performance runs. +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$REPO_ROOT" + +MODEL=${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct} +TRACE=${TRACE:-outputs/real-ali-kvc-iter/samples-balanced/ali-kvc-fit-smallappend.jsonl} +OUT_ROOT=${OUT_ROOT:-outputs/real-ali-kvc-iter/runs} +TIME_SCALE=${TIME_SCALE:-1} +CONCURRENCY=${CONCURRENCY:-32} +REQUEST_TIMEOUT_S=${REQUEST_TIMEOUT_S:-300} +STACK_TIMEOUT_S=${STACK_TIMEOUT_S:-1200} +RUNS=${RUNS:-"dp kvc_bp"} +EXTRA_SERVER_ARGS=${EXTRA_SERVER_ARGS:-} +PREFILL_EXTRA_SERVER_ARGS=${PREFILL_EXTRA_SERVER_ARGS:-} +DECODE_EXTRA_SERVER_ARGS=${DECODE_EXTRA_SERVER_ARGS:-} +KVC_SEED_MIN_TURN_ID=${KVC_SEED_MIN_TURN_ID:-1} +KVC_SEED_ONLY_MULTITURN=${KVC_SEED_ONLY_MULTITURN:-0} + +mkdir -p "$OUT_ROOT" +LOG="$OUT_ROOT/sweep.log" + +log() { + echo "[$(date '+%F %T')] $*" | tee -a "$LOG" +} + +common_args=( + --trace "$TRACE" + --model-path "$MODEL" + --output-root "$OUT_ROOT" + --use-trace-as-sample + --time-scale "$TIME_SCALE" + --concurrency-limit "$CONCURRENCY" + --timeout-s "$STACK_TIMEOUT_S" + --request-timeout-s "$REQUEST_TIMEOUT_S" +) +if [[ -n "$EXTRA_SERVER_ARGS" ]]; then + common_args+=(--extra-server-args "$EXTRA_SERVER_ARGS") +fi +if [[ -n "$PREFILL_EXTRA_SERVER_ARGS" ]]; then + common_args+=(--prefill-extra-server-args "$PREFILL_EXTRA_SERVER_ARGS") +fi +if [[ -n "$DECODE_EXTRA_SERVER_ARGS" ]]; then + common_args+=(--decode-extra-server-args "$DECODE_EXTRA_SERVER_ARGS") +fi + +kvc_args=( + "${common_args[@]}" + --mechanism kvcache-centric + --policy kv-aware + --prefill-workers 2 + --decode-workers 6 + --prefill-tp-size 1 + --decode-tp-size 1 + --prefill-gpu-ids 0,1 + --decode-gpu-ids 2,3,4,5,6,7 + --transfer-backend mooncake + --gpu-budget 8 + --kvcache-admission-mode worker + --kvcache-seed-min-turn-id "$KVC_SEED_MIN_TURN_ID" + --kvcache-seed-max-inflight-decode -1 + --kvcache-prefill-backup-policy release-after-transfer + --kvcache-prefill-priority-eviction +) +if [[ "$KVC_SEED_ONLY_MULTITURN" == "1" ]]; then + kvc_args+=(--kvcache-seed-only-multiturn-sessions) +fi + +run_dp() { + log "=== DP cache-aware baseline: 8 direct workers ===" + uv run agentic-pd-hybrid benchmark-live \ + "${common_args[@]}" \ + --mechanism pd-colo \ + --policy kv-aware \ + --prefill-workers 0 \ + --decode-workers 0 \ + --direct-workers 8 \ + --direct-tp-size 1 \ + --direct-gpu-ids 0,1,2,3,4,5,6,7 \ + --gpu-budget 8 +} + +run_pd_disagg() { + log "=== PD-disaggregation baseline: 2P6D ===" + uv run agentic-pd-hybrid benchmark-live \ + "${common_args[@]}" \ + --mechanism pd-disaggregation \ + --policy kv-aware \ + --prefill-workers 2 \ + --decode-workers 6 \ + --prefill-tp-size 1 \ + --decode-tp-size 1 \ + --prefill-gpu-ids 0,1 \ + --decode-gpu-ids 2,3,4,5,6,7 \ + --transfer-backend mooncake \ + --gpu-budget 8 +} + +run_pd_sticky() { + log "=== PD-disaggregation sticky baseline: 2P6D ===" + uv run agentic-pd-hybrid benchmark-live \ + "${common_args[@]}" \ + --mechanism pd-disaggregation \ + --policy sticky \ + --prefill-workers 2 \ + --decode-workers 6 \ + --prefill-tp-size 1 \ + --decode-tp-size 1 \ + --prefill-gpu-ids 0,1 \ + --decode-gpu-ids 2,3,4,5,6,7 \ + --transfer-backend mooncake \ + --gpu-budget 8 +} + +run_kvc() { + log "=== KVC baseline: 2P6D worker admission, no backpressure ===" + uv run agentic-pd-hybrid benchmark-live "${kvc_args[@]}" +} + +run_kvc_bp() { + log "=== KVC candidate: 2P6D worker admission + backpressure ===" + uv run agentic-pd-hybrid benchmark-live \ + "${kvc_args[@]}" \ + --enable-backpressure \ + --backpressure-max-pause-s 2.0 +} + +summarize_latest() { + log "=== Latest summaries ===" + find "$OUT_ROOT" -maxdepth 2 -name 'request-metrics.jsonl.summary.json' -print \ + | sort \ + | while read -r summary; do + python - "$summary" <<'PY' +import json, sys +p=sys.argv[1] +d=json.load(open(p)) +lat=d.get("latency_stats_s") or {} +tt=d.get("ttft_stats_s") or {} +em=d.get("execution_modes") or {} +print(p) +print(" reqs", d.get("request_count"), "errors", d.get("error_count"), "trunc", d.get("truncated_request_count")) +print(" lat mean/p50/p90/p99", lat.get("mean"), lat.get("p50"), lat.get("p90"), lat.get("p99")) +print(" ttft mean/p50/p90", tt.get("mean"), tt.get("p50"), tt.get("p90")) +print(" modes", em) +PY + done | tee -a "$LOG" +} + +log "Trace: $TRACE" +log "Model: $MODEL" +log "Runs: $RUNS | time-scale=$TIME_SCALE concurrency=$CONCURRENCY | kvc-seed-min-turn-id=$KVC_SEED_MIN_TURN_ID | kvc-seed-only-multiturn=$KVC_SEED_ONLY_MULTITURN" + +for run in $RUNS; do + case "$run" in + dp) run_dp ;; + pd) run_pd_disagg ;; + pd_sticky) run_pd_sticky ;; + kvc) run_kvc ;; + kvc_bp) run_kvc_bp ;; + *) log "Unknown run name: $run"; exit 2 ;; + esac +done + +summarize_latest +log "DONE" diff --git a/src/agentic_pd_hybrid/benchmark.py b/src/agentic_pd_hybrid/benchmark.py index c43de8d..d8fe5a3 100644 --- a/src/agentic_pd_hybrid/benchmark.py +++ b/src/agentic_pd_hybrid/benchmark.py @@ -3,13 +3,20 @@ from __future__ import annotations import asyncio import json import signal +import shutil +from collections import Counter 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.sampling import ( + SessionSampleConfig, + SessionSampleSummary, + sample_trace_sessions, +) from agentic_pd_hybrid.stack import ManagedPdStack, launch_pd_stack +from agentic_pd_hybrid.trace import load_trace from agentic_pd_hybrid.topology import SingleNodeTopology @@ -47,12 +54,14 @@ class BenchmarkConfig: pool_poll_include_sessions: bool = True enable_backpressure: bool = False backpressure_max_pause_s: float = 2.0 + progress_interval_s: float = 30.0 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 + use_trace_as_sample: bool = False launch_stack: bool = True @@ -94,22 +103,37 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts: ) 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, + if config.use_trace_as_sample: + shutil.copyfile(config.trace_path, sampled_trace_path) + sample_summary = _summarize_trace_sample( + input_trace_path=config.trace_path, + sampled_trace_path=sampled_trace_path, + profile=config.sample_profile, 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, ) - ) + else: + 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) @@ -198,6 +222,7 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts: pool_poll_include_sessions=config.pool_poll_include_sessions, enable_backpressure=config.enable_backpressure, backpressure_max_pause_s=config.backpressure_max_pause_s, + progress_interval_s=config.progress_interval_s, ) if config.request_timeout_s is not None: replay_config = replace( @@ -258,12 +283,14 @@ def run_live_benchmark(config: BenchmarkConfig) -> BenchmarkArtifacts: "pool_poll_include_sessions": config.pool_poll_include_sessions, "enable_backpressure": config.enable_backpressure, "backpressure_max_pause_s": config.backpressure_max_pause_s, + "progress_interval_s": config.progress_interval_s, "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, + "use_trace_as_sample": config.use_trace_as_sample, "sample_summary": asdict(sample_summary), "topology": { "model_path": config.topology.model_path, @@ -310,3 +337,44 @@ def _header_mode_for(policy_name: str) -> str: if policy_name == "kv-aware": return "target-worker" return "none" + + +def _summarize_trace_sample( + *, + input_trace_path: Path, + sampled_trace_path: Path, + profile: str, + session_sample_rate: float, + min_turns: int, + min_initial_input_tokens: int | None, + max_initial_input_tokens: int | None, + max_append_input_tokens: int | None, + max_output_tokens: int | None, + min_overlap_ratio: float | None, +) -> SessionSampleSummary: + requests = load_trace(sampled_trace_path) + if not requests: + raise ValueError(f"Trace sample is empty: {sampled_trace_path}") + session_turns = Counter(request.session_id for request in requests) + start_time_s = requests[0].timestamp_s + end_time_s = requests[-1].timestamp_s + return SessionSampleSummary( + input_trace_path=str(input_trace_path), + output_trace_path=str(sampled_trace_path), + request_count=len(requests), + session_count=len(session_turns), + multi_turn_session_count=sum(1 for turns in session_turns.values() if turns > 1), + start_time_s=start_time_s, + end_time_s=end_time_s, + sampled_duration_s=end_time_s - start_time_s, + session_sample_rate=session_sample_rate, + min_turns=min_turns, + profile=profile, + min_initial_input_tokens=min_initial_input_tokens, + max_initial_input_tokens=max_initial_input_tokens, + max_append_input_tokens=max_append_input_tokens, + max_output_tokens=max_output_tokens, + min_overlap_ratio=min_overlap_ratio, + mean_append_input_tokens=None, + mean_turn_overlap_ratio=None, + ) diff --git a/src/agentic_pd_hybrid/cli.py b/src/agentic_pd_hybrid/cli.py index c3d0858..86e20a3 100644 --- a/src/agentic_pd_hybrid/cli.py +++ b/src/agentic_pd_hybrid/cli.py @@ -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 /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 /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)), + ), ) diff --git a/src/agentic_pd_hybrid/replay.py b/src/agentic_pd_hybrid/replay.py index f8b51ce..bf2ad70 100644 --- a/src/agentic_pd_hybrid/replay.py +++ b/src/agentic_pd_hybrid/replay.py @@ -107,6 +107,7 @@ class ReplayConfig: enable_backpressure: bool = False backpressure_max_pause_s: float = 2.0 structural_log_dir: Path | None = None + progress_interval_s: float = 30.0 @dataclass @@ -174,6 +175,62 @@ class ExecutionResult: finish_reason: str | None = None +@dataclass +class ReplayProgress: + total_requests: int + output_path: Path + interval_s: float + start_time_s: float + submitted_count: int = 0 + completed_count: int = 0 + error_count: int = 0 + truncated_count: int = 0 + last_request_id: str | None = None + last_session_id: str | None = None + last_trace_timestamp_s: float | None = None + execution_modes: Counter[str] = field(default_factory=Counter) + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + async def record_submitted(self, request: TraceRequest) -> None: + async with self.lock: + self.submitted_count += 1 + self.last_request_id = request.request_id + self.last_session_id = request.session_id + self.last_trace_timestamp_s = request.timestamp_s + + async def record_completed(self, row: RequestMetrics) -> None: + async with self.lock: + self.completed_count += 1 + if row.error is not None: + self.error_count += 1 + if _is_truncated(row): + self.truncated_count += 1 + self.execution_modes[row.execution_mode] += 1 + self.last_request_id = row.request_id + self.last_session_id = row.session_id + self.last_trace_timestamp_s = row.trace_timestamp_s + + async def emit(self, phase: str) -> None: + async with self.lock: + event = { + "phase": phase, + "elapsed_s": round(time.perf_counter() - self.start_time_s, 3), + "total_requests": self.total_requests, + "submitted_count": self.submitted_count, + "completed_count": self.completed_count, + "inflight_count": self.submitted_count - self.completed_count, + "error_count": self.error_count, + "truncated_count": self.truncated_count, + "last_request_id": self.last_request_id, + "last_session_id": self.last_session_id, + "last_trace_timestamp_s": self.last_trace_timestamp_s, + "execution_modes": dict(sorted(self.execution_modes.items())), + } + self.output_path.parent.mkdir(parents=True, exist_ok=True) + with self.output_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(event, sort_keys=True) + "\n") + + async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]: structural_dir = config.structural_log_dir if structural_dir is None and config.output_path is not None: @@ -199,6 +256,23 @@ async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]: session_tail_tasks: dict[str, asyncio.Task[RequestMetrics]] = {} direct_sessions: dict[str, DirectSessionState] = {} direct_session_lock = asyncio.Lock() + progress = ( + ReplayProgress( + total_requests=len(requests), + output_path=config.output_path.parent / "replay-progress.jsonl", + interval_s=config.progress_interval_s, + start_time_s=start_time, + ) + if config.progress_interval_s > 0 + else None + ) + progress_stop = asyncio.Event() + progress_task: asyncio.Task[None] | None = None + if progress is not None: + await progress.emit("start") + progress_task = asyncio.create_task( + _progress_heartbeat(progress, progress_stop) + ) async with httpx.AsyncClient(timeout=config.timeout_s, trust_env=False) as client: decode_residency = await _discover_decode_residency( client=client, @@ -230,6 +304,8 @@ async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]: sleep_s = target_offset - (time.perf_counter() - start_time) if sleep_s > 0: await asyncio.sleep(sleep_s) + if progress is not None: + await progress.record_submitted(request) tasks.append( asyncio.create_task( _run_request( @@ -244,12 +320,15 @@ async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]: direct_session_lock=direct_session_lock, decode_residency=decode_residency, depends_on=session_tail_tasks.get(request.session_id), + progress=progress, ) ) ) session_tail_tasks[request.session_id] = tasks[-1] results = await asyncio.gather(*tasks) + if progress is not None: + await progress.emit("requests-complete") if poll_task is not None: poll_task.cancel() try: @@ -285,6 +364,14 @@ async def replay_trace(config: ReplayConfig) -> list[RequestMetrics]: trace_path=config.trace_path, router_url=config.router_url, ) + if progress is not None: + await progress.emit("final") + progress_stop.set() + if progress_task is not None: + try: + await progress_task + except asyncio.CancelledError: + pass _structural_close() return results @@ -302,6 +389,7 @@ async def _run_request( direct_session_lock: asyncio.Lock, decode_residency: DecodeResidencyState, depends_on: asyncio.Task[RequestMetrics] | None, + progress: ReplayProgress | None = None, ) -> RequestMetrics: if depends_on is not None: await depends_on @@ -351,7 +439,7 @@ async def _run_request( async with state_lock: state.finish(request, decision) - return RequestMetrics.from_decision( + row = RequestMetrics.from_decision( request, decision, mechanism_name=config.mechanism_name, @@ -371,6 +459,29 @@ async def _run_request( requested_output_tokens=execution.requested_output_tokens, finish_reason=execution.finish_reason, ) + if progress is not None: + await progress.record_completed(row) + return row + + +async def _progress_heartbeat( + progress: ReplayProgress, + stop_event: asyncio.Event, +) -> None: + while not stop_event.is_set(): + try: + await asyncio.wait_for(stop_event.wait(), timeout=progress.interval_s) + except asyncio.TimeoutError: + await progress.emit("heartbeat") + + +def _is_truncated(row: RequestMetrics) -> bool: + return ( + row.actual_output_tokens is not None + and row.requested_output_tokens is not None + and row.requested_output_tokens > 1 + and row.actual_output_tokens < row.requested_output_tokens * 0.5 + ) async def _invoke_router( @@ -643,16 +754,41 @@ async def _open_streaming_session( request.input_length * 16, (request.input_length + request.output_length) * 16, ) - response = await client.post( - f"{server_url.rstrip('/')}/open_session", - json={ - "capacity_of_str_len": capacity, - "session_id": session_id, - "streaming": True, - }, - timeout=_ADMISSION_PROBE_TIMEOUT_S, - ) - response.raise_for_status() + payload = { + "capacity_of_str_len": capacity, + "session_id": session_id, + "streaming": True, + } + url = f"{server_url.rstrip('/')}/open_session" + response = await client.post(url, json=payload, timeout=_ADMISSION_PROBE_TIMEOUT_S) + try: + response.raise_for_status() + except httpx.HTTPStatusError: + if response.status_code != 400: + raise + await _structural_emit( + "session-lifecycle.jsonl", + { + "event": "open-session-400-retry", + "server_url": server_url, + "session_id": session_id, + "request_id": request.request_id, + "turn_id": request.turn_id, + "response_text": response.text[:512], + }, + ) + await _close_streaming_session( + client=client, + server_url=server_url, + session_id=session_id, + allow_missing=True, + ) + response = await client.post( + url, + json=payload, + timeout=_ADMISSION_PROBE_TIMEOUT_S, + ) + response.raise_for_status() opened_session_id = response.json() if opened_session_id != session_id: raise ValueError(