Python Shared Memory Multiprocessing: Batch Features Without Pickle Blowups

Feature jobs that Pool.map multi-gigabyte NumPy arrays die on pickle: every worker gets a full copy, RSS explodes, and the parent stalls serializing. The unfair advantage is multiprocessing.shared_memory (or RawArray) with explicit lifetime and dtype contracts so workers see the same bytes without a second heap.

⚡ TL;DR: Put dense float32/int64 matrices in SharedMemory; pass only (name, shape, dtype) handles; never pickle the array body; unlink aggressively; pin worker count to cores × memory budget. Pair with NumPy Operator Fusion and Python Prefork Pitfalls.

Why pickle melts Batch

# ❌ Classic blow-up — each task pickles X
from multiprocessing import Pool
import numpy as np

def score_row(i_x):
    i, x = i_x
    return float(x[i].sum())

X = np.random.rand(2_000_000, 128).astype(np.float32)  # ~1 GB
with Pool(16) as pool:
    # pickle copies X into every task payload → disaster
    pool.map(score_row, [(i, X) for i in range(0, len(X), 10_000)])

SharedMemory pattern that survives

# ✅ Parent allocates once; workers attach by name
from multiprocessing import Pool, shared_memory
import numpy as np

def _init_worker(shm_name: str, shape: tuple[int, ...], dtype: str):
    global SHM, ARR
    SHM = shared_memory.SharedMemory(name=shm_name)
    ARR = np.ndarray(shape, dtype=np.dtype(dtype), buffer=SHM.buf)

def score_block(start_end: tuple[int, int]) -> float:
    start, end = start_end
    # ✅ read-only view into shared buffer — no copy
    return float(ARR[start:end].sum())

def run(X: np.ndarray, n_workers: int = 8) -> list[float]:
    shm = shared_memory.SharedMemory(create=True, size=X.nbytes)
    arr = np.ndarray(X.shape, dtype=X.dtype, buffer=shm.buf)
    arr[:] = X  # one copy into SHM
    blocks = [(i, min(i + 50_000, len(X))) for i in range(0, len(X), 50_000)]
    try:
        with Pool(
            n_workers,
            initializer=_init_worker,
            initargs=(shm.name, X.shape, X.dtype.str),
        ) as pool:
            return pool.map(score_block, blocks)
    finally:
        shm.close()
        shm.unlink()  # ✅ always unlink or you leak /dev/shm

Ownership and failure modes

# ✅ Partition writes: each worker owns a disjoint row range
def fill_features(start_end: tuple[int, int]) -> None:
    start, end = start_end
    ARR[start:end, 0] = ARR[start:end, 1:5].mean(axis=1)

❌ Sharing Python list/dict graphs via managers for numeric work — you reintroduce serialization and lock contention.

AWS Batch sizing

Size tasks so X.nbytes + worker_overhead * n fits in the instance. Prefer one shared matrix per job over shipping shards through S3 for every subprocess when the matrix already fits in RAM.

Closing checklist

✅ Dos
– ✅ Pass SHM name + shape + dtype only
– ✅ unlink in finally and on SIGTERM handlers
– ✅ Use np.ndarray(..., buffer=shm.buf) views
– ✅ Prefer spawn start method for clarity on Linux containers
– ✅ Cap workers by memory, not just vCPU count

❌ Don’ts
– ❌ Don’t pickle multi-GB arrays into Pool.map args
– ❌ Don’t mutate overlapping ranges without barriers
– ❌ Don’t leave SHM segments for “debug later”
– ❌ Don’t mix managed dicts with dense numeric pipelines

Related reading

• NumPy Operator Fusion: Stop Temporary Blowups in AWS Batch ETL
• Python Prefork Pitfalls
• Python Free-Threading
• N-API Addons for Node (same idea: move hot loops)

Last updated on September 11, 2026