Context Navigation

sara_nodes.py.in @ 367

Last change on this file since 367 was 367, checked in by martijk, 6 years ago
rendering fix
Property svn:executable set to ``*
File size: 23.8 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 line['ticket'].strip():
326	show_line_fields.append([w_ticket,'#' + line['ticket']])
327	elif w_ticket > 0:
328	show_line_fields.append([w_ticket,''])
329	show_line_fields.append([w_note,line['note']])
330
331	print(' %s' % ' \| '.join(print_create_list(show_line_fields)))
332
333	def print_overview_format(hosts=None, format=None):
334	'''Print the information in a certain format, when you want to use it in a
335	different program'''
336
337	matched, rest = print_get_nodes(hosts)
338	print_list = print_process_dict(rest)
339	print_list.extend(print_process_dict(matched))
340
341	for line in print_list:
342	print(format % line)
343	## END functions for printing
344	####
345
346	class SaraNodes(object):
347	'''This class is used to communicate with the batch server'''
348
349	ticket = None
350
351	def _get_current_notes(self, nodes):
352	'''A function to retrieve the current message'''
353	if ARGS_VERBOSE:
354	print('class:SaraNodes func:_get_current_notes input:%s' % str(nodes), file=sys.stderr)
355
356	pbsq = PBSQuery.PBSQuery()
357	rdict = dict()
358
359	## We are only intereseted in the note
360	for node, properties in pbsq.getnodes(['note']).items():
361	if node in nodes and properties.has_key('note'):
362	rdict[node] = properties['note']
363	return rdict
364
365	def _get_curdate(self):
366	'''Returns the current time'''
367	if ARGS_VERBOSE:
368	print('class:SaraNodes func:_get_curdate', file=sys.stderr)
369	return time.strftime('%d-%m %H:%M', time.localtime())
370
371	def _get_uid(self, prev_uid=None):
372	'''Get the username'''
373	if ARGS_VERBOSE:
374	print('class:SaraNodes func:_get_uid input:%s' % prev_uid, file=sys.stderr)
375	cur_uid = getpass.getuser()
376	if prev_uid and cur_uid == 'root':
377	return prev_uid
378	return cur_uid
379
380	def _get_ticket(self, prev_ticket=None):
381	'''Check if we already have a ticket number'''
382	if ARGS_VERBOSE:
383	print('class:SaraNodes func:_get_ticket input:%s' % prev_ticket, file=sys.stderr)
384	cur_ticket = '%s' % self.ticket
385	if prev_ticket and cur_ticket == prev_ticket:
386	return prev_ticket
387	elif self.ticket:
388	return cur_ticket
389	elif self.ticket in ['c','clear','N',]:
390	return ''
391	elif prev_ticket:
392	return prev_ticket
393	return ''
394
395	def _generate_note(self, nodes=None, note=None, append=True):
396	'''Generates the node in a specific format'''
397	if ARGS_VERBOSE:
398	print('class:SaraNodes func:_generate_note input:%s,%s,%s' % (str(nodes), note, str(append)), file=sys.stderr)
399
400	## First step, is to get the current info of a host
401	cur_data = self._get_current_notes(nodes)
402	rdict = dict()
403
404	for node in nodes:
405	date_add = date_edit = username = ticket = nnote = None
406	if node in cur_data.keys():
407	date_add = cur_data[node][0]
408	date_edit = self._get_curdate()
409	username = self._get_uid(cur_data[node][2])
410	ticket = self._get_ticket(cur_data[node][3])
411	nnote = ",".join(cur_data[node][4:])
412	else:
413	date_add = date_edit = self._get_curdate()
414	username = self._get_uid()
415	ticket = self._get_ticket()
416	nnote = None
417
418	if nnote and append and note:
419	nnote = '%s, %s' % (nnote, note)
420	elif note:
421	nnote = note
422
423	rdict[node] = '%s,%s,%s,%s,%s' % (date_add, date_edit, username, ticket, nnote)
424	return rdict
425
426	def do_offline(self, nodes, note):
427	'''Change the state of node(s) to offline with a specific note'''
428
429	if ARGS_VERBOSE:
430	print('class:SaraNodes func:do_offline input:%s,%s' % (str(nodes), note), file=sys.stderr)
431	attributes = pbs.new_attropl(2)
432	attributes[0].name = pbs.ATTR_NODE_state
433	attributes[0].value = 'offline'
434	attributes[0].op = pbs.SET
435	attributes[1].name = pbs.ATTR_NODE_note
436	attributes[1].op = pbs.SET
437
438	batch_list = list()
439
440	## again a loop, now create the attrib dict list
441	for node, note in self._generate_note(nodes, note).items():
442	attributes[1].value = note
443	batch_list.append(tuple([node, attributes]))
444
445	self._process(batch_list)
446
447	def do_clear(self, nodes):
448	'''Clear the state on a node(s) to down, also clear the note'''
449
450	if ARGS_VERBOSE:
451	print('class:SaraNodes func:do_clear input:%s' % str(nodes), file=sys.stderr)
452	attributes = pbs.new_attropl(2)
453	attributes[0].name = pbs.ATTR_NODE_state
454	attributes[0].value = 'down'
455	attributes[0].op = pbs.SET
456	attributes[1].name = pbs.ATTR_NODE_note
457	attributes[1].op = pbs.SET
458	attributes[1].value = ''
459
460	batch_list = list()
461
462	## again a loop, now create the attrib dict list
463	for node in nodes:
464	batch_list.append(tuple([node, attributes]))
465
466	self._process(batch_list)
467
468	def do_modify(self, nodes, note):
469	'''Modify the note on a node, override the previous note'''
470
471	if ARGS_VERBOSE:
472	print('class:SaraNodes func:do_modify input:%s,%s' % (str(nodes), note), file=sys.stderr)
473	attributes = pbs.new_attropl(1)
474	attributes[0].name = pbs.ATTR_NODE_note
475	attributes[0].op = pbs.SET
476
477	batch_list = list()
478
479	## again a loop, now create the attrib dict list
480	for node, note in self._generate_note(nodes, note, append=False).items():
481	attributes[0].value = note
482	batch_list.append(tuple([node, attributes]))
483
484	self._process(batch_list)
485
486	def do_clear_note(self, nodes):
487	'''Clear the note on the node(s)'''
488
489	if ARGS_VERBOSE:
490	print('class:SaraNodes func:do_clear_note input:%s' % str(nodes), file=sys.stderr)
491	attributes = pbs.new_attropl(1)
492	attributes[0].name = pbs.ATTR_NODE_note
493	attributes[0].op = pbs.SET
494	attributes[0].value = ''
495
496	batch_list = list()
497
498	## again a loop, now create the attrib dict list
499	for node in nodes:
500	batch_list.append(tuple([node, attributes]))
501
502	self._process(batch_list)
503
504	def _process(self, batch_list):
505	'''This function execute the change to the batch server'''
506
507	if ARGS_VERBOSE:
508	print('class:SaraNodes func:_process input:%s' % str(batch_list), file=sys.stderr)
509
510	## Always get the pbs_server name, even in dry-run mode
511	pbs_server = pbs.pbs_default()
512	if not pbs_server:
513	print('Could not locate a pbs server', file=sys.stderr)
514	sys.exit(1)
515
516	if ARGS_VERBOSE:
517	print('class:SaraNodes func:_process pbs_server:%s' % pbs_server, file=sys.stderr)
518
519	## If dry-run is not specified create a connection
520	if not ARGS_DRYRUN:
521	pbs_connection = pbs.pbs_connect(pbs_server)
522
523	## Execute the changes
524	for node in batch_list:
525	if not ARGS_DRYRUN:
526	pbs_connection = pbs.pbs_connect(pbs_server)
527	rcode = pbs.pbs_manager(pbs_connection, pbs.MGR_CMD_SET, pbs.MGR_OBJ_NODE, node[0], node[1], 'NULL')
528	if rcode > 0:
529	errno, text = pbs.error()
530	print('PBS error for node \'%s\': %s (%s)' % (node[0], text, errno), file=sys.stderr)
531	else:
532	print("pbs.pbs_manager(pbs_connection, pbs.MGR_CMD_SET, pbs.MGR_OBJ_NODE, %s, %s, 'NULL')" % (node[0], str(node[1])))
533
534	## Close the connection with the batch system
535	if not ARGS_DRYRUN:
536	pbs.pbs_disconnect(pbs_connection)
537
538	if __name__ == '__main__':
539	## The arguments of sara_nodes
540	parser = argparse.ArgumentParser(
541	description=HELP_DESCRIPTION,
542	epilog=HELP_EPILOG,
543	)
544	parser.add_argument('nodes', metavar='NODES', nargs='*', type=str)
545	parser.add_argument('-j', '--jobs', action='store_true', help='use job id\'s instead of nodenames')
546	parser.add_argument('-v','--verbose', action='store_true', help='enables verbose mode')
547	parser.add_argument('-n','--dry-run', action='store_true', help='enables dry-run mode')
548	parser.add_argument('-q','--quiet', action='store_true', help='enables to supress all feedback')
549	parser.add_argument('-o','--offline', metavar='NOTE', help='change state to offline with message')
550	parser.add_argument('-m','--modify', metavar='NOTE', help='change the message of a node')
551	parser.add_argument('-c','--clear', action='store_true', help='change the state to down')
552	parser.add_argument('-N','--clear-note', action='store_true', help='clear the message of a node')
553	parser.add_argument('-f','--format', metavar='FORMAT', help='change the output of sara_nodes (see footer of --help)')
554	parser.add_argument('-t','--ticket', metavar='TICKET', type=int, help='add a ticket number to a node')
555	parser.add_argument('--version', action='version', version=pbs.version)
556
557	## Parse the arguments
558	args = parser.parse_args()
559
560	## The options quiet, verbose and dry-run are processed first
561	if args.quiet:
562	ARGS_QUIET = True
563	if args.verbose:
564	ARGS_VERBOSE = True
565	if args.dry_run:
566	ARGS_DRYRUN = ARGS_VERBOSE = True
567
568	if ARGS_VERBOSE:
569	print('func:__main__ checking type of operation', file=sys.stderr)
570
571	if args.nodes and args.jobs:
572	args.nodes = get_nodes(args.nodes)
573	elif args.nodes:
574	args.nodes = parse_args(args.nodes)
575
576	## If offline, modify, clear, clear_note or ticket then initiate the SaraNodes class
577	if args.offline or args.modify or args.clear or args.clear_note or args.ticket:
578	if not args.nodes:
579	print('You did not specify any nodes, see --help', file=sys.stderr)
580	sys.exit(1)
581
582	sn = SaraNodes()
583	if args.ticket:
584	sn.ticket = args.ticket
585
586	if args.offline:
587	if ARGS_VERBOSE:
588	print('func:__main__ call sn.do_offline', file=sys.stderr)
589	sn.do_offline(args.nodes, args.offline)
590	elif args.modify:
591	if ARGS_VERBOSE:
592	print('func:__main__ call sn.do_modify', file=sys.stderr)
593	sn.do_modify(args.nodes, args.modify)
594	elif args.clear:
595	if ARGS_VERBOSE:
596	print('func:__main__ call sn.do_clear', file=sys.stderr)
597	sn.do_clear(args.nodes)
598	elif args.clear_note:
599	if ARGS_VERBOSE:
600	print('func:__main__ call sn.do_clear_note', file=sys.stderr)
601	sn.do_clear_note(args.nodes)
602	elif args.ticket:
603	if ARGS_VERBOSE:
604	print('func:__main__ call sn.do_modify')
605	sn.do_offline(args.nodes, '')
606	else:
607	if ARGS_DRYRUN:
608	print('Dry-run is not available when we use PBSQuery', file=sys.stderr)
609
610	if args.format:
611	if ARGS_VERBOSE:
612	print('func:__main__ call print_overview_format', file=sys.stderr)
613	print_overview_format(args.nodes, args.format)
614	else:
615	if ARGS_VERBOSE:
616	print('func:__main__ call print_overview_normal', file=sys.stderr)
617	print_overview_normal(args.nodes)
618
619	if ARGS_VERBOSE:
620	print('func:__main__ exit', file=sys.stderr)

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

Download in other formats:

Original Format