Merge pull request #4 from vanderwb/subjob-support

Add proper support for array flags

Author: vanderwb

README.md

@@ -30,10 +30,12 @@ Options:
     -f              display full output for a job
     -Fjson          display full output in JSON format (use with -f)
     -H              job output regardless of state or all finished jobs
+    -J              only show information for jobs (or subjobs with -t)
     -l              disable labels (no header)
     -n              display a list of nodes at the end of the line
     -s              display administrator comment on the next line
     --status        filter jobs by specific single-character status code
+    -t              show information for both jobs and array subjobs
     -u              filter jobs by the submitting user
     -w              use wide format output (120 columns)
     -x              include recently finished jobs in output

gen_data.sh

@@ -60,12 +60,12 @@ function main_gen {
 
     # Get data from PBS
     QSS_TIME=$SECONDS
-    $PBSPREFIX $QSTATBIN -x | sed '/^[0-9]/,$!d' | sed 's/\([0-9]\+\) b/ \1b/' > newlist-default.dat &
-    $PBSPREFIX $QSTATBIN -1 -n -s -x | sed '/^[0-9]/,$!d' | sed 's/\([0-9]\+\) b/ \1b/' > newlist-info.dat &
-    $PBSPREFIX $QSTATBIN -a -1 -n -s -w -x | sed '/^[0-9]/,$!d' | sed 's/\([0-9]\+\) b/ \1b/' > newlist-wide.dat &
+    $PBSPREFIX $QSTATBIN -t -x | sed '/^[0-9]/,$!d' | sed 's/\([0-9]\+\) b/ \1b/' > newlist-default.dat &
+    $PBSPREFIX $QSTATBIN -t -1 -n -s -x | sed '/^[0-9]/,$!d' | sed 's/\([0-9]\+\) b/ \1b/' > newlist-info.dat &
+    $PBSPREFIX $QSTATBIN -t -a -1 -n -s -w -x | sed '/^[0-9]/,$!d' | sed 's/\([0-9]\+\) b/ \1b/' > newlist-wide.dat &
 
     if [[ " $CACHEFLAGS " == *" f "* ]]; then
-        $PBSPREFIX $QSTATBIN -f > joblist-full.dat &
+        $PBSPREFIX $QSTATBIN -t -f > joblist-full.dat &
     else
         rm -f joblist-full.dat
     fi
@@ -75,7 +75,7 @@ function main_gen {
         #   1. Numbers after a number 0 (octal) that aren't strings
         #   2. Trailing decimal points in numbers
         #   3. Numbers that begin with a decimal point
-        $PBSPREFIX $QSTATBIN -f -F json | sed 's/":\(0[0-9][^,]*\)/":"\1"/; s/":\([0-9]*\)\.,/":"\1\.",/; s/":\(\.[^,]*\)/":"\1"/' > joblist-fulljson.dat &
+        $PBSPREFIX $QSTATBIN -t -f -F json | sed 's/":\(0[0-9][^,]*\)/":"\1"/; s/":\([0-9]*\)\.,/":"\1\.",/; s/":\(\.[^,]*\)/":"\1"/' > joblist-fulljson.dat &
     else
         rm -f joblist-fulljson.dat
     fi

qstat

@@ -21,10 +21,12 @@ Options:
     -f              display full output for a job
     -Fjson          display full output in JSON format (use with -f)
     -H              job output regardless of state or all finished jobs
+    -J              only show information for jobs (or subjobs with -t)
     -l              disable labels (no header)
     -n              display a list of nodes at the end of the line
     -s              display administrator comment on the next line
     --status        filter jobs by specific single-character status code
+    -t              show information for both jobs and array subjobs
     -u              filter jobs by the submitting user
     -w              use wide format output (120 columns)
     -x              include recently finished jobs in output
@@ -241,6 +243,10 @@ while [[ $# -gt 0 ]]; do
                         HMODE=true
                         DATAFILE=${DATAFILE/default/info}
                         ;;
+                    J)
+                        JMODE=true
+                        JSOPTS="-J $JSOPTS"
+                        ;;
                     l)
                         NOLABELS=true
                         ;;
@@ -253,6 +259,10 @@ while [[ $# -gt 0 ]]; do
                         DATAFILE=${DATAFILE/job/comm}
                         DATAFILE=${DATAFILE/default/info}
                         ;;
+                    t)
+                        SUBJOBS=true
+                        JSOPTS="-t $JSOPTS"
+                        ;;
                     u)
                         DATAFILE=${DATAFILE/default/info}
 
@@ -288,7 +298,11 @@ while [[ $# -gt 0 ]]; do
             find_server
             ;;&
 		[0-9]*)
-            JOBID=${1%%@*}
+            if [[ $1 == *\[[0-9]*\]* ]]; then
+                JOBID=$(sed 's/\[/\\\[/; s/\]/\\\]/' <<< ${1%%@*})
+            else
+                JOBID=$(sed 's/\[\]/\\[\[0-9\]\*\\]/' <<< ${1%%@*})
+            fi
 
             if [[ -z $CUSTOMSERVER ]] && [[ $JOBID == *.* ]]; then
                 CUSTOMSERVER=${JOBID#*.} find_server
@@ -358,7 +372,7 @@ if [[ $FULLMODE == true ]]; then
             DATAFILE=joblist-fulljson.dat
 
             if [[ -n $JOBLIST ]]; then
-                JSOPTS="-j $JOBLIST"
+                JSOPTS="$JSOPTS -j $JOBLIST"
             else
                 if [[ -n $DESTLIST ]]; then
                     JSOPTS="$JSOPTS -q $DESTLIST"
@@ -397,7 +411,7 @@ if [[ $FULLMODE == true ]]; then
         fi
 
         if [[ -n $JOBLIST ]]; then
-            CMD="awk -vRS='' '/Job Id: ($JOBLIST)./{print \$0\"\n\"}' $DATAPATH/$DATAFILE"
+            CMD="awk -vRS='' '/Job Id: ($JOBLIST)\./{print \$0\"\n\"}' $DATAPATH/$DATAFILE"
         else
             if [[ $SHOWALL == true ]]; then
                 CMD="cat $DATAPATH/$DATAFILE"
@@ -414,6 +428,14 @@ if [[ $FULLMODE == true ]]; then
             CMD="$CMD | awk -vRS='' '/job_state = $STATCODE/{print \$0\"\n\"}'"
         fi
 
+        if [[ $SUBJOBS == true ]]; then
+            if [[ $JMODE == true ]]; then
+                CMD="$CMD | awk -vRS='' '/Job Id:.*\[[0-9]+\]/{print \$0\"\n\"}'"
+            fi
+        else
+            CMD="$CMD | awk -vRS='' '!/Job Id:.*\[[0-9]+\]/{print \$0\"\n\"}'"
+        fi
+
         # Get requested output
         query_cache
 
@@ -453,19 +475,32 @@ else
     # Attach various filters to output
     if [[ ! -z "$STATCODE" ]]; then
         CMD="$CMD | awk '\$$SCOL && \$$SCOL $COPER \"$STATCODE\" && \$1 ~ /^[0-9]/ { $AWKOPT }'"
+        AWKOPT=print
     fi
 
     if [[ ! -z "$USERSTR" ]] && [[ ${USERSTR^^} != ALL ]]; then
         CMD="$CMD | awk '\$$UCOL == \"$USERSTR\" { $AWKOPT }'"
+        AWKOPT=print
     fi
 
     if [[ ! -z "$JOBLIST" ]]; then
-        JOBLIST=${JOBLIST/\[\]/\\\[\\\]}
         CMD="$CMD | awk '\$1 ~ /^($JOBLIST)\./ { $AWKOPT }'"
+        AWKOPT=print
     fi
 
     if [[ ! -z "$DESTLIST" ]]; then
         CMD="$CMD | awk '\$$DCOL ~ /^($DESTLIST)$/ { $AWKOPT }'"
+        AWKOPT=print
+    fi
+
+    if [[ $SUBJOBS == true ]]; then
+        if [[ $JMODE == true ]]; then
+            CMD="$CMD | awk '\$1 ~ /\[[0-9]+\]/ { $AWKOPT }'"
+            AWKOPT=print
+        fi
+    else
+        CMD="$CMD | awk '\$1 !~ /\[[0-9]+\]/ { $AWKOPT }'"
+        AWKOPT=print
     fi
 
     # Get requested output

qstat-json.py

@@ -1,21 +1,23 @@
 #!/usr/bin/env python3
 
-import json, collections, argparse, sys
+import json, collections, argparse, sys, re
 
 from signal import signal, SIGPIPE, SIG_DFL
 signal(SIGPIPE,SIG_DFL) 
 
 parser = argparse.ArgumentParser()
 parser.add_argument("input_file")
 parser.add_argument("--jobs", "-j", nargs='*')
+parser.add_argument("--filter", "-J", action = "store_true")
 parser.add_argument("--queues", "-q", nargs='*')
 parser.add_argument("--status", "-s")
+parser.add_argument("--arrays", "-t", action = "store_true")
 parser.add_argument("--user", "-u")
 args = parser.parse_args()
 
 with open(args.input_file, 'r') as f:
     try:
-        data = json.load(f, object_pairs_hook=collections.OrderedDict)
+        data = json.load(f, object_pairs_hook=collections.OrderedDict, strict = False)
     except json.decoder.JSONDecodeError:
         sys.exit(3)
 
@@ -25,9 +27,12 @@
         if "." in job:
             jobs.append(job)
         else:
-            jobs.append("{}.".format(job))
+            jobs.append("{}\.".format(job))
 
-    data['Jobs'] = { job : data['Jobs'][job] for job in data['Jobs'].keys() if job.startswith(tuple(jobs)) }
+    try:
+        data['Jobs'] = { job : data['Jobs'][job] for job in data['Jobs'].keys() if any([re.match(jid, job) for jid in jobs]) }
+    except:
+        sys.exit(3)
 
 if args.queues:
     data['Jobs'] = { job : data['Jobs'][job] for job in data['Jobs'].keys() if data['Jobs'][job]['queue'] in args.queues }
@@ -38,4 +43,10 @@
 if args.status:
     data['Jobs'] = { job : data['Jobs'][job] for job in data['Jobs'].keys() if data['Jobs'][job]['job_state'] == args.status }
 
+if args.arrays:
+    if args.filter:
+        data['Jobs'] = { job : data['Jobs'][job] for job in data['Jobs'].keys() if re.search("\[[0-9]+\]", job) }
+else:
+    data['Jobs'] = { job : data['Jobs'][job] for job in data['Jobs'].keys() if not re.search("\[[0-9]+\]", job) }
+
 print(json.dumps(data, indent = 4))