File size: 7,773 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """Coordinator-only bounded profiling traffic; --execute required for any network.
Reuses the existing bench2 prompt construction, sampler settings and stream accounting.
No generated text is written. Profile all TP ranks via vLLM's existing API; the local_rank
field is not a global-rank filter. B4 decode capture begins after ALL streams emit text.
"""
from __future__ import annotations
import argparse
import importlib
import json
import sys
import threading
import time
import urllib.request
import uuid
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Callable, Protocol, cast
from urllib.parse import urlparse
from offload.gb10_budget.evidence import JSON, Object
class Bench(Protocol):
URL: str
MODEL: str
DATA_DIR: Path
def models(self) -> Object: ...
def scrape(self, url: str | None = None) -> str: ...
def parse_metrics(self, text: str) -> dict[str, float]: ...
def require_idle(self, values: dict[str, float]) -> None: ...
def make_prompt(self, ctx: int, workload: str, filler: str, seed: int,
template: Object) -> tuple[str, int]: ...
def request_body(self, prompt: str, args: argparse.Namespace, salt: str) -> Object: ...
def stream_request(self, body: Object, started: threading.Event | None = None) -> Object: ...
def load_bench() -> Bench:
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "bench"))
return cast(Bench, importlib.import_module("bench2"))
def control(url: str, verb: str) -> None:
req = urllib.request.Request(url.rstrip("/") + "/" + verb, data=b"", method="POST")
with urllib.request.urlopen(req, timeout=120) as response:
if response.status != 200:
raise RuntimeError(f"{verb} returned HTTP {response.status}")
def run(bench: Bench, args: argparse.Namespace,
command: Callable[[str, str], None] = control) -> Object:
bench.URL, bench.DATA_DIR = args.url, args.data_dir
bench.require_idle(bench.parse_metrics(bench.scrape()))
info = bench.models()
models = info.get("data")
if not isinstance(models, list) or not models or not isinstance(models[0], dict):
raise ValueError("No served model in /v1/models")
bench.MODEL = str(models[0]["id"])
cap = models[0].get("max_model_len")
if not isinstance(cap, int):
raise ValueError("Missing served max_model_len")
template: Object = {"thinking": True, "reasoning_effort": "high"}
workload = "synth" if args.ctx == 150000 else "summary"
prompts = [bench.make_prompt(args.ctx, workload, "real", args.seed + i, template)
for i in range(args.batch)]
if any(n + args.max_tokens > cap for _, n in prompts):
raise ValueError("Occupied prompt + generation exceeds live model cap")
bench.require_idle(bench.parse_metrics(bench.scrape()))
starts = [threading.Event() for _ in prompts]
barrier = threading.Barrier(args.batch)
armed = False
mark = time.monotonic()
def request(i: int) -> Object:
body = bench.request_body(prompts[i][0], args, uuid.uuid4().hex)
barrier.wait(timeout=30)
result = bench.stream_request(body, starts[i])
result.pop("reasoning_text", None)
result.pop("answer_text", None)
return result
prefill_delay = float(getattr(args, "prefill_delay", 0))
if args.phase == "prefill" and prefill_delay == 0:
command(args.url, "start_profile")
armed = True
with ThreadPoolExecutor(max_workers=args.batch) as pool:
futures = [pool.submit(request, i) for i in range(args.batch)]
try:
if args.phase == "prefill" and prefill_delay > 0:
time.sleep(prefill_delay)
for future in futures:
if future.done():
future.result()
raise RuntimeError("Stream ended before delayed prefill capture")
if any(event.is_set() for event in starts):
raise RuntimeError("Decode already started; delayed prefill capture would be mislabeled")
command(args.url, "start_profile")
armed = True
if args.phase == "decode":
deadline = time.monotonic() + args.ttft_timeout
while not all(event.is_set() for event in starts):
for future in futures:
if future.done():
future.result() # Propagate a stream failure instead of timing out.
raise RuntimeError("Stream ended before all streams reached decode")
if time.monotonic() >= deadline:
raise TimeoutError("Not all streams reached decode before TTFT deadline")
time.sleep(0.02)
command(args.url, "start_profile")
armed = True
if armed:
time.sleep(args.seconds)
finally:
if armed:
command(args.url, "stop_profile")
results: list[JSON] = [future.result() for future in futures]
return {"status": "MEASURED profile-affected client observations, not promotion timings",
"phase": args.phase, "context": args.ctx, "batch": args.batch,
"prefill_delay_s": prefill_delay, "workload": workload,
"occupied_prompt_tokens": [n for _, n in prompts],
"elapsed_s": time.monotonic() - mark, "requests": results}
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--execute", action="store_true")
ap.add_argument("--url", default="http://127.0.0.1:18001")
ap.add_argument("--data-dir", type=Path, default=Path("/opt/dsv41/data/bench-data"))
ap.add_argument("--ctx", type=int, choices=(32000, 150000), default=32000)
ap.add_argument("--batch", type=int, choices=(1, 4), default=1)
ap.add_argument("--phase", choices=("none", "prefill", "decode"), default="decode")
ap.add_argument("--seconds", type=float, default=3)
ap.add_argument("--prefill-delay", type=float, default=0, help="Capture later chunks; fail if decode already started")
ap.add_argument("--ttft-timeout", type=float, default=600)
ap.add_argument("--max-tokens", type=int, default=1500)
ap.add_argument("--output", type=Path, required=True)
args = ap.parse_args()
if not 0 < args.seconds <= 30 or args.ttft_timeout <= 0 or args.max_tokens < 1:
ap.error("Require seconds in (0,30], positive timeout and positive generation budget")
if not 0 <= args.prefill_delay <= 60 or (args.prefill_delay and args.phase != "prefill"):
ap.error("Prefill delay must be 0..60 seconds and only used with --phase prefill")
if urlparse(args.url).hostname not in ("127.0.0.1", "localhost") or urlparse(args.url).port != 18001:
ap.error("Use loopback SSH forwarding 18001 -> reserved experiment port 8001")
if not args.execute:
print("Preview only: no network. Add --execute in the coordinator's reserved window.")
return
if args.output.exists():
ap.error("Output exists; use a fresh artifact path")
for filename in ("opinions_cl.jsonl", "filings_recap.jsonl"):
if not (args.data_dir / filename).is_file():
ap.error(f"Missing existing benchmark corpus: {filename}")
args.seed, args.priority, args.temperature, args.top_p = 42, 10, 1, 1
args.thinking, args.reasoning_effort = True, "high"
result = run(load_bench(), args)
with args.output.open("x") as target:
target.write(json.dumps(result, indent=2, allow_nan=False) + "\n")
print(f"Wrote aggregate profile observations: {args.output}")
if __name__ == "__main__":
main()
|