Home | History | Annotate | Download | only in scheduler
      1 #pylint: disable-msg=C0111
      2 
      3 import cPickle
      4 import logging
      5 import os
      6 import time
      7 
      8 import common
      9 from autotest_lib.scheduler import drone_utility, email_manager
     10 from autotest_lib.client.bin import local_host
     11 from autotest_lib.client.common_lib import error, global_config, utils
     12 from autotest_lib.client.common_lib.cros.graphite import autotest_stats
     13 
     14 
     15 AUTOTEST_INSTALL_DIR = global_config.global_config.get_config_value('SCHEDULER',
     16                                                  'drone_installation_directory')
     17 DEFAULT_CONTAINER_PATH = global_config.global_config.get_config_value(
     18         'AUTOSERV', 'container_path')
     19 
     20 class DroneUnreachable(Exception):
     21     """The drone is non-sshable."""
     22     pass
     23 
     24 
     25 class _BaseAbstractDrone(object):
     26     """
     27     Attributes:
     28     * allowed_users: set of usernames allowed to use this drone.  if None,
     29             any user can use this drone.
     30     """
     31     def __init__(self, timestamp_remote_calls=True):
     32         """Instantiate an abstract drone.
     33 
     34         @param timestamp_remote_calls: If true, drone_utility is invoked with
     35             the --call_time option and the current time. Currently this is only
     36             used for testing.
     37         """
     38         self._calls = []
     39         self.hostname = None
     40         self.enabled = True
     41         self.max_processes = 0
     42         self.active_processes = 0
     43         self.allowed_users = None
     44         self._autotest_install_dir = AUTOTEST_INSTALL_DIR
     45         self._host = None
     46         self.timestamp_remote_calls = timestamp_remote_calls
     47         # If drone supports server-side packaging. The property support_ssp will
     48         # init self._support_ssp later.
     49         self._support_ssp = None
     50 
     51 
     52     def shutdown(self):
     53         pass
     54 
     55 
     56     @property
     57     def _drone_utility_path(self):
     58         return os.path.join(self._autotest_install_dir,
     59                             'scheduler', 'drone_utility.py')
     60 
     61 
     62     def used_capacity(self):
     63         """Gets the capacity used by this drone
     64 
     65         Returns a tuple of (percentage_full, -max_capacity). This is to aid
     66         direct comparisons, so that a 0/10 drone is considered less heavily
     67         loaded than a 0/2 drone.
     68 
     69         This value should never be used directly. It should only be used in
     70         direct comparisons using the basic comparison operators, or using the
     71         cmp() function.
     72         """
     73         if self.max_processes == 0:
     74             return (1.0, 0)
     75         return (float(self.active_processes) / self.max_processes,
     76                 -self.max_processes)
     77 
     78 
     79     def usable_by(self, user):
     80         if self.allowed_users is None:
     81             return True
     82         return user in self.allowed_users
     83 
     84 
     85     def _execute_calls_impl(self, calls):
     86         if not self._host:
     87             raise ValueError('Drone cannot execute calls without a host.')
     88         drone_utility_cmd = self._drone_utility_path
     89         if self.timestamp_remote_calls:
     90             drone_utility_cmd = '%s --call_time %s' % (
     91                     drone_utility_cmd, time.time())
     92         logging.info("Running drone_utility on %s", self.hostname)
     93         result = self._host.run('python %s' % drone_utility_cmd,
     94                                 stdin=cPickle.dumps(calls), stdout_tee=None,
     95                                 connect_timeout=300)
     96         try:
     97             return cPickle.loads(result.stdout)
     98         except Exception: # cPickle.loads can throw all kinds of exceptions
     99             logging.critical('Invalid response:\n---\n%s\n---', result.stdout)
    100             raise
    101 
    102 
    103     def _execute_calls(self, calls):
    104         autotest_stats.Gauge('drone_execute_call_count').send(
    105                     self.hostname.replace('.', '_'), len(calls))
    106         return_message = self._execute_calls_impl(calls)
    107         for warning in return_message['warnings']:
    108             subject = 'Warning from drone %s' % self.hostname
    109             logging.warning(subject + '\n' + warning)
    110             email_manager.manager.enqueue_notify_email(subject, warning)
    111         return return_message['results']
    112 
    113 
    114     def get_calls(self):
    115         """Returns the calls queued against this drone.
    116 
    117         @return: A list of calls queued against the drone.
    118         """
    119         return self._calls
    120 
    121 
    122     def call(self, method, *args, **kwargs):
    123         return self._execute_calls(
    124             [drone_utility.call(method, *args, **kwargs)])
    125 
    126 
    127     def queue_call(self, method, *args, **kwargs):
    128         self._calls.append(drone_utility.call(method, *args, **kwargs))
    129 
    130 
    131     def clear_call_queue(self):
    132         self._calls = []
    133 
    134 
    135     def execute_queued_calls(self):
    136         if not self._calls:
    137             return
    138         results = self._execute_calls(self._calls)
    139         self.clear_call_queue()
    140         return results
    141 
    142 
    143     def set_autotest_install_dir(self, path):
    144         pass
    145 
    146 
    147     @property
    148     def support_ssp(self):
    149         """Check if the drone supports server-side packaging with container.
    150 
    151         @return: True if the drone supports server-side packaging with container
    152         """
    153         if not self._host:
    154             raise ValueError('Can not determine if drone supports server-side '
    155                              'packaging before host is set.')
    156         if self._support_ssp is None:
    157             try:
    158                 # TODO(crbug.com/471316): We need a better way to check if drone
    159                 # supports container, and install/upgrade base container. The
    160                 # check of base container folder is not reliable and shall be
    161                 # obsoleted once that bug is fixed.
    162                 self._host.run('which lxc-start')
    163                 # Test if base container is setup.
    164                 base_container = os.path.join(DEFAULT_CONTAINER_PATH, 'base')
    165                 # SSP uses privileged containers, sudo access is required. If
    166                 # the process can't run sudo command without password, SSP can't
    167                 # work properly. sudo command option -n will avoid user input.
    168                 # If password is required, the command will fail and raise
    169                 # AutoservRunError exception.
    170                 self._host.run('sudo -n ls "%s"' %  base_container)
    171                 self._support_ssp = True
    172             except (error.AutoservRunError, error.AutotestHostRunError):
    173                 # Local drone raises AutotestHostRunError, while remote drone
    174                 # raises AutoservRunError.
    175                 self._support_ssp = False
    176         return self._support_ssp
    177 
    178 
    179 SiteDrone = utils.import_site_class(
    180    __file__, 'autotest_lib.scheduler.site_drones',
    181    '_SiteAbstractDrone', _BaseAbstractDrone)
    182 
    183 
    184 class _AbstractDrone(SiteDrone):
    185     pass
    186 
    187 
    188 class _LocalDrone(_AbstractDrone):
    189     def __init__(self, timestamp_remote_calls=True):
    190         super(_LocalDrone, self).__init__(
    191                 timestamp_remote_calls=timestamp_remote_calls)
    192         self.hostname = 'localhost'
    193         self._host = local_host.LocalHost()
    194         self._drone_utility = drone_utility.DroneUtility()
    195 
    196 
    197     def send_file_to(self, drone, source_path, destination_path,
    198                      can_fail=False):
    199         if drone.hostname == self.hostname:
    200             self.queue_call('copy_file_or_directory', source_path,
    201                             destination_path)
    202         else:
    203             self.queue_call('send_file_to', drone.hostname, source_path,
    204                             destination_path, can_fail)
    205 
    206 
    207 class _RemoteDrone(_AbstractDrone):
    208     def __init__(self, hostname, timestamp_remote_calls=True):
    209         super(_RemoteDrone, self).__init__(
    210                 timestamp_remote_calls=timestamp_remote_calls)
    211         self.hostname = hostname
    212         self._host = drone_utility.create_host(hostname)
    213         if not self._host.is_up():
    214             logging.error('Drone %s is unpingable, kicking out', hostname)
    215             raise DroneUnreachable
    216 
    217 
    218     def set_autotest_install_dir(self, path):
    219         self._autotest_install_dir = path
    220 
    221 
    222     def shutdown(self):
    223         super(_RemoteDrone, self).shutdown()
    224         self._host.close()
    225 
    226 
    227     def send_file_to(self, drone, source_path, destination_path,
    228                      can_fail=False):
    229         if drone.hostname == self.hostname:
    230             self.queue_call('copy_file_or_directory', source_path,
    231                             destination_path)
    232         elif isinstance(drone, _LocalDrone):
    233             drone.queue_call('get_file_from', self.hostname, source_path,
    234                              destination_path)
    235         else:
    236             self.queue_call('send_file_to', drone.hostname, source_path,
    237                             destination_path, can_fail)
    238 
    239 
    240 def get_drone(hostname):
    241     """
    242     Use this factory method to get drone objects.
    243     """
    244     if hostname == 'localhost':
    245         return _LocalDrone()
    246     try:
    247         return _RemoteDrone(hostname)
    248     except DroneUnreachable:
    249         return None
    250