Testing Multiprocessing Code

import concurrent.futures as cf
from typing import Callable, Iterable


def chunk_doc(doc: str, size: int = 4) -> list:
    words = doc.split()
    return [" ".join(words[i:i + size]) for i in range(0, len(words), size)]


def ingest(docs: Iterable[str], make_executor: Callable[[], cf.Executor]) -> list:
    with make_executor() as ex:
        return [c for chunks in ex.map(chunk_doc, docs) for c in chunks]


class InlineExecutor(cf.Executor):
    """Runs tasks synchronously in the caller: deterministic, debuggable, no pickling."""

    def submit(self, fn, /, *args, **kwargs):
        fut = cf.Future()
        try:
            fut.set_result(fn(*args, **kwargs))
        except BaseException as exc:
            fut.set_exception(exc)
        return fut


def test_ingest_inline():
    assert ingest(["a b c d e"], InlineExecutor) == ["a b c d", "e"]


test_ingest_inline()
print("ok")   # ok

import functools
import pickle


class ParseError(Exception):
    def __init__(self, doc_id: str, reason: str) -> None:
        super().__init__(doc_id, reason)
        self.doc_id, self.reason = doc_id, reason


def normalize(doc: str, *, lower: bool) -> str:
    return doc.lower() if lower else doc


def roundtrip(obj):
    return pickle.loads(pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL))


def test_task_callable_pickles():
    fn = roundtrip(functools.partial(normalize, lower=True))
    assert fn("ABC") == "abc"


def test_exceptions_pickle():
    err = roundtrip(ParseError("doc-1", "encrypted"))
    assert (err.doc_id, err.reason) == ("doc-1", "encrypted")


def test_payload_is_small():
    args = ("s3://bucket/doc-1.pdf", 1)
    assert len(pickle.dumps(args)) < 1024          # paths, not contents


for t in (test_task_callable_pickles, test_exceptions_pickle, test_payload_is_small):
    t()
print("ok")   # ok

# test_pool_orchestration.py
import concurrent.futures as cf
import multiprocessing as mp

SPAWN = mp.get_context("spawn")


def square(x: int) -> int:
    return x * x


def test_map_preserves_order():
    with cf.ProcessPoolExecutor(2, mp_context=SPAWN) as ex:
        assert list(ex.map(square, range(10), chunksize=3)) == [x * x for x in range(10)]


def test_imap_unordered_returns_everything():
    with SPAWN.Pool(2) as pool:
        assert sorted(pool.imap_unordered(square, range(10), chunksize=2)) == [x * x for x in range(10)]


if __name__ == "__main__":
    test_map_preserves_order()
    test_imap_unordered_returns_everything()
    print("ok")   # ok

import multiprocessing as mp

import pytest


@pytest.fixture(scope="module")
def pool():
    with mp.get_context("spawn").Pool(2) as p:
        yield p

import faulthandler
import sys

faulthandler.dump_traceback_later(600, exit=True, file=sys.stderr)   # conftest.py: hard ceiling
faulthandler.cancel_dump_traceback_later()                            # (cancelled here so the example exits)
print("armed")   # armed

import concurrent.futures as cf
import multiprocessing as mp

SPAWN = mp.get_context("spawn")


def parse(doc: str) -> str:
    if doc == "bad":
        raise ValueError("corrupt document: bad")
    return doc.upper()


def test_task_error_propagates_with_worker_traceback():
    with cf.ProcessPoolExecutor(1, mp_context=SPAWN) as ex:
        fut = ex.submit(parse, "bad")
        try:
            fut.result(timeout=30)
        except ValueError as exc:
            assert "corrupt document" in str(exc)
            assert "in parse" in str(exc.__cause__)        # remote traceback text
        else:
            raise AssertionError("expected ValueError")


def test_error_callback_called():
    errors = []
    with SPAWN.Pool(1) as pool:
        r = pool.apply_async(parse, ("bad",), error_callback=errors.append)
        r.wait(timeout=30)
    assert len(errors) == 1 and isinstance(errors[0], ValueError)


if __name__ == "__main__":
    test_task_error_propagates_with_worker_traceback()
    test_error_callback_called()
    print("ok")   # ok

import concurrent.futures as cf
import multiprocessing as mp
import os
from concurrent.futures.process import BrokenProcessPool


def crash_on(x: int) -> int:
    if x == 2:
        os._exit(13)                                   # simulates an OOM kill / segfault
    return x


def test_crash_surfaces_as_broken_pool():
    with cf.ProcessPoolExecutor(2, mp_context=mp.get_context("spawn")) as ex:
        futs = [ex.submit(crash_on, i) for i in range(4)]
        outcomes = []
        for f in futs:
            try:
                outcomes.append(f.result(timeout=30))
            except BrokenProcessPool:
                outcomes.append("broken")
    assert "broken" in outcomes


if __name__ == "__main__":
    test_crash_surfaces_as_broken_pool()
    print("ok")   # ok

import os
import signal
import subprocess
import sys
import tempfile
import textwrap
import time

PROGRAM = textwrap.dedent("""
    import multiprocessing as mp, signal, sys, time

    def init():
        signal.signal(signal.SIGINT, signal.SIG_IGN)

    def work(x):
        time.sleep(0.2)
        return x

    if __name__ == "__main__":
        pool = mp.get_context("spawn").Pool(2, initializer=init)
        try:
            print("started", flush=True)
            pool.map(work, range(1000))
        except KeyboardInterrupt:
            pool.terminate()
            pool.join()
            sys.exit(130)
""")

if sys.platform != "win32":
    path = os.path.join(tempfile.mkdtemp(), "prog.py")
    with open(path, "w") as fh:
        fh.write(PROGRAM)
    proc = subprocess.Popen([sys.executable, path], stdout=subprocess.PIPE, text=True)
    assert proc.stdout.readline().strip() == "started"
    time.sleep(1.0)                                     # let workers start
    t0 = time.monotonic()
    proc.send_signal(signal.SIGINT)
    code = proc.wait(timeout=30)
    print(code, time.monotonic() - t0 < 10)   # 130 True

import hashlib
import multiprocessing as mp
import random


def augment(item: str) -> float:
    seed = int.from_bytes(hashlib.blake2b(item.encode(), digest_size=8).digest(), "little")
    return random.Random(seed).random()


def run(workers: int, chunksize: int) -> list:
    with mp.get_context("spawn").Pool(workers) as pool:
        return pool.map(augment, [f"doc{i}" for i in range(50)], chunksize=chunksize)


if __name__ == "__main__":
    baseline = [augment(f"doc{i}") for i in range(50)]
    assert run(1, 50) == run(2, 1) == run(3, 7) == baseline
    print("deterministic")   # deterministic

import multiprocessing as mp
import os


def square(x: int) -> int:
    return x * x


def test_no_leaked_children():
    with mp.get_context("spawn").Pool(2) as pool:
        pool.map(square, range(4))
    assert mp.active_children() == []


def open_fds() -> int:
    return len(os.listdir("/proc/self/fd")) if os.path.isdir("/proc/self/fd") else -1


if __name__ == "__main__":
    before = open_fds()
    for _ in range(3):
        test_no_leaked_children()
    print(mp.active_children(), open_fds() - before <= 2)   # [] True