qtl-histogram

#!/usr/bin/env bash

set -e
set -u
set -o pipefail
source $(dirname $0)/../libexec/environ.sh

# constants ############################################################
# slurm settings
SLURM_MEMORY=2000 # must be more than max heap size in $TIMELINE_JAVA_OPTS
SLURM_DISK=32M
SLURM_TIME=10:00:00
SLURM_LOG=/farm_out/%u/%x-%A_%a
########################################################################


# default options
dataset=test_v0
declare -A modes
for key in findhipo rundir eachdir flatdir single series submit check-cache fast-ls swifjob focus-detectors focus-physics help; do
  modes[$key]=false
done
outputDir=""

# usage
sep="================================================================"
usageTerse() {
  echo """
  Run the timeline histogramming jobs.
  NOTE: chefs typically run this program using the CLAS12 workflow

  $sep
  USAGE: qtl histogram -d [DATASET_NAME] [RUN_DIRECTORY]
  $sep

     -d [DATASET_NAME]   unique dataset name, defined by the user, used for organization

     [RUN_DIRECTORY]     input data directory, with subdirectories organized by run,
                         **must be specified last**

  After running this script, run either or both of the suggested 'sbatch' command(s);
  one is for detector QA timelines and the other is for physics QA timelines.
  Output files will appear in ./outfiles/[DATASET_NAME]

  For more options, run with '--help'
  """ >&2
}
usageVerbose() {
  echo """
  $sep
  USAGE: qtl histogram  [OPTIONS]...  [RUN_DIRECTORY]...
  $sep

  REQUIRED ARGUMENTS:

    [RUN_DIRECTORY]...   One or more directories, each directory corresponds to
                         one run and should contain reconstructed HIPO files
                         - See \"INPUT FINDING OPTIONS\" below for more control,
                           so that you don't have to specify each run's directory
                         - A regexp or globbing (wildcards) can be used to
                           specify the list of directories as well, if your shell
                           supports it
                         - for a directory of files, with one run per file (e.g.,
                           SKIM files), use the option \`--flatdir\`

  $sep

  OPTIONS:

     -d [DATASET_NAME]      unique dataset name, defined by the user, used for organization
                            default = '$dataset'

     -o [OUTPUT_DIR]        custom output directory
                            default = ./outfiles/[DATASET_NAME]

     *** INPUT FINDING OPTIONS: control how the input HIPO files are found;
         choose only one:

       --findhipo     use \`find\` to find all HIPO files in each
                      [RUN_DIRECTORY]; this is useful if you have a
                      directory tree, e.g., runs grouped by target

       --rundir       assume each specified [RUN_DIRECTORY] contains
                      subdirectories named as just run numbers;
                      - it is not recommended to use wildcards for this option
                      - see also the \`--fast-ls\` option
                      **this is the DEFAULT option**

       --eachdir      assume each specified [RUN_DIRECTORY] is a single
                      run's directory full of HIPO files (e.g., DST files)

       --flatdir      assume each specified [RUN_DIRECTORY] contains HIPO
                      files, with one run per file; use this option for SKIM files

       --check-cache  cross check /cache directories with tape stub directories
                      (/mss) and exit without creating or running any jobs; this is
                      useful if you are running QA on older DSTs which may no longer be
                      fully pinned on /cache

     *** EXECUTION CONTROL OPTIONS: choose only one, or the default will generate a
         Slurm job description and print out the suggested \`sbatch\` command

       --single    run only the first job, locally; useful for
                   testing before submitting jobs to slurm

       --series    run all jobs locally, one at a time; useful
                   for testing on systems without slurm

       --submit    submit the slurm jobs, rather than just
                   printing the \`sbatch\` command

       --swifjob   run this on a workflow runner, where the input
                   HIPO files are found in ./ and specifying [RUN_DIRECTORIES] is
                   not required; overrides some other settings; this is NOT meant
                   to be used interactively, but rather as a part of a workflow

     *** FOCUS OPTIONS: these options allow for running single types of jobs,
         rather than the default of running everything; you may specify more
         than one

       --focus-detectors   run histogramming for detector QA timelines

       --focus-physics     run histogramming for physics QA timelines

     *** MISCELLANEOUS OPTIONS

       --fast-ls           use \`jasmine list-files\` to get lists of files, rather
                           than commands such as \`find\` or \`ls\`, since \`jasmine\`
                           may be much faster than the latter commands for tape dirs.
                           - WARNING: assumes you already tested \`--check-cache\`,
                                      since \`jasmine\` uses \`/mss\` stub files
                           - useful for options such as \`--rundir\` for DST files

  $sep

  EXAMPLES:

  $  $0 -v v1.0.0 --submit --rundir /volatile/mon
       -> submit slurm jobs for all numerical subdirectories of /volatile/mon/,
          where each subdirectory should be a run number; this is the most common usage

  $  $0 -v v1.0.0 --eachdir /volatile/mon/*
       -> generate the slurm script to run on all subdirectories of
          /volatile/mon/ no matter their name

  $  $0 -v v1.0.0 --single /volatile/mon/run*
       -> run on the first directory named run[RUNNUM], where [RUNNUM] is a run number

  """ # stream to stdout to permit grepping
}
if [ $# -lt 1 ]; then
  usageTerse
  exit 101
fi

# parse options
while getopts "d:o:h-:" opt; do
  case $opt in
    d) dataset=$OPTARG;;
    o) outputDir=$OPTARG;;
    h) modes['help']=true;;
    -)
      for key in "${!modes[@]}"; do
        [ "$key" == "$OPTARG" ] && modes[$OPTARG]=true && break
      done
      [ -z "${modes[$OPTARG]-}" ] && printError "unknown option --$OPTARG" && exit 100
      ;;
    *) exit 100;;
  esac
done
shift $((OPTIND - 1))

# print full usage guide
if ${modes['help']}; then
  usageVerbose
  exit 101
fi

# set the input-finding method; use the DEFAULT one, if not set by user
numTrueInputOpts=0
for key in findhipo rundir eachdir flatdir; do
  if ${modes[$key]}; then
    numTrueInputOpts=$((numTrueInputOpts+1))
  fi
done
if [ $numTrueInputOpts -eq 0 ]; then
  modes['rundir']=true # set the DEFAULT option
elif [ $numTrueInputOpts -gt 1 ]; then
  printError "more than one input-finding option set"
  exit 100
fi

# parse input directories
rdirs=()
if ${modes['swifjob']}; then
  rdirs=(.) # all input files reside in ./ on a workflow runner
else
  [ $# == 0 ] && printError "no run directories specified" && exit 100
  rdirsArgs="$@"
  for topdir in ${rdirsArgs[@]}; do
    [[ "$topdir" =~ ^- ]] && printError "option '$topdir' must be specified before run directories" && exit 100
  done
  if ${modes['rundir']}; then
    for topdir in ${rdirsArgs[@]}; do
      if [ -d $topdir ]; then
        for subdir in $(ls $topdir | grep -E "[0-9]+"); do
          rdirs+=($(echo "$topdir/$subdir " | sed 's;//;/;g'))
        done
      else
        printError "run directory '$topdir' does not exist"
        exit 100
      fi
    done
  elif ${modes['eachdir']}; then
    rdirs=$@
  elif ${modes['flatdir']}; then
    rdirs=$@
  elif ${modes['findhipo']}; then
    for topdir in ${rdirsArgs[@]}; do
      echo "finding .hipo files in $topdir ....."
      fileList=$(find -L $topdir -type f -name "*.hipo")
      if [ -z "$fileList" ]; then
        printWarning "run directory '$topdir' has no HIPO files"
      else
        rdirs+=($(echo $fileList | xargs dirname | sort -u))
      fi
    done
  else
    printError "unknown input option"
    exit 100
  fi
fi
[ ${#rdirs[@]} -eq 0 ] && printError "no run directories found" && exit 100
echo "done finding input run directories"

# set and make output directory
if ${modes['swifjob']}; then
  outputDir=$(pwd -P)/outfiles
else
  [ -z "$outputDir" ] && outputDir=$(pwd -P)/outfiles/$dataset
fi
mkdir -p $outputDir

# check focus options
modes['focus-all']=true
for key in focus-detectors focus-physics; do
  if ${modes[$key]}; then modes['focus-all']=false; fi
done
if ${modes['swifjob']} && ${modes['focus-all']}; then
  printError "option --swifjob must be used with either --focus-detectors or --focus-physics"
  exit 100
fi

# print arguments
echo """
Settings:
$sep
DATASET_NAME = $dataset
OUTPUT_DIR   = $outputDir
OPTIONS = {"""
for key in "${!modes[@]}"; do printf "%20s => %s,\n" $key ${modes[$key]}; done
echo """}
RUN_DIRECTORIES = ["""
for rdir in ${rdirs[@]}; do echo "  $rdir,"; done
echo """]
$sep
"""

# check cache (and exit), if requested
if ${modes['check-cache']}; then
  echo "Cross-checking /cache and /mss..."
  $TIMELINESRC/libexec/check-cache.sh ${rdirs[@]}
  exit $?
fi

# if `flatdir` mode, populate `rdirs` with the list of files, since our job loop will be over `rdirs` elements
if ${modes['flatdir']}; then
  rdirsIn=("${rdirs[@]}") # make a copy, since it will be overwritten
  rdirs=()
  for rdir in ${rdirsIn[@]}; do
    for hipofile in $(ls $(realpath $rdir)/*.hipo); do
      rdirs+=($hipofile)
    done
  done
fi

# initial checks and preparations
echo $dataset | grep -q "/" && printError "dataset name must not contain '/' " && echo && exit 100
[ -z "$dataset" ] && printError "dataset name must not be empty" && echo && exit 100
slurmJobName=clas12-timeline--$dataset

# start job lists
slurmDir=./slurm
mkdir -p $slurmDir/scripts
jobkeys=()
for key in detectors physics; do
  if ${modes['focus-all']} || ${modes['focus-'$key]}; then
    jobkeys+=($key)
  fi
done
declare -A joblists
for key in ${jobkeys[@]}; do
  joblists[$key]=$slurmDir/job.$dataset.$key.list
  > ${joblists[$key]}
done

# get run numbers for each input
echo "Determining run number list..."
declare -A runnumHash # `rdirs` element -> run number
for rdir in ${rdirs[@]}; do
  # get the run number, either from `rdir` basename (fast), or from `RUN::config` (slow)
  [[ ! -e $rdir ]] && printError "the run file/directory '$rdir' does not exist" && continue
  runnum=$(basename $rdir | grep -m1 -o -E "[0-9]+" | tail -n1 || echo '') # first, try from run directory (or file) basename
  if [ -z "$runnum" ] || ${modes['swifjob']}; then # otherwise, use RUN::config from a HIPO file (NOTE: assumes all HIPO files have the same run number)
    if ${modes['flatdir']}; then
      $TIMELINESRC/libexec/hipo-check.sh $rdir
      runnum=$($TIMELINESRC/libexec/run-groovy-timeline.sh $TIMELINESRC/libexec/get-run-number.groovy $rdir | tail -n1 | grep -m1 -o -E "[0-9]+" || echo '')
    else
      firstHipo=$(find $rdir -name "*.hipo" | head -n1)
      [ -z "$firstHipo" ] && printError "no HIPO files in run directory '$rdir'; cannot get run number or create job" && continue
      echo "using HIPO file $firstHipo to get run number for run directory '$rdir'"
      $TIMELINESRC/libexec/hipo-check.sh $firstHipo
      runnum=$($TIMELINESRC/libexec/run-groovy-timeline.sh $TIMELINESRC/libexec/get-run-number.groovy $firstHipo | tail -n1 | grep -m1 -o -E "[0-9]+" || echo '')
    fi
  fi
  [ -z "$runnum" ] &&  printError "unknown run number for '$rdir'; ignoring it!" && continue
  [ $runnum -eq 0 ] && printError "unknown run number for '$rdir'; ignoring it!" && continue
  runnum=$((10#$runnum))
  runnumHash[$rdir]=$runnum
done
echo "RUN NUMBERS = {"
for rdir in "${!runnumHash[@]}"; do echo "  $rdir => ${runnumHash[$rdir]},"; done
echo "}"
echo $sep

# get beam energy for each run number
echo "Retrieving beam energy from RCDB..."
declare -A ebeamHash  # run number -> beam energy
while read -r runnum ebeam; do
  ebeamHash[$runnum]=$ebeam
done < <($TIMELINESRC/libexec/get-beam-energy.sh "${runnumHash[@]}")
echo "BEAM ENERGY = {"
# override beam energy, for cases where RCDB is incorrect
for runnum in "${!ebeamHash[@]}"; do # loop over run numbers
  ebeamOverride=""
  if [ $runnum -ge 11620 -a $runnum -le 11657 ]; then ebeamOverride=2.182 # RG-F
  elif [ $runnum -ge 12389 -a $runnum -le 12443 ]; then ebeamOverride=2.182 # RG-F
  elif [ $runnum -ge 12444 -a $runnum -le 12951 ]; then ebeamOverride=10.389 # RG-F
  fi
  if [ -n "$ebeamOverride" ]; then
    printWarning "overriding RCDB beam energy for run $runnum: ${ebeamHash[$runnum]} -> $ebeamOverride"
    ebeamHash[$runnum]=$ebeamOverride
  fi
  echo "  $runnum => ${ebeamHash[$runnum]} GeV,"
done
echo "}"
echo $sep

# define backup directory (used only if the output files already exist; not used `if ${modes['swifjob']}`)
backupDir=$(pwd -P)/tmp/backup.$dataset.$(date +%s) # use unixtime for uniqueness

# loop over input directories, building the job lists
echo """
Generating job scripts..."""
for rdir in ${rdirs[@]}; do

  echo "processing '$rdir'..."
  runnum=${runnumHash[$rdir]:-unknown}
  [ "$runnum" = "unknown" ] && printError "unknown run number, skipping!" && continue
  ebeam=${ebeamHash[$runnum]:-unknown}
  [ "$ebeam" = "unknown" ] && printError "unknown beam energy, skipping!" && continue

  # get list of input files, and append prefix for SWIF
  echo "..... getting its input files ....."
  inputListFile=$slurmDir/files.$dataset.$runnum.inputs.list
  if ${modes['flatdir']}; then
    realpath $rdir > $inputListFile
  else
    [[ "$(realpath $rdir)" =~ /mss/ ]] && swifPrefix="mss:" || swifPrefix="file:"
    if ${modes['fast-ls']}; then
      rdir_mss=$(realpath $rdir | sed 's;^.*/cache/;/mss/;')
      rdir_cache=$(realpath $rdir)
      # fast-ls mode will use `jasmine list-files`, but requires the stub files on `mss` to exist; if a cook is ongoing,
      # it's possible the `/mss` directory won't be populated yet, although files are already on `/cache`; for that case,
      # fall back to "slow" ls (i.e., the command that runs when not using `--fast-ls`)
      jasmine list-files -display json -under $rdir_mss | jq -r '.[].mss' | xargs -n1 basename | sed "s;^;${swifPrefix}${rdir_cache}/;" > $inputListFile \
        || { printWarning "falling back to slow ls..."; realpath $rdir/*.hipo | sed "s;^;$swifPrefix;" > $inputListFile; }
    else
      echo "      (use \`--fast-ls\` if this seems too slow)"
      realpath $rdir/*.hipo | sed "s;^;$swifPrefix;" > $inputListFile
    fi
  fi

  # generate job scripts
  echo "..... generating its job scripts ....."
  for key in ${jobkeys[@]}; do

    # preparation: make output subdirectory and backup old one, if it exists
    outputSubDir=$outputDir/timeline_$key/$runnum
    if ${modes['swifjob']}; then
      outputSubDir=$outputDir # no need for run subdirectory or backup on swif runner
    else
      if [ -d $outputSubDir ]; then
        mkdir -p $backupDir/timeline_$key
        mv -v $outputSubDir $backupDir/timeline_$key/
      fi
    fi
    mkdir -p $outputSubDir

    # make job scripts for each $key
    jobscript=$slurmDir/scripts/$key.$dataset.$runnum.sh
    case $key in

      detectors)
        cat > $jobscript << EOF
#!/usr/bin/env bash
set -e
set -u
set -o pipefail
echo "RUN $runnum"

# set env vars
source $TIMELINESRC/libexec/environ.sh

# produce histograms
java $TIMELINE_JAVA_OPTS \\
  org.jlab.clas.timeline.histograms.run_histograms \\
    $runnum \\
    $outputSubDir \\
    $inputListFile \\
    $ebeam

# check output HIPO files
$TIMELINESRC/libexec/hipo-check.sh \$(find $outputSubDir -name "*.hipo")
EOF
        ;;

      physics)
        if ${modes['flatdir']}; then
          monitorReadType=skim
        else
          monitorReadType=dst
        fi
        cat > $jobscript << EOF
#!/usr/bin/env bash
set -e
set -u
set -o pipefail
echo "RUN $runnum"

# set env vars
source $TIMELINESRC/libexec/environ.sh

# produce histograms
$TIMELINESRC/libexec/run-groovy-timeline.sh \\
  $TIMELINESRC/qa-physics/monitorRead.groovy \\
    $(realpath $rdir) \\
    $outputSubDir \\
    $monitorReadType \\
    $runnum \\
    $ebeam

# check output HIPO files
$TIMELINESRC/libexec/hipo-check.sh \$(find $outputSubDir -name "*.hipo")
EOF
        ;;

    esac

    # grant permission and add it `joblists`
    chmod u+x $jobscript
    echo $jobscript >> ${joblists[$key]}

  done # loop over `jobkeys`

done # loop over `rdirs`


# now generate slurm descriptions and/or local scripts
echo """
Generating batch scripts..."""
exelist=()
for key in ${jobkeys[@]}; do

  # check if we have any jobs to run
  joblist=${joblists[$key]}
  [ ! -s $joblist ] && printError "there are no $key timeline jobs to run" && continue
  slurm=$(echo $joblist | sed 's;.list$;.slurm;')

  # either generate single/sequential run scripts
  if ${modes['single']} || ${modes['series']} || ${modes['swifjob']}; then
    localScript=$(echo $joblist | sed 's;.list$;.local.sh;')
    echo "#!/usr/bin/env bash" > $localScript
    echo "set -e" >> $localScript
    if ${modes['single']}; then
      head -n1 $joblist >> $localScript
    else # ${modes['series']} || ${modes['swifjob']}
      cat $joblist >> $localScript
    fi
    chmod u+x $localScript
    exelist+=($localScript)

  # otherwise generate slurm description
  else
    cat > $slurm << EOF
#!/bin/sh
#SBATCH --ntasks=1
#SBATCH --job-name=$slurmJobName--$key
#SBATCH --output=$SLURM_LOG.out
#SBATCH --error=$SLURM_LOG.err
#SBATCH --partition=production
#SBATCH --account=clas12

#SBATCH --mem-per-cpu=$SLURM_MEMORY
#SBATCH --gres=disk:$SLURM_DISK
#SBATCH --time=$SLURM_TIME

#SBATCH --array=1-$(cat $joblist | wc -l)
#SBATCH --ntasks=1

srun \$(head -n\$SLURM_ARRAY_TASK_ID $joblist | tail -n1)
EOF
    exelist+=($slurm)
  fi
done


# execution
[ ${#exelist[@]} -eq 0 ] && printError "no jobs were created at all; check errors and warnings above" && exit 100
echo """
$sep
"""
if ${modes['single']} || ${modes['series']} || ${modes['swifjob']}; then
  if ${modes['single']}; then
    echo "RUNNING ONE SINGLE JOB LOCALLY:"
  elif ${modes['series']}; then
    echo "RUNNING ALL JOBS SEQUENTIALLY, LOCALLY:"
  fi
  for exe in ${exelist[@]}; do
    echo """
    $sep
    EXECUTING: $exe
    $sep"""
    $exe
  done
elif ${modes['submit']}; then
  echo "SUBMITTING JOBS TO SLURM"
  echo $sep
  for exe in ${exelist[@]}; do sbatch $exe; done
  echo $sep
  echo "JOBS SUBMITTED!"
else
  echo """  SLURM JOB DESCRIPTIONS GENERATED
  - Slurm job name prefix will be: $slurmJobName
  - To submit all jobs to slurm, run:
    ------------------------------------------"""
  for exe in ${exelist[@]}; do echo "    sbatch $exe"; done
  echo """    ------------------------------------------
  """
fi