|
| 1 | +"""Benchmark the parallel read_csv path (PR #64347) across max_threads values. |
| 2 | +
|
| 3 | +Investigation-only helper for the Windows perf discussion. Generates a CSV of |
| 4 | +large integers (mirroring the bench_csv file used in the PR thread), then times |
| 5 | +``pd.read_csv`` at each requested ``mode.max_threads`` value. |
| 6 | +
|
| 7 | +IMPORTANT METHODOLOGY NOTE |
| 8 | +-------------------------- |
| 9 | +These are **warm-cache** numbers: the file is read straight after being |
| 10 | +written, so it is resident in the OS page cache, and there is no portable way |
| 11 | +to drop the page cache on Windows. That means this script isolates the |
| 12 | +**CPU / lock** behaviour, not disk I/O: |
| 13 | +
|
| 14 | +* If T=2 is still *slower* than T=1 here (warm cache), the regression is a |
| 15 | + genuine in-process contention point (the shared-mmap working-set lock |
| 16 | + hypothesis), and a code change is warranted. |
| 17 | +* If T=2 is *not* slower warm, then the inversion seen on a laptop is disk |
| 18 | + bound, and the fix is to gate/document rather than rewrite the tokenizer |
| 19 | + path. Settling the disk question needs the AWS storage-tier run, not this. |
| 20 | +
|
| 21 | +Run via the ``bench-read-csv-windows.yml`` workflow, or locally: |
| 22 | +
|
| 23 | + BENCH_N_ROWS=2000000 BENCH_THREADS="1 2 4" python scripts/bench_read_csv_parallel.py |
| 24 | +""" |
| 25 | + |
| 26 | +from __future__ import annotations |
| 27 | + |
| 28 | +import os |
| 29 | +import statistics |
| 30 | +import time |
| 31 | + |
| 32 | +import numpy as np |
| 33 | + |
| 34 | +import pandas as pd |
| 35 | + |
| 36 | +# The parallel read_csv path only engages above a size threshold (currently |
| 37 | +# 50 MB; see _PARALLEL_READ_MIN_BYTES in pandas/io/parsers/readers.py). Every |
| 38 | +# other eligibility gate (C engine, uncompressed local path, integer header) is |
| 39 | +# satisfied by this benchmark's call, so a plain size check is a sufficient |
| 40 | +# guard against silently measuring the serial fallback -- and avoids importing |
| 41 | +# a private pandas helper. |
| 42 | +PARALLEL_MIN_BYTES = 50 * 1024 * 1024 |
| 43 | + |
| 44 | +N_COLS = 8 |
| 45 | +N_ROWS = int(os.environ.get("BENCH_N_ROWS", "2000000")) |
| 46 | +N_RUNS = int(os.environ.get("BENCH_N_RUNS", "7")) |
| 47 | +N_WARMUP = int(os.environ.get("BENCH_N_WARMUP", "2")) |
| 48 | +THREADS = [int(tok) for tok in os.environ.get("BENCH_THREADS", "1 2 4").split()] |
| 49 | +PATH = os.environ.get("BENCH_PATH", "bench_data.csv") |
| 50 | + |
| 51 | + |
| 52 | +def generate_file() -> float: |
| 53 | + """Write an 8-column large-integer CSV; return its size in MB.""" |
| 54 | + rng = np.random.default_rng(0) |
| 55 | + data = rng.integers(10**15, 10**18, size=(N_ROWS, N_COLS), dtype=np.int64) |
| 56 | + pd.DataFrame(data).to_csv(PATH, index=False, header=False) |
| 57 | + return os.path.getsize(PATH) / 1e6 |
| 58 | + |
| 59 | + |
| 60 | +def time_read(max_threads: int) -> list[float]: |
| 61 | + """Return per-run wall times (ms) for read_csv at *max_threads*.""" |
| 62 | + pd.options.mode.max_threads = max_threads |
| 63 | + |
| 64 | + def one_read() -> float: |
| 65 | + start = time.perf_counter() |
| 66 | + pd.read_csv(PATH, header=None) |
| 67 | + return (time.perf_counter() - start) * 1000 |
| 68 | + |
| 69 | + for _ in range(N_WARMUP): |
| 70 | + one_read() |
| 71 | + return [one_read() for _ in range(N_RUNS)] |
| 72 | + |
| 73 | + |
| 74 | +def main() -> None: |
| 75 | + print(f"os.cpu_count() = {os.cpu_count()}") |
| 76 | + size_mb = generate_file() |
| 77 | + print(f"generated {PATH}: {size_mb:.1f} MB, {N_ROWS:,} rows x {N_COLS} cols") |
| 78 | + |
| 79 | + # Guard against silently benchmarking the serial fallback path. |
| 80 | + if os.path.getsize(PATH) < PARALLEL_MIN_BYTES: |
| 81 | + raise SystemExit( |
| 82 | + "ERROR: generated file is below the parallel-path size threshold -- " |
| 83 | + "the benchmark would measure the serial path. Increase BENCH_N_ROWS." |
| 84 | + ) |
| 85 | + print("parallel path eligible: OK\n") |
| 86 | + |
| 87 | + results = [] |
| 88 | + baseline = None |
| 89 | + for n_threads in THREADS: |
| 90 | + times = time_read(n_threads) |
| 91 | + best = min(times) |
| 92 | + median = statistics.median(times) |
| 93 | + if baseline is None: |
| 94 | + baseline = best |
| 95 | + speedup = baseline / best |
| 96 | + results.append((n_threads, best, median, speedup)) |
| 97 | + print( |
| 98 | + f"T={n_threads:>2} best={best:7.1f}ms " |
| 99 | + f"median={median:7.1f}ms speedup={speedup:.2f}x" |
| 100 | + ) |
| 101 | + |
| 102 | + table = [ |
| 103 | + "| Threads | best (ms) | median (ms) | speedup vs T=1 |", |
| 104 | + "|--:|--:|--:|--:|", |
| 105 | + ] |
| 106 | + for n_threads, best, median, speedup in results: |
| 107 | + table.append(f"| {n_threads} | {best:.1f} | {median:.1f} | {speedup:.2f}x |") |
| 108 | + markdown = "\n".join(table) |
| 109 | + |
| 110 | + summary_path = os.environ.get("GITHUB_STEP_SUMMARY") |
| 111 | + if summary_path: |
| 112 | + with open(summary_path, "a", encoding="utf-8") as handle: |
| 113 | + handle.write("## read_csv parallel benchmark (warm cache)\n\n") |
| 114 | + handle.write(f"- runner cores: `{os.cpu_count()}`\n") |
| 115 | + handle.write(f"- file: {size_mb:.1f} MB, {N_ROWS:,} rows x {N_COLS} cols\n") |
| 116 | + handle.write(f"- Python env: `{os.environ.get('BENCH_ENV', 'n/a')}`\n\n") |
| 117 | + handle.write(markdown + "\n\n") |
| 118 | + handle.write( |
| 119 | + "_Warm-cache numbers: isolates CPU/lock behaviour, not disk I/O. " |
| 120 | + "A T=2 inversion here points to in-process contention; no " |
| 121 | + "inversion means the laptop result is disk bound._\n" |
| 122 | + ) |
| 123 | + |
| 124 | + |
| 125 | +if __name__ == "__main__": |
| 126 | + main() |
0 commit comments