snakemakeslurm

#!/bin/bash
#
# launch snakemake to run jobs on UAB CHEAHA via SLURM
#
# 20220918 curtish; add gres support for GPUs
#
# debugging
EXEC="exec"
if [[ "$1" == -*debug ]]; then
    # echo commandline instead of running it
    EXEC="echo CMD: "
    DEBUG=1
    shift 1
fi

# auto-flag conda & singularity
AUTO_USE=0
AUTO_USE_ARGS=""
if [[ "$1" == -*use ]]; then
    # automatically generate --use-conda and --use-singularity, if directives present in Snaekfile
    AUTO_USE=1
    shift 1
fi

#
# check for conda: sections, but no --use-conda
#
snakefile_has_conda=$(egrep -c "^[\t ]+conda:" Snakefile)
if [[ snakefile_has_conda -gt 0 ]]; then 
    if [[ "$*" != *--use-conda* ]]; then
	if [[ "$AUTO_USE" == 0 ]]; then
            echo "# WARNING: Snakefile has $snakefile_has_conda 'conda:' directives, but you did not specify --use-conda" >&2
        else
	    echo "# Snkefile has 'conda:' directive: adding --use-conda"
            AUTO_USE_ARGS="$AUTO_USE_ARGS --use-conda"
        fi
    fi
fi

#
# check for singularity: sections, but no --use-singularity
#
snakefile_has_sing=$(egrep -c "^[\t ]+singularity:" Snakefile)
if [[ snakefile_has_sing -gt 0 ]]; then 
    if [[ "$*" != *--use-singularity* ]]; then
	if [[ "$AUTO_USE" == 0 ]]; then
            echo "# WARNING: Snakefile has $snakefile_has_sing 'singularity:' directives, but you did not specify --use-singularity" >&2
        else
	    echo "# Snkefile has 'singularity:' directive: adding --use-singularity"
            AUTO_USE_ARGS="$AUTO_USE_ARGS --use-singularity"
        fi
    fi
fi

#
# SLURM same name
#
# list of SLURM parameters that will come straight from cluster.slurm.cheaha.json/cluster.json with SAME NAME
#
SM_PARAMS="job-name ntasks partition time mail-type error output gres
"

#
# SLURM mapped name
#
# list of SLURM parameters that will come from Snakefile and cluster.slurm.cheaha.json/cluster.json
# with a name mapping
#
SM_ARGS=""

# "threads: 4" attribute directly on the rule, so works correctly for non-cluster case
SM_ARGS="$SM_ARGS --cpus-per-task {threads}"

# "mem-per-cpu-mb" - add units to the name to cut down on errors.  
SM_ARGS="$SM_ARGS --mem-per-cpu {cluster.mem-per-cpu-mb}"

# "mail-user" - set it to our blazerID
SM_ARGS="$SM_ARGS --mail-user $USER@uab.edu"
#
# build SLURM parameter mapping
#
for P in ${SM_PARAMS}; do
    SM_ARGS="$SM_ARGS --$P {cluster.$P}";
done

#
# validate we have a Snakefile
#
if [ ! -e "./Snakefile" ]; then
    echo "ERROR: $0 can not find $PWD/Snakefile"
    exit 1
fi

#
# look for cluster config files
#
# list in order, later files override earlier ones
#

# global config lives with this script

#
# check if we have a script to get "cause of death" from SLURM
# w/o this, if SLURM kills a job for over-time or over-mem, snakemake wont know
# 
# See: https://snakemake.readthedocs.io/en/stable/tutorial/additional_features.html?highlight=--cluster-status%20#using-cluster-status
#
SLURM_STATUS_FLAG=
SLURM_STATUS_SCRIPT=$(dirname $0)/slurm-status.py
if [ -e "$SLURM_STATUS_SCRIPT" ]; then 
    if [ "$DEBUG" == "1" ]; then echo "# found $SLURM_STATUS_SCRIPT"; fi; 
    SLURM_STATUS_FLAG=" --cluster-status $SLURM_STATUS_SCRIPT"
else
    echo "# WARNING: missing $SLURM_STATUS_SCRIPT "
fi

# global config lives with this script
CCONFIGS="--cluster-config $(dirname $0)/cluster.slurm.cheaha.json"

# look for local ones to merge
for LOCAL_CCONFIG in ./cluster.slurm.cheaha.json ./cluster.json ./slurm-config.yaml ./cluster.yaml; do
    if [ -e "$LOCAL_CCONFIG" ]; then 
	CCONFIGS="$CCONFIGS --cluster-config $LOCAL_CCONFIG"; 
	if [ "$DEBUG" == "1" ]; then echo "# found $LOCAL_CCONFIG"; fi;  
    fi
done

#
# our __default__ SLURM error/output paths expect a logs/ subdir in PWD
#
mkdir -p ./logs/
RC=$?;
if [ "$RC" != "0" ]; then exit $RC; fi

#
# actually run snakemake
#
# --latency-wait SECONDS
#   We need this for files created on compute nodes to become reliably visible on the login node where snakemake runs.
#
# --jobs CONCURRENT_JOB_COUNT
#   Number of jobs snakemake can submit to slurm at one time.
#
# --cluster-config FILE.json
#   SLURM parameters to use when submitting jobs for each rule.
#
$EXEC snakemake \
    $AUTO_USE_ARGS \
    --latency-wait 45 \
    --jobs 999 \
    $CCONFIGS \
    $SLURM_STATUS_FLAG \
    --cluster "sbatch $SM_ARGS" \
    $*