MPI Support

bamboost supports running simulation scripts in parallel using MPI via mpi4py. The MPI layer is fully transparent — the same Python code works both with and without MPI.

The MPI layer is subject to change as we gain more experience with it.

How MPI detection works

bamboost uses automatic detection to decide whether to activate MPI support. On startup, it checks:

Whether config.options.mpi is True (default: auto-detected based on mpi4py being installed).
Whether the environment variable BAMBOOST_MPI is not set to "0".
Whether any of the standard MPI launcher environment variables are set (OMPI_COMM_WORLD_SIZE, PMI_SIZE, SLURM_PROCID, etc.).

If all three checks pass, the real mpi4py.MPI module is used. Otherwise, a lightweight mock MPI implementation is used, so that single-process code requires no changes.

The mock communicator has rank == 0 and size == 1, which means all rank-checks and barriers are satisfied transparently when running without MPI.

Prerequisites

To use bamboost with MPI you need:

A working MPI installation (Open MPI, MPICH, Intel MPI, etc.)
mpi4py (pip install mpi4py)
h5py compiled with Parallel HDF5 support

Check whether your h5py supports MPI:

import h5py
print(h5py.get_config().mpi)  # True if compiled with MPI

Many Linux distributions provide a parallel HDF5 package. Alternatively, follow the h5py MPI build guide.

Writing MPI-aware simulation scripts

No special code is needed for read-only access. For writing, bamboost serialises operations that must run only on rank 0 (attribute writes, metadata updates) through a single-process queue, while datasets can be written from all ranks using the mpio file driver.

import os
import numpy as np
from bamboost import SimulationWriter

with SimulationWriter.from_uid(os.environ["SIMULATION_ID"]) as sim:
    # Metadata writes are queued and executed on rank 0 only
    sim.metadata["description"] = "MPI run"

    for n, t in enumerate(solver.times()):
        state = solver.step(state)
        step = sim.data.require_step(value=t, step=n)

        # field data is written in parallel — each rank contributes its
        # local portion; bamboost concatenates along the first axis
        step.add_field("phi", state.local_phi)
        step.add_scalar("energy", state.energy)   # only rank 0 writes

Launching with MPI

Use your MPI launcher to start the run script:

mpirun -n 4 python run_script.py
# or via SLURM
srun --ntasks=4 python run_script.py

The create_run_script / SimulationWriter.submit helpers automatically forward the correct SIMULATION_ID environment variable to the launched process.

Disabling MPI

To force single-process mode even when mpi4py is installed:

bamboost.toml

[options]
mpi = false

Or at runtime via an environment variable:

BAMBOOST_MPI=0 python script.py