Ctrl+K

Introduction to MPI

Contents

Introduction to MPI#

This section adapts mpitutorial.com materials using IPython Parallel and mpi4py to run MPI code in Jupyter notebooks. We won’t go into detail in using IPython Parallel, but cover the key bits for getting started.

mpitutorial.com materials are used under the MIT License.

import ipyparallel as ipp

# create a cluster
cluster = ipp.Cluster(engines="mpi", n=2)
# start that cluster and connect to it
rc = cluster.start_and_connect_sync()

Starting 2 engines with <class 'ipyparallel.cluster.launcher.MPIEngineSetLauncher'>

What did that do?

mpiexec -n 2 python -m ipyparallel.engine --mpi

cluster.engine_set.args

['mpiexec',
 '-n',
 '2',
 '/__w/mpi-tutorial/mpi-tutorial/3/envs/mpi-tutorial/bin/python',
 '-m',
 'ipyparallel.engine',
 '--mpi']

If we ‘activate’ the client, it registers magics with IPython, so we can use %%px to run cells on the engines instead of in the local notebook.

rc.activate()

<DirectView all>

Now we have %%px available:

%%px
import os
pid = os.getpid()
pid

Out[0:1]: 4053

Out[1:1]: 4054

A cell passed to %%px is run on all engines at once. This is the equivalent of mpiexec myscript.py, when running noninteractive MPI.

From now on, notebooks will start with a brief boilerplate to start and register the cluster, so we can use %%px.

Rank and size#

Our very first MPI code, to test %%px. We are going to get the “MPI World communicator”.

The rank is the integer id of the current process, while the size is the number of processes in the communicator.

%%px
# Find out rank, size
from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.rank
size = comm.size

print(f"I am rank {rank} / {size}")

[stdout:0] 
I am rank 0 / 2
[stdout:1] 
I am rank 1 / 2

In IPython Parallel, the state is persistent. This means the rank and size variables can be used in subsequent cells:

%%px
print(f"Rank {rank} has PID {pid}")

[stdout:0] 
Rank 0 has PID 4053
[stdout:1] 
Rank 1 has PID 4054

To translate a notebook written with %%px to a script for mpiexec, you would concatenate all the %%px cells into a single .py file.

Now we can stop the cluster if we want. It should get cleaned up automatically when the notebook exits, but it’s good to be explicit.

cluster.stop_cluster_sync()

Stopping controller

Controller stopped: {'exit_code': 0, 'pid': 4014, 'identifier': 'ipcontroller-1704465289-hilm-3993'}

Stopping engine(s): 1704465290

engine set stopped 1704465290: {'exit_code': 0, 'pid': 4048, 'identifier': 'ipengine-1704465289-hilm-1704465290-3993'}
Contents