#!/bin/bash

# determine STUBL install location
# Copied from Apache Ant:
# https://git-wip-us.apache.org/repos/asf?p=ant.git;a=blob;f=src/script/ant;h=b5ed5be6a8fe3a08d26dea53ea0fb3f5fab45e3f
if [ -z "$STUBL_HOME" -o ! -d "$STUBL_HOME" ] ; then
  ## resolve links - $0 may be a link to stubl's home
  PRG="$0"
  progname=`basename "$0"`

  # need this for relative symlinks
  while [ -h "$PRG" ] ; do
    ls=`ls -ld "$PRG"`
    link=`expr "$ls" : '.*-> \(.*\)$'`
    if expr "$link" : '/.*' > /dev/null; then
    PRG="$link"
    else
    PRG=`dirname "$PRG"`"/$link"
    fi
  done

  STUBL_HOME=`dirname "$PRG"`/..

  # make it fully qualified
  STUBL_HOME=`cd "$STUBL_HOME" > /dev/null && pwd`
fi

# setup STUBL environment
. $STUBL_HOME/conf/stubl 

BODY=$STUBL_HOME/bin/body

if [ "$1" == "--help" ]; then
  echo "======================================================="
  echo ""
  echo "  slogs - A script that retrieves resource usage and   "
  echo "  accounting information for a user or list of users.  "
  echo "  For each job run after the given start date, the     "
  echo "  following information is gathered from the SLURM     "
  echo "  accounting logs:  " 
  echo "     number of CPUS, wait time, start time, elapsed    "
  echo "     time, Amount of RAM Requested, Avg RAM Used, and  "
  echo "     Max RAM Used "
  echo ""
  echo "  Usage:  slogs [start_date] [users] [sacct_args]      "
  echo ""
  echo "  Notes: start_date should be in MMDDYY format.        "
  echo "         Use a comma-separated list for multiple users."
  echo "         All sacct_args passed along to sacct command. "
  echo ""
  echo " Defaults:                                             "
  echo "    start_date : 010113 if no start date is given.     "
  echo ""
  echo "    users : Current user if no user name is given.     "
  echo "            \"all\" selects all users                  "
  echo "======================================================="

  exit
fi

# try to figure out if users are in first or second arg
if [ "$1" != "" -a "`echo $1 | tr ',' '\n' | head -n1 | xargs id 2>/dev/null | wc -l`" != 0 ]; then
  dol1=$2
  dol2=$1
else
  dol1=$1
  dol2=$2
fi

dol1Test=`echo $dol1 | grep '\--' | wc -l`
if [ "$dol1" == "" -o "$dol1Test" == "1" ]; then
  StartDate=010113
else
  StartDate=$dol1
  shift
fi

if [ "$dol2" == "" ]; then
  TheGroup=`whoami`
else
  TheGroup=$dol2
  shift
fi

if [ "$TheGroup" == "all" ]; then
    TheGroup=`sacct --start=${StartDate} --allusers $@ -ouser -n | \
             sort -u | \
             sed 's/ //g' | \
             tr '\n' ',' | \
             sed 's/^,//g' | \
             sed 's/,$//g'`
fi
TheList=`echo "$TheGroup" | tr ',' ' '`

EndDate=`date +%m%d%y`

export SACCT_FORMAT="jobid%20,User,ncpus,start,elapsed,ReqMem,AveRSS,MaxRSS"

TMPFILE=${TheGroup}_usage_logs_${StartDate}_to_${EndDate}.tmp
OUTFILE=${TheGroup}_usage_logs_${StartDate}_to_${EndDate}.txt

echo "JobID                User      NCPUS   Start Date Time         Elapsed     ReqMem    AveRSS     MaxRSS" > $TMPFILE

for i in $TheList; do
  echo "Retrieving accouning data for user $i ..."
  sacct --start=${StartDate} --user=$i $@ | \
  sed "s/$i//g" |
  sed "s/T/ /g" | \
  awk -vuser=$i '{ 
         printf("%-20s %-8s  %5s   %10s %8s   %10s   %8s  %9s   %9s\n", $1, user, $2, $3, $4, $5, $6, $7, $8);
       }' \
   >> $TMPFILE
done

cat $TMPFILE | $BODY sort -k1 -V | tee $OUTFILE
rm $TMPFILE

