[BUG] importing a simsopt python module blocks MPI execution

[jons-pf]

@mbkumar

As soon as a simsopt python module is imported, execution of a MPI program in a subprocess is not possible anymore. This prevents running physics codes, which rely on MPI, from within the same scripts that use simsopt.

A minimal non-working example:

import subprocess

# if this line is commented out, can execute FIELDLINES
from simsopt.mhd import Vmec

conf = "test"
subprocess.run(f"mpirun xfieldlines -vmec {conf} -mgrid mgrid_{conf}.nc -vac | tee fieldlines.log", shell=True)

If the line from simsopt.mhd import Vmec is commented out, FIELDLINES get executed correctly. If the line is not commented out, the subprocess exits immediately with a return code of 0 and FIELDLINES does not get executed.

FIELDLINES is chosen here only as an example, since it is freely available via STELLOPT. This error occurs also for other programs using MPI.

Desired behavior would be that MPI slots (?) are only blocked when, e.g., VMEC gets executed via Simsopt and not until eternity once a Simsopt module is imported.

[mbkumar]

@jons-pf

I'll look into it. Thanks for bringing this to our attention.

[mbkumar]

@jons-pf

In the original code, you are trying to use two MPI programs at once, because importing Vmec initializes MPI. And so you can not run another MPI program in the same process.

We can hide the importing of MPI from mpi4py to inside Vmec class to make your code work. However, if you create a Vmec object, then you can't use subprocess to launch another MPI program. So it won't be a good solution.

The correct way to run your example is as follows.

from simsopt.mhd import Vmec
from mpi4py import MPI

# Do whatever you want with the imported Vmec class

conf = "test"
MPI.COMM_SELF.Spawn("xfieldlines",
                    args=["-vmec", f"{conf}", "-mgrid",  f"mgrid_{conf}.nc", "-vac", "|", "tee", "fieldlines.log"], 
                    maxprocs=<your_choice_of_cores>)

The modified example imports MPI from mpi4py and uses the initialized MPI to spawn a child MPI process. There may be minor bugs in the code snippet above, but you should get the gist.