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