diff --git a/bin/exec.py b/bin/exec.py new file mode 100644 index 0000000..f665102 --- /dev/null +++ b/bin/exec.py @@ -0,0 +1,23 @@ +# coding: utf-8 + +import argparse +import sys + +from common import manager + + +parser = argparse.ArgumentParser( + prog='mpienv exec', + description='Call mpiexec with appropriate arguments') +parser.add_argument('args', nargs=argparse.REMAINDER, + default=[]) + + +def main(): + args = sys.argv[1:] + + manager.exec_(args) + + +if __name__ == "__main__": + main() diff --git a/bin/info.py b/bin/info.py index 1f15bbc..5d62be5 100644 --- a/bin/info.py +++ b/bin/info.py @@ -6,6 +6,7 @@ import sys from common import manager +from common import UnknownMPI parser = argparse.ArgumentParser( prog='mpienv info', @@ -17,13 +18,19 @@ args = parser.parse_args() - name = args.name or manager.get_current_name() + try: + name = manager.get_current_name() + except UnknownMPI: + sys.stderr.write("Error: the current MPI is not under control\n") + exit(-1) - if name in manager: + name = args.name or name + + if name not in manager: + sys.stderr.write("Error: '{}' is unknown.\n".format(name)) + else: if args.json: print(json.dumps(manager[name])) else: print(name) pprint.pprint(manager[name]) - else: - sys.stderr.write("Error: {} is not installed.\n".format(name)) diff --git a/common.py b/common.py index 19c92c8..8867c08 100644 --- a/common.py +++ b/common.py @@ -26,6 +26,10 @@ import builtins +class UnknownMPI(RuntimeError): + pass + + def yes_no_input(msg): if hasattr(__builtin__, 'raw_input'): input = __builtin__.raw_input @@ -258,31 +262,11 @@ def _load_info(self): info['name'] = name self._installed[name] = info - def get_info(self, name): - """Obtain information of the MPI installed under prefix.""" - prefix = os.path.join(self._mpi_dir, name) + def get_info_from_prefix(self, prefix): + info = {} mpiexec = os.path.join(prefix, 'bin', 'mpiexec') mpi_h = os.path.join(prefix, 'include', 'mpi.h') - info = {} - - if is_broken_symlink(prefix): - # This means the symlink under versions/ directory - # is broken. - # (The installed MPI has been removed after registration) - return { - 'name': name, - 'broken': True, - } - elif not os.path.exists(mpiexec): - # If `name` does not exist - return None - else: - # If `name` exists (would be the most cases) - info['broken'] = False - - info['symlink'] = os.path.islink(prefix) - p = Popen([mpiexec, '--version'], stderr=PIPE, stdout=PIPE) out, err = p.communicate() ver_str = decode(out + err) @@ -318,6 +302,34 @@ def get_info(self, name): return info + def prefix(self, name): + return os.path.join(self._mpi_dir, name) + + def get_info(self, name): + """Obtain information of the MPI installed under prefix.""" + info = {} + + mpiexec = os.path.join(self.prefix(name), 'bin', 'mpiexec') + if is_broken_symlink(self.prefix(name)): + # This means the symlink under versions/ directory + # is broken. + # (The installed MPI has been removed after registration) + return { + 'name': name, + 'broken': True, + } + elif not os.path.exists(mpiexec): + # If `name` does not exist + return None + else: + # If `name` exists (would be the most cases) + info['broken'] = False + + info['symlink'] = os.path.islink(self.prefix(name)) + + info.update(self.get_info_from_prefix(self.prefix(name))) + return info + def items(self): return self._installed.items() @@ -352,7 +364,10 @@ def is_installed(self, path): return None def get_current_name(self): - return next(name for name, info in self.items() if info['active']) + try: + return next(name for name, info in self.items() if info['active']) + except StopIteration: + raise UnknownMPI() def add(self, prefix, name=None): info = self.get_info(prefix) @@ -456,6 +471,56 @@ def use(self, name, mpi4py=False): mpi4py.install() mpi4py.use() + def exec_(self, cmds): + envs = os.environ.copy() + + try: + name = self.get_current_name() + + mpi4py = MPI4Py(self, name) + if mpi4py.is_installed(): + envs['PYTHONPATH'] = mpi4py.pylib_dir() + + info = self.get_info(name) + + # TODO(keisukefukuda): if hostfile is given, convert it + + except UnknownMPI: + raise RuntimeError("Internal Error: Unknown MPI") + + if info['broken']: + sys.stderr.write("Error: the current MPI is broken\n") + exit(-1) + + if info['type'] == 'Open MPI': + pref = self.prefix(name) + if os.path.islink(pref): + pref = os.readlink(pref) + + cmds[:0] = ['--prefix', pref] + cmds[:0] = ['-x', 'PYTHONPATH'] + # Transfer some environ vars + vars = ['PATH', 'LD_LIBRARY_PATH'] # vars to be transferred + vars += [v for v in os.environ if v.startswith('OMPI_')] + for var in vars: + if var in envs: + cmds[:0] = ['-x', var] + + elif info['type'] in ['MPICH', 'MVAPICH']: + cmds[:0] = ['-genvlist', 'PATH,LD_LIBRARY_PATH,PYTHONPATH'] + + # sys.stderr.write("{}\n".format(info['type'])) + + mpiexec = os.path.realpath( + os.path.join(self.prefix(name), 'bin', 'mpiexec')) + + cmds[:0] = [mpiexec] + + # sys.stderr.write(' '.join(cmds) + "\n") + p = Popen(cmds, env=envs) + p.wait() + exit(p.returncode) + def _mirror_file(self, f, dst_dir): dst = os.path.join(dst_dir, os.path.basename(f)) diff --git a/init b/init index 7944895..9c6683e 100755 --- a/init +++ b/init @@ -39,8 +39,8 @@ function mpienv() { { eval $(env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ python $root/bin/use.py $*) - env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/use.py $* + #env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ + # python $root/bin/use.py $* if [ -z "${BASH_VERSION:-}" -a ! -z "${ZSH_VERSION:-}" ]; then rehash fi @@ -49,25 +49,25 @@ function mpienv() { "configure" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/configure.py $* + python $root/bin/configure.py "$@" } ;; "build" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/build.py $* + python $root/bin/build.py "$@" } ;; "install" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/install.py $* + python $root/bin/install.py "$@" } ;; "clean" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/clean.py $* + python $root/bin/clean.py "$@" } ;; "add" ) @@ -79,43 +79,49 @@ function mpienv() { "rm" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/rm.py $* + python $root/bin/rm.py "$@" } ;; "rename" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/rename.py $* + python $root/bin/rename.py "$@" } ;; "list" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/list.py $* + python $root/bin/list.py "$@" } ;; "info" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/info.py $* + python $root/bin/info.py "$@" } ;; "autodiscover" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/autodiscover.py $* + python $root/bin/autodiscover.py "$@" } ;; "prefix" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/prefix.py $* + python $root/bin/prefix.py "$@" + } + ;; + "exec" ) + { + env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ + python $root/bin/exec.py "$@" } ;; "help" ) { env PYTHONPATH=$MPIENV_ROOT:${PYTHONPATH:-} \ - python $root/bin/help.py $* + python $root/bin/help.py "$@" } ;; * ) diff --git a/mpienv/py.py b/mpienv/py.py index 48eed74..6e7b15a 100644 --- a/mpienv/py.py +++ b/mpienv/py.py @@ -68,6 +68,9 @@ def use(self): print("export PYTHONPATH={}".format(pypath)) + def pylib_dir(self): + return self._pylib_dir + class MPI4Py(PyModule): def __init__(self, mgr, name): diff --git a/tests/test_main.sh b/tests/test_main.sh index c271661..7be5750 100755 --- a/tests/test_main.sh +++ b/tests/test_main.sh @@ -39,12 +39,12 @@ oneTimeTearDown() { install_mpich() { export MPIENV_CONFIGURE_OPTS="--disable-fortran" - mpienv install mpich-3.2 # >/dev/null 2>&1 + mpienv install mpich-3.2 >/dev/null 2>&1 } install_ompi() { export MPIENV_CONFIGURE_OPTS="--disable-mpi-fortran" - mpienv install openmpi-2.1.1 # >/dev/null 2>&1 + mpienv install openmpi-2.1.1 >/dev/null 2>&1 } #----------------------------------------------------------- @@ -155,19 +155,11 @@ test_info() { test_mpi4py() { export TMPDIR=/tmp - #install_mpich - #mpienv use --mpi4py mpich-3.2 - install_ompi - mpienv use --mpi4py openmpi-2.1.1 - - mpiexec -n 2 python -c "from mpi4py import MPI" - assertTrue $? - local SCRIPT=$(mktemp) cat <$SCRIPT from mpi4py import MPI import sys -print(MPI.__file__) +#print(MPI.__file__) comm = MPI.COMM_WORLD rank = comm.Get_rank() for i in range(0, comm.Get_size()): @@ -175,10 +167,30 @@ for i in range(0, comm.Get_size()): sys.stdout.write(str(rank)) sys.stdout.flush() comm.barrier() - EOF - mpiexec --prefix $(mpienv prefix) \ - -x PYTHONPATH -x PATH -x LD_LIBRARY_PATH -n 2 python $SCRIPT + # test mpich + install_mpich + echo "------------------" + mpienv use --mpi4py mpich-3.2 + echo "------------------" + mpienv exec -n 2 python -c "from mpi4py import MPI" + echo "------------------" + assertTrue $? + mpienv exec -n 2 python $SCRIPT + assertTrue $? + OUT=$(mpienv use --mpi4py mpich-3.2; mpienv exec -n 2 python $SCRIPT) + assertEquals "01" "$OUT" + + # test Open MPI + install_ompi + mpienv use --mpi4py openmpi-2.1.1 + mpienv exec -n 2 python -c "from mpi4py import MPI" + assertTrue $? + mpienv exec -n 2 python $SCRIPT + assertTrue $? + OUT=$(mpienv use --mpi4py openmpi-2.1.1; mpienv exec -n 2 python $SCRIPT) + assertEquals "01" "$OUT" + rm -f ${SCRIPT} }