|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +# pyre-strict |
| 8 | + |
| 9 | +"""Benchmark: SPDL pipeline handoff latency with and without thread output queue. |
| 10 | +
|
| 11 | +Compares main-thread ``get_item`` latency between the default asyncio-based |
| 12 | +handoff (``run_coroutine_threadsafe`` + ``Future.result`` polling) and the |
| 13 | +``queue.Queue``-based handoff (direct ``queue.Queue.get``). |
| 14 | +
|
| 15 | +The benchmark emulates a realistic training loop: each iteration the |
| 16 | +foreground thread receives a batch, does simulated work (busy-wait to |
| 17 | +mimic a GPU forward/backward pass), and then times how long the *next* |
| 18 | +``get_item`` call takes. When the simulated work is long enough the |
| 19 | +producer has time to pre-fill the queue, so the handoff latency |
| 20 | +isolates the cross-thread scheduling overhead. |
| 21 | +
|
| 22 | +Foreground work is swept from 0ms to 30ms in 3ms steps. |
| 23 | +
|
| 24 | +Usage:: |
| 25 | +
|
| 26 | + buck2 run //spdl/examples:benchmark_thread_output_queue |
| 27 | +
|
| 28 | +Example results (devserver, no GPU, 500 lightweight int items):: |
| 29 | +
|
| 30 | + FG work default (p50) default (p99) thread_q (p50) thread_q (p99) |
| 31 | + ------- ------------- ------------- -------------- -------------- |
| 32 | + 0ms 199us 625us 116us 313us |
| 33 | + 3ms 246us 642us 223us 646us |
| 34 | + 6ms 232us 532us 190us 555us |
| 35 | + 9ms 287us 485us 153us 549us |
| 36 | + 12ms 280us 822us 14us 464us |
| 37 | + 15ms 221us 396us 8us 385us |
| 38 | + 18ms 227us 431us 9us 70us |
| 39 | + 21ms 224us 450us 11us 41us |
| 40 | + 24ms 240us 533us 12us 34us |
| 41 | + 27ms 227us 470us 12us 27us |
| 42 | + 30ms 242us 512us 12us 26us |
| 43 | +
|
| 44 | +The default asyncio path stays flat at ~220-280us regardless of overlap |
| 45 | +time — that is the fixed ``run_coroutine_threadsafe`` scheduling tax. |
| 46 | +The thread output queue drops to ~8-14us once the foreground work is |
| 47 | +long enough (>= ~12ms on this machine) for the producer to pre-fill it. |
| 48 | +""" |
| 49 | + |
| 50 | +from __future__ import annotations |
| 51 | + |
| 52 | +import statistics |
| 53 | +import time |
| 54 | +from dataclasses import dataclass, field |
| 55 | + |
| 56 | +from spdl.pipeline import PipelineBuilder |
| 57 | + |
| 58 | +__all__ = ["BenchResult", "main"] |
| 59 | + |
| 60 | + |
| 61 | +@dataclass |
| 62 | +class BenchResult: |
| 63 | + """Collected latency measurements for a single benchmark run.""" |
| 64 | + |
| 65 | + name: str |
| 66 | + latencies_us: list[float] = field(default_factory=list) |
| 67 | + |
| 68 | + @property |
| 69 | + def p50(self) -> float: |
| 70 | + s = sorted(self.latencies_us) |
| 71 | + return s[len(s) // 2] |
| 72 | + |
| 73 | + @property |
| 74 | + def p95(self) -> float: |
| 75 | + s = sorted(self.latencies_us) |
| 76 | + return s[int(len(s) * 0.95)] |
| 77 | + |
| 78 | + @property |
| 79 | + def p99(self) -> float: |
| 80 | + s = sorted(self.latencies_us) |
| 81 | + return s[int(len(s) * 0.99)] |
| 82 | + |
| 83 | + @property |
| 84 | + def mean(self) -> float: |
| 85 | + return statistics.mean(self.latencies_us) |
| 86 | + |
| 87 | + def summary(self) -> str: |
| 88 | + """Return a one-line summary string with mean/p50/p95/p99.""" |
| 89 | + return ( |
| 90 | + f"{self.name:<50s} " |
| 91 | + f"mean={self.mean:>9.1f}us " |
| 92 | + f"p50={self.p50:>9.1f}us " |
| 93 | + f"p95={self.p95:>9.1f}us " |
| 94 | + f"p99={self.p99:>9.1f}us " |
| 95 | + f"({len(self.latencies_us)} meas)" |
| 96 | + ) |
| 97 | + |
| 98 | + |
| 99 | +def _busy_wait_us(us: float) -> None: |
| 100 | + """Busy-wait for *us* microseconds (spin-loop).""" |
| 101 | + deadline = time.perf_counter() + us / 1e6 |
| 102 | + while time.perf_counter() < deadline: |
| 103 | + pass |
| 104 | + |
| 105 | + |
| 106 | +def _run_bench( |
| 107 | + name: str, |
| 108 | + n_items: int, |
| 109 | + buffer_size: int, |
| 110 | + use_thread_output_queue: bool, |
| 111 | + work_us: float, |
| 112 | + warmup: int = 20, |
| 113 | +) -> BenchResult: |
| 114 | + """Run a single benchmark: source -> sink pipeline with foreground work. |
| 115 | +
|
| 116 | + Builds a trivial pipeline (source of *n_items* integers -> sink) and |
| 117 | + iterates it. Between each consumed item the foreground thread |
| 118 | + busy-waits for *work_us* microseconds to simulate consumer-side |
| 119 | + compute (e.g. a GPU training step). After warmup, the time spent |
| 120 | + inside each ``get_item`` call is recorded. |
| 121 | +
|
| 122 | + Args: |
| 123 | + name: Human-readable label for the result. |
| 124 | + n_items: Total number of items the source produces. |
| 125 | + buffer_size: Sink buffer size (number of slots). |
| 126 | + use_thread_output_queue: Whether to use the thread output queue handoff. |
| 127 | + work_us: Duration of simulated foreground work in microseconds. |
| 128 | + warmup: Number of items to consume before recording latencies. |
| 129 | + """ |
| 130 | + items = list(range(n_items)) |
| 131 | + pipeline = ( |
| 132 | + PipelineBuilder() |
| 133 | + .add_source(iter(items)) |
| 134 | + .add_sink(buffer_size=buffer_size) |
| 135 | + .build(num_threads=2, use_thread_output_queue=use_thread_output_queue) |
| 136 | + ) |
| 137 | + |
| 138 | + result = BenchResult(name=name) |
| 139 | + consumed = 0 |
| 140 | + with pipeline.auto_stop(): |
| 141 | + for _ in pipeline.get_iterator(timeout=60): |
| 142 | + consumed += 1 |
| 143 | + |
| 144 | + # Simulate foreground work (e.g. GPU fwd+bwd). |
| 145 | + # This gives the producer time to pre-fill the queue |
| 146 | + # so the next get_item measures pure handoff overhead. |
| 147 | + _busy_wait_us(work_us) |
| 148 | + |
| 149 | + if consumed <= warmup: |
| 150 | + continue |
| 151 | + |
| 152 | + # Time the next get_item call — this is what we're measuring. |
| 153 | + t0 = time.perf_counter() |
| 154 | + try: |
| 155 | + next(pipeline.get_iterator(timeout=60)) |
| 156 | + except StopIteration: |
| 157 | + break |
| 158 | + result.latencies_us.append((time.perf_counter() - t0) * 1e6) |
| 159 | + consumed += 1 |
| 160 | + |
| 161 | + # Do work after the timed item too, so the *next* iteration's |
| 162 | + # measurement also has overlap time. |
| 163 | + _busy_wait_us(work_us) |
| 164 | + |
| 165 | + return result |
| 166 | + |
| 167 | + |
| 168 | +def main() -> None: |
| 169 | + """Run the full benchmark sweep and print results.""" |
| 170 | + n_items = 500 |
| 171 | + buffer_size = 8 |
| 172 | + |
| 173 | + print("SPDL Pipeline Thread Output Queue Handoff Benchmark") |
| 174 | + print("=" * 90) |
| 175 | + |
| 176 | + results: list[BenchResult] = [] |
| 177 | + |
| 178 | + work_ms_values = list(range(0, 31, 3)) |
| 179 | + for work_ms in work_ms_values: |
| 180 | + work_us = work_ms * 1000.0 |
| 181 | + label = f"{work_ms}ms foreground work" |
| 182 | + print(f"\n--- {label} ({n_items} items, buffer_size={buffer_size}) ---") |
| 183 | + |
| 184 | + for use_toq, tag in [ |
| 185 | + (False, "default asyncio"), |
| 186 | + (True, "thread output queue"), |
| 187 | + ]: |
| 188 | + r = _run_bench( |
| 189 | + f"{tag}, {label}", |
| 190 | + n_items, |
| 191 | + buffer_size, |
| 192 | + use_thread_output_queue=use_toq, |
| 193 | + work_us=work_us, |
| 194 | + ) |
| 195 | + results.append(r) |
| 196 | + print(f" {r.summary()}") |
| 197 | + |
| 198 | + print(f"\n{'=' * 90}") |
| 199 | + print("SUMMARY — main thread get_item() latency (lower is better)") |
| 200 | + print(f"{'=' * 90}") |
| 201 | + for r in results: |
| 202 | + if r.latencies_us: |
| 203 | + print(f" {r.summary()}") |
| 204 | + |
| 205 | + |
| 206 | +if __name__ == "__main__": |
| 207 | + main() |
0 commit comments