Skip to content
This repository has been archived by the owner on Jul 19, 2023. It is now read-only.

Commit

Permalink
Merge pull request #57 from keisukefukuda/add-exec-command
Browse files Browse the repository at this point in the history 
close #21
  • Loading branch information
@keisukefukuda
keisukefukuda committed Jul 13, 2017
2 parents a74c661 + 6688e5c commit 58de8eb
Show file tree
Hide file tree
Showing 6 changed files with 170 additions and 54 deletions.
23 changes: 23 additions & 0 deletions bin/exec.py
View file
Original file line number Diff line number Diff line change
@@ -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()
15 changes: 11 additions & 4 deletions bin/info.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import sys

from common import manager
from common import UnknownMPI

parser = argparse.ArgumentParser(
prog='mpienv info',
Expand All @@ -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))
111 changes: 88 additions & 23 deletions common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,10 @@
import builtins


class UnknownMPI(RuntimeError):
pass


def yes_no_input(msg):
if hasattr(__builtin__, 'raw_input'):
input = __builtin__.raw_input
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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))

Expand Down
32 changes: 19 additions & 13 deletions init
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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" )
Expand All @@ -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 "$@"
}
;;
* )
Expand Down
3 changes: 3 additions & 0 deletions mpienv/py.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
40 changes: 26 additions & 14 deletions tests/test_main.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

#-----------------------------------------------------------
Expand Down Expand Up @@ -155,30 +155,42 @@ 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 <<EOF >$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()):
if i == rank:
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}
}

Expand Down

0 comments on commit 58de8eb

Please sign in to comment.