Testing and Benchmarking

import timeit

setup = "import re; rx = re.compile(r'\\s+'); text = 'a  b   c ' * 2000"
stmts = {
    "re.sub": "rx.sub(' ', text)",
    "split/join": "' '.join(text.split())",
}
for name, stmt in stmts.items():
    runs = timeit.repeat(stmt, setup=setup, repeat=5, number=200)
    best = min(runs) / 200
    print(f"{name:<11} best {best * 1e6:8.1f} us/call   spread {(max(runs) - min(runs)) / min(runs):.0%}")

python -m timeit -s "import json; s = json.dumps(list(range(1000)))" "json.loads(s)"

import statistics
import time


def bench(fn, repeat=15, warmup=2):
    for _ in range(warmup):
        fn()
    samples = []
    for _ in range(repeat):
        t0 = time.perf_counter()
        fn()
        samples.append(time.perf_counter() - t0)
    q = statistics.quantiles(samples, n=4)
    return {"min": min(samples), "median": statistics.median(samples), "iqr": q[2] - q[0]}


r = bench(lambda: sorted(range(200_000), key=lambda x: -x))
print({k: f"{v * 1000:.2f} ms" for k, v in r.items()})

pip install pyperf
python -m pyperf system tune                       # needs root: fixed frequency, isolated CPUs (undo: system reset)
python -m pyperf timeit -s "import json; s = json.dumps(list(range(1000)))" "json.loads(s)" -o json.json
python -m pyperf timeit -s "import orjson; s = b'[' + b','.join(b'%d' % i for i in range(1000)) + b']'" \
    "orjson.loads(s)" -o orjson.json
python -m pyperf compare_to json.json orjson.json  # "Mean +- std dev: ... -> ...: 5.2x faster"

# test_bench_clean.py
import re

import pytest

WS = re.compile(r"\s+")
DOC = "Some   scraped\ttext with   extra  spaces. " * 500


def clean_regex(s: str) -> str:
    return WS.sub(" ", s).strip()


def clean_split(s: str) -> str:
    return " ".join(s.split())


@pytest.mark.parametrize("fn", [clean_regex, clean_split], ids=["regex", "split"])
def test_clean_speed(benchmark, fn):
    out = benchmark(fn, DOC)
    assert out == clean_split(DOC)          # correctness is asserted in the benchmark too

pytest test_bench_clean.py --benchmark-only --benchmark-autosave           # saves .benchmarks/
pytest test_bench_clean.py --benchmark-only --benchmark-compare --benchmark-compare-fail=median:10%
pytest --benchmark-disable                                                # in normal test runs

import os
import time
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor

import pytest


def cpu_task(n: int) -> int:
    return sum(i % 7 for i in range(n))


def wall(fn) -> float:
    t0 = time.perf_counter()
    fn()
    return time.perf_counter() - t0


@pytest.mark.skipif(len(os.sched_getaffinity(0)) < 4, reason="needs 4 usable CPUs")
def test_processes_scale_pure_python():
    n, workers = 400_000, 4
    serial = min(wall(lambda: [cpu_task(n) for _ in range(workers)]) for _ in range(2))
    with ProcessPoolExecutor(workers) as ex:
        list(ex.map(cpu_task, [10] * workers))                     # warm up workers
        parallel = min(wall(lambda: list(ex.map(cpu_task, [n] * workers))) for _ in range(2))
    assert serial / parallel > 1.8, f"speedup only {serial / parallel:.2f}x"


def test_threads_do_not_help_pure_python_on_gil_build():
    import sys
    if not getattr(sys, "_is_gil_enabled", lambda: True)():
        pytest.skip("free-threaded interpreter: threads are expected to scale")
    n, workers = 300_000, 4
    serial = wall(lambda: [cpu_task(n) for _ in range(workers)])
    with ThreadPoolExecutor(workers) as ex:
        threaded = wall(lambda: list(ex.map(cpu_task, [n] * workers)))
    assert serial / threaded < 1.5                                 # GIL-bound: no meaningful speedup

python scripts/gil_probe.py --target ingest.clean:normalize_shard --arg 2000 --workers 4 --repeat 3 --json \
  > probe.json
python -c "import json; r = json.load(open('probe.json')); assert 'releases the GIL' in r['classification'], r['classification']"
python scripts/gil_probe.py --target ingest.clean:normalize_shard --fail-if-gil-bound   # exit 1 if GIL-bound

import asyncio
import time
from concurrent.futures import ProcessPoolExecutor


def cpu_heavy(n: int) -> int:
    return sum(i * i for i in range(n))


async def max_loop_lag(during, interval: float = 0.01) -> float:
    lags: list[float] = []
    done = asyncio.Event()

    async def ticker():
        while not done.is_set():
            t0 = time.perf_counter()
            await asyncio.sleep(interval)
            lags.append(time.perf_counter() - t0 - interval)

    tick = asyncio.create_task(ticker())
    await asyncio.sleep(interval * 2)       # let the ticker start sleeping before the work begins
    await during()
    done.set()
    await tick
    return max(lags, default=0.0)


async def main() -> None:
    async def blocking():                   # CPU on the loop thread: lag ~ the whole computation
        cpu_heavy(3_000_000)

    with ProcessPoolExecutor(1) as pool:
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(pool, cpu_heavy, 10)            # warm up the worker

        async def offloaded():
            await loop.run_in_executor(pool, cpu_heavy, 3_000_000)

        print(f"blocking lag  {await max_loop_lag(blocking) * 1000:7.1f} ms")
        print(f"offloaded lag {await max_loop_lag(offloaded) * 1000:7.1f} ms")


if __name__ == "__main__":
    asyncio.run(main())

import os
import sys


def cap_threads(n: str) -> None:            # module-level: picklable initializer
    os.environ["OMP_NUM_THREADS"] = n


def read_env(name: str):
    return os.environ.get(name)


def test_worker_thread_caps():
    from concurrent.futures import ProcessPoolExecutor

    with ProcessPoolExecutor(1, initializer=cap_threads, initargs=("1",)) as ex:
        assert ex.submit(read_env, "OMP_NUM_THREADS").result() == "1"


def test_gil_state_matches_deployment():
    expected_free_threaded = os.environ.get("EXPECT_FREE_THREADED") == "1"
    check = getattr(sys, "_is_gil_enabled", None)
    gil_on = check() if check else True
    assert gil_on is not expected_free_threaded, "an extension re-enabled the GIL (see RuntimeWarning at import)"

import threading


class TokenCounter:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self.total = 0

    def add(self, n: int) -> None:
        with self._lock:                     # remove the lock and this test fails on 3.14t (and sometimes on GIL builds)
            self.total += n


def hammer(counter: TokenCounter, threads: int = 8, iters: int = 20_000) -> None:
    barrier = threading.Barrier(threads)

    def worker() -> None:
        barrier.wait()                       # maximise overlap
        for _ in range(iters):
            counter.add(1)

    ts = [threading.Thread(target=worker) for _ in range(threads)]
    for t in ts:
        t.start()
    for t in ts:
        t.join()


c = TokenCounter()
hammer(c)
assert c.total == 8 * 20_000, c.total
print("ok", c.total)