-
Notifications
You must be signed in to change notification settings - Fork 2
/
Gdel
executable file
·187 lines (147 loc) · 6.29 KB
/
Gdel
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
#!/usr/bin/env python3
'''Gdel
Stop running jobs.
Usage:
Gdel [options]
Gdel [options] <JobId>...
Arguments:
ID-number(s) of the job(s) to delete. (default: all user's jobs)
Options:
--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]
--sep=<NAME> Set column separator. [default: ] (space)
--debug=<FILE> Debug. Output 'squeue -o "%all"' provided from file.
-h, --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
# ------------------------------------ prompt user confirmation ------------------------------------
def confirm(message='Proceed [y/n]?\n'):
while True:
user = input(message)
if not user : print('Please enter y or n.'); continue
if user not in ['y','Y','n','N']: print('Please enter y or n.'); continue
if user in ['y','Y' ]: return True
if user in ['n','N' ]: return False
# ---------------------------------- 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}
# - remove "<...>"
args = {re.sub(r'(<)(.*)(>)',r'\2',key): args[key] for key in args}
# --------------------------------- field-names and print settings ---------------------------------
# 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':'CPUS' , 'width':4 , 'align':'>', 'priority': True },
{'key':'CPUS_R' , 'width':6 , 'align':'>', 'priority': True },
{'key':'CPUS_PD' , 'width':6 , 'align':'>', 'priority': True },
]
# 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,
)
# ------------------------------------- get/check running jobs -------------------------------------
# get user-name
user = pwd.getpwuid(os.getuid())[0]
# filter to current user
lines = [line for line in lines if re.match(user,str(line['USER']))]
# filter to running/queued jobs (other jobs cannot be cancelled)
lines = [line for line in lines if str(line['ST']) in ['R','PD']]
# list with fields on which filters are applied
filters = ['USER', 'ST']
# filter to input list of job-ids
if args['JOBID']:
# - convert to set
jobs = set(args['JOBID'])
# - filter running jobs
lines = [line for line in lines if str(line['JOBID']) in jobs]
# - set selection-color
filters += ['JOBID']
# set selection-color
for key in filters:
for line in lines:
line[key].color = theme['selection']
# no jobs -> quit
if len(lines) == 0:
sys.exit(0)
# -------------------------------- prompt confirmation, cancel jobs --------------------------------
# print header
print('Delete jobs : ')
# print selected jobs
gs.table.print_columns(lines, columns, header, args['no_truncate'], args['sep'], args['width'])
# prompt response
if not confirm():
sys.exit(1)
# construct command
cmd = 'scancel '+' '.join([str(line['JOBID']) for line in lines])
# run/print command
if not args['debug']:
print(subprocess.check_output(cmd,shell=True).decode('utf-8'),end='')
else:
print(cmd)