[252] | 1 | #!/usr/bin/env python |
---|
| 2 | # |
---|
| 3 | # Author: Dennis Stam |
---|
[294] | 4 | # Date : 6th of September 2012 |
---|
[252] | 5 | # |
---|
[294] | 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 |
---|
[252] | 8 | |
---|
[294] | 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' ''' |
---|
[252] | 12 | |
---|
[294] | 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 |
---|
[252] | 21 | import PBSQuery |
---|
[294] | 22 | import re |
---|
[252] | 23 | import sys |
---|
| 24 | import time |
---|
[294] | 25 | import getpass |
---|
| 26 | import types |
---|
[252] | 27 | |
---|
[294] | 28 | ## Use the cli arguments to change the values |
---|
| 29 | ARGS_QUIET = False |
---|
| 30 | ARGS_VERBOSE = False |
---|
| 31 | ARGS_DRYRUN = False |
---|
[252] | 32 | |
---|
[294] | 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''' |
---|
[252] | 37 | |
---|
[294] | 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')) |
---|
[252] | 61 | |
---|
[294] | 62 | ## Import argparse here, as I need the _print function |
---|
| 63 | try: |
---|
| 64 | import argparse |
---|
| 65 | except ImportError: |
---|
[295] | 66 | _print('Cannot find argparse module', file=sys.stderr) |
---|
[294] | 67 | sys.exit(1) |
---|
[252] | 68 | |
---|
[294] | 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() |
---|
[252] | 76 | |
---|
[294] | 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 |
---|
[252] | 87 | |
---|
[294] | 88 | def return_range(string): |
---|
| 89 | '''This function will return the possible values for the given ranges''' |
---|
[252] | 90 | |
---|
[294] | 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') |
---|
[252] | 94 | |
---|
[294] | 95 | numbers_chars = list() |
---|
| 96 | equal_width_length = 0 |
---|
[252] | 97 | |
---|
[294] | 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) |
---|
[252] | 125 | |
---|
[294] | 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))) |
---|
[252] | 133 | else: |
---|
[294] | 134 | tmp_list.append(number_char) |
---|
| 135 | numbers_chars = tmp_list |
---|
[252] | 136 | |
---|
[294] | 137 | return numbers_chars |
---|
[252] | 138 | |
---|
[294] | 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) |
---|
[252] | 148 | |
---|
[294] | 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() |
---|
[252] | 157 | |
---|
[294] | 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 |
---|
[252] | 168 | |
---|
[294] | 169 | ## END functions for hostrange parsing |
---|
| 170 | #### |
---|
[278] | 171 | |
---|
[294] | 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) ] |
---|
[252] | 177 | |
---|
[294] | 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) |
---|
[252] | 182 | |
---|
[294] | 183 | ## there are 2 possible filters, by hostname, or by state |
---|
| 184 | pbsq = PBSQuery.PBSQuery() |
---|
| 185 | split_1 = dict() |
---|
| 186 | split_2 = dict() |
---|
[252] | 187 | |
---|
[294] | 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 |
---|
[252] | 201 | |
---|
[294] | 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 |
---|
[252] | 213 | |
---|
[294] | 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) |
---|
[252] | 218 | |
---|
[294] | 219 | line_print = list() |
---|
| 220 | if ARGS_VERBOSE: |
---|
| 221 | _print('func:print_process_dict processing data', file=sys.stderr) |
---|
[252] | 222 | |
---|
[294] | 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 |
---|
[252] | 225 | |
---|
[294] | 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() |
---|
[252] | 229 | |
---|
[294] | 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 []) |
---|
[278] | 233 | |
---|
[294] | 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:]) |
---|
[252] | 242 | |
---|
[294] | 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 |
---|
[252] | 250 | |
---|
[294] | 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'] = '' |
---|
[252] | 254 | |
---|
[294] | 255 | line_print.append(add_dict) |
---|
[252] | 256 | |
---|
[294] | 257 | if ARGS_VERBOSE: |
---|
| 258 | _print('func:print_process_dict returning values', file=sys.stderr) |
---|
| 259 | return line_print |
---|
[270] | 260 | |
---|
[294] | 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 |
---|
[252] | 266 | |
---|
[294] | 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) |
---|
[252] | 271 | |
---|
[294] | 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 |
---|
[252] | 276 | |
---|
[294] | 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)) |
---|
[252] | 281 | |
---|
[294] | 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']) |
---|
[252] | 300 | |
---|
[294] | 301 | ## The length of the full note |
---|
| 302 | w_notefull = w_date + w_username + w_ticket + w_note |
---|
[252] | 303 | |
---|
[294] | 304 | if not ARGS_QUIET: |
---|
| 305 | show_fields = [ |
---|
| 306 | [w_nodename, 'Nodename'], |
---|
| 307 | [w_state, 'State'], |
---|
| 308 | ] |
---|
| 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']) |
---|
[252] | 318 | |
---|
[294] | 319 | _print(' %s' % ' | '.join(print_create_list(show_fields))) |
---|
| 320 | _print('+'.join([ '-' * (show_field[0]+2) for show_field in show_fields ])) |
---|
[252] | 321 | |
---|
[294] | 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']]) |
---|
[252] | 335 | |
---|
[294] | 336 | _print(' %s' % ' | '.join(print_create_list(show_line_fields))) |
---|
[252] | 337 | |
---|
[294] | 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''' |
---|
[252] | 341 | |
---|
[294] | 342 | matched, rest = print_get_nodes(hosts) |
---|
| 343 | print_list = print_process_dict(rest) |
---|
| 344 | print_list.extend(print_process_dict(matched)) |
---|
[252] | 345 | |
---|
[294] | 346 | for line in print_list: |
---|
| 347 | _print(format % line) |
---|
| 348 | ## END functions for printing |
---|
| 349 | #### |
---|
[252] | 350 | |
---|
[294] | 351 | class SaraNodes(object): |
---|
| 352 | '''This class is used to communicate with the batch server''' |
---|
[252] | 353 | |
---|
[294] | 354 | ticket = None |
---|
[252] | 355 | |
---|
[294] | 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) |
---|
[252] | 360 | |
---|
[294] | 361 | pbsq = PBSQuery.PBSQuery() |
---|
| 362 | rdict = dict() |
---|
[252] | 363 | |
---|
[294] | 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 |
---|
[252] | 369 | |
---|
[294] | 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()) |
---|
[252] | 375 | |
---|
[294] | 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 |
---|
[252] | 384 | |
---|
[294] | 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 '' |
---|
[252] | 399 | |
---|
[294] | 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) |
---|
[252] | 404 | |
---|
[294] | 405 | ## First step, is to get the current info of a host |
---|
| 406 | cur_data = self._get_current_notes(nodes) |
---|
| 407 | rdict = dict() |
---|
[252] | 408 | |
---|
[294] | 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 |
---|
[252] | 422 | |
---|
[294] | 423 | if nnote and append and note: |
---|
| 424 | nnote = '%s, %s' % (nnote, note) |
---|
| 425 | elif note: |
---|
| 426 | nnote = note |
---|
[252] | 427 | |
---|
[294] | 428 | rdict[node] = '%s,%s,%s,%s,%s' % (date_add, date_edit, username, ticket, nnote) |
---|
| 429 | return rdict |
---|
[252] | 430 | |
---|
[294] | 431 | def do_offline(self, nodes, note): |
---|
| 432 | '''Change the state of node(s) to offline with a specific note''' |
---|
[252] | 433 | |
---|
[294] | 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 |
---|
[252] | 442 | |
---|
[294] | 443 | batch_list = list() |
---|
[252] | 444 | |
---|
[294] | 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])) |
---|
[252] | 449 | |
---|
[294] | 450 | self._process(batch_list) |
---|
[252] | 451 | |
---|
[294] | 452 | def do_clear(self, nodes): |
---|
| 453 | '''Clear the state on a node(s) to down, also clear the note''' |
---|
[252] | 454 | |
---|
[294] | 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 = '' |
---|
[252] | 464 | |
---|
[294] | 465 | batch_list = list() |
---|
[252] | 466 | |
---|
[294] | 467 | ## again a loop, now create the attrib dict list |
---|
| 468 | for node in nodes: |
---|
| 469 | batch_list.append(tuple([node, attributes])) |
---|
[252] | 470 | |
---|
[294] | 471 | self._process(batch_list) |
---|
| 472 | |
---|
| 473 | def do_modify(self, nodes, note): |
---|
| 474 | '''Modify the note on a node, override the previous note''' |
---|
[252] | 475 | |
---|
[294] | 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 |
---|
[252] | 481 | |
---|
[294] | 482 | batch_list = list() |
---|
[252] | 483 | |
---|
[294] | 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])) |
---|
[252] | 488 | |
---|
[294] | 489 | self._process(batch_list) |
---|
[252] | 490 | |
---|
[294] | 491 | def do_clear_note(self, nodes): |
---|
| 492 | '''Clear the note on the node(s)''' |
---|
[252] | 493 | |
---|
[294] | 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 = '' |
---|
[252] | 500 | |
---|
[294] | 501 | batch_list = list() |
---|
[252] | 502 | |
---|
[294] | 503 | ## again a loop, now create the attrib dict list |
---|
| 504 | for node in nodes: |
---|
| 505 | batch_list.append(tuple([node, attributes])) |
---|
[252] | 506 | |
---|
[294] | 507 | self._process(batch_list) |
---|
[252] | 508 | |
---|
[294] | 509 | def _process(self, batch_list): |
---|
| 510 | '''This function execute the change to the batch server''' |
---|
[252] | 511 | |
---|
[294] | 512 | if ARGS_VERBOSE: |
---|
| 513 | _print('class:SaraNodes func:_process input:%s' % str(batch_list), file=sys.stderr) |
---|
[252] | 514 | |
---|
[294] | 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) |
---|
[252] | 520 | |
---|
[294] | 521 | if ARGS_VERBOSE: |
---|
| 522 | _print('class:SaraNodes func:_process pbs_server:%s' % pbs_server, file=sys.stderr) |
---|
[252] | 523 | |
---|
[294] | 524 | ## If dry-run is not specified create a connection |
---|
| 525 | if not ARGS_DRYRUN: |
---|
| 526 | pbs_connection = pbs.pbs_connect(pbs_server) |
---|
[252] | 527 | |
---|
[294] | 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]))) |
---|
[252] | 538 | |
---|
[294] | 539 | ## Close the connection with the batch system |
---|
| 540 | if not ARGS_DRYRUN: |
---|
| 541 | pbs.pbs_disconnect(pbs_connection) |
---|
[252] | 542 | |
---|
[294] | 543 | if __name__ == '__main__': |
---|
| 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) |
---|
[252] | 560 | |
---|
[294] | 561 | ## Parse the arguments |
---|
| 562 | args = parser.parse_args() |
---|
[252] | 563 | |
---|
[294] | 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 |
---|
[252] | 571 | |
---|
[294] | 572 | if ARGS_VERBOSE: |
---|
| 573 | _print('func:__main__ checking type of operation', file=sys.stderr) |
---|
[252] | 574 | |
---|
[294] | 575 | if args.nodes: |
---|
| 576 | args.nodes = parse_args(args.nodes) |
---|
[252] | 577 | |
---|
[294] | 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) |
---|
[252] | 583 | |
---|
[294] | 584 | sn = SaraNodes() |
---|
| 585 | if args.ticket: |
---|
| 586 | sn.ticket = args.ticket |
---|
[252] | 587 | |
---|
[294] | 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) |
---|
[252] | 611 | |
---|
[294] | 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) |
---|