jobenv

#!/bin/bash
#
# For each argument JobID, prints environment variables for the job.
# Exit status of 0 if succeeded for all JobIDs, 1 otherwise.
#
# This script relies on Slurm's built-in feature of storing job environment
# in the accounting database (this is not on by default, so need to set
# 'AccountingStoreFlags = job_env' in your slurm.conf).
#
# Caveats:
#   * Requires Slurm 21.08 or newer (but currently no checking);
#   * Must be enabled by a parameter in slurm.conf (but currently no checking).
#
# Lev Gorenstein <lev@purdue.edu>, 2021

VERSION="0.2.0" 			# Increment me!
PROGNAME=${BASH_SOURCE##*/}

# Safety and sanity
set -o pipefail
export PATH="/bin:/usr/bin:/usr/local/bin"

# Helper functions
function usage() {
    echo "Usage: ${PROGNAME} [-h|--help] [-r|--raw] JobID(s)"
}

function help_message() {
    # Yes, I don't like the '<<-' form, so indentation is somewhat off
    cat << _EOF_
$PROGNAME ver. $VERSION - show Slurm job's environment variables.

For each argument JobID, prints environment variables for that job.

$(usage)

Options:
   -r, --raw      Show underlying sacct output as is (default is to trim
                  extra separator lines sacct adds). This is better for
                  looking up multiple jobs at once.

   -v, --verbose  Be verbose (show underlying 'sacct' commands on stderr).
   -V, --version  Print program version and exit.
   -h, --help     Display this help message and exit.

Essentially, this program is just a glorified wrapper for:
     sacct --env-vars -j JobID
in a loop (because it's just easier to type '$PROGNAME NNNN').

For interactive jobs returned "environment" value is 'NONE' (not my choice).
Exit status of 0 if succeeded for all JobIDs, 1 otherwise.

Note: for array jobs (not explicit array elements XXX_YY), underlying
'sacct' returns *separate* environments for each array element.  This
makes sense (those environments do differ by at least the SLURM_ARRAY_*
variables), but most likely this is not what the user typically wants.
This script will only show one such environment by default (but do show
everyting in '--raw' mode).

Note that this script relies on Slurm's built-in feature of storing
job environment variables in the accounting database, so:
  a) Requires Slurm 21.08 or newer;
  b) The feature must be enabled with 'AccountingStoreFlags = job_env'
     in your slurm.conf (it is not on by default).

Author:
  Lev Gorenstein (Purdue University Research Computing) <lev@purdue.edu>, 2021
_EOF_

    echo ""
}


function warn() {
    echo -e "$*" 1>&2
}

# Parse command line
# Note: relies on 'getopt' (a drop-in powerful replacement for 'getopts').
# AFAIK, it is part of any modern Linux distribution, so no further checking.
raw=0
verbo=0
SHORTOPTS="rvVh"
LONGOPTS="raw,verbose,version,help"
if ! TEMP=$(getopt -o "$SHORTOPTS" --long "$LONGOPTS" -- "$@"); then
    usage; exit 1
fi
eval set -- "$TEMP"
while true; do
    case "$1" in
        -r|--raw)
            raw=1; shift ;;
        -v|--verbose)
            verbo=1; shift ;;
        -V | --version)
            echo "$PROGNAME ver. $VERSION"; exit 0 ;;
        -h|--help)
            help_message; exit 0 ;;
        --) shift ; break ;;
        *) warn "$PROGNAME: Internal error ('$1' encountered)"; exit 1 ;;
    esac
done

# Actual work
rc_all=0
declare -a ARRAY_JOBS
for JOB in "$@"; do
    # Query accounting database and save output into an array.
    rc=0
    if [[ $verbo -gt 0 ]]; then
        warn "# /usr/bin/sacct -j $JOB --env-vars"
    fi
    if ! readarray -t RAW < <(/usr/bin/sacct -j $JOB --env-vars); then
        rc=1
        rc_all=1
    fi

    # Normally the above output would be along the lines of:
    #    Environment used for 55 (must be batch to display)
    #    ----------------------------------------------------------
    #    SHELL=/bin/bash
    #    [....]
    # Or for interactive jobs:
    #    Environment used for 55 (must be batch to display)
    #    ----------------------------------------------------------
    #    NONE
    # And a trailing empty line afterwards. So we should get at least 4 lines.
    nlines=${#RAW[@]}
    if [[ $nlines -lt 4 || $rc -ne 0 ]]; then
	echo "No environment found for jobid $JOB"
        rc_all=1
        continue
    fi

    # Catch: for array jobs (not explicit array elements XXX_YY), above
    # 'sacct' returns individual environments for each array element.  While
    # this makes sense (environments differ by respective SLURM_ARRAY_*
    # variables), this is hardly what the user typically wants, so filter out
    # duplicates by default (but do show everyting in '--raw' mode).
    # And save job IDs for the warning.
    pattern="^Environment used for ${JOB}"
    readarray -t FILTERED < <(printf "%s\n" "${RAW[@]}"  \
                                | sed -ne "1,/${pattern}/p" )
    if [[ $raw -eq 0 && $nlines -gt ${#FILTERED[@]} ]]; then
        # Looks like there were duplicates in sacct output (and we got
        # first copy in the FILTERED array).  Now need to remove the last
        # line from FILTERED (because the pattern matched at the tail).
        unset FILTERED[-1]
        ARRAY_JOBS+=( "$JOB" )
    fi
    nfilt=${#FILTERED[@]}                       # Do it here (for final count)

    # Finally, print what we got. Failure mode already handled.
    if [[ $raw -ne 0 ]]; then
        printf "%s\n" "${RAW[@]}"
    else
        printf "%s\n" "${FILTERED[@]:2:$((nfilt - 3))}"
    fi

done

# Special case: throw an array jobs warning if necessary.
if [[ $raw -eq 0 && ${#ARRAY_JOBS[@]} -gt 0 && $verbo -ge 0 ]]; then
    warn "# $PROGNAME: the following array jobs were detected:  ${ARRAY_JOBS[@]}"
    warn "#    For array jobs by default this script only shows the very first array"
    warn "#    element.  If you need to see all elements, use '--raw'.  For specific"
    warn "#    elements, request them explicitly (e.g. '<JOB>_123'). See '-h' for more."
fi
exit $rc_all