richpool

Versions
Status
Tools
License

richpool is a kind of fork that mixes both p_tqdm and schwimmbad into one standalone library, with progress bars rendered natively by rich instead of tqdm. It has no dependency on tqdm, p_tqdm, or schwimmbad.

Since, both p_tqdm and schwimmbad are effectively unmaintained, richpool reimplements what both projects offer, aiming to be the successor to both, with the added functionality of a native rich progress bar.

It gives you two ways to run parallel work, both with a rich progress bar by default:

1. A schwimmbad-style pool interface: SerialPool, MultiPool, JoblibPool, MPIPool, selected via choose_pool(), each with a uniform .map() method. This mirrors all four pool types in schwimmbad.
2. A p_tqdm-style functional interface: p_map, p_imap, p_umap, p_uimap, t_map, t_imap. No pool object to create or manage.

Installation

uv pip install richpool

MPIPool needs mpi4py (also needs a system MPI, e.g. OpenMPI/MPICH) installed too. Everything else (SerialPool, MultiPool, JoblibPool, the functional API) works out of the box:

uv pip install "richpool[mpi]"

Pool interface (schwimmbad)

from richpool import choose_pool


def square(x):
    return x * x


with choose_pool(processes=4) as pool:
    results = pool.map(square, range(20), desc="squaring")
# results == [0, 1, 4, ..., 361]

choose_pool(mpi=False, processes=1, pool=None, **kwargs) picks a pool:

• mpi=True picks MPIPool
• processes != 1 picks MultiPool (backed by pathos.multiprocessing.ProcessPool), or JoblibPool if pool="joblib"
• otherwise picks SerialPool

pool is a richpool extension on top of schwimmbad's choose_pool(mpi, processes, **kwargs) signature: it defaults to None, which keeps the original schwimmbad-compatible behavior (MultiPool), so existing calls keep working unchanged. Pass pool="joblib" to opt into JoblibPool instead, and pass any joblib.Parallel kwarg (e.g. backend="loky" or backend="threading") alongside it. backend is only meaningful for JoblibPool, so choose_pool() forwards it only when pool="joblib" and silently drops it for SerialPool/MultiPool/MPIPool.

All four pool classes share the same .map() interface:

pool.map(func, iterable, callback=None, desc="", total=None, disable=False)

• callback: optional, called on the master process with each individual result as it completes (same contract as schwimmbad's callback).
• desc: progress bar description.
• total: override the progress bar total (inferred from len(iterable) when omitted).
• disable: suppress the progress bar.

You can also instantiate pools directly:

from richpool import SerialPool, MultiPool, JoblibPool, MPIPool

SerialPool()
MultiPool(processes=8)
JoblibPool(processes=8, backend="loky")  # any joblib.Parallel kwargs
MPIPool()  # run script with: mpiexec -n 4 python script.py

MPIPool notes

MPIPool distributes tasks across MPI ranks using mpi4py. Only the master process (rank 0) returns from map() with results and renders the progress bar; other ranks block in a worker loop and map() returns None for them. Launch scripts with mpiexec/mpirun:

from richpool import choose_pool


def square(x):
    return x * x


with choose_pool(mpi=True) as pool:
    if pool.is_master():
        results = pool.map(square, range(20), desc="squaring")
        print(results)

mpiexec -n 4 python script.py

Needs at least 2 MPI ranks (1 master + >=1 worker).

Progress bar rendering under mpiexec

• inplace (default on a terminal): each line is prefixed with a cursor-up-and-erase escape, so the bar redraws over itself and you see a single, live bar.
• lines (default otherwise): plain lines with no escapes, throttled to roughly one print per worker, so logs and CI output stay readable. This is what you get when the bar is redirected to a file.

The mode is autodetected from whether the output stream is a terminal, falling back to TERM. Override it with an environment variable:

RICHPOOL_MPI_PROGRESS=inplace mpiexec -n 4 python script.py   # force a single redrawn bar
RICHPOOL_MPI_PROGRESS=lines   mpiexec -n 4 python script.py   # force the plain scrolling log

Functional interface (p_tqdm)

No pool to create, just call the function:

from richpool import p_map, p_umap, t_map


def add(a, b):
    return a + b


p_map(add, [1, 2, 3], [10, 20, 30])  # parallel, ordered:   [11, 22, 33]
p_umap(add, [1, 2, 3], [10, 20, 30])  # parallel, unordered: e.g. [22, 11, 33]
t_map(add, [1, 2, 3], [10, 20, 30])  # sequential, ordered: [11, 22, 33]

• p_map / p_imap: parallel ordered map / iterator.
• p_umap / p_uimap: parallel unordered map / iterator (results as they complete).
• t_map / t_imap: sequential map / iterator.

All accept num_cpus (int, or float as a proportion of available CPUs), total, desc, disable, and chunksize keyword arguments. The p_* functions also accept kind ("process" or "thread", default "process") to pick the underlying pool.

from functools import partial
from richpool import p_map


def add(a, b, c=0):
    return a + b + c


p_map(partial(add, c=1), [1, 2, 3], [10, 20, 30], num_cpus=0.5, desc="adding")

Examples

Every pool and every functional map has a minimal, runnable example under examples/:

python examples/serial_pool_example.py
python examples/multi_pool_example.py
python examples/joblib_pool_example.py
mpiexec -n 4 python examples/mpi_pool_example.py

python examples/functional_pmap_example.py
python examples/functional_pimap_example.py
python examples/functional_pumap_example.py
python examples/functional_puimap_example.py
python examples/functional_tmap_example.py
python examples/functional_timap_example.py

See them run with their real output, progress bar included, on the examples page of the docs.

Credits

Parts of richpool's code and docs are adapted directly from schwimmbad (Copyright (c) 2016 Adrian Price-Whelan) and p_tqdm (Copyright (c) 2024 Kyle Swanson), both MIT licensed. See THIRD_PARTY_NOTICES.md for their original license text.

License

MIT, see LICENSE.