From bba6f8de496690525cdc3ededbd7c59ca2e8fc4c Mon Sep 17 00:00:00 2001
From: Sergey Koren <sergek@umd.edu>
Date: Thu, 25 Jul 2024 18:18:03 -0400
Subject: [PATCH] update LSF support, add PBS support (untested) #88

---
 src/profiles/slurm-sge-status.sh | 22 ++++++++++++++++++++++
 src/profiles/slurm-sge-submit.sh | 28 ++++++++++++++++++++++++++--
 src/verkko.sh                    |  9 +++------
 3 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/src/profiles/slurm-sge-status.sh b/src/profiles/slurm-sge-status.sh
index daf5bfa6..66ca6c76 100755
--- a/src/profiles/slurm-sge-status.sh
+++ b/src/profiles/slurm-sge-status.sh
@@ -16,6 +16,8 @@ sge=$SGE_ROOT
 #  Test for LSF: if LSF_ENVDIR exists in the environment, assume LSF works.
 lsf=$LSF_ENVDIR
 
+# Test for PBS; if we have pbsnodes binary
+pbs=$(which pbsnodes 2> /dev/null)
 
 #  Check Slurm status.
 #
@@ -126,6 +128,26 @@ elif [ "x$lsf" != "x" ] ; then
    } \
    END { print stat }')
 
+elif [ "x$pbs" != "x" ]; then
+  jobstatus=$(qstat -f "$jobid" | \
+  awk \
+  'BEGIN { stat="running" } \
+   /job_state/ {state=$NF} /exit_status/ {exit_status=$NF} \
+   END { \
+    if      (state == "R")        { stat="running" } \
+    else if (state == "Q")        { stat="running" } \
+    else if (state == "H")        { stat="running" } \
+    else if (state == "T")        { stat="running" } \
+    else if (state == "W")        { stat="running" } \
+    else if (state == "S")        { stat="running" } \
+    else if (state == "B")        { stat="running" } \
+    else if (state == "E")        { stat="running" } \
+    else if (state == "C")        {                  \
+       if (exit_status == 0)      { stat="success" } \
+       else                       { stat="failed"  } \
+    }                                                \
+    else                          { stat="failed"  } \
+    print stat }')
 
 #  Otherwise, do what?  Fail!
 else
diff --git a/src/profiles/slurm-sge-submit.sh b/src/profiles/slurm-sge-submit.sh
index 8d53158d..177e3bb3 100755
--- a/src/profiles/slurm-sge-submit.sh
+++ b/src/profiles/slurm-sge-submit.sh
@@ -52,10 +52,12 @@ mkdir -p batch-scripts/
 #    Slurm: if sinfo is present, assume slurm works.
 #    SGE:   if SGE_ROOT exists in the environment, assume SGE works.
 #    LSF:   if LSF_ENVDIR exists in the environment, assume LSF works.
+#    PBS:   if pbsnodes present, assume PBS works
 #
 slurm=$(which sinfo 2> /dev/null)
 sge=$SGE_ROOT
 lsf=$LSF_ENVDIR
+pbs=$(which pbsnodes 2> /dev/null)
 
 ##########
 #
@@ -141,17 +143,39 @@ elif [ "x$lsf" != "x" ] ; then
   else
     mem=$(dc -e "3 k ${mem_gb} 1048576 * p")
   fi
+  mem_per_thread=$(dc -e "0 k ${mem} ${n_cpus} / p")
 
-  jobid=$(bsub -R "span[hosts=1] rusage[mem=${mem}]" -n ${n_cpus} -oo batch-scripts/${jobid}.${rule_n}.${jobidx}.out "$@" | grep -oE "Job <[0-9]+>" | awk '{print $2}' | tr -d '<>')
+  jobid=$(bsub -R \"span[hosts=1] rusage[mem=${mem_per_thread}]\" -n ${n_cpus} -oo batch-scripts/${jobid}.${rule_n}.${jobidx}.out "$@" | sed -E 's/.*<([0-9]+)>.*/\1/')
+  if [ "x$jobid" = "x" ]; then
+     exit 1
+  fi
+
+  echo > batch-scripts/${jobid}.${rule_n}.${jobidx}.submit \
+          bsub -R \"span[hosts=1] rusage[mem=${mem_per_thread}]\" -n ${n_cpus} -oo batch-scripts/${jobid}.${rule_n}.${jobidx}.out "$@"
+
+#  Submit to PBS.
+#  Other options:
+#    -A account
+#    -q queue
+#    -N job name
+#    -e err-out
+#    -o out-out
+#
+elif [ "x$pbs" != "x" ] ; then
+  jobid=$(qsub -j oe -l select=1:ncpus=${n_cpus}:mem=${mem_gb}gb:walltime=${time_h}:00:00 -o batch-scripts/${jobid}.${rule_n}.${jobidx}.out "$@" | cut -d. -f1)
+  if [ "x$jobid" = "x" ]; then
+     exit 1
+  fi
 
   echo > batch-scripts/${jobid}.${rule_n}.${jobidx}.submit \
-          bsub -R "span[hosts=1] rusage[mem=${mem}]" -n ${n_cpus} -oo batch-scripts/${jobid}.${rule_n}.${jobidx}.out "$@"
+          qsub -j oe -l select=1:ncpus=${n_cpus}:mem=${mem_gb}gb:walltime=${time_h}:00:00 -o batch-scripts/${jobid}.${rule_n}.${jobidx}.out "$@"
 
 ##########
 #
 #  Otherwise, fail.
 #
 else
+  echo "Error: unknown grid, only Slurm, SGE, LSF, and PBS are supported. Please check your environment or use --local instead"
   exit 1
 fi
 
diff --git a/src/verkko.sh b/src/verkko.sh
index ce9074cd..ec542a00 100755
--- a/src/verkko.sh
+++ b/src/verkko.sh
@@ -326,9 +326,8 @@ while [ $# -gt 0 ] ; do
     elif [ "$opt" = "--bwa" ] ;                then bwa=$arg;           shift
     elif [ "$opt" = "--samtools" ] ;           then samtools=$arg;      shift
     elif [ "$opt" = "--local" ] ;              then grid="local";
-    elif [ "$opt" = "--sge" ] ;                then grid="slurm-sge";
-    elif [ "$opt" = "--slurm" ] ;              then grid="slurm-sge";
-    elif [ "$opt" = "--lsf" ] ;                then grid="lsf";
+    elif [ "$opt" = "--grid" ] ;               then grid="slurm-sge-lsf-pbs";
+    elif [ "$opt" = "--slurm" ] ;              then grid="slurm-sge-lsf-pbs";
     elif [ "$opt" = "--local-memory" ] ;       then local_mem=$arg;     shift
     elif [ "$opt" = "--local-cpus" ] ;         then local_cpus=$arg;    shift
     elif [ "$opt" = "--snakeopts" ] ;          then snakeopts=$arg;     shift
@@ -888,9 +887,7 @@ if [ "x$help" = "xhelp" -o "x$errors" != "x" ] ; then
     echo "    --local-memory           Specify the upper limit on memory to use, in GB, default 64"
     echo "    --local-cpus             Specify the number of CPUs to use, default 'all'"
     echo ""
-    echo "    --sge                    Enable Sun Grid Engine support."
-    echo "    --slurm                  Enable Slurm support."
-    echo "    --lsf                    Enable IBM Spectrum LSF support."
+    echo "    --grid                   Enable Grid support (SGE, Slurm, PBS, and LSF are supported)."
     echo ""
     echo "    --snakeopts <string>     Append snakemake options in \"string\" to the"
     echo "                             snakemake command.  Options MUST be quoted."