mpi4pytools 0.2.0

Simple decorators and utilities for MPI parallel computing

mpitools

⚠️ Development Notice: This package is in active development. The API may change significantly between versions until v1.0.0. Use in production environments is not recommended.

A Python package providing simple decorators and utilities for MPI (Message Passing Interface) parallel computing. Built on top of mpi4py, mpitools makes it easy to write parallel code with minimal boilerplate.

Features

• Work distribution decorators: Execute functions on specific ranks or groups of processes
• Communication decorators: Collective communications and reduce operations made simple
• Error handling: Graceful error handling across all MPI processes
• Task queue system: Distributed task processing with the queue submodule

Installation

pip install mpi4pytools

Requirements:

• Python 3.7+
• mpi4py
• An MPI implementation (OpenMPI, MPICH, etc.)

Quick Start

Basic Usage

from mpitools import setup_mpi, broadcast_from_main, gather_to_main, eval_on_main

# Initialize MPI environment
comm, rank, size = setup_mpi()

# Execute only on rank 0, broadcast result to all processes
@broadcast_from_main()
def load_config():
    return {"num_iterations": 1000, "tolerance": 1e-6}

# Execute on all processes, gather results to rank 0
@gather_to_main()
def compute_partial_sum():
    return sum(range(rank * 100, (rank + 1) * 100))

# Execute only on rank 0
@eval_on_main()
def save_results(data):
    with open("results.txt", "w") as f:
        f.write(str(data))

# Usage
config = load_config()  # Same config on all processes
partial_sums = compute_partial_sum()  # List of sums on rank 0, None elsewhere
save_results(partial_sums)  # Only saves on rank 0

Task Queue System

from mpitools import setup_mpi
from mpitools.queue import MPIQueue, Task

# Initialize MPI environment
comm, rank, size = setup_mpi()

# Define a task class
class MyTask(Task):
    def __init__(self, task_id: str, data: int):
        super().__init__(task_id)
        self.data = data

    def execute(self):
        # Perform some computation
        result = self.data * 2  # Example computation
        return result

# Create a distributed task queue
queue = MPIQueue()

# Add tasks to the queue
if rank == 0:
    tasks = [MyTask(f"task_{i}", i) for i in range(10)]
    queue.add_tasks(tasks)

# Run the task queue
results = queue.run()

Error Handling

from mpitools import abort_on_error

@abort_on_error()  # Aborts all processes if any process encounters an error
def risky_computation():
    # If this fails on any process, all processes will terminate
    result = 1 / some_calculation()
    return result

Core Decorators

Error Handling

• @abort_on_error() - Abort all processes if any process raises an exception

Work Distribution

• @eval_on_main() - Execute only on rank 0
• @eval_on_workers() - Execute only on worker ranks (1, 2, ...)
• @eval_on_single(rank) - Execute only on specified rank
• @eval_on_select([ranks]) - Execute only on specified ranks

Collective Communication

• @broadcast_from_main() - Execute on rank 0, broadcast result to all processes
• @broadcast_from_process(rank) - Execute on specified rank, broadcast to all processes
• @scatter_from_main() - Execute on rank 0, scatter data to all processes
• @scatter_from_process(rank) - Execute on specified rank, scatter data to all processes
• @gather_to_main() - Execute on all processes, gather results to rank 0
• @gather_to_process(rank) - Execute on all processes, gather results to specified rank
• @gather_to_all() - Execute on all processes, gather results to all processes
• @all_to_all() - Execute on all processes, exchange data between all processes

Reduction Operations

• @reduce_to_main(op='sum') - Execute on all processes, reduce to rank 0
• @reduce_to_process(rank, op='sum') - Execute on all processes, reduce to specified rank
• @reduce_to_all(op='sum') - Execute on all processes, reduce to all processes

Supported reduction operations: 'sum', 'prod', 'max', 'min', 'land', 'band', 'lor', 'bor', 'lxor', 'bxor', 'maxloc', 'minloc'

Decorator Variants

• @buffered_* - Buffered versions of collective communication and reduction operations for improved performance.
• @variable_* - Variable-sized versions of buffered scatter, gather and all_to_all communications for handling dynamic data sizes.
• Currently, only numpy arrays are supported.

Running MPI Programs

# Run with 4 processes
mpirun -n 4 python your_script.py

# Run with specific hosts
mpirun -n 4 -H host1,host2 python your_script.py

Documentation

• Full API Reference - Complete documentation of all functions and classes
• MPI4PY Documentation - Underlying MPI library

Contributing

Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

• Built on mpi4py