scheduler/drones.py

   1 import cPickle, os, tempfile, logging
   2 import common
   3 from autotest_lib.scheduler import drone_utility, email_manager
   4 from autotest_lib.client.common_lib import error, global_config
   5
   6
   7 AUTOTEST_INSTALL_DIR = global_config.global_config.get_config_value('SCHEDULER',
   8                                                  'drone_installation_directory')
   9
  10 class DroneUnreachable(Exception):
  11     """The drone is non-sshable."""
  12     pass
  13
  14
  15 class _AbstractDrone(object):
  16     """
  17     Attributes:
  18     * allowed_users: set of usernames allowed to use this drone.  if None,
  19             any user can use this drone.
  20     """
  21     def __init__(self):
  22         self._calls = []
  23         self.hostname = None
  24         self.enabled = True
  25         self.max_processes = 0
  26         self.active_processes = 0
  27         self.allowed_users = None
  28
  29
  30     def shutdown(self):
  31         pass
  32
  33
  34     def used_capacity(self):
  35         """Gets the capacity used by this drone
  36
  37         Returns a tuple of (percentage_full, -max_capacity). This is to aid
  38         direct comparisons, so that a 0/10 drone is considered less heavily
  39         loaded than a 0/2 drone.
  40
  41         This value should never be used directly. It should only be used in
  42         direct comparisons using the basic comparison operators, or using the
  43         cmp() function.
  44         """
  45         if self.max_processes == 0:
  46             return (1.0, 0)
  47         return (float(self.active_processes) / self.max_processes,
  48                 -self.max_processes)
  49
  50
  51     def usable_by(self, user):
  52         if self.allowed_users is None:
  53             return True
  54         return user in self.allowed_users
  55
  56
  57     def _execute_calls_impl(self, calls):
  58         raise NotImplementedError
  59
  60
  61     def _execute_calls(self, calls):
  62         return_message = self._execute_calls_impl(calls)
  63         for warning in return_message['warnings']:
  64             subject = 'Warning from drone %s' % self.hostname
  65             logging.warn(subject + '\n' + warning)
  66             email_manager.manager.enqueue_notify_email(subject, warning)
  67         return return_message['results']
  68
  69
  70     def call(self, method, *args, **kwargs):
  71         return self._execute_calls(
  72             [drone_utility.call(method, *args, **kwargs)])
  73
  74
  75     def queue_call(self, method, *args, **kwargs):
  76         self._calls.append(drone_utility.call(method, *args, **kwargs))
  77
  78     def clear_call_queue(self):
  79         self._calls = []
  80
  81
  82     def execute_queued_calls(self):
  83         if not self._calls:
  84             return
  85         self._execute_calls(self._calls)
  86         self.clear_call_queue()
  87
  88
  89     def set_autotest_install_dir(self, path):
  90         pass
  91
  92
  93 class _LocalDrone(_AbstractDrone):
  94     def __init__(self):
  95         super(_LocalDrone, self).__init__()
  96         self.hostname = 'localhost'
  97         self._drone_utility = drone_utility.DroneUtility()
  98
  99
 100     def _execute_calls_impl(self, calls):
 101         return self._drone_utility.execute_calls(calls)
 102
 103
 104     def send_file_to(self, drone, source_path, destination_path,
 105                      can_fail=False):
 106         if drone.hostname == self.hostname:
 107             self.queue_call('copy_file_or_directory', source_path,
 108                             destination_path)
 109         else:
 110             self.queue_call('send_file_to', drone.hostname, source_path,
 111                             destination_path, can_fail)
 112
 113
 114 class _RemoteDrone(_AbstractDrone):
 115     def __init__(self, hostname):
 116         super(_RemoteDrone, self).__init__()
 117         self.hostname = hostname
 118         self._host = drone_utility.create_host(hostname)
 119         if not self._host.is_up():
 120             logging.error('Drone %s is unpingable, kicking out', hostname)
 121             raise DroneUnreachable
 122         self._autotest_install_dir = AUTOTEST_INSTALL_DIR
 123
 124
 125     @property
 126     def _drone_utility_path(self):
 127         return os.path.join(self._autotest_install_dir,
 128                             'scheduler', 'drone_utility.py')
 129
 130
 131     def set_autotest_install_dir(self, path):
 132         self._autotest_install_dir = path
 133
 134
 135     def shutdown(self):
 136         super(_RemoteDrone, self).shutdown()
 137         self._host.close()
 138
 139
 140     def _execute_calls_impl(self, calls):
 141         logging.info("Running drone_utility on %s", self.hostname)
 142         result = self._host.run('python %s' % self._drone_utility_path,
 143                                 stdin=cPickle.dumps(calls), stdout_tee=None,
 144                                 connect_timeout=300)
 145         try:
 146             return cPickle.loads(result.stdout)
 147         except Exception: # cPickle.loads can throw all kinds of exceptions
 148             logging.critical('Invalid response:\n---\n%s\n---', result.stdout)
 149             raise
 150
 151
 152     def send_file_to(self, drone, source_path, destination_path,
 153                      can_fail=False):
 154         if drone.hostname == self.hostname:
 155             self.queue_call('copy_file_or_directory', source_path,
 156                             destination_path)
 157         elif isinstance(drone, _LocalDrone):
 158             drone.queue_call('get_file_from', self.hostname, source_path,
 159                              destination_path)
 160         else:
 161             self.queue_call('send_file_to', drone.hostname, source_path,
 162                             destination_path, can_fail)
 163
 164
 165 def get_drone(hostname):
 166     """
 167     Use this factory method to get drone objects.
 168     """
 169     if hostname == 'localhost':
 170         return _LocalDrone()
 171     try:
 172         return _RemoteDrone(hostname)
 173     except DroneUnreachable:
 174         return None