source: trunk/examples/sara_nodes.py.in @ 368

Last change on this file since 368 was 368, checked in by bas, 5 years ago

sara_nodes also displays NHC offline notes,

  • Property svn:executable set to *
File size: 24.1 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
10from __future__ import print_function
11
12HELP_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.'''
13HELP_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
16SPLIT_SORT = r'r\d+n\d+'
17## This RE pattern is used for the hostrange
18HOSTRANGE = r'\[([0-9az\-,]+)\]'
19## Which states are allowed to show in the print_overview
20ALLOWED_STATES = set(['down', 'offline', 'unknown'])
21
22import pbs
23import PBSQuery
24import re
25import sys
26import time
27import getpass
28import types
29
30## Use the cli arguments to change the values
31ARGS_QUIET      = False
32ARGS_VERBOSE    = False
33ARGS_DRYRUN     = False
34
35## Import argparse here, as I need the print function
36try:
37    import argparse
38except ImportError:
39    print('Cannot find argparse module', file=sys.stderr)
40    sys.exit(1)
41
42####
43## BEGIN functions for hostrange parsing
44def 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
61def 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
112def 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
122def 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
142def 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
162def _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
166def 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
202def 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        elif len(note) == 1:
245            ## There is a note but not set by admin
246            add_dict['date'] = add_dict['date_add'] = add_dict['date_edit'] = add_dict['username'] = add_dict['ticket'] = ''
247            add_dict['note'] = ",".join(note)
248        else:
249            ## If there is no note, just set the variables with a empty string
250            add_dict['date'] = add_dict['date_add'] = add_dict['date_edit'] = add_dict['username'] = add_dict['ticket'] = add_dict['note'] = ''
251
252        lineprint.append(add_dict)
253
254    if ARGS_VERBOSE:
255        print('func:print_process_dict returning values', file=sys.stderr)
256    return lineprint
257
258def print_create_list(values):
259    tmp_list = list()
260    for pair in values:
261        tmp_list.append('%-*s' % tuple(pair))
262    return tmp_list
263
264def print_overview_normal(hosts=None):
265    '''Print the default overview'''
266    if ARGS_VERBOSE:
267        print('func:print_overview_normal input:%s' % str(hosts), file=sys.stderr)
268
269    ## Determine some default values for the column width
270    w_nodename = 8
271    w_state = 5
272    w_date = w_username = w_ticket = w_note = w_date_add = w_date_edit = 0
273
274    ## Get the data, make it one list, the rest first then the matched
275    matched, rest = print_get_nodes(hosts)
276    print_list = print_process_dict(rest)
277    print_list.extend(print_process_dict(matched))
278
279    ## Detect the max width for the columns
280    for line in print_list:
281        if line['nodename'] and len(line['nodename']) > w_nodename:
282            w_nodename = len(line['nodename'])
283        if line['state'] and len(line['state']) > w_state:
284            w_state = len(line['state'])
285        if line['date'] and len(line['date']) > w_date:
286            w_date = len(line['date'])
287        if line['date_add'] and len(line['date_add']) > w_date_add:
288            w_date_add = len(line['date_add'])
289        if line['date_edit'] and len(line['date_edit']) > w_date_edit:
290            w_date_edit = len(line['date_edit'])
291        if line['username'] and len(line['username']) > w_username:
292            w_username = len(line['username'])
293        if line['ticket'] and len(line['ticket']) > w_ticket:
294            w_ticket = len(line['ticket'].strip())
295        if line['note'] and len(line['note']) > w_note:
296            w_note = len(line['note'])
297
298    ## The length of the full note
299    w_notefull  = w_date + w_username + w_ticket + w_note
300
301    if not ARGS_QUIET:
302        show_fields = [
303            [w_nodename, 'Nodename'],
304            [w_state, 'State'],
305        ]
306        if w_date > 0:
307            show_fields.append([w_date_add,'Added'])
308            show_fields.append([w_date_edit,'Modified'])
309            show_fields.append([w_username,'User'])
310            if w_ticket > 0:
311                if w_ticket < 6:
312                    w_ticket = 6
313                show_fields.append([w_ticket,'Ticket'])
314        if w_date > 0 or w_note > 0:
315            show_fields.append([w_note,'Note'])
316
317        print(' %s' % ' | '.join(print_create_list(show_fields)))
318        print('+'.join([ '-' * (show_field[0]+2) for show_field in show_fields ]))
319
320    ## Show the information to the user
321    for line in print_list:
322        show_line_fields = [
323            [w_nodename, line['nodename']],
324            [w_state, line['state']],
325        ]
326        if w_date > 0:
327            show_line_fields.append([w_date_add,line['date_add']])
328            show_line_fields.append([w_date_edit,line['date_edit']])
329            show_line_fields.append([w_username,line['username']])
330            if line['ticket'].strip():
331                show_line_fields.append([w_ticket,'#' + line['ticket']])
332            elif w_ticket > 0:
333                show_line_fields.append([w_ticket,''])
334        if w_date > 0 or w_note > 0:
335            show_line_fields.append([w_note,line['note']])
336
337        print(' %s' % ' | '.join(print_create_list(show_line_fields)))
338
339def print_overview_format(hosts=None, format=None):
340    '''Print the information in a certain format, when you want to use it in a
341    different program'''
342
343    matched, rest = print_get_nodes(hosts)
344    print_list = print_process_dict(rest)
345    print_list.extend(print_process_dict(matched))
346
347    for line in print_list:
348        print(format % line)
349## END functions for printing
350####
351
352class SaraNodes(object):
353    '''This class is used to communicate with the batch server'''
354
355    ticket      = None
356
357    def _get_current_notes(self, nodes):
358        '''A function to retrieve the current message'''
359        if ARGS_VERBOSE:
360            print('class:SaraNodes func:_get_current_notes input:%s' % str(nodes), file=sys.stderr)
361
362        pbsq = PBSQuery.PBSQuery()
363        rdict = dict()
364
365        ## We are only intereseted in the note
366        for node, properties in pbsq.getnodes(['note']).items():
367            if node in nodes and properties.has_key('note'):
368                rdict[node] = properties['note']
369        return rdict
370
371    def _get_curdate(self):
372        '''Returns the current time'''
373        if ARGS_VERBOSE:
374            print('class:SaraNodes func:_get_curdate', file=sys.stderr)
375        return time.strftime('%d-%m %H:%M', time.localtime())
376
377    def _get_uid(self, prev_uid=None):
378        '''Get the username'''
379        if ARGS_VERBOSE:
380            print('class:SaraNodes func:_get_uid input:%s' % prev_uid, file=sys.stderr)
381        cur_uid = getpass.getuser()
382        if prev_uid and cur_uid == 'root':
383            return prev_uid
384        return cur_uid
385
386    def _get_ticket(self, prev_ticket=None):       
387        '''Check if we already have a ticket number'''
388        if ARGS_VERBOSE:
389            print('class:SaraNodes func:_get_ticket input:%s' % prev_ticket, file=sys.stderr)
390        cur_ticket = '%s' % self.ticket
391        if prev_ticket and cur_ticket == prev_ticket:
392            return prev_ticket
393        elif self.ticket:
394            return cur_ticket
395        elif self.ticket in ['c','clear','N',]:
396            return ''
397        elif prev_ticket:
398            return prev_ticket
399        return ''
400
401    def _generate_note(self, nodes=None, note=None, append=True):
402        '''Generates the node in a specific format'''
403        if ARGS_VERBOSE:
404            print('class:SaraNodes func:_generate_note input:%s,%s,%s' % (str(nodes), note, str(append)), file=sys.stderr)
405
406        ## First step, is to get the current info of a host
407        cur_data = self._get_current_notes(nodes)
408        rdict = dict()
409
410        for node in nodes:
411            date_add = date_edit = username = ticket = nnote = None
412            if node in cur_data.keys():
413                date_add    = cur_data[node][0]
414                date_edit   = self._get_curdate()
415                username    = self._get_uid(cur_data[node][2])
416                ticket      = self._get_ticket(cur_data[node][3])
417                nnote       = ",".join(cur_data[node][4:])
418            else:
419                date_add = date_edit = self._get_curdate()
420                username = self._get_uid()
421                ticket   = self._get_ticket()
422                nnote    = None
423
424            if nnote and append and note:
425                nnote = '%s, %s' % (nnote, note)
426            elif note:
427                nnote = note
428
429            rdict[node] = '%s,%s,%s,%s,%s' % (date_add, date_edit, username, ticket, nnote)
430        return rdict
431
432    def do_offline(self, nodes, note):
433        '''Change the state of node(s) to offline with a specific note'''
434
435        if ARGS_VERBOSE:
436            print('class:SaraNodes func:do_offline input:%s,%s' % (str(nodes), note), file=sys.stderr)
437        attributes          = pbs.new_attropl(2)
438        attributes[0].name  = pbs.ATTR_NODE_state
439        attributes[0].value = 'offline'
440        attributes[0].op    = pbs.SET
441        attributes[1].name  = pbs.ATTR_NODE_note
442        attributes[1].op    = pbs.SET
443
444        batch_list = list()
445
446        ## again a loop, now create the attrib dict list
447        for node, note in self._generate_note(nodes, note).items():
448            attributes[1].value = note
449            batch_list.append(tuple([node, attributes]))
450
451        self._process(batch_list)
452
453    def do_clear(self, nodes):
454        '''Clear the state on a node(s) to down, also clear the note'''
455
456        if ARGS_VERBOSE:
457            print('class:SaraNodes func:do_clear input:%s' % str(nodes), file=sys.stderr)
458        attributes          = pbs.new_attropl(2)
459        attributes[0].name  = pbs.ATTR_NODE_state
460        attributes[0].value = 'down'
461        attributes[0].op    = pbs.SET
462        attributes[1].name  = pbs.ATTR_NODE_note
463        attributes[1].op    = pbs.SET
464        attributes[1].value = ''
465
466        batch_list = list()
467
468        ## again a loop, now create the attrib dict list
469        for node in nodes:
470            batch_list.append(tuple([node, attributes]))
471
472        self._process(batch_list)
473 
474    def do_modify(self, nodes, note):
475        '''Modify the note on a node, override the previous note'''
476
477        if ARGS_VERBOSE:
478            print('class:SaraNodes func:do_modify input:%s,%s' % (str(nodes), note), file=sys.stderr)
479        attributes          = pbs.new_attropl(1)
480        attributes[0].name  = pbs.ATTR_NODE_note
481        attributes[0].op    = pbs.SET
482
483        batch_list = list()
484
485        ## again a loop, now create the attrib dict list
486        for node, note in self._generate_note(nodes, note, append=False).items():
487            attributes[0].value = note
488            batch_list.append(tuple([node, attributes]))
489
490        self._process(batch_list)
491
492    def do_clear_note(self, nodes):
493        '''Clear the  note on the node(s)'''
494
495        if ARGS_VERBOSE:
496            print('class:SaraNodes func:do_clear_note input:%s' % str(nodes), file=sys.stderr)
497        attributes          = pbs.new_attropl(1)
498        attributes[0].name  = pbs.ATTR_NODE_note
499        attributes[0].op    = pbs.SET
500        attributes[0].value = ''
501
502        batch_list = list()
503
504        ## again a loop, now create the attrib dict list
505        for node in nodes:
506            batch_list.append(tuple([node, attributes]))
507
508        self._process(batch_list)
509
510    def _process(self, batch_list):
511        '''This function execute the change to the batch server'''
512
513        if ARGS_VERBOSE:
514            print('class:SaraNodes func:_process input:%s' % str(batch_list), file=sys.stderr)
515
516        ## Always get the pbs_server name, even in dry-run mode
517        pbs_server = pbs.pbs_default()
518        if not pbs_server:
519            print('Could not locate a pbs server', file=sys.stderr)
520            sys.exit(1)
521
522        if ARGS_VERBOSE:
523            print('class:SaraNodes func:_process pbs_server:%s' % pbs_server, file=sys.stderr)
524
525        ## If dry-run is not specified create a connection
526        if not ARGS_DRYRUN:
527            pbs_connection = pbs.pbs_connect(pbs_server)
528
529        ## Execute the changes
530        for node in batch_list:
531            if not ARGS_DRYRUN:
532                pbs_connection = pbs.pbs_connect(pbs_server)
533                rcode = pbs.pbs_manager(pbs_connection, pbs.MGR_CMD_SET, pbs.MGR_OBJ_NODE, node[0], node[1], 'NULL')
534                if rcode > 0:
535                    errno, text = pbs.error()
536                    print('PBS error for node \'%s\': %s (%s)' % (node[0], text, errno), file=sys.stderr)
537            else:
538                print("pbs.pbs_manager(pbs_connection, pbs.MGR_CMD_SET, pbs.MGR_OBJ_NODE, %s, %s, 'NULL')" % (node[0], str(node[1])))
539
540        ## Close the connection with the batch system
541        if not ARGS_DRYRUN:
542            pbs.pbs_disconnect(pbs_connection)
543
544if __name__ == '__main__':
545    ## The arguments of sara_nodes
546    parser = argparse.ArgumentParser(
547        description=HELP_DESCRIPTION,
548        epilog=HELP_EPILOG,
549    )
550    parser.add_argument('nodes', metavar='NODES', nargs='*', type=str)
551    parser.add_argument('-j', '--jobs', action='store_true', help='use job id\'s instead of nodenames')
552    parser.add_argument('-v','--verbose', action='store_true', help='enables verbose mode')
553    parser.add_argument('-n','--dry-run', action='store_true', help='enables dry-run mode')
554    parser.add_argument('-q','--quiet', action='store_true', help='enables to supress all feedback')
555    parser.add_argument('-o','--offline', metavar='NOTE', help='change state to offline with message')
556    parser.add_argument('-m','--modify', metavar='NOTE', help='change the message of a node')
557    parser.add_argument('-c','--clear', action='store_true', help='change the state to down')
558    parser.add_argument('-N','--clear-note', action='store_true', help='clear the message of a node')
559    parser.add_argument('-f','--format', metavar='FORMAT', help='change the output of sara_nodes (see footer of --help)')
560    parser.add_argument('-t','--ticket', metavar='TICKET', type=int, help='add a ticket number to a node')
561    parser.add_argument('--version', action='version', version=pbs.version)
562
563    ## Parse the arguments
564    args = parser.parse_args()
565
566    ## The options quiet, verbose and dry-run are processed first
567    if args.quiet:
568        ARGS_QUIET = True
569    if args.verbose:
570        ARGS_VERBOSE = True
571    if args.dry_run:
572        ARGS_DRYRUN = ARGS_VERBOSE = True
573
574    if ARGS_VERBOSE:
575        print('func:__main__ checking type of operation', file=sys.stderr)
576
577    if args.nodes and args.jobs:
578        args.nodes = get_nodes(args.nodes)
579    elif args.nodes:
580        args.nodes = parse_args(args.nodes)
581
582    ## If offline, modify, clear, clear_note or ticket then initiate the SaraNodes class
583    if args.offline or args.modify or args.clear or args.clear_note or args.ticket:
584        if not args.nodes:
585            print('You did not specify any nodes, see --help', file=sys.stderr)
586            sys.exit(1)
587
588        sn = SaraNodes()
589        if args.ticket:
590            sn.ticket = args.ticket
591
592        if args.offline:
593            if ARGS_VERBOSE:
594                print('func:__main__ call sn.do_offline', file=sys.stderr)
595            sn.do_offline(args.nodes, args.offline)
596        elif args.modify:
597            if ARGS_VERBOSE:
598                print('func:__main__ call sn.do_modify', file=sys.stderr)
599            sn.do_modify(args.nodes, args.modify)
600        elif args.clear:
601            if ARGS_VERBOSE:
602                print('func:__main__ call sn.do_clear', file=sys.stderr)
603            sn.do_clear(args.nodes)
604        elif args.clear_note:
605            if ARGS_VERBOSE:
606                print('func:__main__ call sn.do_clear_note', file=sys.stderr)
607            sn.do_clear_note(args.nodes)
608        elif args.ticket:
609            if ARGS_VERBOSE:
610                print('func:__main__ call sn.do_modify')
611            sn.do_offline(args.nodes, '')
612    else:
613        if ARGS_DRYRUN:
614            print('Dry-run is not available when we use PBSQuery', file=sys.stderr)
615
616        if args.format:
617            if ARGS_VERBOSE:
618                print('func:__main__ call print_overview_format', file=sys.stderr)
619            print_overview_format(args.nodes, args.format)
620        else:
621            if ARGS_VERBOSE:
622                print('func:__main__ call print_overview_normal', file=sys.stderr)
623            print_overview_normal(args.nodes)
624   
625    if ARGS_VERBOSE:
626        print('func:__main__ exit', file=sys.stderr)
Note: See TracBrowser for help on using the repository browser.