From 44e4af9adcaef7efe18a027ce2be739566b47444 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Wed, 28 May 2014 13:19:42 -0500
Subject: [PATCH] Rework run_with_mpi_ranks
---
pytools/mpi.py | 38 +++++++++++++++++++++++++-------------
1 file changed, 25 insertions(+), 13 deletions(-)
diff --git a/pytools/mpi.py b/pytools/mpi.py
index 51ce7ff..a575c07 100644
--- a/pytools/mpi.py
+++ b/pytools/mpi.py
@@ -1,16 +1,28 @@
-def in_mpi_relaunch():
- import os
- return "PYTOOLS_RUN_WITHIN_MPI" in os.environ
+def check_for_mpi_relaunch(argv):
+ if argv[1] != "--mpi-relaunch":
+ return
+
+ from pickle import loads
+ f, args, kwargs = loads(argv[2])
+
+ f(*args, **kwargs)
+ import sys
+ sys.exit()
-def run_with_mpi_ranks(py_script, ranks, callable, *args, **kwargs):
- if in_mpi_relaunch():
- callable(*args, **kwargs)
- else:
- import sys, os
- newenv = os.environ.copy()
- newenv["PYTOOLS_RUN_WITHIN_MPI"] = "1"
- from subprocess import check_call
- check_call(["mpirun", "-np", str(ranks),
- sys.executable, py_script], env=newenv)
+def run_with_mpi_ranks(py_script, ranks, callable, args=(), kwargs=None):
+ if kwargs is None:
+ kwargs = {}
+
+ import sys
+ import os
+ newenv = os.environ.copy()
+ newenv["PYTOOLS_RUN_WITHIN_MPI"] = "1"
+
+ from pickle import dumps
+ callable_and_args = dumps((callable, args, kwargs))
+ from subprocess import check_call
+ check_call(["mpirun", "-np", str(ranks),
+ sys.executable, py_script, "--mpi-relaunch", callable_and_args],
+ env=newenv)
--
GitLab