forked from dguest/pandamonium
/
pandamon
executable file
·164 lines (144 loc) · 4.98 KB
/
pandamon
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
#!/usr/bin/env python
"""Script to retreve datasets (input or output) in submitted jobs
With a specified search string, will search for datasets with that
name. If the name doesn't end in `*` or `/`, append a wildcard.
Without a specified search string, the datasets can be piped.
The user name can be specified via environment variable to reduce
clutter.
"""
# help strings
_h_taskname='initial search string'
_h_user='full user name, or blank for all'
_h_stream='stream name fragment to filter for'
_h_days='only look back this many days'
_h_state='prefix dataset name with state'
_h_clean='clean output: better for cut or grep'
_h_taskid='output only taskid. useful for piping.'
# defaults
_def_user='GRID_USER_NAME'
_def_stream='OUT'
import urllib2
import urllib
import json
import sys, os
import re
import argparse
_headers = {'Accept': 'application/json',
'Content-Type':'application/json',
'User-Agent':'User-Agent: curl/7.43.0'}
def get_args():
d = ' (default: %(default)s)'
c = ' (default: %(const)s)'
user=os.environ.get(_def_user, '')
if not user:
de = ' (please set {} environment variable)'.format(_def_user)
else:
de = d
parser = argparse.ArgumentParser(
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('taskname', help=_h_taskname, nargs='?', default="user.%s"%(os.environ["USER"]) )
parser.add_argument('-u','--user', help=_h_user + de, default=user)
parser.add_argument('-d','--days', help=_h_days, type=int)
addinfo = parser.add_mutually_exclusive_group()
addinfo.add_argument('-s','--stream', help=_h_stream + c, nargs='?',
default=None, const=_def_stream)
parser.add_argument('-c','--clean', action='store_true', help=_h_clean)
parser.add_argument('-t','--taskid', action='store_true', help=_h_taskid)
args = parser.parse_args()
if not args.taskname and sys.stdin.isatty():
parser.print_usage()
sys.exit('ERROR: need to pipe datasets or specify a search string')
return args
def get_datasets(taskname, user, days=None):
pars = {
'taskname': taskname,
'username': user,
}
if days is not None:
pars['days'] = days
url = 'http://bigpanda.cern.ch/tasks/?' + urllib.urlencode(pars)
req = urllib2.Request(url, headers=_headers)
return json.loads(urllib2.urlopen(req).read().decode('utf-8'))
def get_ds(taskid, streamfrag):
params = urllib.urlencode({'jeditaskid': taskid})
url = 'http://bigpanda.cern.ch/tasks/?' + params
req = urllib2.Request(url, headers=_headers)
stuff = urllib2.urlopen(req).read().decode('utf-8')
# containers = set()
for entry in json.loads(stuff):
datasets = entry['datasets']
streams = {ds['streamname']: ds for ds in datasets}
for stream, ds in streams.items():
if streamfrag.upper() in stream:
yield streams[stream]['containername']
RED = '\033[0;91m'
GREEN = '\033[92m'
YELLOW = '\033[93m'
BLUE = '\033[94m'
MAGENTA = '\033[95m'
CYAN = '\033[96m'
ENDC = '\033[0m'
BOLD = '\033[1m'
UNDERLINE = '\033[4m'
BLINK = '\033[5m'
_color_dic = {
'running': BLUE + BOLD,
'submitting': CYAN,
'registered': MAGENTA,
'ready': MAGENTA,
'done': GREEN,
'finished': YELLOW,
'broken': RED + BOLD,
'aborted': RED,
'failed': RED,
}
def getstatus(task, args):
if args.clean:
fmt_string = '{s} {i} {p}% {t} '
elif args.taskid:
fmt_string = '{i}'
else:
fmt_string = '{s:<{l}} {i:<9} {p:<6} {t} '
if sys.stdout.isatty() and not args.clean:
color = _color_dic.get(task['superstatus'], ENDC)
status_color = color + task['status'] + ENDC
nonprlen = len(color) + len(ENDC)
else:
status_color = task['status']
nonprlen = 0
if not args.taskid:
outputString = fmt_string.format(
s=status_color, t=task['taskname'], i=task['reqid'],
l=(11 + nonprlen), p="%i%%"%task['dsinfo']["pctfinished"])
if args.taskid:
outputString = fmt_string.format( i=task['reqid'] )
return outputString
def stdin_iter(args):
for line in sys.stdin:
task = line.strip()
if task[-1] not in '/*':
task = task + '*'
for ds in get_datasets(task, args.user, args.days):
yield ds
def run():
args = get_args()
taskname = args.taskname
# try to search
if taskname is not None:
# append a wildcard if I forgot
if args.taskname[-1] not in '/*':
taskname = taskname + '*'
datasets = get_datasets(taskname, args.user, args.days)
else:
# otherwise read from stdin
datasets = stdin_iter(args)
# loop over tasks
for task in datasets:
if args.stream:
for ds in get_ds(task['jeditaskid'], args.stream):
sys.stdout.write(ds + '\n')
else:
sys.stdout.write(getstatus(task, args) + '\n')
if __name__ == '__main__':
run()