| 1 | #!/usr/bin/env python3 |
| 2 | """Run on the target through SSH stdin; use synthetic prompts and disposable files.""" |
| 3 | import argparse |
| 4 | from concurrent.futures import ThreadPoolExecutor |
| 5 | import http.client |
| 6 | import json |
| 7 | import mmap |
| 8 | import os |
| 9 | from pathlib import Path |
| 10 | import random |
| 11 | import socket |
| 12 | import tempfile |
| 13 | import time |
| 14 | |
| 15 | |
| 16 | def request(path, key, data=None): |
| 17 | connection = http.client.HTTPConnection("127.0.0.1", 8081, timeout=600) |
| 18 | try: |
| 19 | body = json.dumps(data) if data is not None else None |
| 20 | connection.request("POST" if body else "GET", path, body, { |
| 21 | "Authorization": "Bearer " + key, "Content-Type": "application/json", |
| 22 | }) |
| 23 | response = connection.getresponse() |
| 24 | payload = response.read() |
| 25 | if response.status != 200: |
| 26 | raise RuntimeError(f"{path}: HTTP {response.status}") |
| 27 | return json.loads(payload) if data is not None else payload.decode() |
| 28 | finally: |
| 29 | connection.close() |
| 30 | |
| 31 | |
| 32 | def metrics(key): |
| 33 | values = {} |
| 34 | for line in request("/metrics", key).splitlines(): |
| 35 | if line.startswith("vllm:"): |
| 36 | name = line.split("{", 1)[0] |
| 37 | values[name] = values.get(name, 0) + float(line.rsplit(" ", 1)[1]) |
| 38 | return values |
| 39 | |
| 40 | |
| 41 | def prefill(pid, repeats): |
| 42 | environment = dict(entry.split(b"=", 1) for entry in Path(f"/proc/{pid}/environ").read_bytes().split(b"\0") if b"=" in entry) |
| 43 | key = environment[b"VLLM_API_KEY"].decode() |
| 44 | vocabulary = request("/tokenize", key, { |
| 45 | "model": "qwen3.8-27b", "prompt": "Synthetic benchmark data: blue green yellow red apple pear stone river cloud mountain 0123456789.\n", |
| 46 | })["tokens"] |
| 47 | for length in (16000, 64000, 148000): |
| 48 | for trial in range(repeats): |
| 49 | before = metrics(key) |
| 50 | if before["vllm:num_requests_running"] or before["vllm:num_requests_waiting"]: |
| 51 | raise RuntimeError("Endpoint is busy; rerun when idle") |
| 52 | prompt = random.choices(vocabulary, k=length) |
| 53 | started = time.perf_counter() |
| 54 | response = request("/v1/completions", key, { |
| 55 | "model": "qwen3.8-27b", "prompt": prompt, "max_tokens": 1, "temperature": 0, |
| 56 | }) |
| 57 | elapsed = time.perf_counter() - started |
| 58 | after = metrics(key) |
| 59 | delta = {name: after.get(name, 0) - value for name, value in before.items()} |
| 60 | print(json.dumps({ |
| 61 | "kind": "prefill", "prompt_tokens": response["usage"]["prompt_tokens"], "trial": trial, |
| 62 | "wall_seconds": elapsed, "prefill_seconds": delta["vllm:request_prefill_time_seconds_sum"], |
| 63 | "queue_seconds": delta["vllm:request_queue_time_seconds_sum"], |
| 64 | "cached_tokens": delta["vllm:prefix_cache_hits_total"], |
| 65 | "completed_requests": delta["vllm:request_success_total"], |
| 66 | }), flush=True) |
| 67 | |
| 68 | |
| 69 | def filesystem(directory, repeats): |
| 70 | os.umask(0o077) |
| 71 | size = 16 * 1024**2 |
| 72 | with tempfile.TemporaryDirectory(prefix="synthetic-ai-cache-", dir=directory) as temporary: |
| 73 | root = Path(temporary) |
| 74 | buffer = mmap.mmap(-1, size) |
| 75 | buffer[:] = os.urandom(size) |
| 76 | paths = [root / f"{random.getrandbits(128):032x}.bin" for _ in range(256)] |
| 77 | started = time.perf_counter() |
| 78 | for path in paths: |
| 79 | descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_DIRECT, 0o600) |
| 80 | try: |
| 81 | if os.write(descriptor, buffer) != size: |
| 82 | raise RuntimeError("Short synthetic write") |
| 83 | os.fsync(descriptor) |
| 84 | finally: |
| 85 | os.close(descriptor) |
| 86 | buffer.close() |
| 87 | print(json.dumps({"kind": "write", "directory": directory, "gib": 4, "seconds": time.perf_counter() - started}), flush=True) |
| 88 | |
| 89 | def read(path): |
| 90 | with mmap.mmap(-1, size) as target: |
| 91 | descriptor = os.open(path, os.O_RDONLY | os.O_DIRECT) |
| 92 | try: |
| 93 | if os.readv(descriptor, [target]) != size: |
| 94 | raise RuntimeError("Short synthetic read") |
| 95 | finally: |
| 96 | os.close(descriptor) |
| 97 | |
| 98 | for gib in (1, 4): |
| 99 | selected = paths[:gib * 64] |
| 100 | for trial in range(repeats): |
| 101 | random.shuffle(selected) |
| 102 | devices = [path for path in Path("/sys/block").iterdir() if (path / "device").exists()] |
| 103 | sectors = {path.name: int((path / "stat").read_text().split()[2]) for path in devices} |
| 104 | started = time.perf_counter() |
| 105 | with ThreadPoolExecutor(max_workers=16) as pool: |
| 106 | list(pool.map(read, selected)) |
| 107 | elapsed = time.perf_counter() - started |
| 108 | disk_reads = {path.name: (int((path / "stat").read_text().split()[2]) - sectors[path.name]) * 512 for path in devices} |
| 109 | print(json.dumps({"kind": "direct-read", "directory": directory, "gib": gib, "trial": trial, |
| 110 | "block_mib": 16, "workers": 16, "seconds": elapsed, |
| 111 | "device_read_bytes": {name: count for name, count in disk_reads.items() if count}}), flush=True) |
| 112 | |
| 113 | |
| 114 | if __name__ == "__main__": |
| 115 | parser = argparse.ArgumentParser(description=__doc__) |
| 116 | parser.add_argument("mode", choices=("prefill", "filesystem")) |
| 117 | parser.add_argument("target", help="Backend PID or disposable parent directory") |
| 118 | parser.add_argument("--hostname", required=True) |
| 119 | parser.add_argument("--repeats", type=int, default=3) |
| 120 | arguments = parser.parse_args() |
| 121 | if socket.gethostname() != arguments.hostname: |
| 122 | raise SystemExit("Unexpected target host") |
| 123 | if arguments.mode == "prefill": |
| 124 | prefill(int(arguments.target), arguments.repeats) |
| 125 | else: |
| 126 | filesystem(arguments.target, arguments.repeats) |