forked from jittor/jittor
118 lines
4.3 KiB
Bash
Executable File
118 lines
4.3 KiB
Bash
Executable File
#!/bin/bash
|
|
|
|
# Copyright 2013 Benedikt Morbach <moben@exherbo.org>
|
|
# Distributed under the terms of the GNU General Public License v2
|
|
|
|
# runs multiple MPI processes as a grid in a new tmux window and multiplexes keyboard input to all of them
|
|
|
|
additional_vars=( LD_LIBRARY_PATH LD_PRELOAD )
|
|
export "${additional_vars[@]}"
|
|
|
|
usage() {
|
|
echo 'tmpi: Run multiple MPI processes as a grid in a new tmux window and multiplex keyboard input to all of them.'
|
|
echo ''
|
|
echo 'Usage:'
|
|
echo ' tmpi [number] [command]'
|
|
echo ''
|
|
echo 'You need to pass at least two arguments.'
|
|
echo 'The first argument is the number of processes to use, every argument after that is the commandline to run.'
|
|
echo 'If you call this script from outside tmux and your command contains important whitespace then you need to appy two levels of quoting to preserve it.'
|
|
echo ''
|
|
echo 'LD_LIBRARY_PATH and LD_PRELOAD are passed through, so you can run it like this:'
|
|
echo 'LD_LIBRARY_PATH="${PWD}/.libs:${LD_LIBRARY_PATH}" tmpi 16 gdb -q bin/.libs/example'
|
|
echo ''
|
|
echo 'The new window is set to remain on exit and has to be closed manually. ("C-b + k" by default)'
|
|
}
|
|
|
|
check_tools() {
|
|
tools=( tmux mpirun )
|
|
|
|
for tool in "${tools[@]}"; do
|
|
if ! which ${tool}; then
|
|
echo "You need to install ${tool} to run this script."
|
|
fi
|
|
done
|
|
}
|
|
|
|
if [[ ${#} -lt 2 ]]; then
|
|
usage
|
|
|
|
exit 1
|
|
fi
|
|
|
|
if [[ -z ${TMUX} ]]; then
|
|
# it seems we aren't in a tmux session.
|
|
# start a new one so that our window doesn't end up in some other session and we have to search it.
|
|
# actually start a new server with '-L' to ensure that our environment carries over.
|
|
socket=$(mktemp --dry-run tmpi.XXXX)
|
|
exec tmux -L ${socket} new-session "${0} ${*}"
|
|
fi
|
|
|
|
if [[ ${1} == runmpi ]] ; then
|
|
# we are being started as one of many processes by mpirun.
|
|
shift
|
|
|
|
# start the processes in the order of their rank.
|
|
# this avoids races, as we have to push the variables in tmux' environment.
|
|
# it has the nice side-effect that the panes are also ordered by rank.
|
|
while [[ $(cat /tmp/tmpi.lock) -ne ${OMPI_COMM_WORLD_RANK} ]] ; do
|
|
sleep 0.02
|
|
done
|
|
|
|
# get all the variables that mpirun starts us with so that we can pass them through.
|
|
mpi_vars=( $( env | grep -e MPI -e OPAL -e PMIX -e PYTHON -e debug | cut -d '=' -f1 ) )
|
|
mpi_vars+=( "${additional_vars[@]}" )
|
|
|
|
# add the variables to tmux' session environment.
|
|
# we can't just export them because the process will be started as a child of tmux, not us.
|
|
for var in "${mpi_vars[@]}"; do
|
|
tmux set-environment -t ${session} "${var}" "${!var}"
|
|
done
|
|
|
|
x=( $(tmux split-window -P -F '#{pane_pid} #{pane_id}' -t ${window} "${*}") )
|
|
pid=${x[0]}
|
|
pane=${x[1]}
|
|
|
|
for var in "${mpi_vars[@]}"; do
|
|
tmux set-environment -t ${session} -u "${var}"
|
|
done
|
|
|
|
# kill the dummy pane that opened the new window
|
|
[[ ${OMPI_COMM_WORLD_RANK} -eq 0 ]] && tmux kill-pane -t ${dummy} &> /dev/null
|
|
|
|
# set the window to tiled mode.
|
|
# have to do this after every new pane is spawned because otherwise the splits get
|
|
# smaller and smaller until tmux refuses to open new panes, despite plenty of space being left.
|
|
tmux select-layout -t ${pane} tiled &> /dev/null
|
|
|
|
# let the next process start
|
|
echo $((${OMPI_COMM_WORLD_RANK}+1)) > /tmp/tmpi.lock
|
|
|
|
# don't exit here as mpirun needs to be kept alive and it would also exit.
|
|
while [[ -d /proc/${pid} ]]; do
|
|
sleep 1
|
|
done
|
|
else
|
|
# we are the parent and set everything up before we start ourselves a bunch of times via mpirun.
|
|
processes=${1}
|
|
self=${0}
|
|
shift
|
|
|
|
# create an empty new dummy window which we sill later split up for the mpi processes.
|
|
x=( $(tmux new-window ${session} -P -F '#{pane_id} #{window_id} #{session_id}') )
|
|
export dummy=${x[0]}
|
|
export window=${x[1]}
|
|
export session=${x[2]}
|
|
|
|
# syncronize input to all panes.
|
|
tmux set-window-option -t ${window} synchronize-panes on &> /dev/null
|
|
tmux set-window-option -t ${window} remain-on-exit on &> /dev/null
|
|
|
|
# always start with rank 0.
|
|
echo 0 > /tmp/tmpi.lock
|
|
|
|
# re-execute ourself to spawn of the processes.
|
|
echo mpirun -np ${processes} ${self} runmpi "${@}"
|
|
mpirun -np ${processes} ${self} runmpi "${@}"
|
|
fi
|