Context Navigation

sara_nodes.py.in @ 366

Last change on this file since 366 was 366, checked in by martijk, 6 years ago
fixed incorrect rendering when no tickets are specified
Property svn:executable set to ``*
File size: 23.7 KB

Line
1	#!@PYTHON@
2	#
3	# Author: Dennis Stam
4	# Date : 6th of September 2012
5	#
6	# Tested with Python 2.5, 2.6, 2.7 (should work for 3.0, 3.1, 3.2)
7	# sara_nodes uses the module argparse
8
9	## The documenation, is shown when you type --help
10	from __future__ import print_function
11
12	HELP_DESCRIPTION = '''This program is a great example what you can achieve with the pbs_python wrapper. You can use sara_nodes to change the state of a machine to offline with a reason. Several information is stored in the nodes note attribute. Information such as; date added, date last change, the username, ticket number (handy when you wan't to reference to an issue in your tracking system) and the message.'''
13	HELP_EPILOG = '''The format argument uses the Python string formatting. Fields that can be used are; nodename, state, date_add, date_edit, username, ticket and note. For example sara_nodes -f '%(nodename)s;%(state)s' '''
14
15	## This variable is used to sort by basename
16	SPLIT_SORT = r'r\d+n\d+'
17	## This RE pattern is used for the hostrange
18	HOSTRANGE = r'\[([0-9az\-,]+)\]'
19	## Which states are allowed to show in the print_overview
20	ALLOWED_STATES = set(['down', 'offline', 'unknown'])
21
22	import pbs
23	import PBSQuery
24	import re
25	import sys
26	import time
27	import getpass
28	import types
29
30	## Use the cli arguments to change the values
31	ARGS_QUIET = False
32	ARGS_VERBOSE = False
33	ARGS_DRYRUN = False
34
35	## Import argparse here, as I need the print function
36	try:
37	import argparse
38	except ImportError:
39	print('Cannot find argparse module', file=sys.stderr)
40	sys.exit(1)
41
42	####
43	## BEGIN functions for hostrange parsing
44	def l_range(start, end):
45	'''The equivalent for the range function, but then with letters, uses the ord function'''
46	start = ord(start)
47	end = ord(end)
48	rlist = list()
49
50	## A ord number must be between 96 (a == 97) and 122 (z == 122)
51	if start < 96 or start > 122 and end < 96 or end > 122:
52	raise Exception('You can only use letters a to z')
53	## If start is greater then end, then the range is invalid
54	elif start > end:
55	raise Exception('The letters must be in alphabetical order')
56	## Just revert the ord number to the char
57	for letter in range(start, end + 1):
58	rlist.append(chr(letter))
59	return rlist
60
61	def return_range(string):
62	'''This function will return the possible values for the given ranges'''
63
64	## First check if the first char is valid
65	if string.startswith(',') or string.startswith('-'):
66	raise Exception('Given pattern is invalid, you can\'t use , and - at the beginning')
67
68	numbers_chars = list()
69	equal_width_length = 0
70
71	## First splitup the sections (divided by ,)
72	for section in string.split(','):
73	## Within a section you can have a range, max two values
74	chars = section.split('-')
75	if len(chars) == 2:
76	## When range is a digit, simply use the range function
77	if chars[0].isdigit() and chars[1].isdigit():
78	## Owke, check for equal_width_length
79	if chars[0][0] == '0' or chars[1][0] == '0':
80	if len(chars[0]) >= len(chars[1]):
81	equal_width_length = len(chars[0])
82	else:
83	equal_width_length = len(chars[1])
84	## Don't forget the +1
85	numbers_chars += range(int(chars[0]), int(chars[1])+1)
86	## If one of the two is a digit, raise an exception
87	elif chars[0].isdigit() or chars[1].isdigit():
88	raise Exception('I can\'t combine integers with letters, change your range please')
89	## Else use the l_range
90	else:
91	numbers_chars += l_range(chars[0], chars[1])
92	else:
93	## If the value of the section is a integer value, check if it has a 0
94	if section.isdigit() and section[0] == '0':
95	if len(section) > equal_width_length:
96	equal_width_length = len(section)
97	numbers_chars.append(section)
98
99	## if the equal_width length is greater then 0, rebuild the list
100	## 01, 02, 03, ... 10
101	if equal_width_length > 0:
102	tmp_list = list()
103	for number_char in numbers_chars:
104	if type(number_char) is types.IntType or number_char.isdigit():
105	tmp_list.append('%0*d' % ( equal_width_length, int(number_char)))
106	else:
107	tmp_list.append(number_char)
108	numbers_chars = tmp_list
109
110	return numbers_chars
111
112	def product(args, *kwargs):
113	'''Taken from the python docs, does the same as itertools.product,
114	but this also works for py2.5'''
115	pools = map(tuple, args) * kwargs.get('repeat', 1)
116	result = [[]]
117	for pool in pools:
118	result = [x+[y] for x in result for y in pool]
119	for prod in result:
120	yield tuple(prod)
121
122	def parse_args(args):
123	rlist = list()
124	for arg in args:
125	parts = re.findall(HOSTRANGE, arg)
126	if parts:
127	## create a formatter string, sub the matched patternn with %s
128	string_format = re.sub(HOSTRANGE, '%s', arg)
129	ranges = list()
130
131	## detect the ranges in the parts
132	for part in parts:
133	ranges.append(return_range(part))
134
135	## produce the hostnames
136	for combination in product(*ranges):
137	rlist.append(string_format % combination)
138	else:
139	rlist.append(arg)
140	return rlist
141
142	def get_nodes(jobs):
143
144	p = PBSQuery.PBSQuery()
145	nodes = list()
146
147	for job in jobs:
148	exec_hosts = p.getjob(job, ['exec_host'])
149	if not exec_hosts or 'exec_host' not in exec_hosts:
150	continue
151	for exec_host in exec_hosts['exec_host'][0].split('+'):
152	hostname = exec_host.split('/')[0]
153	if hostname not in nodes:
154	nodes.append(hostname)
155	return nodes
156
157	## END functions for hostrange parsing
158	####
159
160	####
161	## BEGIN functions for printing
162	def _generate_index(string):
163	'''Is used to generate a index, this way we can also sort nummeric values in a string'''
164	return [ int(y) if y.isdigit() else y for y in re.split(r'(\d+)', string) ]
165
166	def print_get_nodes(hosts=None):
167	'''This function retrieves the information from your batch environment'''
168	if ARGS_VERBOSE:
169	print('func:print_get_nodes input:%s' % str(hosts), file=sys.stderr)
170
171	## there are 2 possible filters, by hostname, or by state
172	pbsq = PBSQuery.PBSQuery()
173	split_1 = dict()
174	split_2 = dict()
175
176	if ARGS_VERBOSE:
177	print('func:print_get_nodes fetching node information', file=sys.stderr)
178	## We ask from the batch all nodes, and with the properties state and note
179	for host, properties in pbsq.getnodes(['state', 'note']).items():
180	do_host = None
181	## Check if the current host matches our criterium (given with the arguments
182	## or has the allowed state)
183	if hosts and host in hosts:
184	do_host = host
185	elif not hosts:
186	## Do a intersection on both set's, if there is a match, then the host is allowed
187	if bool(ALLOWED_STATES.intersection(set(properties.state))):
188	do_host = host
189
190	## when we have a do_host (means matches our criterium) then sort
191	## them by basename
192	if do_host:
193	if SPLIT_SORT and re.findall(SPLIT_SORT, do_host):
194	split_1[host] = properties
195	else:
196	split_2[host] = properties
197
198	if ARGS_VERBOSE:
199	print('func:print_get_nodes returning values', file=sys.stderr)
200	return split_1, split_2
201
202	def print_process_dict(dictin):
203	'''This function processes the data from the batch system and make it for all hosts the same layout'''
204	if ARGS_VERBOSE:
205	print('func:print_process_dict input:%s' % str(dictin), file=sys.stderr)
206
207	lineprint = list()
208	if ARGS_VERBOSE:
209	print('func:print_process_dict processing data', file=sys.stderr)
210
211	## Generate a list containing a dictionary, so we can use the stringformatting functionality
212	## of Python, fieldnames are: nodename, date_edit, date_add, username, ticket, note
213
214	## Replaced real_sort with sorted, this means from 50 lines of code to 3
215	for host in sorted(dictin.keys(), key=_generate_index):
216	add_dict = dict()
217
218	add_dict['nodename'] = host
219	## Glue the state list with a ,
220	add_dict['state'] = ", ".join(dictin[host]['state'] if dictin[host].has_key('state') else [])
221
222	## Check if the given host has a note
223	note = dictin[host]['note'] if dictin[host].has_key('note') else []
224	if len(note) > 3:
225	add_dict['date_add'] = note[0]
226	add_dict['date_edit'] = note[1]
227	add_dict['username'] = note[2]
228
229	if note[3].isdigit():
230	add_dict['ticket'] = note[3]
231	add_dict['note'] = ",".join(note[4:])
232	else:
233	add_dict['ticket'] = ""
234	add_dict['note'] = ",".join(note[3:])
235
236	## Create an extra date field, combined for date_edit and date_add
237	if add_dict['date_add'] and add_dict['date_edit']:
238	add_dict['date'] = '%s, %s' % (add_dict['date_add'], add_dict['date_edit'])
239	elif add_dict['date_add']:
240	add_dict['date'] = add_dict['date_add']
241	else:
242	add_dict['date'] = None
243
244	else:
245	## If there is no note, just set the variables with a empty string
246	add_dict['date'] = add_dict['date_add'] = add_dict['date_edit'] = add_dict['username'] = add_dict['ticket'] = add_dict['note'] = ''
247
248	lineprint.append(add_dict)
249
250	if ARGS_VERBOSE:
251	print('func:print_process_dict returning values', file=sys.stderr)
252	return lineprint
253
254	def print_create_list(values):
255	tmp_list = list()
256	for pair in values:
257	tmp_list.append('%-*s' % tuple(pair))
258	return tmp_list
259
260	def print_overview_normal(hosts=None):
261	'''Print the default overview'''
262	if ARGS_VERBOSE:
263	print('func:print_overview_normal input:%s' % str(hosts), file=sys.stderr)
264
265	## Determine some default values for the column width
266	w_nodename = 8
267	w_state = 5
268	w_date = w_username = w_ticket = w_note = w_date_add = w_date_edit = 0
269
270	## Get the data, make it one list, the rest first then the matched
271	matched, rest = print_get_nodes(hosts)
272	print_list = print_process_dict(rest)
273	print_list.extend(print_process_dict(matched))
274
275	## Detect the max width for the columns
276	for line in print_list:
277	if line['nodename'] and len(line['nodename']) > w_nodename:
278	w_nodename = len(line['nodename'])
279	if line['state'] and len(line['state']) > w_state:
280	w_state = len(line['state'])
281	if line['date'] and len(line['date']) > w_date:
282	w_date = len(line['date'])
283	if line['date_add'] and len(line['date_add']) > w_date_add:
284	w_date_add = len(line['date_add'])
285	if line['date_edit'] and len(line['date_edit']) > w_date_edit:
286	w_date_edit = len(line['date_edit'])
287	if line['username'] and len(line['username']) > w_username:
288	w_username = len(line['username'])
289	if line['ticket'] and len(line['ticket']) > w_ticket:
290	w_ticket = len(line['ticket'].strip())
291	if line['note'] and len(line['note']) > w_note:
292	w_note = len(line['note'])
293
294	## The length of the full note
295	w_notefull = w_date + w_username + w_ticket + w_note
296
297	if not ARGS_QUIET:
298	show_fields = [
299	[w_nodename, 'Nodename'],
300	[w_state, 'State'],
301	]
302	if w_date > 0:
303	show_fields.append([w_date_add,'Added'])
304	show_fields.append([w_date_edit,'Modified'])
305	show_fields.append([w_username,'User'])
306	if w_ticket > 0:
307	if w_ticket < 6:
308	w_ticket = 6
309	show_fields.append([w_ticket,'Ticket'])
310	show_fields.append([w_note,'Note'])
311
312	print(' %s' % ' \| '.join(print_create_list(show_fields)))
313	print('+'.join([ '-' * (show_field[0]+2) for show_field in show_fields ]))
314
315	## Show the information to the user
316	for line in print_list:
317	show_line_fields = [
318	[w_nodename, line['nodename']],
319	[w_state, line['state']],
320	]
321	if w_date > 0:
322	show_line_fields.append([w_date_add,line['date_add']])
323	show_line_fields.append([w_date_edit,line['date_edit']])
324	show_line_fields.append([w_username,line['username']])
325	if w_ticket > 0 and line['ticket'].strip():
326	show_line_fields.append([w_ticket,'#' + line['ticket']])
327	show_line_fields.append([w_note,line['note']])
328
329	print(' %s' % ' \| '.join(print_create_list(show_line_fields)))
330
331	def print_overview_format(hosts=None, format=None):
332	'''Print the information in a certain format, when you want to use it in a
333	different program'''
334
335	matched, rest = print_get_nodes(hosts)
336	print_list = print_process_dict(rest)
337	print_list.extend(print_process_dict(matched))
338
339	for line in print_list:
340	print(format % line)
341	## END functions for printing
342	####
343
344	class SaraNodes(object):
345	'''This class is used to communicate with the batch server'''
346
347	ticket = None
348
349	def _get_current_notes(self, nodes):
350	'''A function to retrieve the current message'''
351	if ARGS_VERBOSE:
352	print('class:SaraNodes func:_get_current_notes input:%s' % str(nodes), file=sys.stderr)
353
354	pbsq = PBSQuery.PBSQuery()
355	rdict = dict()
356
357	## We are only intereseted in the note
358	for node, properties in pbsq.getnodes(['note']).items():
359	if node in nodes and properties.has_key('note'):
360	rdict[node] = properties['note']
361	return rdict
362
363	def _get_curdate(self):
364	'''Returns the current time'''
365	if ARGS_VERBOSE:
366	print('class:SaraNodes func:_get_curdate', file=sys.stderr)
367	return time.strftime('%d-%m %H:%M', time.localtime())
368
369	def _get_uid(self, prev_uid=None):
370	'''Get the username'''
371	if ARGS_VERBOSE:
372	print('class:SaraNodes func:_get_uid input:%s' % prev_uid, file=sys.stderr)
373	cur_uid = getpass.getuser()
374	if prev_uid and cur_uid == 'root':
375	return prev_uid
376	return cur_uid
377
378	def _get_ticket(self, prev_ticket=None):
379	'''Check if we already have a ticket number'''
380	if ARGS_VERBOSE:
381	print('class:SaraNodes func:_get_ticket input:%s' % prev_ticket, file=sys.stderr)
382	cur_ticket = '%s' % self.ticket
383	if prev_ticket and cur_ticket == prev_ticket:
384	return prev_ticket
385	elif self.ticket:
386	return cur_ticket
387	elif self.ticket in ['c','clear','N',]:
388	return ''
389	elif prev_ticket:
390	return prev_ticket
391	return ''
392
393	def _generate_note(self, nodes=None, note=None, append=True):
394	'''Generates the node in a specific format'''
395	if ARGS_VERBOSE:
396	print('class:SaraNodes func:_generate_note input:%s,%s,%s' % (str(nodes), note, str(append)), file=sys.stderr)
397
398	## First step, is to get the current info of a host
399	cur_data = self._get_current_notes(nodes)
400	rdict = dict()
401
402	for node in nodes:
403	date_add = date_edit = username = ticket = nnote = None
404	if node in cur_data.keys():
405	date_add = cur_data[node][0]
406	date_edit = self._get_curdate()
407	username = self._get_uid(cur_data[node][2])
408	ticket = self._get_ticket(cur_data[node][3])
409	nnote = ",".join(cur_data[node][4:])
410	else:
411	date_add = date_edit = self._get_curdate()
412	username = self._get_uid()
413	ticket = self._get_ticket()
414	nnote = None
415
416	if nnote and append and note:
417	nnote = '%s, %s' % (nnote, note)
418	elif note:
419	nnote = note
420
421	rdict[node] = '%s,%s,%s,%s,%s' % (date_add, date_edit, username, ticket, nnote)
422	return rdict
423
424	def do_offline(self, nodes, note):
425	'''Change the state of node(s) to offline with a specific note'''
426
427	if ARGS_VERBOSE:
428	print('class:SaraNodes func:do_offline input:%s,%s' % (str(nodes), note), file=sys.stderr)
429	attributes = pbs.new_attropl(2)
430	attributes[0].name = pbs.ATTR_NODE_state
431	attributes[0].value = 'offline'
432	attributes[0].op = pbs.SET
433	attributes[1].name = pbs.ATTR_NODE_note
434	attributes[1].op = pbs.SET
435
436	batch_list = list()
437
438	## again a loop, now create the attrib dict list
439	for node, note in self._generate_note(nodes, note).items():
440	attributes[1].value = note
441	batch_list.append(tuple([node, attributes]))
442
443	self._process(batch_list)
444
445	def do_clear(self, nodes):
446	'''Clear the state on a node(s) to down, also clear the note'''
447
448	if ARGS_VERBOSE:
449	print('class:SaraNodes func:do_clear input:%s' % str(nodes), file=sys.stderr)
450	attributes = pbs.new_attropl(2)
451	attributes[0].name = pbs.ATTR_NODE_state
452	attributes[0].value = 'down'
453	attributes[0].op = pbs.SET
454	attributes[1].name = pbs.ATTR_NODE_note
455	attributes[1].op = pbs.SET
456	attributes[1].value = ''
457
458	batch_list = list()
459
460	## again a loop, now create the attrib dict list
461	for node in nodes:
462	batch_list.append(tuple([node, attributes]))
463
464	self._process(batch_list)
465
466	def do_modify(self, nodes, note):
467	'''Modify the note on a node, override the previous note'''
468
469	if ARGS_VERBOSE:
470	print('class:SaraNodes func:do_modify input:%s,%s' % (str(nodes), note), file=sys.stderr)
471	attributes = pbs.new_attropl(1)
472	attributes[0].name = pbs.ATTR_NODE_note
473	attributes[0].op = pbs.SET
474
475	batch_list = list()
476
477	## again a loop, now create the attrib dict list
478	for node, note in self._generate_note(nodes, note, append=False).items():
479	attributes[0].value = note
480	batch_list.append(tuple([node, attributes]))
481
482	self._process(batch_list)
483
484	def do_clear_note(self, nodes):
485	'''Clear the note on the node(s)'''
486
487	if ARGS_VERBOSE:
488	print('class:SaraNodes func:do_clear_note input:%s' % str(nodes), file=sys.stderr)
489	attributes = pbs.new_attropl(1)
490	attributes[0].name = pbs.ATTR_NODE_note
491	attributes[0].op = pbs.SET
492	attributes[0].value = ''
493
494	batch_list = list()
495
496	## again a loop, now create the attrib dict list
497	for node in nodes:
498	batch_list.append(tuple([node, attributes]))
499
500	self._process(batch_list)
501
502	def _process(self, batch_list):
503	'''This function execute the change to the batch server'''
504
505	if ARGS_VERBOSE:
506	print('class:SaraNodes func:_process input:%s' % str(batch_list), file=sys.stderr)
507
508	## Always get the pbs_server name, even in dry-run mode
509	pbs_server = pbs.pbs_default()
510	if not pbs_server:
511	print('Could not locate a pbs server', file=sys.stderr)
512	sys.exit(1)
513
514	if ARGS_VERBOSE:
515	print('class:SaraNodes func:_process pbs_server:%s' % pbs_server, file=sys.stderr)
516
517	## If dry-run is not specified create a connection
518	if not ARGS_DRYRUN:
519	pbs_connection = pbs.pbs_connect(pbs_server)
520
521	## Execute the changes
522	for node in batch_list:
523	if not ARGS_DRYRUN:
524	pbs_connection = pbs.pbs_connect(pbs_server)
525	rcode = pbs.pbs_manager(pbs_connection, pbs.MGR_CMD_SET, pbs.MGR_OBJ_NODE, node[0], node[1], 'NULL')
526	if rcode > 0:
527	errno, text = pbs.error()
528	print('PBS error for node \'%s\': %s (%s)' % (node[0], text, errno), file=sys.stderr)
529	else:
530	print("pbs.pbs_manager(pbs_connection, pbs.MGR_CMD_SET, pbs.MGR_OBJ_NODE, %s, %s, 'NULL')" % (node[0], str(node[1])))
531
532	## Close the connection with the batch system
533	if not ARGS_DRYRUN:
534	pbs.pbs_disconnect(pbs_connection)
535
536	if __name__ == '__main__':
537	## The arguments of sara_nodes
538	parser = argparse.ArgumentParser(
539	description=HELP_DESCRIPTION,
540	epilog=HELP_EPILOG,
541	)
542	parser.add_argument('nodes', metavar='NODES', nargs='*', type=str)
543	parser.add_argument('-j', '--jobs', action='store_true', help='use job id\'s instead of nodenames')
544	parser.add_argument('-v','--verbose', action='store_true', help='enables verbose mode')
545	parser.add_argument('-n','--dry-run', action='store_true', help='enables dry-run mode')
546	parser.add_argument('-q','--quiet', action='store_true', help='enables to supress all feedback')
547	parser.add_argument('-o','--offline', metavar='NOTE', help='change state to offline with message')
548	parser.add_argument('-m','--modify', metavar='NOTE', help='change the message of a node')
549	parser.add_argument('-c','--clear', action='store_true', help='change the state to down')
550	parser.add_argument('-N','--clear-note', action='store_true', help='clear the message of a node')
551	parser.add_argument('-f','--format', metavar='FORMAT', help='change the output of sara_nodes (see footer of --help)')
552	parser.add_argument('-t','--ticket', metavar='TICKET', type=int, help='add a ticket number to a node')
553	parser.add_argument('--version', action='version', version=pbs.version)
554
555	## Parse the arguments
556	args = parser.parse_args()
557
558	## The options quiet, verbose and dry-run are processed first
559	if args.quiet:
560	ARGS_QUIET = True
561	if args.verbose:
562	ARGS_VERBOSE = True
563	if args.dry_run:
564	ARGS_DRYRUN = ARGS_VERBOSE = True
565
566	if ARGS_VERBOSE:
567	print('func:__main__ checking type of operation', file=sys.stderr)
568
569	if args.nodes and args.jobs:
570	args.nodes = get_nodes(args.nodes)
571	elif args.nodes:
572	args.nodes = parse_args(args.nodes)
573
574	## If offline, modify, clear, clear_note or ticket then initiate the SaraNodes class
575	if args.offline or args.modify or args.clear or args.clear_note or args.ticket:
576	if not args.nodes:
577	print('You did not specify any nodes, see --help', file=sys.stderr)
578	sys.exit(1)
579
580	sn = SaraNodes()
581	if args.ticket:
582	sn.ticket = args.ticket
583
584	if args.offline:
585	if ARGS_VERBOSE:
586	print('func:__main__ call sn.do_offline', file=sys.stderr)
587	sn.do_offline(args.nodes, args.offline)
588	elif args.modify:
589	if ARGS_VERBOSE:
590	print('func:__main__ call sn.do_modify', file=sys.stderr)
591	sn.do_modify(args.nodes, args.modify)
592	elif args.clear:
593	if ARGS_VERBOSE:
594	print('func:__main__ call sn.do_clear', file=sys.stderr)
595	sn.do_clear(args.nodes)
596	elif args.clear_note:
597	if ARGS_VERBOSE:
598	print('func:__main__ call sn.do_clear_note', file=sys.stderr)
599	sn.do_clear_note(args.nodes)
600	elif args.ticket:
601	if ARGS_VERBOSE:
602	print('func:__main__ call sn.do_modify')
603	sn.do_offline(args.nodes, '')
604	else:
605	if ARGS_DRYRUN:
606	print('Dry-run is not available when we use PBSQuery', file=sys.stderr)
607
608	if args.format:
609	if ARGS_VERBOSE:
610	print('func:__main__ call print_overview_format', file=sys.stderr)
611	print_overview_format(args.nodes, args.format)
612	else:
613	if ARGS_VERBOSE:
614	print('func:__main__ call print_overview_normal', file=sys.stderr)
615	print_overview_normal(args.nodes)
616
617	if ARGS_VERBOSE:
618	print('func:__main__ exit', file=sys.stderr)

Note: See TracBrowser for help on using the repository browser.

Download in other formats:

Original Format