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
     12 
     13 CONFIG = global_config.global_config
     14 AUTOTEST_INSTALL_DIR = CONFIG.get_config_value('SCHEDULER',
     15                                                'drone_installation_directory')
     16 DEFAULT_CONTAINER_PATH = CONFIG.get_config_value('AUTOSERV', 'container_path')
     17 
     18 SSP_REQUIRED = CONFIG.get_config_value('SCHEDULER', 'exit_on_failed_ssp_setup',
     19                                        default=False)
     20 
     21 class DroneUnreachable(Exception):
     22     """The drone is non-sshable."""
     23     pass
     24 
     25 
     26 class SiteDrone(object):
     27     """
     28     Attributes:
     29     * allowed_users: set of usernames allowed to use this drone.  if None,
     30             any user can use this drone.
     31     """
     32     def __init__(self, timestamp_remote_calls=True):
     33         """Instantiate an abstract drone.
     34 
     35         @param timestamp_remote_calls: If true, drone_utility is invoked with
     36             the --call_time option and the current time. Currently this is only
     37             used for testing.
     38         """
     39         self._calls = []
     40         self.hostname = None
     41         self.enabled = True
     42         self.max_processes = 0
     43         self.active_processes = 0
     44         self.allowed_users = None
     45         self._autotest_install_dir = AUTOTEST_INSTALL_DIR
     46         self._host = None
     47         self.timestamp_remote_calls = timestamp_remote_calls
     48         # If drone supports server-side packaging. The property support_ssp will
     49         # init self._support_ssp later.
     50         self._support_ssp = None
     51         self._processes_to_kill = []
     52 
     53 
     54     def shutdown(self):
     55         pass
     56 
     57 
     58     @property
     59     def _drone_utility_path(self):
     60         return os.path.join(self._autotest_install_dir,
     61                             'scheduler', 'drone_utility.py')
     62 
     63 
     64     def used_capacity(self):
     65         """Gets the capacity used by this drone
     66 
     67         Returns a tuple of (percentage_full, -max_capacity). This is to aid
     68         direct comparisons, so that a 0/10 drone is considered less heavily
     69         loaded than a 0/2 drone.
     70 
     71         This value should never be used directly. It should only be used in
     72         direct comparisons using the basic comparison operators, or using the
     73         cmp() function.
     74         """
     75         if self.max_processes == 0:
     76             return (1.0, 0)
     77         return (float(self.active_processes) / self.max_processes,
     78                 -self.max_processes)
     79 
     80 
     81     def usable_by(self, user):
     82         if self.allowed_users is None:
     83             return True
     84         return user in self.allowed_users
     85 
     86 
     87     def _execute_calls_impl(self, calls):
     88         if not self._host:
     89             raise ValueError('Drone cannot execute calls without a host.')
     90         drone_utility_cmd = self._drone_utility_path
     91         if self.timestamp_remote_calls:
     92             drone_utility_cmd = '%s --call_time %s' % (
     93                     drone_utility_cmd, time.time())
     94         logging.info("Running drone_utility on %s", self.hostname)
     95         result = self._host.run('python %s' % drone_utility_cmd,
     96                                 stdin=cPickle.dumps(calls), stdout_tee=None,
     97                                 connect_timeout=300)
     98         try:
     99             return cPickle.loads(result.stdout)
    100         except Exception: # cPickle.loads can throw all kinds of exceptions
    101             logging.critical('Invalid response:\n---\n%s\n---', result.stdout)
    102             raise
    103 
    104 
    105     def _execute_calls(self, 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         """Execute queued calls.
    137 
    138         If there are any processes queued to kill, kill them then process the
    139         remaining queued up calls.
    140         """
    141         if self._processes_to_kill:
    142             self.queue_call('kill_processes', self._processes_to_kill)
    143         self.clear_processes_to_kill()
    144 
    145         if not self._calls:
    146             return
    147         results = self._execute_calls(self._calls)
    148         self.clear_call_queue()
    149         return results
    150 
    151 
    152     def set_autotest_install_dir(self, path):
    153         pass
    154 
    155 
    156     @property
    157     def support_ssp(self):
    158         """Check if the drone supports server-side packaging with container.
    159 
    160         @return: True if the drone supports server-side packaging with container
    161         """
    162         if not self._host:
    163             raise ValueError('Can not determine if drone supports server-side '
    164                              'packaging before host is set.')
    165         if self._support_ssp is None:
    166             try:
    167                 # TODO(crbug.com/471316): We need a better way to check if drone
    168                 # supports container, and install/upgrade base container. The
    169                 # check of base container folder is not reliable and shall be
    170                 # obsoleted once that bug is fixed.
    171                 self._host.run('which lxc-start')
    172                 # Test if base container is setup.
    173                 base_container_name = CONFIG.get_config_value(
    174                         'AUTOSERV', 'container_base_name')
    175                 base_container = os.path.join(DEFAULT_CONTAINER_PATH,
    176                                               base_container_name)
    177                 # SSP uses privileged containers, sudo access is required. If
    178                 # the process can't run sudo command without password, SSP can't
    179                 # work properly. sudo command option -n will avoid user input.
    180                 # If password is required, the command will fail and raise
    181                 # AutoservRunError exception.
    182                 self._host.run('sudo -n ls "%s"' %  base_container)
    183                 self._support_ssp = True
    184             except (error.AutoservRunError, error.AutotestHostRunError):
    185                 # Local drone raises AutotestHostRunError, while remote drone
    186                 # raises AutoservRunError.
    187                 logging.exception('Drone %s does not support server-side '
    188                                   'packaging.', self.hostname)
    189                 self._support_ssp = False
    190                 if SSP_REQUIRED:
    191                   raise
    192         return self._support_ssp
    193 
    194 
    195     def queue_kill_process(self, process):
    196         """Queue a process to kill/abort.
    197 
    198         @param process: Process to kill/abort.
    199         """
    200         self._processes_to_kill.append(process)
    201 
    202 
    203     def clear_processes_to_kill(self):
    204         """Reset the list of processes to kill for this tick."""
    205         self._processes_to_kill = []
    206 
    207 
    208 class _AbstractDrone(SiteDrone):
    209     pass
    210 
    211 
    212 class _LocalDrone(_AbstractDrone):
    213     def __init__(self, timestamp_remote_calls=True):
    214         super(_LocalDrone, self).__init__(
    215                 timestamp_remote_calls=timestamp_remote_calls)
    216         self.hostname = 'localhost'
    217         self._host = local_host.LocalHost()
    218 
    219 
    220     def send_file_to(self, drone, source_path, destination_path,
    221                      can_fail=False):
    222         if drone.hostname == self.hostname:
    223             self.queue_call('copy_file_or_directory', source_path,
    224                             destination_path)
    225         else:
    226             self.queue_call('send_file_to', drone.hostname, source_path,
    227                             destination_path, can_fail)
    228 
    229 
    230 class _RemoteDrone(_AbstractDrone):
    231     def __init__(self, hostname, timestamp_remote_calls=True):
    232         super(_RemoteDrone, self).__init__(
    233                 timestamp_remote_calls=timestamp_remote_calls)
    234         self.hostname = hostname
    235         self._host = drone_utility.create_host(hostname)
    236         if not self._host.is_up():
    237             logging.error('Drone %s is unpingable, kicking out', hostname)
    238             raise DroneUnreachable
    239 
    240 
    241     def set_autotest_install_dir(self, path):
    242         self._autotest_install_dir = path
    243 
    244 
    245     def shutdown(self):
    246         super(_RemoteDrone, self).shutdown()
    247         self._host.close()
    248 
    249 
    250     def send_file_to(self, drone, source_path, destination_path,
    251                      can_fail=False):
    252         if drone.hostname == self.hostname:
    253             self.queue_call('copy_file_or_directory', source_path,
    254                             destination_path)
    255         elif isinstance(drone, _LocalDrone):
    256             drone.queue_call('get_file_from', self.hostname, source_path,
    257                              destination_path)
    258         else:
    259             self.queue_call('send_file_to', drone.hostname, source_path,
    260                             destination_path, can_fail)
    261 
    262 
    263 def get_drone(hostname):
    264     """
    265     Use this factory method to get drone objects.
    266     """
    267     if hostname == 'localhost':
    268         return _LocalDrone()
    269     try:
    270         return _RemoteDrone(hostname)
    271     except DroneUnreachable:
    272         return None
    273