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