Factoring out base class for trigger_multiple_dimensions.py.
This will be used by perf team for intercepting trigger call
to trigger on specific bot ids.
Bug: 758630
Cq-Include-Trybots: master.tryserver.chromium.android:android_optional_gpu_tests_rel;master.tryserver.chromium.linux:linux_optional_gpu_tests_rel;master.tryserver.chromium.mac:mac_optional_gpu_tests_rel;master.tryserver.chromium.win:win_optional_gpu_tests_rel
Change-Id: Ib5152748169ccd50d01dccc1f61979bcb967a481
Reviewed-on: https://chromium-review.googlesource.com/896365
Reviewed-by: Emily Hanley <[email protected]>
Reviewed-by: Ashley Enstad <[email protected]>
Reviewed-by: Kenneth Russell <[email protected]>
Commit-Queue: Emily Hanley <[email protected]>
Cr-Commit-Position: refs/heads/master@{#535001}diff --git a/testing/trigger_scripts/base_test_triggerer.py b/testing/trigger_scripts/base_test_triggerer.py
new file mode 100755
index 0000000..98f19e0
--- /dev/null
+++ b/testing/trigger_scripts/base_test_triggerer.py
@@ -0,0 +1,287 @@
+#!/usr/bin/env python
+# Copyright 2018 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+"""Custom swarming base trigger class.
+
+This base class consolidates custom swarming triggering logic, to allow one bot
+to conceptually span multiple Swarming configurations, while lumping all trigger
+calls under one logical step. It also gives the subclasses the ability to
+define their own logic for pruning the configurations they want to trigger
+jobs on and what configurations to use.
+
+See trigger_multiple_dimensions.py for an example of how to use this base class.
+
+"""
+
+import argparse
+import copy
+import json
+import os
+import random
+import subprocess
+import sys
+import tempfile
+import urllib
+
+
+SRC_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(
+ __file__))))
+
+SWARMING_PY = os.path.join(SRC_DIR, 'tools', 'swarming_client', 'swarming.py')
+
+
+def strip_unicode(obj):
+ """Recursively re-encodes strings as utf-8 inside |obj|. Returns the result.
+ """
+ if isinstance(obj, unicode):
+ return obj.encode('utf-8', 'replace')
+
+ if isinstance(obj, list):
+ return list(map(strip_unicode, obj))
+
+ if isinstance(obj, dict):
+ new_obj = type(obj)(
+ (strip_unicode(k), strip_unicode(v)) for k, v in obj.iteritems() )
+ return new_obj
+
+ return obj
+
+
+class BaseTestTriggerer(object):
+ def __init__(self):
+ self._bot_configs = None
+ self._bot_statuses = []
+ self._total_bots = 0
+
+
+ def modify_args(self, all_args, bot_index, shard_index, total_shards,
+ temp_file):
+ """Modifies the given argument list.
+
+ Specifically, it does the following:
+ * Adds a --dump_json argument, to read in the results of the
+ individual trigger command.
+ * Adds the dimensions associated with the bot config at the given index.
+ * If the number of shards is greater than one, adds --env
+ arguments to set the GTEST_SHARD_INDEX and GTEST_TOTAL_SHARDS
+ environment variables to _shard_index_ and _total_shards_,
+ respectively.
+
+ The arguments are structured like this:
+ <args to swarming.py trigger> -- <args to bot running isolate>
+ This means we have to add arguments to specific locations in the argument
+ list, to either affect the trigger command, or what the bot runs.
+
+ """
+
+
+
+ assert '--' in all_args, (
+ 'Malformed trigger command; -- argument expected but not found')
+ dash_ind = all_args.index('--')
+ bot_args = ['--dump-json', temp_file]
+ if total_shards > 1:
+ bot_args.append('--env')
+ bot_args.append('GTEST_SHARD_INDEX')
+ bot_args.append(str(shard_index))
+ bot_args.append('--env')
+ bot_args.append('GTEST_TOTAL_SHARDS')
+ bot_args.append(str(total_shards))
+ for key, val in sorted(self._bot_configs[bot_index].iteritems()):
+ bot_args.append('--dimension')
+ bot_args.append(key)
+ bot_args.append(val)
+ return self.append_additional_args(
+ all_args[:dash_ind] + bot_args + all_args[dash_ind:])
+
+ def append_additional_args(self, args):
+ """ Gives subclasses ability to append additional args if necessary
+
+ Base class just returns given get."""
+ return args
+
+ def parse_bot_configs(self, args):
+ try:
+ self._bot_configs = strip_unicode(json.loads(
+ args.multiple_trigger_configs))
+ except Exception as e:
+ raise ValueError('Error while parsing JSON from bot config string %s: %s'
+ % (args.multiple_trigger_configs, str(e)))
+ # Validate the input.
+ if not isinstance(self._bot_configs, list):
+ raise ValueError('Bot configurations must be a list, were: %s' %
+ args.multiple_trigger_configs)
+ if len(self._bot_configs) < 1:
+ raise ValueError('Bot configuration list must have at least one entry')
+ if not all(isinstance(entry, dict) for entry in self._bot_configs):
+ raise ValueError('Bot configurations must all be dictionaries')
+
+ def query_swarming_for_bot_configs(self, verbose):
+ # Query Swarming to figure out which bots are available.
+ for config in self._bot_configs:
+ values = []
+ for key, value in sorted(config.iteritems()):
+ values.append(('dimensions', '%s:%s' % (key, value)))
+ # Ignore dead and quarantined bots.
+ values.append(('is_dead', 'FALSE'))
+ values.append(('quarantined', 'FALSE'))
+ query_arg = urllib.urlencode(values)
+
+ temp_file = self.make_temp_file(prefix='base_trigger_dimensions',
+ suffix='.json')
+ try:
+ ret = self.run_swarming(['query',
+ '-S',
+ 'chromium-swarm.appspot.com',
+ '--limit',
+ '0',
+ '--json',
+ temp_file,
+ ('bots/count?%s' % query_arg)],
+ verbose)
+ if ret:
+ raise Exception('Error running swarming.py')
+ with open(temp_file) as fp:
+ query_result = strip_unicode(json.load(fp))
+ # Summarize number of available bots per configuration.
+ count = int(query_result['count'])
+ # Be robust against errors in computation.
+ available = max(0, count - int(query_result['busy']))
+ self._bot_statuses.append({'total': count, 'available': available})
+ if verbose:
+ idx = len(self._bot_statuses) - 1
+ print 'Bot config %d: %s' % (idx, str(self._bot_statuses[idx]))
+ finally:
+ self.delete_temp_file(temp_file)
+ # Sum up the total count of all bots.
+ self._total_bots = sum(x['total'] for x in self._bot_statuses)
+ if verbose:
+ print 'Total bots: %d' % (self._total_bots)
+
+ def remove_swarming_dimension(self, args, dimension):
+ for i in xrange(len(args)):
+ if args[i] == '--dimension' and args[i+1] == dimension:
+ return args[:i] + args[i+3:]
+ return args
+
+ def make_temp_file(self, prefix=None, suffix=None):
+ # This trick of closing the file handle is needed on Windows in order to
+ # make the file writeable.
+ h, temp_file = tempfile.mkstemp(prefix=prefix, suffix=suffix)
+ os.close(h)
+ return temp_file
+
+ def delete_temp_file(self, temp_file):
+ os.remove(temp_file)
+
+ def read_json_from_temp_file(self, temp_file):
+ with open(temp_file) as f:
+ return json.load(f)
+
+ def write_json_to_file(self, merged_json, output_file):
+ with open(output_file, 'w') as f:
+ json.dump(merged_json, f)
+
+ def run_swarming(self, args, verbose):
+ if verbose:
+ print 'Running Swarming with args:'
+ print str(args)
+ return subprocess.call([sys.executable, SWARMING_PY] + args)
+
+ def prune_test_specific_configs(self, args, verbose):
+ # Ability for base class to further prune configs to
+ # run tests on.
+ pass
+
+ def select_config_indices(self, args, verbose):
+ # Main implementation for base class to determine what
+ # configs to trigger jobs on from self._bot_configs.
+ # Returns a list of indices into the self._bot_configs and
+ # len(args.shards) == len(selected_indices).
+ pass
+
+ def trigger_tasks(self, args, remaining):
+ """Triggers tasks for each bot.
+
+ Args:
+ args: Parsed arguments which we need to use.
+ remaining: The remainder of the arguments, which should be passed to
+ swarming.py calls.
+
+ Returns:
+ Exit code for the script.
+ """
+ verbose = args.multiple_dimension_script_verbose
+ self.parse_bot_configs(args)
+ # Prunes config list to the exact set of configurations to trigger jobs on.
+ # This logic is specific to the base class if they want to prune list
+ # further.
+ self.prune_test_specific_configs(args, verbose)
+
+ # In the remaining arguments, find the Swarming dimensions that are
+ # specified by the bot configs and remove them, because for each shard,
+ # we're going to select one of the bot configs and put all of its Swarming
+ # dimensions on the command line.
+ filtered_remaining_args = copy.deepcopy(remaining)
+ for config in self._bot_configs:
+ for k in config.iterkeys():
+ filtered_remaining_args = self.remove_swarming_dimension(
+ filtered_remaining_args, k)
+
+ merged_json = {}
+
+ # Choose selected configs for this run of the test suite.
+ selected_configs = self.select_config_indices(args, verbose)
+ for i in xrange(args.shards):
+ # For each shard that we're going to distribute, do the following:
+ # 1. Pick which bot configuration to use.
+ # 2. Insert that bot configuration's dimensions as command line
+ # arguments, and invoke "swarming.py trigger".
+ bot_index = selected_configs[i]
+ # Holds the results of the swarming.py trigger call.
+ try:
+ json_temp = self.make_temp_file(prefix='base_trigger_dimensions',
+ suffix='.json')
+ args_to_pass = self.modify_args(filtered_remaining_args, bot_index, i,
+ args.shards, json_temp)
+ ret = self.run_swarming(args_to_pass, verbose)
+ if ret:
+ sys.stderr.write('Failed to trigger a task, aborting\n')
+ return ret
+ result_json = self.read_json_from_temp_file(json_temp)
+ if i == 0:
+ # Copy the entire JSON -- in particular, the "request"
+ # dictionary -- from shard 0. "swarming.py collect" uses
+ # some keys from this dictionary, in particular related to
+ # expiration. It also contains useful debugging information.
+ merged_json = copy.deepcopy(result_json)
+ # However, reset the "tasks" entry to an empty dictionary,
+ # which will be handled specially.
+ merged_json['tasks'] = {}
+ for k, v in result_json['tasks'].items():
+ v['shard_index'] = i
+ merged_json['tasks'][k + ':%d:%d' % (i, args.shards)] = v
+ finally:
+ self.delete_temp_file(json_temp)
+ self.write_json_to_file(merged_json, args.dump_json)
+ return 0
+
+
+ def setup_parser_contract(self, parser):
+ parser.add_argument('--multiple-trigger-configs', type=str, required=True,
+ help='The Swarming configurations to trigger tasks on, '
+ 'in the form of a JSON array of dictionaries (these are'
+ ' Swarming dimension_sets). At least one entry in this '
+ 'dictionary is required.')
+ parser.add_argument('--multiple-dimension-script-verbose', type=bool,
+ default=False, help='Turn on verbose logging')
+ parser.add_argument('--dump-json', required=True,
+ help='(Swarming Trigger Script API) Where to dump the'
+ ' resulting json which indicates which tasks were'
+ ' triggered for which shards.')
+ parser.add_argument('--shards', type=int, default=1,
+ help='How many shards to trigger. Duplicated from the'
+ ' `swarming.py trigger` command.')
+ return parser
+
