qase-test

#!/usr/bin/env bash

# Read in the arguments
WD=`pwd`

VASP="/home/ehermes/local/bin/vasp"

NPROC=0
NODES=0
PPN=0
QUEUE=""
TIME=""
EXCLUDE=""
LOCAL=true
DEPENDS=""
NODELIST=""
SENDMAIL=true

while [ $# -gt 0 ]
do
    case "$1" in
    (-np) NPROC="$2"; shift;;
    (-nodes) NODES="$2"; shift;;
    (-ppn) PPN="$2"; shift;;
    (-*gamma) VASP="/home/ehermes/local/bin/vasp.gamma";;
    (-*test) VASP="/home/ehermes/local/bin/vasp.test";;
    (-q) QUEUE="$2"; shift;;
    (-time) TIME="$2"; shift;;
    (-exclude) EXCLUDE="$2"; shift;;
    (-*noncollinear) VASP="/home/ehermes/local/bin/vasp.noncollinear";;
    (-*nolocal) LOCAL=false;;
    (-depends) DEPENDS="$2"; shift;;
    (-target) NODELIST="$2"; shift;;
    (-*nomail) SENDMAIL=false;;
    (-*) echo "$0: error - unrecognized option $1" 1>&2; exit 1;;
    (*)  break;;
    esac
    shift
done

if [ -z $QUEUE ]
then
    if [ $PPN -gt 16 ]
    then
        QUEUE='univ2'
    else
        QUEUE='univ'
    fi
fi

if [ $PPN -eq 0 ] && [ $NPROC -eq 0 ] && [ $NODES -eq 0 ]
then
    PPN=1
    NPROC=1
    NODES=1
elif [ $PPN -eq 0 ] && [ $NPROC -eq 0 ] && [ $NODES -ne 0 ]
then
    if [ "$QUEUE" = "univ2" ]
    then
        PPN=20
    else
        PPN=16
    fi
    NPROC=$(($NODES * $PPN))
elif [ $PPN -eq 0 ] && [ $NPROC -ne 0 ] && [ $NODES -eq 0 ]
then
    if [ "$QUEUE" = "univ2" ]
    then
        MAX_PPN=20
    else
        MAX_PPN=16
    fi

    PPN_REMAINDER=$(($NPROC % $MAX_PPN))

    if [ $PPN_REMAINDER -eq 0 ]
    then
        PPN=$MAX_PPN
        NODES=$(($NPROC / $PPN))
    else
        NODES=$((($NPROC / $MAX_PPN) + 1))
        NODES_REMAINDER=$(($NPROC / $NODES))
        if [ $NODES_REMAINDER -ne 0 ]
        then
            NPROC=$(($NPROC + $NODES - $NODES_REMAINDER))
            echo "WARNING: Rounding number of processors to $NPROC"
        fi
        PPN=$(($NPROC / $NODES))
    fi
elif [ $PPN -eq 0 ] && [ $NPROC -ne 0 ] && [ $NODES -ne 0 ]
then
    PPN=$(($NODES / $NPROC))
    REMAINDER=$(($NODES % $NPROC))
    if [ $REMAINDER -ne 0 ]
    then
        PPN=$(($PPN + 1))
        NPROC=$(($PPN * $NODES))
        echo "WARNING: Rounding number of processors to $NPROC"
    fi
elif [ $PPN -ne 0 ] && [ $NPROC -eq 0 ] && [ $NODES -eq 0 ]
then
    NPROC=$PPN
    NODES=1
elif [ $PPN -ne 0 ] && [ $NPROC -eq 0 ] && [ $NODES -ne 0 ]
then
    NPROC=$(($PPN * $NODES))
elif [ $PPN -ne 0 ] && [ $NPROC -ne 0 ] && [ $NODES -eq 0 ]
then
    NODES=$(($NPROC / $PPN))
    NEW_NPROC=$(($NODES * $PPN))
    if [ $NPROC != $NEW_NPROC ]
    then
        echo "WARNING: Rounding number of processors to $NEW_NPROC"
        NPROC=$NEW_NPROC
    fi
elif [ $PPN -ne 0 ] && [ $NPROC -ne 0 ] && [ $NODES -ne 0 ]
then
    NEW_NPROC=$(($PPN * $NODES))
    if [ $NEW_NPROC != $NPROC ]
    then
        echo "ERROR: Number of processors (-np) is not consistent with the number of nodes (-nodes) and processors per node (-ppn)" 1>&2
        exit 1
    fi
fi


if [ -z $TIME ]
then
    if [ $QUEUE = "univ" ] || [ $QUEUE = "univ2" ]
    then
        if [ -z $TIME ]
        then
            TIME="7-00:00:00"
        fi
    elif [ $QUEUE = "pre" ]
    then
        if [ -z $TIME]
        then
            TIME="1-00:00:00"
        fi
#        EXCLUDE="$EXCLUDE,aci-[001-002,014-048]"
    else
        echo "error - unrecognized queue $QUEUE" 1>&2
        exit 1
    fi
fi

NAME=$1

# Create the input file
TARGET=run.$$.sh

PREAMBLE="#!/bin/bash
#SBATCH --job-name=$NAME
#SBATCH --ntasks-per-node=$PPN
#SBATCH --nodes=$NODES
#SBATCH --partition=$QUEUE
#SBATCH --time=$TIME
#SBATCH --error=$WD/$NAME.%J.err
#SBATCH --output=$WD/$NAME.%J.out
#SBATCH --exclusive
#SBATCH --hint=nomultithread
#SBATCH --exclude=$EXCLUDE
#SBATCH --dependency=$DEPENDS"

if [ -n "$NODELIST" ]
then
    PREAMBLE="$PREAMBLE
#SBATCH --target=$NODELIST"
fi

if [ $SENDMAIL ]
then
    PREAMBLE="$PREAMBLE
#SBATCH --mail-type END
#SBATCH --mail-type FAIL
#SBATCH --mail-user=$USER"
fi

cat >> $TARGET << EOF
$PREAMBLE

cd $WD

export OMP_NUM_THREADS=1
export TMPDIR=/scratch/local/\$USER/\$SLURM_JOB_ID

if [ $LOCAL = true ]
then
    NODES=\${SLURM_NODELIST:4}

    if [ "\${NODES:0:1}" == "[" ]
    then
        LEN=\`expr length \$NODES\`
        IFS=","
        NODES=\${NODES:1:\$((\$LEN - 2))}
        read -ra NODES <<< "\$NODES"
        NNODES=0
        NODELIST=()
        for NODE in "\${NODES[@]}"
        do
            SUBLEN=\`expr length \$NODE\`
            if [ "\$SUBLEN" -eq "7" ]
            then
                for SUBNODE in \`seq -w -s ',' \${NODE:0:3} \${NODE:4:3}\`
                do
                    NODELIST+=(\$SUBNODE)
                    NNODES=\$((\$NNODES + 1))
                done
            elif [ "\$SUBLEN" -eq "3" ]
            then
                NODELIST+=(\$NODE)
                NNODES=\$((\$NNODES + 1))
            fi
        done
    else
        NNODES=1
        NODELIST=(\$NODES)
    fi

    for NODE in "\${NODELIST[@]}"
    do
        ssh aci-\$NODE mkdir -p \$TMPDIR
    done

    export NODELIST

    export VASP_COMMAND="
        rsync -raz $WD/* \$TMPDIR
        IFS=\",\"
        for NODE in \${NODELIST[@]}
        do
            rsync -raz $WD/* aci-\$NODE:\$TMPDIR &
        done
        wait
        cd \$TMPDIR

        srun --mpi=none $VASP

        for NODE in \${NODELIST[@]}
        do
            rsync -raz aci-\$NODE:\$TMPDIR/* $WD &
        done
        wait
        cd $WD"

else
    export VASP_COMMAND="srun --mpi=none $VASP"
fi

export VASP_PP_PATH="/home/ehermes/local/apps/vasp"

sleep 30
/home/ehermes/local/bin/python $NAME.py >& $NAME.out

EOF

# Submit the input file to PBS and clean up
sbatch $TARGET
rm $TARGET
sleep 1