source: trunk/examples/sara_nodes.py @ 270

Last change on this file since 270 was 270, checked in by bas, 13 years ago

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