#!/usr/bin/env python3
# Copyright 2017 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Custom swarming triggering script.
This script does custom swarming triggering logic, to enable device affinity
for our bots, while lumping all trigger calls under one logical step.
For the perf use case of device affinity, this script now enables soft device
affinity. This means that it tries to smartly allocate jobs to bots based
on what is currently alive and what bot the task was last triggered on,
preferring that last triggered bot if available. If the
--multiple-trigger-configs flag is specified than this script overrides
the soft device affinity functionality in favor of the provided ids.
The algorithm is roughly the following:
Find eligible bots, healthy or not.
* Query swarming for eligible bots based on the dimensions passed in
on the swarming call. Determine their health status based on
is not quarantied and is not is_dead
Of the eligible bots determine what bot id to run the shard on.
(Implementation in _select_config_indices_with_soft_affinity)
* First query swarming for the last task that ran that shard with
given dimensions. Assuming they are returned with most recent first.
* Check if the bot id that ran that task is alive, if so trigger
on that bot again.
* If that bot isn't alive, allocate to another alive bot or if no
other alive bots exist, trigger on the same dead one.
Scripts inheriting must have roughly the same command line interface as
swarming.py trigger. It modifies it in the following ways:
* Intercepts the dump-json argument, and creates its own by combining the
results from each trigger call.
* Intercepts the dimensions from the swarming call and determines what bots
are healthy based on the above device affinity algorithm, and triggers
* Adds a tag to the swarming trigger job with the shard so we know the last
bot that ran this shard.
This script is normally called from the swarming recipe module in tools/build.
"""
import argparse
import copy
import os
import subprocess
import sys
import logging
import random
import base_test_triggerer
SRC_DIR = os.path.dirname(
os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(os.path.join(SRC_DIR, 'tools', 'perf'))
# //tools/perf imports.
import generate_perf_sharding
from core import bot_platforms
class Bot(object): # pylint: disable=useless-object-inheritance
"""Eligible bots to run the task."""
def __init__(self, bot_id, is_alive):
self._bot_id = bot_id
self._is_alive = is_alive
def id(self):
return self._bot_id
def is_alive(self):
return self._is_alive
def as_json_config(self):
return {'id': self._bot_id}
class PerfDeviceTriggerer(base_test_triggerer.BaseTestTriggerer):
def __init__(self, args, swarming_args):
# pylint: disable=super-with-arguments
super(PerfDeviceTriggerer, self).__init__()
# pylint: enable=super-with-arguments
self._sharded_query_failed = False
if not args.multiple_trigger_configs:
# Represents the list of current dimensions requested
# by the parent swarming job.
self._dimensions = self._get_swarming_dimensions(swarming_args)
# Store what swarming server we need and whether or not we need
# to send down authentication with it
self._swarming_server = self._get_swarming_server(swarming_args)
# Map of all existing bots in swarming that satisfy the current
# set of dimensions indexed by bot id.
# Note: this assumes perf bot dimensions are unique between
# configurations.
self._eligible_bots_by_ids = (
self._query_swarming_for_eligible_bot_configs(self._dimensions))
if args.multiple_dimension_script_verbose:
logging.basicConfig(level=logging.DEBUG)
def generate_shard_map(self, args, buildername, selected_config):
shard_map = None
num_of_shards = len(selected_config)
builder = bot_platforms.find_bot_platform(buildername)
if args.use_dynamic_shards and builder and num_of_shards:
logging.info(
'Generating dynamic shardmap for builder: %s with %d shards',
buildername, num_of_shards)
shard_map = generate_perf_sharding.GenerateShardMap(
builder=builder, num_of_shards=num_of_shards)
for shard_index, bot_index in selected_config:
bot_id = self._bot_configs[bot_index]['id']
shard_map['extra_infos']['bot #%s' % shard_index] = bot_id
return shard_map
def append_additional_args(self, args, shard_index):
# Append a tag to the swarming task with the shard number
# so we can query for the last bot that ran a specific shard.
tag = 'shard:%d' % shard_index
shard_tag = ['--tag', tag]
# Need to append this before the dash if present so it gets fed to
# the swarming task itself.
if '--' in args:
dash_ind = args.index('--')
return args[:dash_ind] + shard_tag + args[dash_ind:]
return args + shard_tag
def parse_bot_configs(self, args):
if args.multiple_trigger_configs:
# pylint: disable=super-with-arguments
super(PerfDeviceTriggerer, self).parse_bot_configs(args)
# pylint: enable=super-with-arguments
else:
self._bot_configs = []
# For each eligible bot, append the dimension
# to the eligible bot_configs
for _, bot in self._eligible_bots_by_ids.items():
self._bot_configs.append(bot.as_json_config())
def select_config_indices(self, args):
if args.multiple_trigger_configs:
configs = []
# If specific bot ids were passed in, we want to trigger a job for
# every valid config regardless of health status since
# each config represents exactly one bot in the perf swarming pool.
for index in range(len(self.indices_to_trigger(args))):
configs.append((index, index))
if args.use_dynamic_shards:
return self._select_config_indices_with_dynamic_sharding()
return self._select_config_indices_with_soft_affinity(args)
def _select_config_indices_with_dynamic_sharding(self):
alive_bot_ids = [
bot_id for bot_id, b in self._eligible_bots_by_ids.items()
if b.is_alive()
]
trigger_count = len(alive_bot_ids)
indexes = list(range(trigger_count))
random.shuffle(indexes)
selected_config = [(indexes[i],
self._find_bot_config_index(alive_bot_ids[i]))
for i in range(trigger_count)]
selected_config.sort()
for shard_index, bot_index in selected_config:
logging.info('Shard %d\n\tBot: %s', shard_index,
self._bot_configs[bot_index]['id'])
return selected_config
def _select_config_indices_with_soft_affinity(self, args):
trigger_count = len(self.indices_to_trigger(args))
# First make sure the number of shards doesn't exceed the
# number of eligible bots. This means there is a config error somewhere.
if trigger_count > len(self._eligible_bots_by_ids):
self._print_device_affinity_info({}, {}, self._eligible_bots_by_ids,
trigger_count)
raise ValueError(
'Not enough available machines exist in swarming '
'pool. Shards requested (%d) exceeds available bots '
'(%d).' % (trigger_count, len(self._eligible_bots_by_ids)))
shard_to_bot_assignment_map = {}
unallocated_bots_by_ids = copy.deepcopy(self._eligible_bots_by_ids)
for shard_index in self.indices_to_trigger(args):
bot_id = self._query_swarming_for_last_shard_id(shard_index)
if bot_id and bot_id in unallocated_bots_by_ids:
bot = unallocated_bots_by_ids[bot_id]
shard_to_bot_assignment_map[shard_index] = bot
unallocated_bots_by_ids.pop(bot_id)
else:
shard_to_bot_assignment_map[shard_index] = None
# Maintain the current map for debugging purposes
existing_shard_bot_to_shard_map = copy.deepcopy(
shard_to_bot_assignment_map)
# Now create sets of remaining healthy and bad bots
unallocated_healthy_bots = {
b
for b in unallocated_bots_by_ids.values() if b.is_alive()
}
unallocated_bad_bots = {
b
for b in unallocated_bots_by_ids.values() if not b.is_alive()
}
# Try assigning healthy bots for new shards first.
for shard_index, bot in sorted(shard_to_bot_assignment_map.items()):
if not bot and unallocated_healthy_bots:
shard_to_bot_assignment_map[shard_index] = \
unallocated_healthy_bots.pop()
logging.info('First time shard %d has been triggered',
shard_index)
elif not bot:
shard_to_bot_assignment_map[
shard_index] = unallocated_bad_bots.pop()
# Handle the rest of shards that were assigned dead bots:
for shard_index, bot in sorted(shard_to_bot_assignment_map.items()):
if not bot.is_alive() and unallocated_healthy_bots:
dead_bot = bot
healthy_bot = unallocated_healthy_bots.pop()
shard_to_bot_assignment_map[shard_index] = healthy_bot
logging.info(
'Device affinity broken for shard #%d. bot %s is dead,'
' new mapping to bot %s', shard_index, dead_bot.id(),
healthy_bot.id())
# Now populate the indices into the bot_configs array
selected_configs = []
for shard_index in self.indices_to_trigger(args):
selected_configs.append(
(shard_index,
self._find_bot_config_index(
shard_to_bot_assignment_map[shard_index].id())))
self._print_device_affinity_info(shard_to_bot_assignment_map,
existing_shard_bot_to_shard_map,
self._eligible_bots_by_ids,
trigger_count)
return selected_configs
def _print_device_affinity_info(self, new_map, existing_map, health_map,
num_shards):
logging.info('')
for shard_index in range(num_shards):
existing = existing_map.get(shard_index, None)
new = new_map.get(shard_index, None)
existing_id = ''
if existing:
existing_id = existing.id()
new_id = ''
if new:
new_id = new.id()
logging.info('Shard %d\n\tprevious: %s\n\tnew: %s', shard_index,
existing_id, new_id)
healthy_bots = []
dead_bots = []
for _, b in health_map.items():
if b.is_alive():
healthy_bots.append(b.id())
else:
dead_bots.append(b.id())
logging.info('Shards needed: %d', num_shards)
logging.info('Total bots (dead + healthy): %d',
len(dead_bots) + len(healthy_bots))
logging.info('Healthy bots, %d: %s', len(healthy_bots), healthy_bots)
logging.info('Dead Bots, %d: %s', len(dead_bots), dead_bots)
logging.info('')
def _query_swarming_for_eligible_bot_configs(self, dimensions):
"""Query Swarming to figure out which bots are available.
Returns: a dictionary in which the keys are the bot id and
the values are Bot object that indicate the health status
of the bots.
"""
query_result = self.list_bots(dimensions, server=self._swarming_server)
perf_bots = {}
for bot in query_result:
# Device maintenance is usually quick, and we can wait for it to
# finish. However, if the device is too hot, it can take a long time
# for it to cool down, so check for 'Device temperature' in
# maintenance_msg.
alive = (not bot.get('is_dead') and not bot.get('quarantined')
and 'Device temperature' not in bot.get(
'maintenance_msg', ''))
perf_bots[bot['bot_id']] = Bot(bot['bot_id'], alive)
return perf_bots
def _find_bot_config_index(self, bot_id):
# Find the index into the bot_config map that
# maps to the bot id in question
for i, dimensions in enumerate(self._bot_configs):
if dimensions['id'] == bot_id:
return i
return None
def _query_swarming_for_last_shard_id(self, shard_index):
"""Per shard, query swarming for the last bot that ran the task.
Example: swarming.py query -S server-url.com --limit 1 \\
'tasks/list?tags=os:Windows&tags=pool:chrome.tests.perf&tags=shard:12'
"""
values = ['%s:%s' % (k, v) for k, v in self._dimensions.items()]
values.sort()
# Append the shard as a tag
values_with_shard = list(values)
values_with_shard.append('%s:%s' % ('shard', str(shard_index)))
values_with_shard.sort()
# TODO(eyaich): For now we are ignoring the state of the returned
# task (ie completed, timed_out, bot_died, etc) as we are just
# answering the question "What bot did we last trigger this shard on?"
# Evaluate if this is the right decision going forward.
# Query for the last task that ran with these dimensions and this shard.
# Query with the shard param first. This will sometimes time out for
# queries we've never done before, so try querying without it if that
# happens.
try:
if not self._sharded_query_failed:
tasks = self.list_tasks(values_with_shard,
limit='1',
server=self._swarming_server)
except subprocess.CalledProcessError:
self._sharded_query_failed = True
if self._sharded_query_failed:
tasks = self.list_tasks(values,
limit='1',
server=self._swarming_server)
if tasks:
# We queried with a limit of 1 so we could only get back
# the most recent which is what we care about.
task = tasks[0]
if 'bot_id' in task:
return task['bot_id']
for tag in task['tags']:
if tag.startswith('id:'):
return tag[len('id:'):]
# No eligible shard for this bot
return None
def _get_swarming_dimensions(self, args):
dimensions = {}
for i in range(len(args) - 2):
if args[i] == '--dimension':
dimensions[args[i + 1]] = args[i + 2]
return dimensions
# pylint: disable=inconsistent-return-statements
def _get_swarming_server(self, args):
for i, argument in enumerate(args):
if '--swarming' in argument:
server = args[i + 1]
slashes_index = server.index('//') + 2
# Strip out the protocol
return server[slashes_index:]
# pylint: enable=inconsistent-return-statements
def main():
logging.basicConfig(level=logging.INFO,
format='(%(levelname)s) %(asctime)s pid=%(process)d'
' %(module)s.%(funcName)s:%(lineno)d %(message)s')
# Setup args for common contract of base class
parser = base_test_triggerer.BaseTestTriggerer.setup_parser_contract(
argparse.ArgumentParser(description=__doc__))
parser.add_argument('--use-dynamic-shards',
action='store_true',
required=False,
help='Ignore --shards and the existing shard map. Will '
'generate a shard map at run time and use as much '
'device as possible.')
args, remaining = parser.parse_known_args()
triggerer = PerfDeviceTriggerer(args, remaining)
return triggerer.trigger_tasks(args, remaining)
if __name__ == '__main__':
sys.exit(main())