Numba-MPI
The Numba-MPI package enables just-in-time (JIT) compiled Python code using Numba to call Message Passing Interface (MPI) routines directly, allowing parallel communication without leaving compiled execution blocks. Designed as a lightweight, NumPy-based wrapper around the MPI C API, it provides cross-platform high-performance communication capabilities, bridging the gap between Python productivity and native MPI efficiency. We have outlined the project design, features and performance in a SoftwareX paper (Derlatka, Manna, Bulenok, Zwicker & Arabas 2024).
Numba-MPI has been developed by an open group of developers from Jagiellonian University in Kraków, from the Max Planck Institute (MPI-DS) in Göttingen, from our team at the AGH University of Krakow, Poland, with much appreciated contributions from the users' community.
Numba-MPI is an optional dependency of py-pde and PyMPDATA-MPI.
Project resources
- Documentation
- PyPI package
- Anaconda package
- AUR package (Arch Linux)
- GitHub repository
- Zenodo release archive
Project overview video (FOSDEM talk)
🎓 Student project opportunities
- Engineering support for
MPI_Type_create_struct(mapping to NumPy structured arrays) - Redesigning the API to allow for custom JIT-compiler flags to be passed to Numba (e.g.,
boundscheck,fastmath,nogil,cacheorerror_model) - Extending the API to support non-default communicators (for subsetting groups of processes)
- Engineering support for
MPI_IN_PLACEin [all]gather/scatterandallreduce - Extending the test suite to include hybrid threading+MPI parallelism with performance checks
- Adding support for decoding MPI status codes by wrapping
MPI_ERROR_CLASS - Enabling support for passing CuPy arrays to MPI calls (key challenge: GPU support in CI)
Contributors from our team
Kacper Derlatka
Oleksii Bulenok
Michał Wroński
