GitHub - shikokuchuo/pymizu: Thread-speed async and parallel Python over shared memory

GitHub

5 min read Original article ↗

pymizu logo

PyPI ci codecov License: MIT Ruff Python 3.11+

Thread-speed async and parallel Python over shared memory.

A task round trip drops from ~100 µs over sockets to < 1 µs, so you can divide work as finely as you would with threads (see Benchmarks). Common data structures like lists, dicts, numpy arrays, and Arrow tables can be shared between Python and R as zero-copy views.

In the default CPython build, the GIL runs CPU-bound threads on one core at a time, so compute parallelism in Python usually means multiple processes. Processes also isolate failures: a worker that crashes does not take down the host with it.

Channels and work-stealing task pools run over POSIX shared memory (Linux, macOS) or Win32 file mappings (Windows): lock-free single-producer single-consumer rings with batched publication, spin-then-park waiting that stays out of the kernel on the hot path, and OS notification when a peer dies.

pymizu is built on libmizu, a C library for lock-free shared-memory IPC. mizu binds the same core for R, so Python and R can talk to each other.

  • Zero-copy numpy arrays and Arrow interchange (pyarrow, polars, pandas, duckdb)
  • Drop-in concurrent.futures.Executor — real Futures, wait, as_completed, asyncio.wrap_future
  • Event-driven asyncio — await tasks and maps, async for and recv_async() on channels — wake-ups as data arrives, no polling threads
  • Reproducible parallel random numbers (seed=) — identical results for any worker count or steal order
  • Nested parallel maps — a task can map over its own pool, deadlock-free

Lifecycle: experimental. The API is not yet stable and we may make changes in response to feedback.

Installation

Install from PyPI:

Or the development version from GitHub:

pip install git+https://github.com/shikokuchuo/pymizu

Requires Python 3.11 or later on a 64-bit platform. The extension compiles the vendored C core, so no system library is necessary.

Optional extras:

  • pymizu[numpy]: zero-copy array views and raw-tier array staging
  • pymizu[cloudpickle]: lambdas, closures, and local functions as pool tasks

To request an extra with the GitHub install, use pip install "pymizu[numpy] @ git+https://github.com/shikokuchuo/pymizu".

Quickstart

import pymizu

with pymizu.Pool.create(4) as pool:
    task = pool.submit(pow, 2, 16)
    print(task.collect(timeout=5))

Benchmarks

Communication overhead against the stdlib concurrent.futures pools (Apple M4 Pro, from benchmarks/mizu-stdlib-bench.py):

Against ProcessPoolExecutor (tasks run in separate processes, with pickled payloads):

Benchmark pymizu ProcessPoolExecutor Speedup
Trivial task round trip 0.6 µs 94.2 µs 157x
Pipelined throughput, 1 worker 2,390,000 tasks/s 18,500 tasks/s 129x
8 KB vector round trip 2.4 µs 100.7 µs 42x
Parallel map of 2,000 ~5 µs tasks, 4 workers 2.9 ms 123 ms 42x

Against ThreadPoolExecutor (tasks share one process, so the GIL caps CPU-bound work at a single core):

Benchmark pymizu ThreadPoolExecutor Speedup
Trivial task round trip 0.6 µs 9.5 µs 16x
Pipelined throughput, 1 worker 2,390,000 tasks/s 405,000 tasks/s 5.9x
8 KB vector round trip 2.4 µs 9.3 µs 3.9x
Parallel map of 2,000 ~5 µs tasks, 4 workers 2.9 ms 48 ms 17x

benchmarks/mizu-bench.py runs the pymizu rows standalone. See how pymizu compares with Ray, Dask, joblib, and modern CPython.

concurrent.futures and asyncio

PoolExecutor adapts a pool to the stdlib concurrent.futures.Executor interface, so code written against ProcessPoolExecutor drops in. submit returns real Future objects, so wait / as_completed / asyncio.wrap_future all work unchanged:

with pymizu.PoolExecutor.create(4) as ex:
    futures = [ex.submit(pow, 2, i) for i in range(4)]
    print([f.result() for f in futures])

Pool.submit() returns an awaitable Task, so native asyncio code needs no executor or asyncio.wrap_future. Batch collectors and map handles are awaitable too:

with pymizu.Pool.create(4) as pool:
    task = pool.submit(pow, 2, 16)
    print(await task)

Parallel map

Pool.map(fn, x) maps fn over x on the pool and returns a Map handle immediately, just as Pool.submit() returns a task. collect() on the handle waits and returns the values in order. The function and data cross into shared memory once. Workers divide the elements among themselves, claiming ranges from a shared cursor.

A map also nests: the mapped function can map over its own pool with pymizu.current_pool(), deadlock-free.

with pymizu.Pool.create(4) as pool:
    m = pool.map(abs, range(-5, 5))
    print(m.collect())

R interop

Pool workers can be R processes that run the mizu package, the R binding of the same core. pymizu.r_launcher() spawns them and pymizu.call() specifies the task. Here, R's built-in mtcars is summarized in R and read back into polars:

import polars as pl
import pymizu

with pymizu.Pool.create(4, launcher=pymizu.r_launcher()) as pool:
    task = pool.submit(
        pymizu.call(source="aggregate(mpg ~ cyl, data = mtcars, FUN = mean)")
    )
    print(pl.DataFrame(task.collect()))
  • A numpy array arrives in R as a numeric vector, and a polars, pyarrow, or pandas frame as a data frame
  • A vector or data frame arrives back as a numpy array or a pymizu.Frame
  • Arrow data (pyarrow tables and record batch readers, or any __arrow_c_stream__ producer) crosses the same way

The R interop page has the full contract: the portable subset, pymizu.Frame, zero-copy frames, and the dtype matrix. The mizu README shows the reverse direction: an R host that drives Python workers with mizu_py_launcher().

Documentation

The documentation site covers the full package: channels and what crosses them, task pools (batch operations, nested tasks, the concurrent.futures executor, watching a running pool), the parallel map (scheduling, reproducible randomness, templates, prepared maps), R interop and the dtype matrix, deployment on Linux, and the API reference.

Contributing

See CONTRIBUTING.md for the repository layout and contributor notes.

License

MIT. LICENSE.note has the RngStreams attribution for the vendored libmizu core.


Please note that this project is released with a Contributor Code of Conduct. By participating in this project you agree to abide by its terms.