1#! /usr/bin/python 2 3# Copyright (c) 2014 The Chromium OS Authors. All rights reserved. 4# Use of this source code is governed by a BSD-style license that can be 5# found in the LICENSE file. 6 7""" 8This script generates a csv file containing the mapping of 9(device_hostname, rpm_hostname, outlet, hydra_hostname) for each 10host in our lab. The csv file is in the following format. 11 12chromeos-rack2-host1,chromeos-rack2-rpm1,.A1,chromeos-197-hydra1.mtv 13chromeos-rack2-host2,chromeos-rack2-rpm1,.A2,chromeos-197-hydra1.mtv 14... 15 16The generated csv file can be used as input to add_host_powerunit_info.py 17 18Workflow: 19 <Generate the csv file> 20 python generate_rpm_mapping.py --csv mapping_file.csv --server cautotest 21 22 <Upload mapping information in csv file to AFE> 23 python add_host_powerunit_info.py --csv mapping_file.csv 24 25""" 26import argparse 27import collections 28import logging 29import re 30import sys 31 32import common 33 34from autotest_lib.client.common_lib import enum 35from autotest_lib.server.cros.dynamic_suite import frontend_wrappers 36 37CHROMEOS_LABS = enum.Enum('OysterBay', 'Atlantis', 'Chaos', 'Destiny', start_value=1) 38HOST_REGX = 'chromeos(\d+)(-row(\d+))*-rack(\d+)-host(\d+)' 39DeviceHostname = collections.namedtuple( 40 'DeviceHostname', ['lab', 'row', 'rack', 'host']) 41 42 43class BaseLabConfig(object): 44 """Base class for a lab configuration.""" 45 RPM_OUTLET_MAP = {} 46 LAB_NUMBER = -1 47 48 @classmethod 49 def get_rpm_hostname(cls, device_hostname): 50 """Get rpm hostname given a device. 51 52 @param device_hostname: A DeviceHostname named tuple. 53 54 @returns: the rpm hostname, default to empty string. 55 56 """ 57 return '' 58 59 60 @classmethod 61 def get_rpm_outlet(cls, device_hostname): 62 """Get rpm outlet given a device. 63 64 @param device_hostname: A DeviceHostname named tuple. 65 66 @returns: the rpm outlet, default to empty string. 67 68 """ 69 return '' 70 71 72 @classmethod 73 def get_hydra_hostname(cls, device_hostname): 74 """Get hydra hostname given a device. 75 76 @param device_hostname: A DeviceHostname named tuple. 77 78 @returns: the hydra hostname, default to empty string. 79 80 """ 81 return '' 82 83 84 @classmethod 85 def is_device_in_the_lab(cls, device_hostname): 86 """Check whether a dut belongs to the lab. 87 88 @param device_hostname: A DeviceHostname named tuple. 89 90 @returns: True if the dut belongs to the lab, 91 False otherwise. 92 93 """ 94 return device_hostname.lab == cls.LAB_NUMBER 95 96 97class OysterBayConfig(BaseLabConfig): 98 """Configuration for OysterBay""" 99 100 LAB_NUMBER = CHROMEOS_LABS.OYSTERBAY 101 102 103 @classmethod 104 def get_rpm_hostname(cls, device_hostname): 105 """Get rpm hostname. 106 107 @param device_hostname: A DeviceHostname named tuple. 108 109 @returns: hostname of the rpm that has the device. 110 111 """ 112 if not device_hostname.row: 113 return '' 114 return 'chromeos%d-row%d-rack%d-rpm1' % ( 115 device_hostname.lab, device_hostname.row, 116 device_hostname.rack) 117 118 119 @classmethod 120 def get_rpm_outlet(cls, device_hostname): 121 """Get rpm outlet. 122 123 @param device_hostname: A DeviceHostname named tuple. 124 125 @returns: rpm outlet, e.g. '.A1' 126 127 """ 128 if not device_hostname.row: 129 return '' 130 return '.A%d' % device_hostname.host 131 132 133class AtlantisConfig(BaseLabConfig): 134 """Configuration for Atlantis lab.""" 135 136 LAB_NUMBER = CHROMEOS_LABS.ATLANTIS 137 # chromeos2, hostX -> outlet 138 RPM_OUTLET_MAP = { 139 1: 1, 140 7: 2, 141 2: 4, 142 8: 5, 143 3: 7, 144 9: 8, 145 4: 9, 146 10: 10, 147 5: 12, 148 11: 13, 149 6: 15, 150 12: 16} 151 152 @classmethod 153 def get_rpm_hostname(cls, device_hostname): 154 """Get rpm hostname. 155 156 @param device_hostname: A DeviceHostname named tuple. 157 158 @returns: hostname of the rpm that has the device. 159 160 """ 161 return 'chromeos%d-row%d-rack%d-rpm1' % ( 162 device_hostname.lab, device_hostname.row, 163 device_hostname.rack) 164 165 166 @classmethod 167 def get_rpm_outlet(cls, device_hostname): 168 """Get rpm outlet. 169 170 @param device_hostname: A DeviceHostname named tuple. 171 172 @returns: rpm outlet, e.g. '.A1' 173 174 """ 175 return '.A%d' % cls.RPM_OUTLET_MAP[device_hostname.host] 176 177 178 @classmethod 179 def get_hydra_hostname(cls, device_hostname): 180 """Get hydra hostname. 181 182 @param device_hostname: A DeviceHostname named tuple. 183 184 @returns: hydra hostname 185 186 """ 187 row = device_hostname.row 188 rack = device_hostname.rack 189 if row >= 1 and row <= 5 and rack >= 1 and rack <= 7: 190 return 'chromeos-197-hydra1.cros' 191 elif row >= 1 and row <= 5 and rack >= 8 and rack <= 11: 192 return 'chromeos-197-hydra2.cros' 193 else: 194 logging.error('Could not determine hydra for %s', 195 device_hostname) 196 return '' 197 198 199class ChaosConfig(BaseLabConfig): 200 """Configuration for Chaos lab.""" 201 202 LAB_NUMBER = CHROMEOS_LABS.CHAOS 203 204 205 @classmethod 206 def get_rpm_hostname(cls, device_hostname): 207 """Get rpm hostname. 208 209 @param device_hostname: A DeviceHostname named tuple. 210 211 @returns: hostname of the rpm that has the device. 212 213 """ 214 return 'chromeos%d-row%d-rack%d-rpm1' % ( 215 device_hostname.lab, device_hostname.row, 216 device_hostname.rack) 217 218 219 @classmethod 220 def get_rpm_outlet(cls, device_hostname): 221 """Get rpm outlet. 222 223 @param device_hostname: A DeviceHostname named tuple. 224 225 @returns: rpm outlet, e.g. '.A1' 226 227 """ 228 return '.A%d' % device_hostname.host 229 230 231class DestinyConfig(BaseLabConfig): 232 """Configuration for Desitny lab.""" 233 234 LAB_NUMBER = CHROMEOS_LABS.DESTINY 235 # None-densified rack: one host per shelf 236 # (rowX % 2, hostY) -> outlet 237 RPM_OUTLET_MAP = { 238 (1, 1): 1, 239 (0, 1): 2, 240 (1, 2): 4, 241 (0, 2): 5, 242 (1, 3): 7, 243 (0, 3): 8, 244 (1, 4): 9, 245 (0, 4): 10, 246 (1, 5): 12, 247 (0, 5): 13, 248 (1, 6): 15, 249 (0, 6): 16, 250 } 251 252 # Densified rack: one shelf can have two chromeboxes or one notebook. 253 # (rowX % 2, hostY) -> outlet 254 DENSIFIED_RPM_OUTLET_MAP = { 255 (1, 2): 1, (1, 1): 1, 256 (0, 1): 2, (0, 2): 2, 257 (1, 4): 3, (1, 3): 3, 258 (0, 3): 4, (0, 4): 4, 259 (1, 6): 5, (1, 5): 5, 260 (0, 5): 6, (0, 6): 6, 261 (1, 8): 7, (1, 7): 7, 262 (0, 7): 8, (0, 8): 8, 263 # outlet 9, 10 are not used 264 (1, 10): 11, (1, 9): 11, 265 (0, 9): 12, (0, 10): 12, 266 (1, 12): 13, (1, 11): 13, 267 (0, 11): 14, (0, 12): 14, 268 (1, 14): 15, (1, 13): 15, 269 (0, 13): 16, (0, 14): 16, 270 (1, 16): 17, (1, 15): 17, 271 (0, 15): 18, (0, 16): 18, 272 (1, 18): 19, (1, 17): 19, 273 (0, 17): 20, (0, 18): 20, 274 (1, 20): 21, (1, 19): 21, 275 (0, 19): 22, (0, 20): 22, 276 (1, 22): 23, (1, 21): 23, 277 (0, 21): 24, (0, 22): 24, 278 } 279 280 281 @classmethod 282 def is_densified(cls, device_hostname): 283 """Whether the host is on a densified rack. 284 285 @param device_hostname: A DeviceHostname named tuple. 286 287 @returns: True if on a densified rack, False otherwise. 288 """ 289 return device_hostname.rack in (0, 12, 13) 290 291 292 @classmethod 293 def get_rpm_hostname(cls, device_hostname): 294 """Get rpm hostname. 295 296 @param device_hostname: A DeviceHostname named tuple. 297 298 @returns: hostname of the rpm that has the device. 299 300 """ 301 row = device_hostname.row 302 if row == 13: 303 logging.warn('Rule not implemented for row 13 in chromeos4') 304 return '' 305 306 # rpm row is like chromeos4-row1_2-rackX-rpmY 307 rpm_row = ('%d_%d' % (row - 1, row) if row % 2 == 0 else 308 '%d_%d' % (row, row + 1)) 309 310 if cls.is_densified(device_hostname): 311 # Densified rack has two rpms, decide which one the host belongs to 312 # Rule: 313 # odd row number, even host number -> rpm1 314 # odd row number, odd host number -> rpm2 315 # even row number, odd host number -> rpm1 316 # even row number, even host number -> rpm2 317 rpm_number = 1 if (row + device_hostname.host) % 2 == 1 else 2 318 else: 319 # Non-densified rack only has one rpm 320 rpm_number = 1 321 return 'chromeos%d-row%s-rack%d-rpm%d' % ( 322 device_hostname.lab, 323 rpm_row, device_hostname.rack, rpm_number) 324 325 326 @classmethod 327 def get_rpm_outlet(cls, device_hostname): 328 """Get rpm outlet. 329 330 @param device_hostname: A DeviceHostname named tuple. 331 332 @returns: rpm outlet, e.g. '.A1' 333 334 """ 335 try: 336 outlet_map = (cls.DENSIFIED_RPM_OUTLET_MAP 337 if cls.is_densified(device_hostname) else 338 cls.RPM_OUTLET_MAP) 339 outlet_number = outlet_map[(device_hostname.row % 2, 340 device_hostname.host)] 341 return '.A%d' % outlet_number 342 except KeyError: 343 logging.error('Could not determine outlet for device %s', 344 device_hostname) 345 return '' 346 347 348 @classmethod 349 def get_hydra_hostname(cls, device_hostname): 350 """Get hydra hostname. 351 352 @param device_hostname: A DeviceHostname named tuple. 353 354 @returns: hydra hostname 355 356 """ 357 row = device_hostname.row 358 rack = device_hostname.rack 359 if row >= 1 and row <= 6 and rack >=1 and rack <= 11: 360 return 'chromeos-destiny-hydra1.cros' 361 elif row >= 7 and row <= 12 and rack >=1 and rack <= 11: 362 return 'chromeos-destiny-hydra2.cros' 363 elif row >= 1 and row <= 10 and rack >=12 and rack <= 13: 364 return 'chromeos-destiny-hydra3.cros' 365 elif row in [3, 4, 5, 6, 9, 10] and rack == 0: 366 return 'chromeos-destiny-hydra3.cros' 367 elif row == 13 and rack >= 0 and rack <= 11: 368 return 'chromeos-destiny-hydra3.cros' 369 else: 370 logging.error('Could not determine hydra hostname for %s', 371 device_hostname) 372 return '' 373 374 375def parse_device_hostname(device_hostname): 376 """Parse device_hostname to DeviceHostname object. 377 378 @param device_hostname: A string, e.g. 'chromeos2-row2-rack4-host3' 379 380 @returns: A DeviceHostname named tuple or None if the 381 the hostname doesn't follow the pattern 382 defined in HOST_REGX. 383 384 """ 385 m = re.match(HOST_REGX, device_hostname.strip()) 386 if m: 387 return DeviceHostname( 388 lab=int(m.group(1)), 389 row=int(m.group(3)) if m.group(3) else None, 390 rack=int(m.group(4)), 391 host=int(m.group(5))) 392 else: 393 logging.error('Could not parse %s', device_hostname) 394 return None 395 396 397def generate_mapping(hosts, lab_configs): 398 """Generate device_hostname-rpm-outlet-hydra mapping. 399 400 @param hosts: hosts objects get from AFE. 401 @param lab_configs: A list of configuration classes, 402 each one for a lab. 403 404 @returns: A dictionary that maps device_hostname to 405 (rpm_hostname, outlet, hydra_hostname) 406 407 """ 408 # device hostname -> (rpm_hostname, outlet, hydra_hostname) 409 rpm_mapping = {} 410 for host in hosts: 411 device_hostname = parse_device_hostname(host.hostname) 412 if not device_hostname: 413 continue 414 for lab in lab_configs: 415 if lab.is_device_in_the_lab(device_hostname): 416 rpm_hostname = lab.get_rpm_hostname(device_hostname) 417 rpm_outlet = lab.get_rpm_outlet(device_hostname) 418 hydra_hostname = lab.get_hydra_hostname(device_hostname) 419 if not rpm_hostname or not rpm_outlet: 420 logging.error( 421 'Skipping device %s: could not determine ' 422 'rpm hostname or outlet.', host.hostname) 423 break 424 rpm_mapping[host.hostname] = ( 425 rpm_hostname, rpm_outlet, hydra_hostname) 426 break 427 else: 428 logging.info( 429 '%s is not in a know lab ' 430 '(oyster bay, atlantis, chaos, destiny)', 431 host.hostname) 432 return rpm_mapping 433 434 435def output_csv(rpm_mapping, csv_file): 436 """Dump the rpm mapping dictionary to csv file. 437 438 @param rpm_mapping: A dictionary that maps device_hostname to 439 (rpm_hostname, outlet, hydra_hostname) 440 @param csv_file: The name of the file to write to. 441 442 """ 443 with open(csv_file, 'w') as f: 444 for hostname, rpm_info in rpm_mapping.iteritems(): 445 line = ','.join(rpm_info) 446 line = ','.join([hostname, line]) 447 f.write(line + '\n') 448 449 450if __name__ == '__main__': 451 logging.basicConfig(level=logging.DEBUG) 452 parser = argparse.ArgumentParser( 453 description='Generate device_hostname-rpm-outlet-hydra mapping ' 454 'file needed by add_host_powerunit_info.py') 455 parser.add_argument('--csv', type=str, dest='csv_file', required=True, 456 help='The path to the csv file where we are going to ' 457 'write the mapping information to.') 458 parser.add_argument('--server', type=str, dest='server', default=None, 459 help='AFE server that the script will be talking to. ' 460 'If not specified, will default to using the ' 461 'server in global_config.ini') 462 options = parser.parse_args() 463 464 AFE = frontend_wrappers.RetryingAFE(timeout_min=5, delay_sec=10, 465 server=options.server) 466 logging.info('Connected to %s', AFE.server) 467 rpm_mapping = generate_mapping( 468 AFE.get_hosts(), 469 [OysterBayConfig, AtlantisConfig, ChaosConfig, DestinyConfig]) 470 output_csv(rpm_mapping, options.csv_file) 471