Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
executable file 261 lines (193 sloc) 9.46 KB
#!/usr/bin/env python3
'''Gstat
Summarize the status of the jobs (wrapper around "squeue").
Usage:
Gstat [options]
Gstat [options] [--jobid=N...] [--host=N...] [--user=N...] [--name=N...] [--account=N...] [--partition=N...] [--sort=N...] [--output=N...]
Options:
-U Limit output to the current user.
-u, --user=<NAME> Limit output to user(s) (may be a regex).
-j, --jobid=<NAME> Limit output to job-id(s) (may be a regex).
-h, --host=<NAME> Limit output to host(s) (may be a regex).
-a, --account=<NAME> Limit output to account(s) (may be a regex).
-n, --name=<NAME> Limit output to job-name(s) (may be a regex).
--status=<NAME> Limit output to status (may be a regex).
-p, --partition=<NAME> Limit output to partition(s) (may be a regex).
-s, --sort=<NAME> Sort by field (selected by the header name).
-r, --reverse Reverse sort.
-o, --output=<NAME> Select output columns (selected by the header name).
--full-name Show full user names.
-S, --summary Print only summary.
--no-header Suppress header.
--no-truncate Print full columns, do not truncate based on screen width.
--width=<N> Set print with.
--colors=<NAME> Select color scheme from: none, dark. [default: dark]
-l, --list Print selected column as list.
--sep=<NAME> Set column separator. [default: ] (space)
--long Print full information (each column is printed as a line).
--debug=<FILE> Debug: read 'squeue -o "%all"' from file.
--help Show help.
--version Show version.
(c - MIT) T.W.J. de Geus | tom@geus.me | www.geus.me | github.com/tdegeus/GooseSLURM
'''
# --------------------------------------------------------------------------------------------------
import os, sys, re, subprocess, docopt, pwd
import GooseSLURM as gs
# ---------------------------------- parse command line arguments ----------------------------------
# parse command-line options
args = docopt.docopt(__doc__,version='0.0.6')
# change keys to simplify implementation:
# - remove leading "-" and "--" from options
args = {re.sub(r'([\-]{1,2})(.*)',r'\2',key): args[key] for key in args}
# - change "-" to "_" to facilitate direct use in print format
args = {key.replace('-','_'): args[key] for key in args}
# --------------------------------- field-names and print settings ---------------------------------
# handle 'alias' options
if args['U']: args['user'] += [pwd.getpwuid(os.getuid())[0]]
# conversion map: default field-names -> custom field-names
alias = {
'JOBID' :'JobID' ,
'USER' :'User' ,
'ACCOUNT' :'Account' ,
'NAME' :'Name' ,
'START_TIME' :'Tstart' ,
'TIME_LEFT' :'Tleft' ,
'NODES' :'#node' ,
'CPUS' :'#CPU' ,
'CPUS_R' :'#CPU(R)' ,
'CPUS_PD' :'#CPU(PD)' ,
'MIN_MEMORY' :'MEM' ,
'ST' :'ST' ,
'NODELIST(REASON)':'Host' ,
'PARTITION' :'Partition',
}
# conversion map: custom field-names -> default field-names
aliasInv = {alias[key].upper():key for key in alias}
# rename command line options -> default field-names
# - add key-names
aliasInv['STATUS'] = 'ST'
# - apply conversion
for key in [key for key in args]:
if key.upper() in aliasInv:
args[aliasInv[key.upper()]] = args.pop(key)
# print settings of all columns
# - "width" : minimum width, adapted to print width (min_width <= width <= real_width)
# - "align" : alignment of the columns (except the header)
# - "priority": priority of column expansing, columns marked "True" are expanded first
columns = [
{'key':'JOBID' , 'width':7 , 'align':'>', 'priority': True },
{'key':'USER' , 'width':7 , 'align':'<', 'priority': True },
{'key':'ACCOUNT' , 'width':7 , 'align':'<', 'priority': True },
{'key':'NAME' , 'width':11, 'align':'<', 'priority': False},
{'key':'START_TIME' , 'width':6 , 'align':'>', 'priority': True },
{'key':'TIME_LEFT' , 'width':5 , 'align':'>', 'priority': True },
{'key':'NODES' , 'width':5 , 'align':'>', 'priority': True },
{'key':'CPUS' , 'width':4 , 'align':'>', 'priority': True },
{'key':'MIN_MEMORY' , 'width':3 , 'align':'>', 'priority': True },
{'key':'ST' , 'width':2 , 'align':'<', 'priority': True },
{'key':'PARTITION' , 'width':9 , 'align':'<', 'priority': False},
{'key':'NODELIST(REASON)', 'width':5 , 'align':'<', 'priority': False},
]
# header
header = {column['key']: gs.rich.String(alias[column['key']],align=column['align'])
for column in columns}
# print settings for the summary
columns_summary = [
{'key':'USER' , 'width':7 , 'align':'<', 'priority': True },
{'key':'ACCOUNT' , 'width':7 , 'align':'<', 'priority': False},
{'key':'CPUS' , 'width':4 , 'align':'>', 'priority': True },
{'key':'CPUS_R' , 'width':6 , 'align':'>', 'priority': True },
{'key':'CPUS_PD' , 'width':6 , 'align':'>', 'priority': True },
{'key':'PARTITION' , 'width':9 , 'align':'<', 'priority': False},
]
# header
header_summary = {column['key']: gs.rich.String(alias[column['key']],align=column['align'])
for column in columns_summary}
# select color theme
theme = gs.squeue.colors(args['colors'].lower())
# ---------------------------------- load the output of "squeue" -----------------------------------
if not args['debug']:
lines = gs.squeue.read_interpret(theme=theme)
else:
lines = gs.squeue.read_interpret(
data = open(args['debug'],'r').read(),
now = os.path.getctime(args['debug']),
theme = theme,
)
# ------------------------------ limit based on command-line options -------------------------------
for key in ['USER','ACCOUNT','NAME','JOBID','ST','NODELIST(REASON)','PARTITION']:
if args[key]:
# limit data
lines = [l for l in lines if sum([1 if re.match(n,str(l[key])) else 0 for n in args[key]])]
# color-highlight selected columns
# - apply to all remaining lines
for line in lines: line[key].color = theme['selection']
# - apply to the header
header[key].color = theme['selection']
# ---------------------------------------------- sort ----------------------------------------------
# default sort
lines.sort(key=lambda line: line['START_TIME'], reverse=not args['reverse'])
# optional: sort by key(s)
if args['sort']:
for key in args['sort']:
lines.sort(key=lambda line: line[aliasInv[key.upper()]], reverse=args['reverse'])
# ----------------------------------------- select columns -----------------------------------------
if args['output']:
keys = [aliasInv[key.upper()] for key in args['output']]
columns = [column for column in columns if column['key'] in keys]
# --------------------------------------------- print ----------------------------------------------
if not args['summary']:
# optional: print all fields and quit
if args['long']:
gs.table.print_long(lines)
sys.exit(0)
# optional: print as list and quit
elif args['list']:
# - only one field can be selected
if len(columns) > 1:
print('Only one field can be selected')
sys.exit(1)
# - print and quit
gs.table.print_list(lines, columns[0]['key'], args['sep'])
sys.exit(0)
# default: print columns
else:
gs.table.print_columns(lines, columns, header, args['no_truncate'], args['sep'], args['width'])
sys.exit(0)
# ------------------------------------- summarize information --------------------------------------
# get names of the different users
users = sorted(set([str(line['USER']) for line in lines]))
# start a new list of "user information", summed on the relevant users
users = [{'USER':gs.rich.String(key)} for key in users]
# loop over users
for user in users:
# - isolate jobs for this user
N = [line for line in lines if str(line['USER']) == str(user['USER'])]
# - get (a list of) partition(s)/account(s)
user['PARTITION'] = gs.rich.String(','.join(list(set([str(line['PARTITION']) for line in N]))))
user['ACCOUNT' ] = gs.rich.String(','.join(list(set([str(line['ACCOUNT' ]) for line in N]))))
# - count used CPU (per category)
user['CPUS' ] = gs.rich.Integer(sum([int(line['CPUS' ]) for line in N]))
user['CPUS_R' ] = gs.rich.Integer(sum([int(line['CPUS_R' ]) for line in N]))
user['CPUS_PD'] = gs.rich.Integer(sum([int(line['CPUS_PD']) for line in N]))
# - remove zeros from output for more intuitive output
for key in ['CPUS_R', 'CPUS_PD']:
if int(user[key]) == 0:
user[key] = gs.rich.Integer('-')
# rename field
lines = users
# ---------------------------------------------- sort ----------------------------------------------
# default sort
lines.sort(key=lambda line: line['USER'], reverse=args['reverse'])
# optional: sort by key(s)
if args['sort']:
# get available keys in the setting with fewer columns
keys = [alias[column['key']].upper() for column in columns_summary]
# filter sort keys that are not available in this mode
args['sort'] = [key for key in args['sort'] if key.upper() in keys]
# apply sort
for key in args['sort']:
lines.sort(key=lambda line: line[aliasInv[key.upper()]], reverse=args['reverse'])
# --------------------------------------------- print ----------------------------------------------
gs.table.print_columns(lines, columns_summary, header_summary,
args['no_truncate'], args['sep'], args['width'])
You can’t perform that action at this time.