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