#!/bin/bash
# This runs on the compute nodes...
# Args: [-a i] wd [-d] [-t nthreads] [-b CMDS] [-i input] CMDS
# See ug4/azure/notes.txt for documentation

N=$SLURM_JOB_NUM_NODES
n=$SLURM_NTASKS
c=$SLURM_CPUS_PER_TASK
nodename=$SLURMD_NODENAME
local=$SLURM_LOCALID
node=$SLURM_NODEID
task=$SLURM_PROCID

if [ "$1" = "-a" ]
then
 shift
 xarg="$1"
 xa="run $1 "
 shift
else
 xarg=
 xa=
fi
 
cd "$1"
shift

if [ "$1" = "-d" ]
then
  shift
  debug=1
fi

if [ "$1" = "-t" ]
then
  shift
  threadsPerJob=$1
  shift
else
  threadsPerJob=2
fi

pjobs=$((c / $threadsPerJob))

echo $(date) ${xa}task $n.$task on $nodename:$N.$node start $pjobs jobs 1>&2

PATH=$W/$USER/bin:$W/shared/bin:$PATH
export task PATH n xarg

if [ "$1" = "-b" ]
then
  shift
  eval "$1"
  shift
fi

input="seq 1 $pjobs" # default to make sure something runs

if [ "$1" = "-i" ]
then
 shift
 input="$1"
 shift
 if [ "$debug" ]
 then
     echo $(date) task $n.$task input "|$input|"
 fi
fi

export cmd="$1"
shift

doit () {
 arg="$1"
 i="$2"
 echo $(date) start $task $i "|$xarg|$arg|"
 eval "$cmd"
 echo $(date) end $task $i
}

export -f doit

if [ "$debug" ]
then
 echo $(date) task $n.$task cmd "|$cmd|" $(type doit)
fi

eval "$input" | \
   if [ "$debug" ]
   then
     tee >(cat 1>&2)
   else
     cat
   fi | \
   if [ $pjobs -le 1 ]
   then
    xargs -I ^ -n 1 bash -lc "doit ^"
   else
    parallel --line-buffer -j $pjobs doit '{}' '{#}'
   fi

echo $(date) task $n.$task on $nodename:$N.$node end 1>&2