perf_device_trigger.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401
  1. #!/usr/bin/env python3
  2. # Copyright 2017 The Chromium 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. """Custom swarming triggering script.
  6. This script does custom swarming triggering logic, to enable device affinity
  7. for our bots, while lumping all trigger calls under one logical step.
  8. For the perf use case of device affinity, this script now enables soft device
  9. affinity. This means that it tries to smartly allocate jobs to bots based
  10. on what is currently alive and what bot the task was last triggered on,
  11. preferring that last triggered bot if available. If the
  12. --multiple-trigger-configs flag is specified than this script overrides
  13. the soft device affinity functionality in favor of the provided ids.
  14. The algorithm is roughly the following:
  15. Find eligible bots, healthy or not.
  16. * Query swarming for eligible bots based on the dimensions passed in
  17. on the swarming call. Determine their health status based on
  18. is not quarantied and is not is_dead
  19. Of the eligible bots determine what bot id to run the shard on.
  20. (Implementation in _select_config_indices_with_soft_affinity)
  21. * First query swarming for the last task that ran that shard with
  22. given dimensions. Assuming they are returned with most recent first.
  23. * Check if the bot id that ran that task is alive, if so trigger
  24. on that bot again.
  25. * If that bot isn't alive, allocate to another alive bot or if no
  26. other alive bots exist, trigger on the same dead one.
  27. Scripts inheriting must have roughly the same command line interface as
  28. swarming.py trigger. It modifies it in the following ways:
  29. * Intercepts the dump-json argument, and creates its own by combining the
  30. results from each trigger call.
  31. * Intercepts the dimensions from the swarming call and determines what bots
  32. are healthy based on the above device affinity algorithm, and triggers
  33. * Adds a tag to the swarming trigger job with the shard so we know the last
  34. bot that ran this shard.
  35. This script is normally called from the swarming recipe module in tools/build.
  36. """
  37. from __future__ import print_function
  38. import argparse
  39. import copy
  40. import os
  41. import sys
  42. import logging
  43. import random
  44. import base_test_triggerer
  45. import six
  46. SRC_DIR = os.path.dirname(
  47. os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
  48. sys.path.append(os.path.join(SRC_DIR, 'tools', 'perf'))
  49. import generate_perf_sharding
  50. from core import bot_platforms
  51. class Bot(object): # pylint: disable=useless-object-inheritance
  52. """Eligible bots to run the task."""
  53. def __init__(self, bot_id, is_alive):
  54. self._bot_id = bot_id
  55. self._is_alive = is_alive
  56. def id(self):
  57. return self._bot_id
  58. def is_alive(self):
  59. return self._is_alive
  60. def as_json_config(self):
  61. return {'id': self._bot_id}
  62. class PerfDeviceTriggerer(base_test_triggerer.BaseTestTriggerer):
  63. def __init__(self, args, swarming_args):
  64. # pylint: disable=super-with-arguments
  65. super(PerfDeviceTriggerer, self).__init__()
  66. # pylint: enable=super-with-arguments
  67. self._sharded_query_failed = False
  68. if not args.multiple_trigger_configs:
  69. # Represents the list of current dimensions requested
  70. # by the parent swarming job.
  71. self._dimensions = self._get_swarming_dimensions(swarming_args)
  72. # Store what swarming server we need and whether or not we need
  73. # to send down authentication with it
  74. self._swarming_server = self._get_swarming_server(swarming_args)
  75. # Map of all existing bots in swarming that satisfy the current
  76. # set of dimensions indexed by bot id.
  77. # Note: this assumes perf bot dimensions are unique between
  78. # configurations.
  79. self._eligible_bots_by_ids = (
  80. self._query_swarming_for_eligible_bot_configs(
  81. self._dimensions))
  82. if args.multiple_dimension_script_verbose:
  83. logging.basicConfig(level=logging.DEBUG)
  84. def generate_shard_map(self, args, buildername, selected_config):
  85. shard_map = None
  86. num_of_shards = len(selected_config)
  87. builder = bot_platforms.find_bot_platform(buildername)
  88. if args.use_dynamic_shards and builder and num_of_shards:
  89. logging.info(
  90. 'Generating dynamic shardmap for builder: %s with %d shards',
  91. buildername, num_of_shards)
  92. shard_map = generate_perf_sharding.GenerateShardMap(
  93. builder=builder, num_of_shards=num_of_shards)
  94. for shard_index, bot_index in selected_config:
  95. bot_id = self._bot_configs[bot_index]['id']
  96. shard_map['extra_infos']['bot #%s' % shard_index] = bot_id
  97. return shard_map
  98. def append_additional_args(self, args, shard_index):
  99. # Append a tag to the swarming task with the shard number
  100. # so we can query for the last bot that ran a specific shard.
  101. tag = 'shard:%d' % shard_index
  102. shard_tag = ['--tag', tag]
  103. # Need to append this before the dash if present so it gets fed to
  104. # the swarming task itself.
  105. if '--' in args:
  106. dash_ind = args.index('--')
  107. return args[:dash_ind] + shard_tag + args[dash_ind:]
  108. return args + shard_tag
  109. def parse_bot_configs(self, args):
  110. if args.multiple_trigger_configs:
  111. # pylint: disable=super-with-arguments
  112. super(PerfDeviceTriggerer, self).parse_bot_configs(args)
  113. # pylint: enable=super-with-arguments
  114. else:
  115. self._bot_configs = []
  116. # For each eligible bot, append the dimension
  117. # to the eligible bot_configs
  118. for _, bot in self._eligible_bots_by_ids.items():
  119. self._bot_configs.append(bot.as_json_config())
  120. def select_config_indices(self, args):
  121. if args.multiple_trigger_configs:
  122. configs = []
  123. # If specific bot ids were passed in, we want to trigger a job for
  124. # every valid config regardless of health status since
  125. # each config represents exactly one bot in the perf swarming pool.
  126. for index in range(len(self.indices_to_trigger(args))):
  127. configs.append((index, index))
  128. if args.use_dynamic_shards:
  129. return self._select_config_indices_with_dynamic_sharding()
  130. return self._select_config_indices_with_soft_affinity(args)
  131. def _select_config_indices_with_dynamic_sharding(self):
  132. alive_bot_ids = [
  133. bot_id for bot_id, b in self._eligible_bots_by_ids.items()
  134. if b.is_alive()
  135. ]
  136. trigger_count = len(alive_bot_ids)
  137. indexes = list(range(trigger_count))
  138. random.shuffle(indexes)
  139. selected_config = [(indexes[i],
  140. self._find_bot_config_index(alive_bot_ids[i]))
  141. for i in range(trigger_count)]
  142. selected_config.sort()
  143. for shard_index, bot_index in selected_config:
  144. logging.info('Shard %d\n\tBot: %s', shard_index,
  145. self._bot_configs[bot_index]['id'])
  146. return selected_config
  147. def _select_config_indices_with_soft_affinity(self, args):
  148. trigger_count = len(self.indices_to_trigger(args))
  149. # First make sure the number of shards doesn't exceed the
  150. # number of eligible bots. This means there is a config error somewhere.
  151. if trigger_count > len(self._eligible_bots_by_ids):
  152. self._print_device_affinity_info({}, {},
  153. self._eligible_bots_by_ids,
  154. trigger_count)
  155. raise ValueError(
  156. 'Not enough available machines exist in swarming '
  157. 'pool. Shards requested (%d) exceeds available bots '
  158. '(%d).' % (trigger_count, len(self._eligible_bots_by_ids)))
  159. shard_to_bot_assignment_map = {}
  160. unallocated_bots_by_ids = copy.deepcopy(self._eligible_bots_by_ids)
  161. for shard_index in self.indices_to_trigger(args):
  162. bot_id = self._query_swarming_for_last_shard_id(shard_index)
  163. if bot_id and bot_id in unallocated_bots_by_ids:
  164. bot = unallocated_bots_by_ids[bot_id]
  165. shard_to_bot_assignment_map[shard_index] = bot
  166. unallocated_bots_by_ids.pop(bot_id)
  167. else:
  168. shard_to_bot_assignment_map[shard_index] = None
  169. # Maintain the current map for debugging purposes
  170. existing_shard_bot_to_shard_map = copy.deepcopy(
  171. shard_to_bot_assignment_map)
  172. # Now create sets of remaining healthy and bad bots
  173. unallocated_healthy_bots = {
  174. b
  175. for b in unallocated_bots_by_ids.values() if b.is_alive()
  176. }
  177. unallocated_bad_bots = {
  178. b
  179. for b in unallocated_bots_by_ids.values() if not b.is_alive()
  180. }
  181. # Try assigning healthy bots for new shards first.
  182. for shard_index, bot in sorted(
  183. shard_to_bot_assignment_map.items()):
  184. if not bot and unallocated_healthy_bots:
  185. shard_to_bot_assignment_map[shard_index] = \
  186. unallocated_healthy_bots.pop()
  187. logging.info('First time shard %d has been triggered',
  188. shard_index)
  189. elif not bot:
  190. shard_to_bot_assignment_map[
  191. shard_index] = unallocated_bad_bots.pop()
  192. # Handle the rest of shards that were assigned dead bots:
  193. for shard_index, bot in sorted(
  194. shard_to_bot_assignment_map.items()):
  195. if not bot.is_alive() and unallocated_healthy_bots:
  196. dead_bot = bot
  197. healthy_bot = unallocated_healthy_bots.pop()
  198. shard_to_bot_assignment_map[shard_index] = healthy_bot
  199. logging.info(
  200. 'Device affinity broken for shard #%d. bot %s is dead,'
  201. ' new mapping to bot %s', shard_index, dead_bot.id(),
  202. healthy_bot.id())
  203. # Now populate the indices into the bot_configs array
  204. selected_configs = []
  205. for shard_index in self.indices_to_trigger(args):
  206. selected_configs.append(
  207. (shard_index,
  208. self._find_bot_config_index(
  209. shard_to_bot_assignment_map[shard_index].id())))
  210. self._print_device_affinity_info(shard_to_bot_assignment_map,
  211. existing_shard_bot_to_shard_map,
  212. self._eligible_bots_by_ids,
  213. trigger_count)
  214. return selected_configs
  215. def _print_device_affinity_info(self, new_map, existing_map, health_map,
  216. num_shards):
  217. logging.info('')
  218. for shard_index in range(num_shards):
  219. existing = existing_map.get(shard_index, None)
  220. new = new_map.get(shard_index, None)
  221. existing_id = ''
  222. if existing:
  223. existing_id = existing.id()
  224. new_id = ''
  225. if new:
  226. new_id = new.id()
  227. logging.info('Shard %d\n\tprevious: %s\n\tnew: %s', shard_index,
  228. existing_id, new_id)
  229. healthy_bots = []
  230. dead_bots = []
  231. for _, b in health_map.items():
  232. if b.is_alive():
  233. healthy_bots.append(b.id())
  234. else:
  235. dead_bots.append(b.id())
  236. logging.info('Shards needed: %d', num_shards)
  237. logging.info('Total bots (dead + healthy): %d',
  238. len(dead_bots) + len(healthy_bots))
  239. logging.info('Healthy bots, %d: %s', len(healthy_bots), healthy_bots)
  240. logging.info('Dead Bots, %d: %s', len(dead_bots), dead_bots)
  241. logging.info('')
  242. def _query_swarming_for_eligible_bot_configs(self, dimensions):
  243. """Query Swarming to figure out which bots are available.
  244. Returns: a dictionary in which the keys are the bot id and
  245. the values are Bot object that indicate the health status
  246. of the bots.
  247. """
  248. query_result = self.list_bots(dimensions,
  249. server=self._swarming_server)
  250. perf_bots = {}
  251. for bot in query_result:
  252. # Device maintenance is usually quick, and we can wait for it to
  253. # finish. However, if the device is too hot, it can take a long time
  254. # for it to cool down, so check for 'Device temperature' in
  255. # maintenance_msg.
  256. alive = (not bot.get('is_dead') and not bot.get('quarantined')
  257. and 'Device temperature' not in bot.get(
  258. 'maintenance_msg', ''))
  259. perf_bots[bot['bot_id']] = Bot(bot['bot_id'], alive)
  260. return perf_bots
  261. def _find_bot_config_index(self, bot_id):
  262. # Find the index into the bot_config map that
  263. # maps to the bot id in question
  264. for i, dimensions in enumerate(self._bot_configs):
  265. if dimensions['id'] == bot_id:
  266. return i
  267. return None
  268. def _query_swarming_for_last_shard_id(self, shard_index):
  269. """Per shard, query swarming for the last bot that ran the task.
  270. Example: swarming.py query -S server-url.com --limit 1 \\
  271. 'tasks/list?tags=os:Windows&tags=pool:chrome.tests.perf&tags=shard:12'
  272. """
  273. values = ['%s:%s' % (k, v) for k, v in self._dimensions.items()]
  274. values.sort()
  275. # Append the shard as a tag
  276. values_with_shard = list(values)
  277. values_with_shard.append('%s:%s' % ('shard', str(shard_index)))
  278. values_with_shard.sort()
  279. # TODO(eyaich): For now we are ignoring the state of the returned
  280. # task (ie completed, timed_out, bot_died, etc) as we are just
  281. # answering the question "What bot did we last trigger this shard on?"
  282. # Evaluate if this is the right decision going forward.
  283. # Query for the last task that ran with these dimensions and this shard.
  284. # Query with the shard param first. This will sometimes time out for
  285. # queries we've never done before, so try querying without it if that
  286. # happens.
  287. try:
  288. if not self._sharded_query_failed:
  289. tasks = self.list_tasks(values_with_shard,
  290. limit='1',
  291. server=self._swarming_server)
  292. except Exception:
  293. self._sharded_query_failed = True
  294. if self._sharded_query_failed:
  295. tasks = self.list_tasks(values,
  296. limit='1',
  297. server=self._swarming_server)
  298. if tasks:
  299. # We queried with a limit of 1 so we could only get back
  300. # the most recent which is what we care about.
  301. task = tasks[0]
  302. if 'bot_id' in task:
  303. return task['bot_id']
  304. for tag in task['tags']:
  305. if tag.startswith('id:'):
  306. return tag[len('id:'):]
  307. # No eligible shard for this bot
  308. return None
  309. def _get_swarming_dimensions(self, args):
  310. dimensions = {}
  311. for i in range(len(args) - 2):
  312. if args[i] == '--dimension':
  313. dimensions[args[i + 1]] = args[i + 2]
  314. return dimensions
  315. # pylint: disable=inconsistent-return-statements
  316. def _get_swarming_server(self, args):
  317. for i in range(len(args)):
  318. if '--swarming' in args[i]:
  319. server = args[i + 1]
  320. slashes_index = server.index('//') + 2
  321. # Strip out the protocol
  322. return server[slashes_index:]
  323. # pylint: enable=inconsistent-return-statements
  324. def main():
  325. logging.basicConfig(level=logging.INFO,
  326. format='(%(levelname)s) %(asctime)s pid=%(process)d'
  327. ' %(module)s.%(funcName)s:%(lineno)d %(message)s')
  328. # Setup args for common contract of base class
  329. parser = base_test_triggerer.BaseTestTriggerer.setup_parser_contract(
  330. argparse.ArgumentParser(description=__doc__))
  331. parser.add_argument(
  332. '--use-dynamic-shards',
  333. action='store_true',
  334. required=False,
  335. help='Ignore --shards and the existing shard map. Will '
  336. 'generate a shard map at run time and use as much '
  337. 'device as possible.')
  338. args, remaining = parser.parse_known_args()
  339. triggerer = PerfDeviceTriggerer(args, remaining)
  340. return triggerer.trigger_tasks(args, remaining)
  341. if __name__ == '__main__':
  342. sys.exit(main())