diff --git a/pytools/mpi.py b/pytools/mpi.py new file mode 100644 index 0000000000000000000000000000000000000000..d67106f12f663c33ffdae60c9d6aaa29c5ac37f5 --- /dev/null +++ b/pytools/mpi.py @@ -0,0 +1,13 @@ +def run_with_mpi_ranks(py_script, ranks, callable, *args, **kwargs): + import os + if "BOOSTMPI_RUN_WITHIN_MPI" in os.environ: + callable(*args, **kwargs) + else: + import sys + newenv = os.environ.copy() + newenv["BOOSTMPI_RUN_WITHIN_MPI"] = "1" + + from subprocess import check_call + check_call(["mpirun", "-np", str(ranks), + sys.executable, py_script], env=newenv) +