1 # Copyright (c) 2014 The Chromium OS Authors. All rights reserved. 2 # Use of this source code is governed by a BSD-style license that can be 3 # found in the LICENSE file. 4 5 """RDB utilities. 6 7 Do not import rdb or autotest modules here to avoid cyclic dependencies. 8 """ 9 10 import collections 11 12 import common 13 from autotest_lib.client.common_lib import priorities 14 from autotest_lib.client.common_lib import utils 15 16 try: 17 from chromite.lib import metrics 18 except ImportError: 19 metrics = utils.metrics_mock 20 21 22 RDB_STATS_KEY = 'rdb' 23 24 class RDBException(Exception): 25 """Generic RDB exception.""" 26 27 def wire_format(self, **kwargs): 28 """Convert the exception to a format better suited to an rpc response. 29 """ 30 return str(self) 31 32 33 class CacheMiss(RDBException): 34 """Generic exception raised for a cache miss in the rdb.""" 35 pass 36 37 38 class LabelIterator(object): 39 """An Iterator for labels. 40 41 Within the rdb any label/dependency comparisons are performed based on label 42 ids. However, the host object returned needs to contain label names instead. 43 This class returns label ids for iteration, but a list of all label names 44 when accessed through get_label_names. 45 """ 46 47 def __init__(self, labels): 48 self.labels = labels 49 50 51 def __iter__(self): 52 return iter(label.id for label in self.labels) 53 54 55 def get_label_names(self): 56 """Get all label names of the labels associated with this class. 57 58 @return: A list of label names. 59 """ 60 return [label.name for label in self.labels] 61 62 63 class RequestAccountant(object): 64 """A helper class that count requests and manages min_duts requirement. 65 66 On initialization, this object takes a list of host requests. 67 It will batch the requests by grouping similar requests together 68 and generate a mapping from unique request-> count of the request. 69 It will also generates a mapping from suite_job_id -> min_duts. 70 71 RDB does a two-round of host aquisition. The first round attempts 72 to get min_duts for each suite. The second round attemps to satisfy 73 the rest requests. RDB calls get_min_duts and get_rest to 74 figure out how many duts it should attempt to get for a unique 75 request in the first and second round respectively. 76 77 Assume we have two distinct requests 78 R1 (parent_job_id: 10, need hosts: 2) 79 R2 (parent_job_id: 10, need hosts: 4) 80 And parent job P (job_id:10) has min dut requirement of 3. So we got 81 requests_to_counts = {R1: 2, R2: 4} 82 min_duts_map = {P: 3} 83 84 First round acquiring: 85 Call get_min_duts(R1) 86 return 2, because P hasn't reach its min dut limit (3) yet 87 requests_to_counts -> {R1: 2-2=0, R2: 4} 88 min_duts_map -> {P: 3-2=1} 89 90 Call get_min_duts(R2) 91 return 1, because although R2 needs 4 duts, P's min dut limit is now 1 92 requests_to_counts -> {R1: 0, R2: 4-1=3} 93 min_duts_map -> {P: 1-1=0} 94 95 Second round acquiring: 96 Call get_rest(R1): 97 return 0, requests_to_counts[R1] 98 Call get_rest(R2): 99 return 3, requests_to_counts[R2] 100 101 Note it is possible that in the first round acquiring, although 102 R1 requested 2 duts, it may only get 1 or None. However get_rest 103 doesn't need to care whether the first round succeeded or not, as 104 in the case when the first round failed, regardless how many duts 105 get_rest requests, it will not be fullfilled anyway. 106 """ 107 108 _host_ratio_metric = metrics.Float( 109 'chromeos/autotest/scheduler/rdb/host_acquisition_ratio') 110 111 112 def __init__(self, host_requests): 113 """Initialize. 114 115 @param host_requests: A list of request to acquire hosts. 116 """ 117 self.requests_to_counts = {} 118 # The order matters, it determines which request got fullfilled first. 119 self.requests = [] 120 for request, count in self._batch_requests(host_requests): 121 self.requests.append(request) 122 self.requests_to_counts[request] = count 123 self.min_duts_map = dict( 124 (r.parent_job_id, r.suite_min_duts) 125 for r in self.requests_to_counts.iterkeys() if r.parent_job_id) 126 127 128 @classmethod 129 def _batch_requests(cls, requests): 130 """ Group similar requests, sort by priority and parent_job_id. 131 132 @param requests: A list or unsorted, unordered requests. 133 134 @return: A list of tuples of the form (request, number of occurances) 135 formed by counting the number of requests with the same acls/deps/ 136 priority in the input list of requests, and sorting by priority. 137 The order of this list ensures against priority inversion. 138 """ 139 sort_function = lambda request: (request[0].priority, 140 -request[0].parent_job_id) 141 return sorted(collections.Counter(requests).items(), key=sort_function, 142 reverse=True) 143 144 145 def get_min_duts(self, host_request): 146 """Given a distinct host request figure out min duts to request for. 147 148 @param host_request: A request. 149 @returns: The minimum duts that should be requested. 150 """ 151 parent_id = host_request.parent_job_id 152 count = self.requests_to_counts[host_request] 153 if parent_id: 154 min_duts = self.min_duts_map.get(parent_id, 0) 155 to_acquire = min(count, min_duts) 156 self.min_duts_map[parent_id] = max(0, min_duts - to_acquire) 157 else: 158 to_acquire = 0 159 self.requests_to_counts[host_request] -= to_acquire 160 return to_acquire 161 162 163 def get_duts(self, host_request): 164 """Return the number of duts host_request still need. 165 166 @param host_request: A request. 167 @returns: The number of duts need to be requested. 168 """ 169 return self.requests_to_counts[host_request] 170 171 172 # TODO(akeshet): Possibly this code is dead, see crbug.com/738508 for 173 # context. 174 def record_acquire_min_duts(cls, host_request, hosts_required, 175 acquired_host_count): 176 """Send stats about host acquisition. 177 178 @param host_request: A request. 179 @param hosts_required: Number of hosts required to satisfy request. 180 @param acquired_host_count: Number of host acquired. 181 """ 182 try: 183 priority = priorities.Priority.get_string(host_request.priority) 184 except ValueError: 185 return 186 cls._host_ratio_metric.set(acquired_host_count/float(hosts_required)) 187
