xref: /aosp_15_r20/external/autotest/site_utils/rpm_control_system/rpm_client.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1#!/usr/bin/python3
2# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6from __future__ import division
7from __future__ import print_function
8
9import argparse
10import logging
11import sys
12import six.moves.xmlrpc_client
13
14import common
15
16from autotest_lib.client.common_lib.cros import retry
17from autotest_lib.site_utils.rpm_control_system import rpm_constants
18
19try:
20    from autotest_lib.utils.frozen_chromite.lib import metrics
21except ImportError:
22    from autotest_lib.client.bin.utils import metrics_mock as metrics
23
24
25class RemotePowerException(Exception):
26    """This is raised when we fail to set the state of the device's outlet."""
27    pass
28
29
30def set_power(host,
31              new_state,
32              timeout_mins=rpm_constants.RPM_CALL_TIMEOUT_MINS):
33    """Sends the power state change request to the RPM Infrastructure.
34
35    @param host: A CrosHost or ServoHost instance.
36    @param new_state: State we want to set the power outlet to.
37    """
38    # servo V3 is handled differently from the rest.
39    # The best heuristic we have to determine servo V3 is the hostname.
40    if host.hostname.endswith('servo'):
41        args_tuple = (host.hostname, new_state)
42    else:
43        info = host.host_info_store.get()
44        try:
45            args_tuple = (
46                    host.hostname,
47                    info.attributes[rpm_constants.POWERUNIT_HOSTNAME_KEY],
48                    info.attributes[rpm_constants.POWERUNIT_OUTLET_KEY],
49                    info.attributes.get(rpm_constants.HYDRA_HOSTNAME_KEY),
50                    new_state)
51        except KeyError as e:
52            logging.warning('Powerunit information not found. Missing:'
53                            ' %s in host_info_store.', e)
54            raise RemotePowerException('Remote power control is not applicable'
55                                       ' for %s, it could be either RPM is not'
56                                       ' supported on the rack or powerunit'
57                                       ' attributes is not configured in'
58                                       ' inventory.' % host.hostname)
59    _set_power(args_tuple, timeout_mins)
60
61
62def _set_power(args_tuple, timeout_mins=rpm_constants.RPM_CALL_TIMEOUT_MINS):
63    """Sends the power state change request to the RPM Infrastructure.
64
65    @param args_tuple: A args tuple for rpc call. See example below:
66        (hostname, powerunit_hostname, outlet, hydra_hostname, new_state)
67    """
68    client = six.moves.xmlrpc_client.ServerProxy(
69            rpm_constants.RPM_FRONTEND_URI, verbose=False, allow_none=True)
70    timeout = None
71    result = None
72    endpoint = (client.set_power_via_poe if len(args_tuple) == 2
73                else client.set_power_via_rpm)
74    try:
75        timeout, result = retry.timeout(endpoint,
76                                        args=args_tuple,
77                                        timeout_sec=timeout_mins * 60,
78                                        default_result=False)
79    except Exception as e:
80        logging.exception(e)
81        raise RemotePowerException('Client call exception (%s): %s' %
82                                   (rpm_constants.RPM_FRONTEND_URI, e))
83    if timeout:
84        raise RemotePowerException(
85                'Call to RPM Infrastructure timed out (%s).' %
86                rpm_constants.RPM_FRONTEND_URI)
87    if not result:
88        error_msg = ('Failed to change outlet status for host: %s to '
89                     'state: %s.' % (args_tuple[0], args_tuple[-1]))
90        logging.error(error_msg)
91        if len(args_tuple) > 2:
92            # Collect failure metrics if we set power via rpm.
93            _send_rpm_failure_metrics(args_tuple[0], args_tuple[1],
94                                      args_tuple[2])
95        raise RemotePowerException(error_msg)
96
97
98def _send_rpm_failure_metrics(hostname, rpm_host, outlet):
99    metrics_fields = {
100            'hostname': hostname,
101            'rpm_host': rpm_host,
102            'outlet': outlet
103    }
104    metrics.Counter('chromeos/autotest/rpm/rpm_failure2').increment(
105            fields=metrics_fields)
106
107
108def parse_options():
109    """Parse the user supplied options."""
110    parser = argparse.ArgumentParser()
111    parser.add_argument('-m', '--machine', dest='machine', required=True,
112                        help='Machine hostname to change outlet state.')
113    parser.add_argument('-s', '--state', dest='state', required=True,
114                        choices=['ON', 'OFF', 'CYCLE'],
115                        help='Power state to set outlet: ON, OFF, CYCLE')
116    parser.add_argument('-p', '--powerunit_hostname', dest='p_hostname',
117                        help='Powerunit hostname of the host.')
118    parser.add_argument('-o', '--outlet', dest='outlet',
119                        help='Outlet of the host.')
120    parser.add_argument('-y', '--hydra_hostname', dest='hydra', default='',
121                        help='Hydra hostname of the host.')
122    parser.add_argument('-d', '--disable_emails', dest='disable_emails',
123                        help='Hours to suspend RPM email notifications.')
124    parser.add_argument('-e', '--enable_emails', dest='enable_emails',
125                        action='store_true',
126                        help='Resume RPM email notifications.')
127    return parser.parse_args()
128
129
130def main():
131    """Entry point for rpm_client script."""
132    options = parse_options()
133    if options.machine.endswith('servo'):
134        args_tuple = (options.machine, options.state)
135    elif not options.p_hostname or not options.outlet:
136        print("Powerunit hostname and outlet info are required for DUT.")
137        return
138    else:
139        args_tuple = (options.machine, options.p_hostname, options.outlet,
140                      options.hydra, options.state)
141    _set_power(args_tuple)
142
143    if options.disable_emails is not None:
144        client = six.moves.xmlrpc_client.ServerProxy(
145                rpm_constants.RPM_FRONTEND_URI, verbose=False)
146        client.suspend_emails(options.disable_emails)
147    if options.enable_emails:
148        client = six.moves.xmlrpc_client.ServerProxy(
149                rpm_constants.RPM_FRONTEND_URI, verbose=False)
150        client.resume_emails()
151
152
153if __name__ == "__main__":
154    sys.exit(main())
155