Home | History | Annotate | Download | only in contrib
      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 """
      8 This script generates a csv file containing the mapping of
      9 (device_hostname, rpm_hostname, outlet, hydra_hostname) for each
     10 host in our lab. The csv file is in the following format.
     11 
     12 chromeos-rack2-host1,chromeos-rack2-rpm1,.A1,chromeos-197-hydra1.mtv
     13 chromeos-rack2-host2,chromeos-rack2-rpm1,.A2,chromeos-197-hydra1.mtv
     14 ...
     15 
     16 The generated csv file can be used as input to add_host_powerunit_info.py
     17 
     18 Workflow:
     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 """
     26 import argparse
     27 import collections
     28 import logging
     29 import re
     30 import sys
     31 
     32 import common
     33 
     34 from autotest_lib.client.common_lib import enum
     35 from autotest_lib.server.cros.dynamic_suite import frontend_wrappers
     36 
     37 CHROMEOS_LABS = enum.Enum('OysterBay', 'Atlantis', 'Chaos', 'Destiny', start_value=1)
     38 HOST_REGX = 'chromeos(\d+)(-row(\d+))*-rack(\d+)-host(\d+)'
     39 DeviceHostname = collections.namedtuple(
     40         'DeviceHostname', ['lab', 'row', 'rack', 'host'])
     41 
     42 
     43 class 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 
     97 class 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 
    133 class 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 
    199 class 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 
    231 class 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 
    375 def 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 
    397 def 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 
    435 def 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 
    450 if __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