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