Python MPI Job Example

This page shows an example of running a python job that uses a MPI for parallelization.

Overview

The example shown is a rather trivial Hello World type program in which each task prints out a "Hello" message identifying the thread number. Although trivial, it shows the basics of submitting a python job using MPI, as well as the basics of using MPI in a python script. The job also displays various other good practices, including:

• setting standard sbatch options within the script
• loading the needed modules within the script
• printing some useful diagnostic information at start of the script
• creating a job specific work directory
• running the code and saving the exit code
• exiting with the exit code from the main application

Many of the practices above are rather overkill for such a simple job --- indeed, the vast majority of lines in the submission scripts is for these "good practices" rather than the running of the intended code, but are included for educational purposes.

The submission script

The submission script can be downloaded as plain text, suitable for you to submit to the scheduler with the sbatch command. We present a copy with line numbers for discussion below.

This script is essentially the same as the hello-umd using MPI submission script, with the main difference being the mpirun command.

Source of submission script

The line numbers in the listing below are to facilitate a line-by-line discussion pf the script; they should not present in the actual script. Downloaded as plain text version of the script.

Example

1  #!/bin/bash
2  # The line above this is the "shebang" line.  It must be first line in script
3  #-----------------------------------------------------
4  #    Example Submission Script: python version of Hello World using MPI
5  #    Runs an MPI-enabled python "Hello World" script
6  #-----------------------------------------------------
7  #
8  # Slurm sbatch parameters section:
9  #       These SBATCH lines should immediately follow the "shebang" line
10 #       They must occur before any non-empty/non-comment lines in the script
11 # 
12 #    Request 60 tasks with one CPU core each
13 #SBATCH -n  60
14 #SBATCH -c 1
15 #    Request 5 minutes of walltime
16 #SBATCH -t 5
17 #    Request 1 GB of memory per CPU core
18 #SBATCH --mem-per-cpu=1024
19 #    Allow other jobs to run on same node
20 #SBATCH --oversubscribe
21 #    Run on debug partition for rapid turnaround.  You will need
22 #    to change this (remove the line) if walltime > 15 minutes
23 #SBATCH --partition=debug

24 # This job will run the code in hello_mpi.py script from the submission dir
25 # We create a working directory on parallel file system and run the job
26 # from there

27 # Process our dot files (and in particular define the "module" command)
28 source /etc/profile

29 # Module load section
30 # First clear our module list 
31 module purge
32 # and reload the standard modules
33 module load hpcc/zaratan
34 # Load the desired compiler, MPI, and python modules
35 # NOTE: You need to use the same compiler and MPI module used
36 # when compiling the python (and its mpi4py package).  Module load-ing
37 # python first, followed by openmpi will ensure you get a compatible set
38 module load python
39 module load openmpi

40 # Section to make a scratch directory for this job
41 # Because different MPI tasks, which might be on different nodes, need
42 # access to it, we put it in scratch.  We include the SLURM jobid in the
43 # directory name to avoid interference if multiple jobs running at same time.
44 TMPWORKDIR="${HOME}/scratch/mpi4py-job.${SLURM_JOBID}"
45 mkdir $TMPWORKDIR
46 cd $TMPWORKDIR

47 # Section to output information identifying the job, etc.
48 echo "Slurm job ${SLURM_JOBID} running on"
49 echo "Slurm job submit-python-mpi.sh ${SLURM_JOBID} running on:"
50 hostname
51 echo "Ntasks=${SLURM_NTASKS}; CPU cores per task=${SLURM_CPUS_PER_TASK}"
52 echo "On ${SLURM_JOB_NUM_NODES} nodes: ${SLURM_JOB_NODELIST}"
53 date
54 pwd
55 echo "Loaded modules are:"
56 module list
57 echo "Running in $TMPWORKDIR"

58 # Setting this variable will suppress the warnings
59 # about lack of CUDA support on non-GPU enabled nodes.  We
60 # are not using CUDA, so warning is harmless.
61 export OMPI_MCA_mpi_cuda_support=0

62 # Get the full path to our python executable.  It is best
63 # to provide full path of our executable to , etc.
64 MYPYTHON=`which python`
65 MYPYSCRIPT="${SLURM_SUBMIT_DIR}/hello-mpi.py"
66 echo "Using python $MYPYTHON"
67 echo "Using python script $MYPYSCRIPT"

68 # Run our script using mpirun
69 # We do not specify the number of tasks here, and instead rely on
70 # it defaulting to the number of tasks requested of Slurm
71 mpirun  ${MYPYTHON} $MYPYSCRIPT > hello.out 2>&1
72 # Save the exit code from the previous command
73 ECODE=$?

74 echo "Job finished with exit code $ECODE.  Working dir is $TMPWORKDIR"
75 date

76 # Exit with the cached exit code
77 exit $ECODE

Discussion of submission script

The following is a detailed, more or less line-by-line discussion of the job submission script:

• Line 1: This shebang line specifies that the script is to be processed using the bash shell.
  For more information on the shebang line.
• Lines 2-9: These are comment lines, which are ignored by the bash shell.
  They are only present to make the script more readable to humans. For more information on comment lines.
• Lines 12-23: These lines specify various flags to be passed to the sbatch command. For more information about #SBATCH directives.
• Lines 13-14: As this is a pure MPI job, we request a 60 MPI takss with each task having a single threads, for a total of 60 CPU cores. The scheduler is free to allocate these cores on any number (from 1 to 60) nodes as it desires. For more information about core, task, and node settings
• Line 16: This is a trivial job, so we only request 5 minutes for it. It should finish in well under 5 minutes. For more information about maximum walltime settings
• Line 18: This is a trivial job which does not need much RAM, so we only request 1 GiB per CPU core (so 60 GiB in total). Normally jobs needing no more than 4 GiB per CPU core can just omit this and get the default value of 4 GiB per core. For more information about memory settings
• Line 20: This #SBATCH directive says that we are allowing other jobs to share the same node(s) as our job. This is the default (so this directive can usually be omitted) and is recommended, especially for jobs using only a fraction of the cores available on a node (128 cores per node on the Zaratan cluster), otherwise a large number of the cores on the node will not be used by your job nor will the scheduler allow other jobs to use them. For more information about exclusive/shared mode.
• Line 23: This tells the scheduler to put our job in the debug partition. Since this is a short tutorial job, the debug partition is fine and generally will allow the job to spend less time in the queue. For more information about selecting partitions.
• Lines 28-39: These lines set up the environment within which to run your job. These are explained in more detail in the main job examples page, but in this case the lines basically:
  • Ensure the module command is defined
  • Clear out any previously loaded modules (inherited from .bashrc or similar)
  • Load the modules for python and the appropriate MPI library.
• Lines 40-46: These lines set up a work directory for our job, in this case since it is an MPI based job and will potentially run on multiple nodes, a shared filesystem is appropriate, so we use a work directory in the scratch. For more information about setting up work directories.
• Lines 47-57: These lines print out various pieces of information to identify the job. The output will go into the file "slurm-JOB_NUMBER.out" in the directory from which you issued the sbatch command (where JOB_NUMBER is the Slurm job number). There are options you can pass to sbatch to change the name of the file.
• Line 61: This line suppresses some uninmportant warnings.
• Line 64: This line uses the Unix which command to determine the full path to the python command. Generally, when using MPI, it is best to provide the full path to the command to run to the mpirun command.
• Line 65: This line sets the full path to the python script.
• Line 71: This line actually does the work. It uses the mpirun command to launch a python command running the python script for each task. python script. All output is redirected to the file hello.out in the working directory created in lines 40-46.
  I.e., this one line is the meat of the job script; everything else basically just exists to get to this point or cleanup afterwards.
  
  The "> hello.out" part means that the standard output from the command will be saved in the file hello.out in the working directory. The "2>&1" part means that any output to the standard error output stream will be redirected to the standard output. (NOTE: if you wish to save both standard output and standard error to the same file, you need to give these output redirection arguments in that order; redirect standard output first, then redirect standard error to standard output.)
• Line 73: This line saves the exit code from the hello-umd command so thatt be used when the script ends.
  For more information about exit codes.
• Lines 74-75: These lines print some additional information just before the job terminates.
  These lines will go in the same file as those from the start of the script, usually "slurm-JOBNUMBER.out".
  This just prints out:
  • The exit code from the hello-umd command (as stashed in the ECODE environment variable)
  • The time when the job finished
• Line 77: This line actually ends the job script, and terminates the job. The script and job will actually end at the end of the file even without an explicit exit command, but in that case the exit code for the job script will be that of the last command which ran, which typically will be part of the cleanup commands, and not the meat of the script.
  
  Here we explicitly set the exit code to the exit code from the hello-umd command (which we previously saved in the variable ECODE).

The python script

The pythonscript can be downloaded as plain text. We present a copy with line numbers for discussion below.

Source of python script

The line numbers in the listing below are to facilitate a line-by-line discussion pf the script; they should not present in the actual script. Downloaded as plain text version of the script.

Example

1  #!/bin/env python
2  #
3  def hello_mpi ( ):

4  #*****************************************************************************
5  #
6  # HELLO_MPI is a simple 'Hello, world!' test of MPI4PY.
7  #  Licensing:
8  #    This code is distributed under the GNU LGPL license.
9  #  Author:
10 #    John Burkardt
11 #    with some modifications by Tom Payerle
12 #  Retrieved from:
13 #     https://people.sc.fsu.edu/~jburkardt/py_src/hello_mpi/hello_mpi.py
14 #

15   # Import some useful packages
16   import sys
17   import platform
18   from mpi4py import MPI

19   # Get the name of the compute node we are running on
20   node = platform.node()

21   # Use the default MPI communicator
22   comm = MPI.COMM_WORLD
23   # Get the size (total number of tasks) and rank (number of our task)
24   id = comm.Get_rank()
25   p = comm.Get_size()

26   if ( id == 0 ):
27     # Print our version number, but only if task # 0
28     print ( 'HELLO_MPI: Version 1.5' )

29   #Print a message from each task
30   print ( 'Hello, world! from MPI task ', id, ' of ', p, ' on node ', node)

31  return

32 # Run the hello_mpi routine when called as a script
33 if ( __name__ == '__main__' ):
34   hello_mpi ( )

Discussion of python script

The following is a detailed, more or less line-by-line discussion of the python script:

• Line 1: This shebang line specifies that the script is to be processed as a python script. The use of /bin/env python means that it will run under the python command found in the PATH of the calling process.
  
  As we will be invoking the script as the argument of the an python command, it is not really necessary (the shebang line really only is used if one tries to execute the script directly). But it is good practice to include it to identify this as a python script.
• Lines 3-31: These lines define a python function which is a Python version of an MPI-enabled Hello World program, printing out a message from each MPI task. These lines just define the function; the lines are not actually executed until the function invocation in line 34.
• Lines 4-13: These lines are comments, describing the Python function, and in particular the license and authorship.
• Lines 15-18: These lines import some useful Python packages, in particular the "platform" package so that we can get our node name, and the "MPI" package from "mpi4py" in order to access the MPI API.
• Line 20: This line obtains the hostname of the node this task is running on and stores it in the Python variable "node".
• Line 22: This line stores the default MPI Communicator, MPI_COMM_WORLD into the Python variable "comm". This communicate consists of all the processes for the job, which is all that we need for this job.
• Line 24: This line sets the Python variable "id" to the rank or the current task.
• Line 25: This line sets the Python variable "p" to the size or the MPI communicator in comm, and since we are using MPI_COMM_WORLD, this is the total number of tasks in the job.
• Lines 26-28: These lines print out the version of the script, but only if we are the task with rank 0.
• Line 30: This line prints out the "Hello World" message for this task, printing out the rank, size, and the name of the node on which the task is running.
• Line 31: This line returns from the python function.
• Lines 33-34: These lines actually invoke the python function hello_mpi defined in lines 31. The if statement we wrap the function call in is a common Python idiom which makes it so that the function does not get invoked if this script is imported, only when it is executed (i.e. called as a script or the argument of the python command, etc).

Submitting the example

To submit the script, you can just: - Download the submission script and Python script as plain text and copy them to one of the cluster login nodes.
Move them to an appropriate directory. - From the directory with the job script, run the command sbatch submit.sh (replacing submit.sh with the name of the script you downloaded). The scheduler should return something like Submitted batch job 11274515 --- the number will vary and is the job number for this job). The job number can be used to reference the job in various Slurm commands, etc. (Please always give the job number(s) when requesting help about a job you submitted).

The job will be queued for partition specified (or defaulted to); in this example we specified the the debug partition so it will be in that partition. It should typically start running within 15 minutes or so, and finish shortly after it starts (as it is a rather trivial job).

When it finishes running, you should be able to find the file "slurm-JOBNUMBER.out" in the directory from which you ran the sbatch command which should contain the output of the diagnostic commands (the job number, when it started, what hosts running on, when it ended and exit code, etc), something like:

Example

Loading python/gcc/11.3.0/nocuda/linux-rhel8-zen2/3.10.10
  Loading requirement: gcc/11.3.0
Slurm job 11381257 running on
Slurm job submit-python-mpi.sh 11381257 running on:
compute-b8-57.zaratan.umd.edu
Ntasks=60; CPU cores per task=1
On 4 nodes: compute-b8-[57-60]
Fri Apr 25 09:37:14 EDT 2025
/home/payerle/scratch/mpi4py-job.11381257
Loaded modules are:
Currently Loaded Modulefiles:
 1) hpcc/zaratan(default)                              
 2) gcc/11.3.0                                         
 3) python/gcc/11.3.0/nocuda/linux-rhel8-zen2/3.10.10  
 4) openmpi/gcc/11.3.0/zen2/4.1.5                      
Running in /home/payerle/scratch/mpi4py-job.11381257
Using python /cvmfs/hpcsw.umd.edu/spack-software/2023.11.20/views/2023/linux-rhel8-zen2/gcc@11.3.0/python-3.10.10/mpi-nocuda/linux-rhel8-zen2/gcc/11.3.0/bin/python
Using python script /home/payerle/example-jobs/python-mpi/hello-mpi.py
Job finished with exit code 0.  Working dir is /home/payerle/scratch/mpi4py-job.11381257
Fri Apr 25 09:37:16 EDT 2025

Similar, in the working directory of the job (normally "~/scratch/mpi4py-job.JOBNUMBER") the file "hello.out" should contain the output of the actual "hello-umd" command, something like:

Example

Hello, world! from MPI task  6  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  7  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  8  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  9  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  10  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  11  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  12  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  13  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  14  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  15  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  16  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  17  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  18  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  19  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  20  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  21  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  22  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  23  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  24  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  25  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  26  of  60  on node  compute-b8-57.zaratan.umd.edu
HELLO_MPI: Version 1.5
Hello, world! from MPI task  0  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  1  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  2  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  3  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  4  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  5  of  60  on node  compute-b8-57.zaratan.umd.edu
Hello, world! from MPI task  42  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  43  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  44  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  45  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  46  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  47  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  48  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  38  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  39  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  40  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  41  of  60  on node  compute-b8-59.zaratan.umd.edu
Hello, world! from MPI task  51  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  52  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  53  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  54  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  55  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  56  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  57  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  58  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  59  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  49  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  50  of  60  on node  compute-b8-60.zaratan.umd.edu
Hello, world! from MPI task  27  of  60  on node  compute-b8-58.zaratan.umd.edu
Hello, world! from MPI task  28  of  60  on node  compute-b8-58.zaratan.umd.edu
Hello, world! from MPI task  29  of  60  on node  compute-b8-58.zaratan.umd.edu
Hello, world! from MPI task  30  of  60  on node  compute-b8-58.zaratan.umd.edu
Hello, world! from MPI task  31  of  60  on node  compute-b8-58.zaratan.umd.edu
Hello, world! from MPI task  32  of  60  on node  compute-b8-58.zaratan.umd.edu
Hello, world! from MPI task  33  of  60  on node  compute-b8-58.zaratan.umd.edu
Hello, world! from MPI task  34  of  60  on node  compute-b8-58.zaratan.umd.edu
Hello, world! from MPI task  35  of  60  on node  compute-b8-58.zaratan.umd.edu
Hello, world! from MPI task  36  of  60  on node  compute-b8-58.zaratan.umd.edu

In both cases, some of the specifics (e.g. job numbers, timestamps, pids, and node names) will differ from run to run. Also note that the ordering of the lines from the various tasks will vary.