Detect if code is run with 'srun'

[Sebastian-Belkner]

Hi dear community,

I am running code on different HPC's and struggle with the following.

Assume I have a file foo.py with a single line of code:

from mpi4py import MPI

Everything works great when I srun python3 foo.py
However, using python3 foo.py breaks, due to the PMI init not working. I don't mind this, as long as I can detect if foo.py was called using srun or not, but I have no idea how I could detect this.

I am thinking of something along

if code_called_with_srun:
    from mpi4py import MPI
    print('OK')
else:
    print('MPI not loaded')

My current workaround is to set my own os.environ['dont_use_MPI'], but I am thinking there should be a better way?

Hope someone can help

[dalcinl]

Looks like your MPI implementation is using PMI for process management. If that is the case, you can use:

import os
if "PMI_SIZE" in os.environ:
    from mpi4py import MPI
else:
    MPI = None

[Sebastian-Belkner]

That's really helpful, thanks! I also found that os.environ['_'] holds the actual name of the program with which the Python command was called, and can therefore also easily catch 'mpirun' calls.

[dalcinl]

I would advise against relying on $_. What if a user of your code use srun to execute a bash script that in turn exec Python code using mpi4py?
Checking for PMI_SIZE or PMIX_SIZE (and maybe OMPI_COMM_WORLD_SIZE) environment variables being defined should be enough for the vast majority of platforms and even MPI implementations.

[Sebastian-Belkner]

Hm.. I hadn't thought about this, thanks. I agree, directly checking for the pmi environ variables is much safer.