gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(kvc): add real Ali replay workflow

Browse Source
This commit is contained in:
Gahow Wang
2026-05-12 05:28:00 +00:00
parent 1d51704dad
commit 4e8f943875
5 changed files with 901 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

450
scripts/prepare_real_ali_samples.py Executable file
View File

@@ -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()

170
scripts/sweep_real_ali_kvc.sh Executable file
View File

@@ -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"

86
src/agentic_pd_hybrid/benchmark.py
View File

@@ -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,
)

58
src/agentic_pd_hybrid/cli.py
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)),
),
)

158
src/agentic_pd_hybrid/replay.py
View File

@@ -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(