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