| | """A toy LLM streaming gateway that reproduces event-loop stalls. | | | | | | Many clients hold open token streams on one asyncio loop. New requests keep | | | arriving, and each one does some setup before its own stream starts. In | | | "naive" mode the setup does the usual per-request things on the loop; in | | | "fixed" mode the same work is built once per process before traffic. | | | | | | Run: python gateway.py naive | python gateway.py fixed | python gateway.py both | | | """ | | | | | | from future import annotations | | | | | | import asyncio | | | import gc | | | import json | | | import os | | | import random | | | import ssl | | | import statistics | | | import sys | | | import tempfile | | | import threading | | | import time | | | import traceback | | | import zlib | | | from collections import Counter | | | | | | import certifi | | | import numpy as np | | | | | | STREAMS = 40 # long-lived client streams | | | CHUNK_EVERY = 0.020 # each stream sends a token chunk every 20 ms | | | ARRIVAL_MEAN = 0.25 # a new request every 250 ms on average | | | DURATION = 20.0 # seconds of traffic per run | | | VOCAB_SIZE = 200_000 | | | DOCS, DIMS = 400, 384 # retrieval candidates scored per request | | | PAYLOAD = 1_000_000 # bytes of already-compressed data per request | | | STARTUP_OBJECTS = 1_500_000 | | | | | | rng = random.Random(7) | | | | | | | | | # ---------- watchdog (same design as the article) ---------- | | | | | | class LoopWatchdog: | | | def init(self, loop, tick=0.010, margin=0.040, poll=0.005, report=None): | | | self.loop, self.report = loop, report | | | self.tick, self.limit, self.poll = tick, tick + margin, poll |
| | self.loop_thread = threading.get_ident() |
| | self.heartbeat = self.seen = time.monotonic() |
| | self.stacks = [] |
| | self._stop = False | | | threading.Thread(target=self._watch, daemon=True).start() | | | | | | async def beat(self): | | | while True: |
| | self.heartbeat = time.monotonic() |
| | await asyncio.sleep(self.tick) |
| | | | | def stop(self): | | | self._stop = True | | | | | | def _watch(self): | | | while not self._stop: | | | time.sleep(self.poll) | | | beat = self.heartbeat | | | if beat != self.seen: | | | gap, self.seen = beat - self.seen, beat |
| | stacks, self.stacks = self.stacks, [] |
| | if gap > self.limit: |
| | blocked_ms = (gap - self.tick) * 1000 |
| | self.loop.call_soon_threadsafe(self.report, blocked_ms, stacks) | | | continue |
| | overdue = time.monotonic() - beat > self.limit |
| | if overdue and len(self.stacks) < 2: |
| | frame = sys._current_frames().get(self.loop_thread) |
| | if self.heartbeat == beat: |
| | self.stacks.append(traceback.extract_stack(frame, limit=15)) |
| | | | | | | | # ---------- the expensive things a request start does ---------- | | | |
| | def make_vocab_file() -> str: |
| | path = os.path.join(tempfile.gettempdir(), "toy_vocab.json") |
| | if not os.path.exists(path): |
| | vocab = {f"tok_{i:06d}_{i * 7919 % 100003}": i for i in range(VOCAB_SIZE)} |
| | with open(path, "w") as f: |
| | json.dump(vocab, f) |
| | return path | | | | | | |
| | def load_tokenizer(path: str) -> dict: |
| | with open(path) as f: |
| | vocab = json.load(f) |
| | return {v: k for k, v in vocab.items()} |
| | | | | |
| | def make_tls_context() -> ssl.SSLContext: |
| | return ssl.create_default_context(cafile=certifi.where()) |
| | | | | |
| | def score_python(query: list[float], docs: list[list[float]]) -> list[float]: |
| | out = [] |
| | for doc in docs: |
| | dot = qn = dn = 0.0 |
| | for q, d in zip(query, doc): |
| | dot += q * d | | | qn += q * q | | | dn += d * d | | | out.append(dot / ((qn ** 0.5) * (dn ** 0.5))) | | | return out | | | | | | |
| | def score_numpy(query: np.ndarray, docs: np.ndarray) -> np.ndarray: |
| | return (docs @ query) / (np.linalg.norm(docs, axis=1) * np.linalg.norm(query)) |
| | | | | |
| | def compress_payload(data: bytes) -> bytes: |
| | return zlib.compress(data, 6) |
| | | | | | | | def make_garbage(): | | | # Per-request bookkeeping: lots of small containers, some cyclic. |
| | items = [] |
| | for i in range(20_000): |
| | d = {"i": i, "meta": [i, str(i)]} |
| | d["self"] = d |
| | items.append(d) |
| | return len(items) |
| | | | | | | | # ---------- the gateway ---------- | | | | | | class Gateway: | | | def init(self, mode: str): | | | self.mode = mode |
| | self.vocab_path = make_vocab_file() |
| | self.docs_list = [[rng.random() for _ in range(DIMS)] for _ in range(DOCS)] |
| | self.docs_np = np.array(self.docs_list) |
| | self.payload = os.urandom(PAYLOAD) |
| | # A big long-lived heap, like a prompt/template catalog loaded at boot. |
| | self.catalog = [(i, str(i), [i]) for i in range(STARTUP_OBJECTS)] |
| | self.setup_ms: list[float] = [] |
| | self.step_ms: dict[str, list[float]] = {} |
| | if mode == "fixed": |
| | self.tls = make_tls_context() |
| | self.tokenizer = load_tokenizer(self.vocab_path) |
| | gc.collect() |
| | gc.freeze() |
| | |
| | def _step(self, name, fn): |
| | t = time.perf_counter() |
| | out = fn() |
| | self.step_ms.setdefault(name, []).append((time.perf_counter() - t) * 1000) |
| | return out | | | |
| | def start_request(self): |
| | t0 = time.perf_counter() |
| | step = self._step |
| | if self.mode == "naive": |
| | clients = step("TLS context", lambda: [make_tls_context() for _ in range(3)]) |
| | tokenizer = step("tokenizer load", lambda: load_tokenizer(self.vocab_path)) |
| | q = [rng.random() for _ in range(DIMS)] |
| | step("pure-Python scoring", lambda: score_python(q, self.docs_list)) |
| | step("compressing incompressible data", lambda: compress_payload(self.payload)) |
| | else: | | | clients = [self.tls] * 3 | | | tokenizer = self.tokenizer |
| | q = np.array([rng.random() for _ in range(DIMS)]) |
| | step("numpy scoring", lambda: score_numpy(q, self.docs_np)) |
| | # payload is already compressed: send as is | | | step("request bookkeeping", make_garbage) | | | self.setup_ms.append((time.perf_counter() - t0) * 1000) | | | return clients, tokenizer | | | | | | | | | TIMELINE: list[tuple[float, float]] = [] # (seconds since start, lateness) for --dump | | | | | | |
| | async def stream(lateness: list[float], stop_at: float): |
| | next_due = time.monotonic() |
| | while True: | | | next_due += CHUNK_EVERY |
| | await asyncio.sleep(max(0.0, next_due - time.monotonic())) |
| | now = time.monotonic() |
| | if now >= stop_at: |
| | return |
| | lateness.append(now - next_due) |
| | TIMELINE.append((now - (stop_at - DURATION), now - next_due)) |
| | if now - next_due > CHUNK_EVERY: # fell behind: resync like a real sender | | | next_due = now | | | | | | |
| | def arrival_offsets() -> list[float]: |
| | r, t, out = random.Random(11), 0.0, [] |
| | while True: |
| | t += r.expovariate(1 / ARRIVAL_MEAN) |
| | if t >= DURATION - 0.5: |
| | return out | | | out.append(t) | | | | | | | | | async def arrivals(gw: Gateway, start: float): | | | # Same arrival times in both modes, even if the loop falls behind. |
| | for offset in arrival_offsets(): |
| | await asyncio.sleep(max(0.0, start + offset - time.monotonic())) |
| | gw.start_request() |
| | | | | |
| | def _bucket_worst() -> list[float]: |
| | buckets = [0.0] * int(DURATION * 10) |
| | for t, late in TIMELINE: |
| | i = min(int(t * 10), len(buckets) - 1) |
| | buckets[i] = max(buckets[i], late * 1000) |
| | TIMELINE.clear() |
| | return [round(b, 1) for b in buckets] |
| | | | | | | | CATEGORIES = [ | | | ("make_tls_context", "TLS context"), | | | ("load_tokenizer", "tokenizer load"), | | | ("score_python", "pure-Python scoring"), | | | ("compress_payload", "compressing incompressible data"), | | | ] | | | | | | |
| | def categorize(stack) -> str: |
| | names = {f.name for f in stack} |
| | for fn, label in CATEGORIES: | | | if fn in names: | | | return label | | | return "other" | | | | | | |
| | async def run(mode: str) -> dict: |
| | gw = Gateway(mode) |
| | gc_s: list[float] = [] |
| | gc_start = [0.0] |
| | |
| | def on_gc(phase, info): |
| | if phase == "start": |
| | gc_start[0] = time.perf_counter() |
| | else: | | | gc_s.append((time.perf_counter() - gc_start[0]) * 1000) | | | |
| | gc.callbacks.append(on_gc) |
| | stalls = [] |
| | |
| | def report(blocked_ms, stacks): |
| | stalls.append((blocked_ms, categorize(stacks[0]) if stacks else "no stack (GIL held)")) |
| | |
| | wd = LoopWatchdog(asyncio.get_running_loop(), report=report) |
| | beat = asyncio.create_task(wd.beat()) |
| | stop_at = time.monotonic() + DURATION |
| | lateness: list[list[float]] = [[] for _ in range(STREAMS)] |
| | await asyncio.gather(arrivals(gw, stop_at - DURATION), *(stream(l, stop_at) for l in lateness)) |
| | wd.stop() |
| | beat.cancel() |
| | gc.callbacks.remove(on_gc) |
| | gc.unfreeze() |
| | | | | all_late = [x * 1000 for l in lateness for x in l] | | | |
| | def share_late(threshold_ms): |
| | return 100 * sum(x > threshold_ms for x in all_late) / len(all_late) |
| | | | | causes = Counter(c for _, c in stalls) | | | return { | | | "mode": mode, |
| | "requests": len(gw.setup_ms), |
| | "setup_p50_ms": statistics.median(gw.setup_ms), |
| | "stalls_per_request": len(stalls) / len(gw.setup_ms), |
| | "stall_p50_ms": statistics.median([s for s, _ in stalls]) if stalls else 0.0, |
| | "stall_max_ms": max([s for s, _ in stalls], default=0.0), |
| | "loop_blocked_pct": 100 * sum(s for s, _ in stalls) / 1000 / DURATION, |
| | "chunk_late_p99_ms": float(np.percentile(all_late, 99)), |
| | "chunk_late_max_ms": max(all_late), |
| | "chunks_late_50ms_pct": share_late(50), |
| | "chunks_late_100ms_pct": share_late(100), |
| | "gc_collections": len(gc_s), |
| | "gc__max_ms": max(gc_s, default=0.0), |
| | "gc__total_ms": sum(gc_s), |
| | "setup_step_p50_ms": {k: round(statistics.median(v), 2) for k, v in gw.step_ms.items()}, |
| | "timeline_100ms_worst_late_ms": _bucket_worst(), |
| | "stall_causes": {k: round(100 * v / len(stalls), 1) for k, v in causes.most_common()} if stalls else {}, |
| | } | | | | | | |
| | def main(): |
| | args = sys.argv[1:] |
| | dump = None |
| | if "--dump" in args: |
| | i = args.index("--dump") |
| | dump = args[i + 1] |
| | del args[i:i + 2] |
| | modes = args or ["both"] |
| | if modes == ["both"]: |
| | modes = ["naive", "fixed"] |
| | results = [] |
| | for mode in modes: | | | result = asyncio.run(run(mode)) | | | timeline = result.pop("timeline_100ms_worst_late_ms") |
| | print(json.dumps(result, indent=2)) |
| | result["timeline_100ms_worst_late_ms"] = timeline |
| | results.append(result) |
| | gc.collect() |
| | if dump: |
| | with open(dump, "w") as f: |
| | json.dump(results, f) |
| | | | | |
| | if __name__ == "__main__": |
| | main() |