source: trunk/examples/sara_nodes.py @ 278

Last change on this file since 278 was 278, checked in by bas, 12 years ago

changed tabs to spaces

  • Property svn:executable set to *
File size: 22.8 KB
Line 
1#!/usr/bin/env python
2#
3# Author: Dennis Stam
4# Date  : 14-05-2009
5# Desc. : This program/module allows you to change
6#         the state of a node to offline or down
7#
8# SVN Info:
9#       $Id: sara_nodes 4552 2010-04-29 12:15:42Z dennis $
10#       $URL: https://subtrac.sara.nl/hpcv/svn/beowulf/trunk/torque/utils/sara_nodes $
11#
12
13
14# imports of the pbs_python module
15import PBSQuery
16import PBSAdvancedParser
17import pbs
18
19# python core modules
20from optparse import make_option
21import types
22import sys
23import re
24import os
25import time
26import string
27
28# GINA also uses python2.3
29if sys.version_info < ( 2,4 ):
30    from set import Set as set
31
32__author__ = 'Dennis Stam'
33
34# Specify here your BATCH name pattern, this is
35# used for sorting when you are using basenames
36RE_PATTERN = '(r\d+n\d+)'
37
38# Cfegine uses : for lists, so a override is needed
39TIME_SEPARATOR = ':'
40
41class sara_nodesException( Exception ):
42
43        def __init__(self, msg='' ):
44                self.msg = msg
45                Exception.__init__( self, msg )
46
47        def __repr__(self):
48                return self.msg
49
50        def islist(self):
51                if type(self.msg) is types.ListType:
52                        return True
53
54                return False
55
56        def getlist(self):
57                return self.msg
58
59class sara_nodesCli:
60        '''
61        This class is the Command Line Interface from here we call the sara_nodes class / module
62        '''
63        option_list = [
64                make_option( '-v', '--verbose', dest='verbose', action='store_true', help='enables verbose mode' ),
65                make_option( '-n', '--dry-run', dest='dryrun', action='store_true', help='enables dry-run mode' ),
66                make_option( '-q', '--quiet', dest='quiet', action='store_true', help='enable this function to supress all feedback'),
67                make_option( '-o', '--offline', dest='offline', help='change state to offline', metavar='NOTE' ),
68                make_option( '-c', '--clear', dest='clear', action='store_true', help='change state to down' ),
69                make_option( '-N', '--clearnote', dest='note', action='store_true', help='clear note of node' ),
70                make_option( '-m', '--modify', dest='modify', help='use this option the modify the note, it will replace the current one!', metavar='NOTE' ),
71                make_option( '-f', '--format', dest='format', nargs=2, help='specify how to display the information, also sets quiet to True' ),
72                make_option( '-t', '--ticket', dest='ticket', help='add/change/remove ticket number, removing use -t c' ),
73        ]
74
75        def __init__(self):
76                '''
77  sara_nodes [ <options> <nodenames> | [nodenames] ]
78
79  -f/--format needs 2 arguments, first the pattern secondly the variables
80        the pattern you specify must be the string format pattern of Python
81        fields: node, state, date_add, date_edit, user, ticket, remark'''
82
83                self.obj_sara_nodes = sara_nodes()
84                self.parser = PBSAdvancedParser.AdvancedParser( 
85                option_list=self.option_list, 
86                version=pbs.version,
87                usage=self.__init__.__doc__
88                )
89
90                self.parser.set_default('verbose', False)
91                self.parser.set_default('dryrun', False)
92                self.parser.set_default('offline', False)
93                self.parser.set_default('clear', False)
94                self.parser.set_default('note', False)
95                self.parser.set_default('quiet', False)
96                self.parser.set_default('ticket', None)
97                self.parser.set_default('modify', False)
98                self.parser.set_default('format', 'default' )
99
100                options, args = self.parser.parse_args()
101
102                if options.format == 'default':
103                    options.format = (' %-10s | %-19s | %-11s, %-11s %-6s %-5s: %s','node,state,date_add,date_edit,user,ticket,remark' )
104                else:
105                    options.quiet = True
106
107                if not options.quiet:
108                        self.obj_sara_nodes.dryrun = options.dryrun
109
110                        if options.dryrun or options.verbose:
111                                self.obj_sara_nodes.verbose = True
112
113                if str( options.offline ).rstrip() == '' or str( options.modify ).rstrip() == '':
114                        sys.stderr.write( 'sara_nodes: error: option requires an argument\n' )
115
116                        sys.exit( 1 )
117               
118                try:
119                        if options.offline and not options.clear:
120                                if args:
121                                        self.obj_sara_nodes.pbs_change_state_offline( args, options.offline, options.ticket )
122                                else:
123                                        raise sara_nodesException, 'No hostnames given'
124                        elif options.clear and not options.offline:
125                                if args:
126                                        self.obj_sara_nodes.pbs_change_state_down( args )
127                                else:
128                                        raise sara_nodesException, 'No hostnames given'
129                        elif options.note:
130                                if args:
131                                        self.obj_sara_nodes.pbs_change_note_clear( args )
132                                else:
133                                        raise sara_nodesException, 'No hostnames given'
134                        elif options.modify:
135                                if args:
136                                        self.obj_sara_nodes.pbs_change_note( args, options.modify, options.ticket )
137                                else:
138                                        raise sara_nodesException, 'No hostnames given'
139                        elif options.ticket:
140                                if args:
141                                        self.obj_sara_nodes.pbs_change_note_ticket( args, options.ticket )
142                                else:
143                                        raise sara_nodesException, 'No hostnames given'
144                        else:
145                                if not options.quiet:
146                                        print '\n Use option --help for help'
147                                self.print_list(args, options)
148
149
150                except sara_nodesException, msg:
151                        if msg.islist():
152                                for item in msg.getlist():
153                                        sys.stderr.write( 'sara_nodes: error: %s\n' % item )
154                        else:
155                                sys.stderr.write( 'sara_nodes: error: %s\n' % str( msg ) )
156
157                        sys.exit( 1 )
158
159        def return_note( self, pre_parts ):
160
161                if type(pre_parts) is types.ListType and len( pre_parts) >= 5:
162                        return { 
163                                'date_add': pre_parts[0].strip(), 
164                                'date_edit': pre_parts[1].strip(), 
165                                'user': pre_parts[2].strip(), 
166                                'ticket': pre_parts[3].strip(),
167                                'remark': ','.join( pre_parts[4:] )
168                        }
169                else:
170                        return { 
171                                'date_add': '', 
172                                'date_edit': '', 
173                                'user': '', 
174                                'ticket': '',
175                                'remark': str( pre_parts ) 
176                        }
177
178        def convert_format( self, format_options ):
179            pattern = r'\%([-|+]){0,1}([0-9]{0,2})([a-z]{1})'
180            parts = re.findall( pattern, format_options[0], re.VERBOSE )
181            line = re.sub( pattern, '%s', format_options[0], re.VERBOSE )
182
183            rlist = list()
184            counter = 0
185            for variable in format_options[1].split( ',' ):
186                rlist.append( '%s(%s)%s%s%s' % (
187                    '%',
188                    variable, 
189                    parts[ counter ][0], 
190                    parts[ counter ][1],
191                    parts[ counter ][2],
192                ) )
193                counter += 1
194            return line % tuple( rlist )
195
196        def print_table( self, node_list, pbs_nodes, format_options ):
197            '''
198            This method prints the rows of a table
199            '''
200
201            try:
202                line_format = self.convert_format( format_options )
203                for node in node_list:
204                    note = '' 
205                    if pbs_nodes[ node ].has_key('note'):
206                        note = pbs_nodes[ node ]['note']
207                               
208                    if self.allowed_state( pbs_nodes[ node ]['state'] ) or note:
209                        node_note = self.return_note( note )
210
211                        fields = { 
212                            'node': node,
213                            'state': ', '.join( pbs_nodes[ node ]['state'] ),
214                            'date_add': node_note['date_add'],
215                            'date_edit': node_note['date_edit'],
216                            'user': node_note['user'],
217                            'ticket': node_note['ticket'],
218                            'remark': node_note['remark'],
219                        }
220
221                       
222                        print line_format % fields
223
224            except KeyError, e:
225                raise sara_nodesException, 'Given host does not exist'
226
227        def print_list(self, args, options):
228                '''
229                A method that is used for collecting all nodes with the state down, offline or unknown
230                '''
231
232               
233                p = PBSQuery.PBSQuery()
234                if pbs.version_info >= ( 4,0,0 ):
235                        if self.obj_sara_nodes.verbose:
236                                print "Enabling new_data_structure for PBSQuery"
237                        p.new_data_structure()
238
239                try:
240                    nodes = p.getnodes( ['state', 'note'] )
241                except PBSQuery.PBSError, detail:
242                    print "PBSQuery error: %s" %detail
243                    sys.exit(1)
244
245                header = ' %-10s | %-19s | %s' % ( 'Nodename', 'State', 'Note' )
246                if not options.quiet:
247                        print '\n%s\n%s' % ( header, ( '-' * 80 ) )
248
249                if args:
250                        args = self.sort_nodes( args )
251                        self.print_table( args[0], nodes, options.format )
252                else:
253                        sorted_nodes, sorted_other = self.sort_nodes( nodes )
254
255                        self.print_table( sorted_other, nodes, options.format )
256                        self.print_table( sorted_nodes, nodes, options.format )
257
258        def real_sort( self, inlist ):
259                '''
260                Use this method instead of the x.sort(), because with x.sort()
261                numeric values in a string are not correctly sorted!
262                '''
263                indices = map(self._generate_index, inlist )
264                decorated = zip( indices, inlist )
265                decorated.sort()
266
267                return [ item for index, item in decorated ]
268
269        def _generate_index( self, str ):
270                '''
271                Spliting a string in aplha and numeric elements
272                '''
273
274                index = []
275
276                def _append( fragment, alist=index ):
277                        if fragment.isdigit():
278                                fragment = int( fragment )
279                        alist.append( fragment )
280
281                prev_isdigit = str[0].isdigit()
282                current_fragment = ''
283
284                for char in str:
285                        curr_isdigit = char.isdigit()
286
287                        if curr_isdigit == prev_isdigit:
288                                current_fragment += char
289                        else:
290                                _append( current_fragment )
291                                current_fragment = char
292                                prev_isdigit = curr_isdigit
293
294                _append( current_fragment )
295
296                return tuple( index )
297
298        def sort_nodes(self, nodes):
299                '''
300                Sorts the nodes list and returns two lists
301                the first the nodes secondly the other machines
302
303                When RE_PATTERN is not supplied then all names
304                will be sorted the same way.
305                '''
306
307                if not globals().has_key('RE_PATTERN'):
308                        global RE_PATTERN
309                        RE_PATTERN = ''
310
311                pattern = re.compile( RE_PATTERN, re.VERBOSE )
312
313                tmplist = list()
314                tmplist_other = list()
315
316                for node in nodes:
317                        match = pattern.findall( node )
318
319                        if match and len( match ) == 1:
320                                tmplist.append( node )
321                        else:
322                                tmplist_other.append( node )
323
324                tmplist = self.real_sort( tmplist )
325                tmplist_other.sort()
326
327                return tmplist, tmplist_other
328
329        def allowed_state(self, state):
330                '''
331                This method checks is a node complies with the following states:
332                down, offline and or unknown
333                '''
334                allowed_list = set( ['down', 'offline', 'unknown'] )
335
336                return bool( allowed_list.intersection( set( state ) ) )
337
338class sara_nodes:
339
340        def __init__(self):
341                '''
342                Just initialize two optional variables
343                '''
344                self.dryrun = False
345                self.verbose = False
346
347        def note_check_ticket( self, ticketno, oldticket ):
348               
349                if ticketno:
350                        try:
351                            return '#%d' % int( ticketno )
352                        except ValueError:
353                                if ticketno == 'c':
354                                        return ''
355
356                return oldticket
357
358        def note_return_username( self, old_username ):
359                try:
360                        username = os.getlogin()
361
362                        if username != 'root':
363                                return username
364                        else:
365                                return old_username
366                except OSError, err:
367                        return 'root'
368
369        def note_create( self, new_note, mode = 'a', old_note = None ):
370                if mode == 'w':
371                        return new_note
372                else:
373                        if old_note and old_note.find( new_note ) < 0:
374                                return '%s, %s' % ( old_note, new_note )
375                        else:
376                                return new_note
377
378        def create_date( self ):
379                if not globals().has_key('TIME_SEPARATOR'):
380                        global TIME_SEPARATOR
381                        TIME_SEPARATOR = ':'
382
383                curtime = time.localtime()
384                day = time.strftime( '%d-%m', curtime )
385                hour = time.strftime( '%H', curtime )
386                minutes = time.strftime( '%M', curtime )
387                return '%s %s%s%s' % ( day, hour, TIME_SEPARATOR, minutes )
388
389        def note_init( self ):
390                current_date = self.create_date()
391                try:
392                        current_username = os.getlogin()
393                except OSError, err:
394                        current_username = 'root'
395
396                return [ current_date, current_date, current_username, '' ]
397
398        def note( self, node, note_attr ):
399                '''
400                This method combines all note methods and returns the new note
401                '''
402                p = PBSQuery.PBSQuery()
403                p.new_data_structure()
404                pbs_info = p.getnode( node )
405               
406                pre_parts = list()
407                old_note = None
408                new_note = None
409
410                if pbs_info.has_key( 'note' ):
411                        pbs_note = pbs_info[ 'note' ]
412                        if len( pbs_note ) > 4:
413                                pre_parts = pbs_note[:4]
414                                old_note = ', '.join( pbs_note[4:] )
415
416                                pre_parts[1] = self.create_date()
417                                pre_parts[2] = self.note_return_username( pre_parts[2] )
418
419                else:
420                        pre_parts = self.note_init()
421
422                if note_attr.has_key( 'ticket' ):
423                        pre_parts[3] = self.note_check_ticket( note_attr['ticket'], pre_parts[3] )
424
425                if note_attr.has_key( 'note' ) and note_attr.has_key( 'mode' ):
426                        if note_attr[ 'note' ] and note_attr[ 'mode' ] in [ 'a','w' ]:
427                                if old_note:
428                                        new_note = self.note_create( note_attr[ 'note' ], note_attr[ 'mode' ], old_note )
429                                else:
430                                        new_note = self.note_create( note_attr[ 'note' ], note_attr[ 'mode' ] )
431                        else:
432                                new_note = old_note
433
434                return '%s,%s' % ( ','.join( pre_parts ), new_note )
435
436        def verbose_print( self, msg ):
437                if self.verbose:
438                        print msg
439
440        def pbs_change_note_clear( self, nodes ):
441                attributes = pbs.new_attropl(1)
442                attributes[0].name = pbs.ATTR_NODE_note
443                attributes[0].value = ''
444                attributes[0].op = pbs.SET
445
446                self.verbose_print( '%*s: cleared' % ( 7, 'Note') )
447                self.pbs_batch( nodes, attributes )
448
449        def pbs_change_note_ticket( self, nodes, ticket ):
450                note_attributes = { 'note': None, 'ticket': ticket, 'mode': 'a' }
451                self.verbose_print( '%*s: %s' % ( 7, 'Ticket', ticket  ) )
452                self.pbs_batch( nodes, None, note_attributes)
453
454        def pbs_change_note( self, nodes, note, ticket=None ):
455                note_attributes = { 'note': note, 'ticket': ticket, 'mode': 'w' }
456
457                self.verbose_print( '%*s: %s' % ( 7, 'Note', note ) )
458                if ticket:
459                        self.verbose_print( '%*s: %s' % ( 7, 'Ticket', ticket ) )
460                self.pbs_batch( nodes, None, note_attributes)
461
462        def pbs_change_state_offline( self, nodes, note, ticket=None ):
463                attributes = pbs.new_attropl(1)
464                attributes[0].name = pbs.ATTR_NODE_state
465                attributes[0].value = 'offline'
466                attributes[0].op = pbs.SET
467
468                note_attributes = { 'note': note, 'ticket': ticket, 'mode': 'a' }
469
470                self.verbose_print( '%*s: offline' % ( 7, 'State') )
471                self.verbose_print( '%*s: %s' % ( 7, 'Note', note ) )
472                if ticket:
473                        self.verbose_print( '%*s: %s' % ( 7, 'Ticket', ticket ) )
474                self.pbs_batch( nodes, attributes, note_attributes )
475
476        def pbs_change_state_down( self, nodes ):
477                attributes = pbs.new_attropl(2)
478                attributes[0].name = pbs.ATTR_NODE_state
479                attributes[0].value = 'down'
480                attributes[0].op = pbs.SET
481
482                attributes[1].name = 'note'
483                attributes[1].value = ''
484
485                self.verbose_print( '%*s: down' % ( 7, 'State') )
486                self.verbose_print( '%*s: cleared' % ( 7, 'Note' ) )
487                self.pbs_batch( nodes, attributes )
488
489        def pbs_batch( self, nodes, attrs=None, note_attributes=None ):
490                nodeserror = list()
491                if not attrs and not note_attributes:
492                        raise sara_nodesException, 'attrs and note_attributes can not be empty together!'
493
494                if not self.dryrun:
495                        if note_attributes and len( note_attributes ) == 3:
496                                if attrs:
497                                        attributes = attrs + pbs.new_attropl(1)
498                                        attributes[1].name = pbs.ATTR_NODE_note
499                                        attributes[1].op = pbs.SET
500                                else:
501                                        attributes = pbs.new_attropl(1)
502                                        attributes[0].name = pbs.ATTR_NODE_note
503                                        attributes[0].op = pbs.SET
504                        else:
505                                attributes = attrs
506                        # Some hacking here because some limitation in the Torque 2.4 version
507                        # fetching note data first for all nodes!
508                        tmp_node_note = dict()
509
510                        for node in nodes:
511                                if note_attributes and len( note_attributes ) == 3:
512                                        tmp_node_note[ node ] = self.note( node, note_attributes )
513
514                        pbs_server = pbs.pbs_default()
515
516                        if not pbs_server:
517                                raise sara_nodesException, 'Default pbs server not found!'
518
519                        pbs_connection = pbs.pbs_connect( pbs_server )
520                        for node in nodes:
521                                if note_attributes and len( note_attributes ) == 3:
522                                        try:
523                                                if attrs:
524                                                        attributes[1].value = tmp_node_note[ node ]
525                                                else:
526                                                        attributes[0].value = tmp_node_note[ node ]
527                                        except KeyError:
528                                                pass
529                                rcode = pbs.pbs_manager( pbs_connection, pbs.MGR_CMD_SET, pbs.MGR_OBJ_NODE, node, attributes, 'NULL' )
530                                if rcode > 0:
531                                        errno, text = pbs.error()
532                                        nodeserror.append( '%s: %s (%s)' % ( node, text, errno ) )
533                else:
534                        p = PBSQuery.PBSQuery()
535                        pbsnodes = p.getnodes().keys()
536
537                        print '%*s:' % ( 7, 'Nodes' ),
538                        firstitem = True
539
540                        for node in nodes:
541
542                                if node in pbsnodes:
543                                        if firstitem:
544                                                print '%s' % node
545                                                firstitem = False
546                                        else:
547                                                print '%*s' % ( 17, node )
548                                else:
549                                        nodeserror.append( '%s: does not exist' % node )
550
551                if len( nodeserror ) > 0:
552                        raise sara_nodesException, nodeserror
553
554if __name__ == '__main__':
555        sara_nodesCli()
Note: See TracBrowser for help on using the repository browser.