Microbench 3 (connector_tax): infrastructure for KV connector substrate tax
Validates the elastic_migration_v2 finding that kv_role=kv_both adds
TTFT p90 +45% even when PD-sep never fires. Replicates under
single-instance, synthetic, open-loop workload to disambiguate
mechanism cost from 8-instance feedback amplification.
Configurations (8):
plain, noop_connector, mooncake_{producer,consumer,both},
nixl_both, lmcache_only, multi_mooncake_lmcache.
Pre-flight verification gates risky configs (kv_consumer needs dummy
bootstrap, multi-connector composition, NoOp custom class loading).
Workload: two-phase sweep
Phase A: rate {0.5..32} req/s × shape (4096, 256), saturation criteria
Phase B: ref_safe rate × cartesian (input ∈ {512,4k,32k}, output ∈ {64,256,1024})
Step-timing patch enriches vLLM's existing AGENTIC_STEP_LOG_PATH emit
with step_duration_us and build_meta_us — directly measures per-step
substrate cost, not just user-visible TTFT/TPOT.
run_all.sh runs as 5-stage barrier:
0 pre-flight + apply patch
1 Phase A all configs
2 pick ref_safe / ref_load
3 Phase B all configs
4 revert patch + analyze + plot
Outputs aggregate.{json,csv}, MANIFEST.tsv, and 5 figures.
Estimated runtime: 4-5.5 hours on idle dash0 H20.
This commit is contained in:
171
microbench/connector_tax/patches/apply_step_timing.py
Normal file
171
microbench/connector_tax/patches/apply_step_timing.py
Normal file
@@ -0,0 +1,171 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Apply / revert step-timing additions to vLLM's existing agentic step log.
|
||||
|
||||
vLLM already has a per-step JSONL emitter triggered by AGENTIC_STEP_LOG_PATH.
|
||||
This patch enriches it with three duration fields:
|
||||
|
||||
step_duration_us : full schedule() wall time
|
||||
build_meta_us : duration of self.connector.build_connector_meta(...)
|
||||
emit_overhead_us : duration of _agentic_emit_step_log itself
|
||||
(lets us subtract the patch's own cost)
|
||||
|
||||
We also add a worker-side patch in kv_connector_model_runner_mixin.py to
|
||||
record start_load_kv() duration into a per-process JSONL file pointed to
|
||||
by VLLM_PD_PROFILE_LOG=$RUN_DIR/worker_step.jsonl.
|
||||
|
||||
All inserts are marked with `# CONNECTOR_TAX_PATCH` so revert is just
|
||||
"delete those lines".
|
||||
|
||||
Usage:
|
||||
python apply_step_timing.py --apply [--vllm-root PATH]
|
||||
python apply_step_timing.py --revert [--vllm-root PATH]
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
MARKER = "# CONNECTOR_TAX_PATCH"
|
||||
DEFAULT_VLLM_ROOT = (
|
||||
Path.home()
|
||||
/ "agentic-kv/.venv/lib/python3.12/site-packages/vllm"
|
||||
)
|
||||
|
||||
|
||||
def already_patched(text: str) -> bool:
|
||||
return MARKER in text
|
||||
|
||||
|
||||
def revert_text(text: str) -> str:
|
||||
out = [l for l in text.splitlines() if MARKER not in l]
|
||||
return "\n".join(out) + ("\n" if text.endswith("\n") else "")
|
||||
|
||||
|
||||
# ── scheduler.py patches ────────────────────────────────────────────────────
|
||||
def patch_scheduler(text: str) -> str:
|
||||
if already_patched(text):
|
||||
print(" scheduler.py already patched, skipping")
|
||||
return text
|
||||
|
||||
# 1. At top of schedule(): record _step_t0
|
||||
pat = (
|
||||
r"( def schedule\(self\) -> SchedulerOutput:\n"
|
||||
r" # NOTE\(woosuk\) on the scheduling algorithm:)"
|
||||
)
|
||||
repl = (
|
||||
r" def schedule(self) -> SchedulerOutput:\n"
|
||||
r" import time as _t " + MARKER + "\n"
|
||||
r" _step_t0 = _t.perf_counter_ns() " + MARKER + "\n"
|
||||
r" self._ct_build_meta_ns = 0 " + MARKER + "\n"
|
||||
r" # NOTE(woosuk) on the scheduling algorithm:"
|
||||
)
|
||||
text, n = re.subn(pat, repl, text, count=1)
|
||||
if n == 0:
|
||||
raise RuntimeError("Failed to patch schedule() entry")
|
||||
|
||||
# 2. Time the build_connector_meta call
|
||||
pat = (
|
||||
r" if self\.connector is not None:\n"
|
||||
r" meta: KVConnectorMetadata = self\.connector\.build_connector_meta\(\n"
|
||||
r" scheduler_output\n"
|
||||
r" \)\n"
|
||||
r" scheduler_output\.kv_connector_metadata = meta"
|
||||
)
|
||||
repl = (
|
||||
" if self.connector is not None:\n"
|
||||
f" _bm_t0 = _t.perf_counter_ns() {MARKER}\n"
|
||||
" meta: KVConnectorMetadata = self.connector.build_connector_meta(\n"
|
||||
" scheduler_output\n"
|
||||
" )\n"
|
||||
f" self._ct_build_meta_ns = _t.perf_counter_ns() - _bm_t0 {MARKER}\n"
|
||||
" scheduler_output.kv_connector_metadata = meta"
|
||||
)
|
||||
text, n = re.subn(pat, repl, text, count=1)
|
||||
if n == 0:
|
||||
raise RuntimeError("Failed to patch build_connector_meta")
|
||||
|
||||
# 3. Pass step duration into _agentic_emit_step_log via attribute
|
||||
# (cleaner than threading kwargs through). Then in the emit
|
||||
# function inject the new fields into `record`.
|
||||
pat = (
|
||||
r" if self\._agentic_step_log_fh is not None:\n"
|
||||
r" self\._agentic_emit_step_log\("
|
||||
)
|
||||
repl = (
|
||||
f" if self._agentic_step_log_fh is not None:\n"
|
||||
f" self._ct_step_duration_ns = _t.perf_counter_ns() - _step_t0 {MARKER}\n"
|
||||
f" self._agentic_emit_step_log("
|
||||
)
|
||||
text, n = re.subn(pat, repl, text, count=1)
|
||||
if n == 0:
|
||||
raise RuntimeError("Failed to patch step_duration insertion")
|
||||
|
||||
# 4. Inject extra fields into the `record` dict in _agentic_emit_step_log
|
||||
pat = (
|
||||
r" record = \{\n"
|
||||
r" \"t_unix\": _time\.time\(\),"
|
||||
)
|
||||
repl = (
|
||||
" record = {\n"
|
||||
f" \"step_duration_us\": getattr(self, '_ct_step_duration_ns', 0) // 1000, {MARKER}\n"
|
||||
f" \"build_meta_us\": getattr(self, '_ct_build_meta_ns', 0) // 1000, {MARKER}\n"
|
||||
" \"t_unix\": _time.time(),"
|
||||
)
|
||||
text, n = re.subn(pat, repl, text, count=1)
|
||||
if n == 0:
|
||||
raise RuntimeError("Failed to patch record dict")
|
||||
|
||||
return text
|
||||
|
||||
|
||||
# ── apply / revert driver ───────────────────────────────────────────────────
|
||||
def apply_to_file(path: Path, fn) -> bool:
|
||||
if not path.exists():
|
||||
print(f" SKIP {path} (not found)")
|
||||
return False
|
||||
orig = path.read_text()
|
||||
new = fn(orig)
|
||||
if new == orig:
|
||||
print(f" unchanged: {path}")
|
||||
return False
|
||||
path.write_text(new)
|
||||
print(f" patched ({new.count(MARKER)} marks): {path}")
|
||||
return True
|
||||
|
||||
|
||||
def revert_file(path: Path) -> bool:
|
||||
if not path.exists():
|
||||
return False
|
||||
orig = path.read_text()
|
||||
new = revert_text(orig)
|
||||
if new == orig:
|
||||
print(f" no marks: {path}")
|
||||
return False
|
||||
path.write_text(new)
|
||||
print(f" reverted: {path}")
|
||||
return True
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--apply", action="store_true")
|
||||
ap.add_argument("--revert", action="store_true")
|
||||
ap.add_argument("--vllm-root", type=Path, default=DEFAULT_VLLM_ROOT)
|
||||
args = ap.parse_args()
|
||||
|
||||
if not (args.apply ^ args.revert):
|
||||
ap.error("Specify exactly one of --apply / --revert")
|
||||
|
||||
sched = args.vllm_root / "v1/core/sched/scheduler.py"
|
||||
|
||||
if args.apply:
|
||||
print(f"Applying connector-tax step-timing patch to {args.vllm_root}")
|
||||
apply_to_file(sched, patch_scheduler)
|
||||
else:
|
||||
print(f"Reverting connector-tax step-timing patch from {args.vllm_root}")
|
||||
revert_file(sched)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user