AI-Engineering Multiprocessing

import hashlib
import multiprocessing as mp
import os
import re
import tempfile

CHUNK_WORDS, OVERLAP = 120, 20


def parse(path: str) -> str:
    with open(path, encoding="utf-8", errors="replace") as fh:   # stand-in for pypdf / trafilatura
        html = fh.read()
    return re.sub(r"<[^>]+>", " ", html)


def chunk_file(path: str) -> list:
    words = parse(path).split()
    step = CHUNK_WORDS - OVERLAP
    out = []
    for i in range(0, max(1, len(words)), step):
        text = " ".join(words[i:i + CHUNK_WORDS])
        if text:
            cid = hashlib.sha1(f"{path}:{i}".encode()).hexdigest()[:12]
            out.append({"id": cid, "source": os.path.basename(path), "text": text})
    return out


def make_corpus(n: int) -> list:
    root = tempfile.mkdtemp()
    paths = []
    for i in range(n):
        p = os.path.join(root, f"page{i}.html")
        with open(p, "w", encoding="utf-8") as fh:
            fh.write("<p>" + " ".join(f"token{j}" for j in range(50 * (i % 7 + 1))) + "</p>")
        paths.append(p)
    return paths


if __name__ == "__main__":
    paths = sorted(make_corpus(40), key=os.path.getsize, reverse=True)   # big first: shorter tail
    records = []
    with mp.get_context("spawn").Pool(4) as pool:
        for chunks in pool.imap_unordered(chunk_file, paths, chunksize=2):
            records.extend(chunks)                       # parent: batch upserts to the vector store here
    print(len(records), len({r["id"] for r in records}))   # 89 89

import multiprocessing as mp
import signal
import sys
import time


class TaskTimeout(Exception):
    pass


def _alarm(signum, frame):
    raise TaskTimeout("document took too long")


def parse_with_deadline(args: tuple) -> tuple:
    doc_id, seconds_of_work = args
    signal.signal(signal.SIGALRM, _alarm)
    signal.alarm(1)                                      # per-task deadline in whole seconds
    try:
        time.sleep(seconds_of_work)                      # stand-in for a parser stuck on a bad PDF
        return doc_id, "ok"
    except TaskTimeout as exc:
        return doc_id, f"timeout: {exc}"
    finally:
        signal.alarm(0)


if __name__ == "__main__":
    if sys.platform != "win32":
        with mp.get_context("spawn").Pool(2) as pool:
            print(pool.map(parse_with_deadline, [("a", 0.1), ("b", 5), ("c", 0.1)]))
        # [('a', 'ok'), ('b', 'timeout: document took too long'), ('c', 'ok')]

import multiprocessing as mp
import os
import tempfile

import numpy as np

_vocab = None


def init_tokenizer(vocab_size: int) -> None:
    global _vocab
    os.environ["TOKENIZERS_PARALLELISM"] = "false"       # HF fast tokenizers: avoid nested threads
    _vocab = vocab_size                                   # stand-in for AutoTokenizer.from_pretrained(...)


def tokenize_shard(job: tuple) -> tuple:
    in_path, out_path = job
    ids = []
    with open(in_path, encoding="utf-8") as fh:
        for line in fh:                                   # real code: tokenizer(batch_of_lines)["input_ids"]
            ids.extend(hash(w) % _vocab for w in line.split())
    np.save(out_path, np.asarray(ids, dtype=np.int32))
    return os.path.basename(out_path), len(ids)


if __name__ == "__main__":
    root = tempfile.mkdtemp()
    jobs = []
    for s in range(8):
        src = os.path.join(root, f"shard{s}.txt")
        with open(src, "w", encoding="utf-8") as fh:
            fh.write("the quick brown fox\n" * 1000)
        jobs.append((src, os.path.join(root, f"shard{s}.npy")))
    with mp.get_context("spawn").Pool(4, initializer=init_tokenizer, initargs=(50_000,)) as pool:
        stats = pool.map(tokenize_shard, jobs, chunksize=1)
    print(sum(n for _, n in stats))   # 32000

import os
from concurrent.futures import ProcessPoolExecutor
import multiprocessing as mp

_model = None


def init_model(model_name: str, threads: int) -> None:
    global _model
    os.environ["OMP_NUM_THREADS"] = str(threads)
    # import torch; torch.set_num_threads(threads)
    # from sentence_transformers import SentenceTransformer
    # _model = SentenceTransformer(model_name, device="cpu")
    _model = lambda texts: [[float(len(t)), float(t.count(" "))] for t in texts]   # stand-in, set in-worker


def embed_batch(texts: list) -> list:
    return _model(texts)                                  # real: _model.encode(texts, batch_size=32).tolist()


if __name__ == "__main__":
    texts = [f"passage number {i} about retrieval" for i in range(1000)]
    batches = [texts[i:i + 64] for i in range(0, len(texts), 64)]
    workers, threads = 2, 1                               # workers x threads <= cores
    ctx = mp.get_context("spawn")
    with ProcessPoolExecutor(workers, mp_context=ctx, initializer=init_model,
                             initargs=("all-MiniLM-L6-v2", threads)) as ex:
        vectors = [v for batch in ex.map(embed_batch, batches) for v in batch]
    print(len(vectors))   # 1000

import multiprocessing as mp
import os
import tempfile

import numpy as np


def preprocess(job: tuple) -> tuple:
    src, dst = job
    img = np.load(src)                                    # real: PIL.Image.open(src).convert("RGB")
    small = img[::2, ::2].astype(np.float32) / 255.0      # real: resize to 224x224, normalize
    np.save(dst, small)
    return os.path.basename(dst), small.shape


if __name__ == "__main__":
    root = tempfile.mkdtemp()
    jobs = []
    rng = np.random.default_rng(0)
    for i in range(16):
        src = os.path.join(root, f"img{i}.npy")
        np.save(src, rng.integers(0, 255, (448, 448, 3), dtype=np.uint8))
        jobs.append((src, os.path.join(root, f"img{i}_224.npy")))
    with mp.get_context("spawn").Pool(4) as pool:
        out = pool.map(preprocess, jobs, chunksize=2)
    print(out[0])   # ('img0_224.npy', (224, 224, 3))

import multiprocessing as mp
import re
from collections import Counter


def f1(pred: str, gold: str) -> float:
    p, g = re.findall(r"\w+", pred.lower()), re.findall(r"\w+", gold.lower())
    common = sum((Counter(p) & Counter(g)).values())
    if not common:
        return 0.0
    prec, rec = common / len(p), common / len(g)
    return 2 * prec * rec / (prec + rec)


def score(row: tuple) -> dict:
    qid, pred, gold = row
    return {"id": qid, "em": float(pred.strip().lower() == gold.strip().lower()), "f1": f1(pred, gold)}


if __name__ == "__main__":
    rows = [(i, "Paris is the capital", "Paris") for i in range(2000)]
    with mp.get_context("spawn").Pool(4) as pool:
        scores = pool.map(score, rows, chunksize=100)
    print(round(sum(s["f1"] for s in scores) / len(scores), 3))   # 0.4

import hashlib
import multiprocessing as mp
import random
from collections import defaultdict

NUM_PERM, BANDS = 64, 16
ROWS = NUM_PERM // BANDS
_seeds = None


def init(seed: int) -> None:
    global _seeds
    rng = random.Random(seed)                                     # same seeds in every worker
    _seeds = [rng.getrandbits(32) for _ in range(NUM_PERM)]


def signature(doc: tuple) -> tuple:
    doc_id, text = doc
    words = text.split()
    shingles = {" ".join(words[i:i + 3]) for i in range(max(1, len(words) - 2))}
    hashes = [int.from_bytes(hashlib.blake2b(s.encode(), digest_size=8).digest(), "little") for s in shingles]
    return doc_id, tuple(min(h ^ seed for h in hashes) for seed in _seeds)


if __name__ == "__main__":
    base = "retrieval augmented generation combines search with a language model to answer questions"
    docs = [(i, base + (" extra" if i % 2 else "")) for i in range(200)] + [(999, "a totally different text here")]
    with mp.get_context("spawn").Pool(4, initializer=init, initargs=(42,)) as pool:
        sigs = pool.map(signature, docs, chunksize=25)
    buckets = defaultdict(set)
    for doc_id, sig in sigs:                                       # LSH banding in the parent
        for b in range(BANDS):
            buckets[(b, sig[b * ROWS:(b + 1) * ROWS])].add(doc_id)
    dupes = {d for ids in buckets.values() if len(ids) > 1 for d in ids}
    print(len(dupes), 999 in dupes)   # 200 False

import multiprocessing as mp
import os


def gpu_worker(device: int, jobs, results) -> None:
    os.environ["CUDA_VISIBLE_DEVICES"] = str(device)     # before importing torch in this process
    # import torch; model = load_model().to("cuda").eval()
    while (batch := jobs.get(timeout=60)) is not None:
        results.put((device, [len(t) for t in batch]))   # real: model(**tokenize(batch)) under no_grad
    results.put(None)


if __name__ == "__main__":
    ctx = mp.get_context("spawn")
    n_gpus = 2                                           # torch.cuda.device_count() in a subprocess, or nvidia-smi -L
    jobs, results = ctx.Queue(maxsize=8), ctx.Queue()
    procs = [ctx.Process(target=gpu_worker, args=(d, jobs, results)) for d in range(n_gpus)]
    for p in procs:
        p.start()
    batches = [[f"text {i}"] * 32 for i in range(20)]
    for b in batches:
        jobs.put(b, timeout=60)
    for _ in procs:
        jobs.put(None)
    done, outputs = 0, []
    while done < n_gpus:
        item = results.get(timeout=120)
        if item is None:
            done += 1
        else:
            outputs.append(item)
    for p in procs:
        p.join(timeout=30)
    print(len(outputs))   # 20

import asyncio
import concurrent.futures as cf
import multiprocessing as mp


def parse(html: str) -> int:                              # CPU-bound, module level
    return len(html.split())


async def fetch(i: int) -> str:
    await asyncio.sleep(0.01)                             # stand-in for httpx.AsyncClient.get
    return "<p>" + "word " * (i + 1) + "</p>"


async def pipeline(n: int, pool: cf.ProcessPoolExecutor) -> int:
    loop = asyncio.get_running_loop()
    sem = asyncio.Semaphore(8)                            # bound fetches and parse jobs in flight

    async def one(i: int) -> int:
        async with sem:
            html = await fetch(i)
            return await loop.run_in_executor(pool, parse, html)

    return sum(await asyncio.gather(*(one(i) for i in range(n))))


if __name__ == "__main__":
    with cf.ProcessPoolExecutor(max_workers=2, mp_context=mp.get_context("spawn")) as pool:
        print(asyncio.run(pipeline(20, pool)))   # 230

import hashlib
import multiprocessing as mp
import random


def stable_seed(base: int, item_id: str) -> int:
    return int.from_bytes(hashlib.blake2b(f"{base}:{item_id}".encode(), digest_size=8).digest(), "little")


def augment(item_id: str) -> float:
    rng = random.Random(stable_seed(1234, item_id))
    return round(rng.random(), 6)


if __name__ == "__main__":
    ids = [f"doc{i}" for i in range(100)]
    ctx = mp.get_context("spawn")
    with ctx.Pool(2) as a, ctx.Pool(4) as b:
        r1 = a.map(augment, ids, chunksize=7)
        r2 = b.map(augment, ids, chunksize=1)
    print(r1 == r2 == [augment(i) for i in ids])   # True