source: trunk/examples/sara_nodes.py @ 271

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

Changed debian dependencies. closes #30
and fixed some conflicts with SARA packages
  • Property svn:executable set to *
File size: 22.1 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                except sara_nodesException, msg:
150                        if msg.islist():
151                                for item in msg.getlist():
152                                        sys.stderr.write( 'sara_nodes: error: %s\n' % item )
153                        else:
154                                sys.stderr.write( 'sara_nodes: error: %s\n' % str( msg ) )
155
156                        sys.exit( 1 )
157
158        def return_note( self, pre_parts ):
159
160                if type(pre_parts) is types.ListType and len( pre_parts) >= 5:
161                        return { 
162                                'date_add': pre_parts[0].strip(), 
163                                'date_edit': pre_parts[1].strip(), 
164                                'user': pre_parts[2].strip(), 
165                                'ticket': pre_parts[3].strip(),
166                                'remark': ','.join( pre_parts[4:] )
167                        }
168                else:
169                        return { 
170                                'date_add': '', 
171                                'date_edit': '', 
172                                'user': '', 
173                                'ticket': '',
174                                'remark': str( pre_parts )
175                        }
176
177        def convert_format( self, format_options ):
178            pattern = r'\%([-|+]){0,1}([0-9]{0,2})([a-z]{1})'
179            parts = re.findall( pattern, format_options[0], re.VERBOSE )
180            line = re.sub( pattern, '%s', format_options[0], re.VERBOSE )
181
182            rlist = list()
183            counter = 0
184            for variable in format_options[1].split( ',' ):
185                rlist.append( '%s(%s)%s%s%s' % (
186                    '%',
187                    variable, 
188                    parts[ counter ][0], 
189                    parts[ counter ][1],
190                    parts[ counter ][2],
191                ) )
192                counter += 1
193            return line % tuple( rlist )
194
195        def print_table( self, node_list, pbs_nodes, format_options ):
196            '''
197            This method prints the rows of a table
198            '''
199
200            try:
201                line_format = self.convert_format( format_options )
202                for node in node_list:
203                    note = '' 
204                    if pbs_nodes[ node ].has_key('note'):
205                        note = pbs_nodes[ node ]['note']
206                               
207                    if self.allowed_state( pbs_nodes[ node ]['state'] ) or note:
208                        node_note = self.return_note( note )
209
210                        fields = { 
211                            'node': node,
212                            'state': ', '.join( pbs_nodes[ node ]['state'] ),
213                            'date_add': node_note['date_add'],
214                            'date_edit': node_note['date_edit'],
215                            'user': node_note['user'],
216                            'ticket': node_note['ticket'],
217                            'remark': node_note['remark'],
218                        }
219
220                       
221                        print line_format % fields
222            except KeyError, e:
223                raise sara_nodesException, 'Given host does not exist'
224
225        def print_list(self, args, options):
226                '''
227                A method that is used for collecting all nodes with the state down, offline or unknown
228                '''
229
230               
231                p = PBSQuery.PBSQuery()
232                if pbs.version_info >= ( 4,0,0 ):
233                        if self.obj_sara_nodes.verbose:
234                                print "Enabling new_data_structure for PBSQuery"
235                        p.new_data_structure()
236
237                try:
238                    nodes = p.getnodes( ['state', 'note'] )
239                except PBSQuery.PBSError, detail:
240                    print "PBSQuery error: %s" %detail
241                    sys.exit(1)
242
243                header = ' %-10s | %-19s | %s' % ( 'Nodename', 'State', 'Note' )
244                if not options.quiet:
245                        print '\n%s\n%s' % ( header, ( '-' * 80 ) )
246
247                if args:
248                        args = self.sort_nodes( args )
249                        self.print_table( args[0], nodes, options.format )
250                else:
251                        sorted_nodes, sorted_other = self.sort_nodes( nodes )
252
253                        self.print_table( sorted_other, nodes, options.format )
254                        self.print_table( sorted_nodes, nodes, options.format )
255
256        def real_sort( self, inlist ):
257                '''
258                Use this method instead of the x.sort(), because with x.sort()
259                numeric values in a string are not correctly sorted!
260                '''
261                indices = map(self._generate_index, inlist )
262                decorated = zip( indices, inlist )
263                decorated.sort()
264
265                return [ item for index, item in decorated ]
266
267        def _generate_index( self, str ):
268                '''
269                Spliting a string in aplha and numeric elements
270                '''
271
272                index = []
273
274                def _append( fragment, alist=index ):
275                        if fragment.isdigit():
276                                fragment = int( fragment )
277                        alist.append( fragment )
278
279                prev_isdigit = str[0].isdigit()
280                current_fragment = ''
281
282                for char in str:
283                        curr_isdigit = char.isdigit()
284
285                        if curr_isdigit == prev_isdigit:
286                                current_fragment += char
287                        else:
288                                _append( current_fragment )
289                                current_fragment = char
290                                prev_isdigit = curr_isdigit
291
292                _append( current_fragment )
293
294                return tuple( index )
295
296        def sort_nodes(self, nodes):
297                '''
298                Sorts the nodes list and returns two lists
299                the first the nodes secondly the other machines
300
301                When RE_PATTERN is not supplied then all names
302                will be sorted the same way.
303                '''
304
305                if not globals().has_key('RE_PATTERN'):
306                        global RE_PATTERN
307                        RE_PATTERN = ''
308
309                pattern = re.compile( RE_PATTERN, re.VERBOSE )
310
311                tmplist = list()
312                tmplist_other = list()
313
314                for node in nodes:
315                        match = pattern.findall( node )
316
317                        if match and len( match ) == 1:
318                                tmplist.append( node )
319                        else:
320                                tmplist_other.append( node )
321
322                tmplist = self.real_sort( tmplist )
323                tmplist_other.sort()
324
325                return tmplist, tmplist_other
326
327        def allowed_state(self, state):
328                '''
329                This method checks is a node complies with the following states:
330                down, offline and or unknown
331                '''
332                allowed_list = set( ['down', 'offline', 'unknown'] )
333
334                return bool( allowed_list.intersection( set( state ) ) )
335
336class sara_nodes:
337
338        def __init__(self):
339                '''
340                Just initialize two optional variables
341                '''
342                self.dryrun = False
343                self.verbose = False
344
345        def note_check_ticket( self, ticketno, oldticket ):
346               
347                if ticketno:
348                        try:
349                            return '#%d' % int( ticketno )
350                        except ValueError:
351                                if ticketno == 'c':
352                                        return ''
353
354                return oldticket
355
356        def note_return_username( self, old_username ):
357                try:
358                        username = os.getlogin()
359
360                        if username != 'root':
361                                return username
362                        else:
363                                return old_username
364                except OSError, err:
365                        return 'root'
366
367        def note_create( self, new_note, mode = 'a', old_note = None ):
368                if mode == 'w':
369                        return new_note
370                else:
371                        if old_note and old_note.find( new_note ) < 0:
372                                return '%s, %s' % ( old_note, new_note )
373                        else:
374                                return new_note
375
376        def create_date( self ):
377                if not globals().has_key('TIME_SEPARATOR'):
378                        global TIME_SEPARATOR
379                        TIME_SEPARATOR = ':'
380
381                curtime = time.localtime()
382                day = time.strftime( '%d-%m', curtime )
383                hour = time.strftime( '%H', curtime )
384                minutes = time.strftime( '%M', curtime )
385                return '%s %s%s%s' % ( day, hour, TIME_SEPARATOR, minutes )
386
387        def note_init( self ):
388                current_date = self.create_date()
389                try:
390                        current_username = os.getlogin()
391                except OSError, err:
392                        current_username = 'root'
393
394                return [ current_date, current_date, current_username, '' ]
395
396        def note( self, node, note_attr ):
397                '''
398                This method combines all note methods and returns the new note
399                '''
400                p = PBSQuery.PBSQuery()
401                p.new_data_structure()
402                pbs_info = p.getnode( node )
403               
404                pre_parts = list()
405                old_note = None
406                new_note = None
407
408                if pbs_info.has_key( 'note' ):
409                        pbs_note = pbs_info[ 'note' ]
410                        if len( pbs_note ) > 4:
411                                pre_parts = pbs_note[:4]
412                                old_note = ', '.join( pbs_note[4:] )
413
414                                pre_parts[1] = self.create_date()
415                                pre_parts[2] = self.note_return_username( pre_parts[2] )
416
417                else:
418                        pre_parts = self.note_init()
419
420                if note_attr.has_key( 'ticket' ):
421                        pre_parts[3] = self.note_check_ticket( note_attr['ticket'], pre_parts[3] )
422
423                if note_attr.has_key( 'note' ) and note_attr.has_key( 'mode' ):
424                        if note_attr[ 'note' ] and note_attr[ 'mode' ] in [ 'a','w' ]:
425                                if old_note:
426                                        new_note = self.note_create( note_attr[ 'note' ], note_attr[ 'mode' ], old_note )
427                                else:
428                                        new_note = self.note_create( note_attr[ 'note' ], note_attr[ 'mode' ] )
429                        else:
430                                new_note = old_note
431
432                return '%s,%s' % ( ','.join( pre_parts ), new_note )
433
434        def verbose_print( self, msg ):
435                if self.verbose:
436                        print msg
437
438        def pbs_change_note_clear( self, nodes ):
439                attributes = pbs.new_attropl(1)
440                attributes[0].name = pbs.ATTR_NODE_note
441                attributes[0].value = ''
442                attributes[0].op = pbs.SET
443
444                self.verbose_print( '%*s: cleared' % ( 7, 'Note') )
445                self.pbs_batch( nodes, attributes )
446
447        def pbs_change_note_ticket( self, nodes, ticket ):
448                note_attributes = { 'note': None, 'ticket': ticket, 'mode': 'a' }
449                self.verbose_print( '%*s: %s' % ( 7, 'Ticket', ticket  ) )
450                self.pbs_batch( nodes, None, note_attributes)
451
452        def pbs_change_note( self, nodes, note, ticket=None ):
453                note_attributes = { 'note': note, 'ticket': ticket, 'mode': 'w' }
454
455                self.verbose_print( '%*s: %s' % ( 7, 'Note', note ) )
456                if ticket:
457                        self.verbose_print( '%*s: %s' % ( 7, 'Ticket', ticket ) )
458                self.pbs_batch( nodes, None, note_attributes)
459
460        def pbs_change_state_offline( self, nodes, note, ticket=None ):
461                attributes = pbs.new_attropl(1)
462                attributes[0].name = pbs.ATTR_NODE_state
463                attributes[0].value = 'offline'
464                attributes[0].op = pbs.SET
465
466                note_attributes = { 'note': note, 'ticket': ticket, 'mode': 'a' }
467
468                self.verbose_print( '%*s: offline' % ( 7, 'State') )
469                self.verbose_print( '%*s: %s' % ( 7, 'Note', note ) )
470                if ticket:
471                        self.verbose_print( '%*s: %s' % ( 7, 'Ticket', ticket ) )
472                self.pbs_batch( nodes, attributes, note_attributes )
473
474        def pbs_change_state_down( self, nodes ):
475                attributes = pbs.new_attropl(2)
476                attributes[0].name = pbs.ATTR_NODE_state
477                attributes[0].value = 'down'
478                attributes[0].op = pbs.SET
479
480                attributes[1].name = 'note'
481                attributes[1].value = ''
482
483                self.verbose_print( '%*s: down' % ( 7, 'State') )
484                self.verbose_print( '%*s: cleared' % ( 7, 'Note' ) )
485                self.pbs_batch( nodes, attributes )
486
487        def pbs_batch( self, nodes, attrs=None, note_attributes=None ):
488                nodeserror = list()
489                if not attrs and not note_attributes:
490                        raise sara_nodesException, 'attrs and note_attributes can not be empty together!'
491
492                if not self.dryrun:
493                        if note_attributes and len( note_attributes ) == 3:
494                                if attrs:
495                                        attributes = attrs + pbs.new_attropl(1)
496                                        attributes[1].name = pbs.ATTR_NODE_note
497                                        attributes[1].op = pbs.SET
498                                else:
499                                        attributes = pbs.new_attropl(1)
500                                        attributes[0].name = pbs.ATTR_NODE_note
501                                        attributes[0].op = pbs.SET
502                        else:
503                                attributes = attrs
504                        # Some hacking here because some limitation in the Torque 2.4 version
505                        # fetching note data first for all nodes!
506                        tmp_node_note = dict()
507
508                        for node in nodes:
509                                if note_attributes and len( note_attributes ) == 3:
510                                            tmp_node_note[ node ] = self.note( node, note_attributes )
511
512                        pbs_server = pbs.pbs_default()
513
514                        if not pbs_server:
515                                raise sara_nodesException, 'Default pbs server not found!'
516
517                        pbs_connection = pbs.pbs_connect( pbs_server )
518                        for node in nodes:
519                                if note_attributes and len( note_attributes ) == 3:
520                                        try:
521                                                if attrs:
522                                                        attributes[1].value = tmp_node_note[ node ]
523                                                else:
524                                                        attributes[0].value = tmp_node_note[ node ]
525                                        except KeyError:
526                                                pass
527                                rcode = pbs.pbs_manager( pbs_connection, pbs.MGR_CMD_SET, pbs.MGR_OBJ_NODE, node, attributes, 'NULL' )
528                                if rcode > 0:
529                                        errno, text = pbs.error()
530                                        nodeserror.append( '%s: %s (%s)' % ( node, text, errno ) )
531                else:
532                        p = PBSQuery.PBSQuery()
533                        pbsnodes = p.getnodes().keys()
534
535                        print '%*s:' % ( 7, 'Nodes' ),
536                        firstitem = True
537
538                        for node in nodes:
539                                if node in pbsnodes:
540                                        if firstitem:
541                                                print '%s' % node
542                                                firstitem = False
543                                        else:
544                                                print '%*s' % ( 17, node )
545                                else:
546                                        nodeserror.append( '%s: does not exist' % node )
547
548                if len( nodeserror ) > 0:
549                        raise sara_nodesException, nodeserror
550
551if __name__ == '__main__':
552        sara_nodesCli()
Note: See TracBrowser for help on using the repository browser.