source: trunk/examples/sara_nodes.py @ 295

Last change on this file since 295 was 295, checked in by dennis, 11 years ago

Bumped version number to 4.3.6, modified debian/control file to ensure a dependicy
for python-argparse when your are using Python 2.6 or lower.

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