Slurm SSH Script

slurm-ssh.sh

#!/bin/bash
# usage: slurm-ssh [host] [shell]
# e.g.: salloc -N2 slurm-ssh
#       n01 (salloc n[01,04]) >
# or:
#   slurm-ssh n01
# or:
#   slurm-ssh n01 zsh

DIR=$(pwd)

if [ -n "$SLURM_NODELIST" ]; then
  first=$(scontrol show hostname $SLURM_NODELIST | head -n1)
fi
if [ "$#" -gt 2 ]; then
  exec ssh -q "cd DIR; $@"
else
  host=${1-$first}
  env=$(printenv | grep SLURM | sed -rn "s/=(.*)/='\1'/p" | paste -d' ' -s)
  exec ssh -q $host -t "cd $DIR; $env ${2-$SHELL} --login"
fi