Home | History | Annotate | Download | only in cli 
      1 #
      2 # Copyright 2008 Google Inc. All Rights Reserved.
      3 
      4 """
      5 The shard module contains the objects and methods used to
      6 manage shards in Autotest.
      7 
      8 The valid actions are:
      9 create:       creates shard
     10 remove:       deletes shard(s)
     11 list:         lists shards with label
     12 add_boards:   add boards to a given shard
     13 remove_board: remove board from a given shard
     14 
     15 See topic_common.py for a High Level Design and Algorithm.
     16 """
     17 
     18 import sys
     19 from autotest_lib.cli import topic_common, action_common
     20 
     21 
     22 class shard(topic_common.atest):
     23     """shard class
     24     atest shard [create|delete|list|add_boards|remove_board] <options>"""
     25     usage_action = '[create|delete|list|add_boards|remove_board]'
     26     topic = msg_topic = 'shard'
     27     msg_items = '<shards>'
     28 
     29     def __init__(self):
     30         """Add to the parser the options common to all the
     31         shard actions"""
     32         super(shard, self).__init__()
     33 
     34         self.topic_parse_info = topic_common.item_parse_info(
     35             attribute_name='shards',
     36             use_leftover=True)
     37 
     38 
     39     def get_items(self):
     40         return self.shards
     41 
     42 
     43 class shard_help(shard):
     44     """Just here to get the atest logic working.
     45     Usage is set by its parent"""
     46     pass
     47 
     48 
     49 class shard_list(action_common.atest_list, shard):
     50     """Class for running atest shard list"""
     51 
     52     def execute(self):
     53         filters = {}
     54         if self.shards:
     55             filters['hostname__in'] = self.shards
     56         return super(shard_list, self).execute(op='get_shards',
     57                                                filters=filters)
     58 
     59 
     60     def warn_if_label_assigned_to_multiple_shards(self, results):
     61         """Prints a warning if one label is assigned to multiple shards.
     62 
     63         This should never happen, but if it does, better be safe.
     64 
     65         @param results: Results as passed to output().
     66         """
     67         assigned_labels = set()
     68         for line in results:
     69             for label in line['labels']:
     70                 if label in assigned_labels:
     71                     sys.stderr.write('WARNING: label %s is assigned to '
     72                                      'multiple shards.\n'
     73                                      'This will lead to unpredictable behavor '
     74                                      'in which hosts and jobs will be assigned '
     75                                      'to which shard.\n')
     76                 assigned_labels.add(label)
     77 
     78 
     79     def output(self, results):
     80         self.warn_if_label_assigned_to_multiple_shards(results)
     81         super(shard_list, self).output(results, ['id', 'hostname', 'labels'])
     82 
     83 
     84 class shard_create(action_common.atest_create, shard):
     85     """Class for running atest shard create -l <label> <shard>"""
     86     def __init__(self):
     87         super(shard_create, self).__init__()
     88         self.parser.add_option('-l', '--labels',
     89                                help=('Assign LABELs to the SHARD. All jobs that '
     90                                      'require one of the labels will be run on  '
     91                                      'the shard. List multiple labels separated '
     92                                      'by a comma.'),
     93                                type='string',
     94                                metavar='LABELS')
     95 
     96 
     97     def parse(self):
     98         (options, leftover) = super(shard_create, self).parse(
     99                 req_items='shards')
    100         self.data_item_key = 'hostname'
    101         self.data['labels'] = options.labels or ''
    102         return (options, leftover)
    103 
    104 
    105 class shard_add_boards(shard_create):
    106     """Class for running atest shard add_boards -l <label> <shard>"""
    107     usage_action = 'add_boards'
    108     op_action = 'add_boards'
    109     msg_done = 'Added boards for'
    110 
    111     def execute(self):
    112         """Running the rpc to add boards to the target shard.
    113 
    114         Returns:
    115           A tuple, 1st element is the target shard. 2nd element is the list of
    116           boards labels to be added to the shard.
    117         """
    118         target_shard = self.shards[0]
    119         self.data[self.data_item_key] = target_shard
    120         super(shard_add_boards, self).execute_rpc(op='add_board_to_shard',
    121                                                   item=target_shard,
    122                                                   **self.data)
    123         return (target_shard, self.data['labels'])
    124 
    125 
    126 class shard_delete(action_common.atest_delete, shard):
    127     """Class for running atest shard delete <shards>"""
    128 
    129     def parse(self):
    130         (options, leftover) = super(shard_delete, self).parse()
    131         self.data_item_key = 'hostname'
    132         return (options, leftover)
    133 
    134 
    135     def execute(self, *args, **kwargs):
    136         print 'Please ensure the shard host is powered off.'
    137         print ('Otherwise DUTs might be used by multiple shards at the same '
    138                'time, which will lead to serious correctness problems.')
    139         return super(shard_delete, self).execute(*args, **kwargs)
    140 
    141 
    142 class shard_remove_board(shard):
    143     """Class for running atest shard remove_board -l <label> <shard>"""
    144     usage_action = 'remove_board'
    145     op_action = 'remove_board'
    146     msg_done = 'Removed board'
    147 
    148     def __init__(self):
    149         super(shard_remove_board, self).__init__()
    150         self.parser.add_option('-l', '--board_label', type='string',
    151                                metavar='BOARD_LABEL',
    152                                help=('Remove the board with the given '
    153                                      'BOARD_LABEL from shard.'))
    154 
    155     def parse(self):
    156         (options, leftover) = super(shard_remove_board, self).parse(
    157               req_items='shards')
    158         self.data['board_label'] = options.board_label
    159         self.data['hostname'] = self.shards[0]
    160         return (options, leftover)
    161 
    162 
    163     def execute(self):
    164         """Validate args and execute the remove_board_from_shard rpc."""
    165         if not self.data.get('board_label'):
    166             self.invalid_syntax('Must provide exactly 1 BOARD_LABEL')
    167             return
    168         if not self.data['board_label'].startswith('board:'):
    169             self.invalid_arg('BOARD_LABEL must begin with "board:"')
    170             return
    171         return super(shard_remove_board, self).execute_rpc(
    172                 op='remove_board_from_shard',
    173                 hostname=self.data['hostname'],
    174                 label=self.data['board_label'])
    175 
    176 
    177     def output(self, results):
    178         """Print command results.
    179 
    180         @param results: Results of rpc execution.
    181         """
    182         print results
    183