Escape Hatches: Getting CPU Work Off the GIL

import time

docs = [f"doc-{i}" for i in range(20_000)]
seen_list = docs[:5_000]

t0 = time.perf_counter()
dupes_slow = [d for d in docs if d in seen_list]        # O(n*m): list membership scans
t1 = time.perf_counter()
seen = set(seen_list)
dupes_fast = [d for d in docs if d in seen]             # O(n): hash lookups
t2 = time.perf_counter()
assert dupes_slow == dupes_fast
print(f"list {t1 - t0:.3f}s  set {t2 - t1:.4f}s")

import numpy as np

rng = np.random.default_rng(0)
emb = rng.standard_normal((10_000, 384)).astype(np.float32)
query = rng.standard_normal(384).astype(np.float32)

# Python loop: holds the GIL, ~1000x slower
def top_k_loop(emb, query, k=5):
    scores = []
    for row in emb.tolist():
        dot = sum(a * b for a, b in zip(row, query.tolist()))
        scores.append(dot)
    return sorted(range(len(scores)), key=scores.__getitem__, reverse=True)[:k]

# Vectorised: normalise, one matrix-vector product, partial sort
def top_k_vec(emb, query, k=5):
    emb_n = emb / np.linalg.norm(emb, axis=1, keepdims=True)
    q_n = query / np.linalg.norm(query)
    scores = emb_n @ q_n
    idx = np.argpartition(-scores, k)[:k]
    return idx[np.argsort(-scores[idx])].tolist()

print(top_k_vec(emb, query))

import numpy as np
import pandas as pd

df = pd.DataFrame({"tokens": [120, 5000, 800], "latency_ms": [200, 2500, 900]})
# slow: df["bucket"] = df.apply(lambda r: "long" if r.tokens > 1000 else "short", axis=1)
df["bucket"] = np.where(df["tokens"] > 1000, "long", "short")
df["tok_per_s"] = df["tokens"] / (df["latency_ms"] / 1000)

import polars as pl

df = pl.DataFrame({"text": ["a b c", "d e"], "source": ["web", "pdf"]})
out = (
    df.with_columns(pl.col("text").str.split(" ").list.len().alias("n_words"))
      .group_by("source")
      .agg(pl.col("n_words").sum())
)

import numpy as np
from concurrent.futures import ThreadPoolExecutor
from numba import njit, prange


@njit(nogil=True, cache=True)
def minhash_row(tokens: np.ndarray, seeds: np.ndarray) -> np.ndarray:
    out = np.empty(seeds.shape[0], dtype=np.uint64)
    for j in range(seeds.shape[0]):
        best = np.uint64(0xFFFFFFFFFFFFFFFF)
        for t in tokens:
            h = (t * np.uint64(0x9E3779B97F4A7C15) + seeds[j]) & np.uint64(0xFFFFFFFFFFFFFFFF)
            if h < best:
                best = h
        out[j] = best
    return out


@njit(parallel=True, cache=True)
def l2_normalise(x: np.ndarray) -> np.ndarray:
    out = np.empty_like(x)
    for i in prange(x.shape[0]):                       # rows split across numba's threads
        out[i] = x[i] / np.sqrt((x[i] ** 2).sum())
    return out


docs = [np.arange(i, i + 500, dtype=np.uint64) for i in range(64)]
seeds = np.arange(128, dtype=np.uint64)
with ThreadPoolExecutor(8) as ex:                       # threads scale because of nogil=True
    sigs = list(ex.map(lambda d: minhash_row(d, seeds), docs))

# cython: language_level=3, boundscheck=False, wraparound=False
# file: fastscore.pyx
from cython.parallel cimport prange

def dot_rows(const float[:, ::1] emb, const float[::1] q, float[::1] out):
    cdef Py_ssize_t i, j, n = emb.shape[0], d = emb.shape[1]
    cdef float acc
    with nogil:                                   # other Python threads run meanwhile
        for i in prange(n, schedule="static"):    # OpenMP: compile with -fopenmp
            acc = 0
            for j in range(d):
                acc = acc + emb[i, j] * q[j]
            out[i] = acc

pip install mypy
mypyc ingest/chunking.py          # builds chunking.<abi>.so next to the source
python -c "import ingest.chunking as c; print(c.__file__)"   # confirm the compiled module is used

use pyo3::prelude::*;

#[pyfunction]
fn count_tokens(py: Python<'_>, texts: Vec<String>) -> Vec<usize> {
    // texts were copied out of Python objects above; no Python access inside the closure
    py.allow_threads(|| texts.iter().map(|t| t.split_whitespace().count()).collect())
}

#[pymodule(gil_used = false)]          // declares free-threading support (recent PyO3)
fn fasttok(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(count_tokens, m)?)?;
    Ok(())
}

from concurrent.futures import ProcessPoolExecutor
import re

WS = re.compile(r"\s+")


def clean_shard(lines: list[str]) -> list[str]:          # module-level: picklable
    return [WS.sub(" ", ln).strip().lower() for ln in lines]


def shards(lines, size):
    for i in range(0, len(lines), size):
        yield lines[i:i + size]


if __name__ == "__main__":                               # required for spawn/forkserver
    corpus = [f"  Line {i}\twith   EXTRA   spaces " for i in range(200_000)]
    with ProcessPoolExecutor(max_workers=4) as ex:
        cleaned = [ln for part in ex.map(clean_shard, shards(corpus, 10_000)) for ln in part]
    print(len(cleaned), cleaned[0])

import hashlib
import zlib
from concurrent.futures import ThreadPoolExecutor

blobs = [bytes([i % 251]) * 2_000_000 for i in range(16)]   # 2 MB "documents"


def fingerprint_and_pack(blob: bytes) -> tuple[str, int]:
    digest = hashlib.sha256(blob).hexdigest()       # GIL released for big buffers
    packed = zlib.compress(blob, 6)                 # GIL released
    return digest, len(packed)


with ThreadPoolExecutor(max_workers=8) as ex:
    results = list(ex.map(fingerprint_and_pack, blobs))
print(len(results), results[0][1])

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
emb = torch.randn(100_000, 768, device=device, dtype=torch.float16 if device == "cuda" else torch.float32)
queries = torch.randn(64, 768, device=device, dtype=emb.dtype)
with torch.inference_mode():
    emb = torch.nn.functional.normalize(emb, dim=1)
    q = torch.nn.functional.normalize(queries, dim=1)
    top = (q @ emb.T).topk(10, dim=1)              # one kernel for 64 x 100k similarities
print(top.indices.shape)

import asyncio


async def call_llm(prompt: str) -> str:
    await asyncio.sleep(0.1)                         # stands in for an HTTP call
    return prompt.upper()


async def main(prompts: list[str]) -> list[str]:
    sem = asyncio.Semaphore(16)                      # concurrency limit

    async def one(p: str) -> str:
        async with sem:
            return await call_llm(p)

    return await asyncio.gather(*(one(p) for p in prompts))


print(len(asyncio.run(main([f"q{i}" for i in range(100)]))))   # ~0.7 s, not 10 s