Add custom perf device trigger script

This will be used by the one buildbot step effort for perf builders.
This CL shouldn't have any effect when it lands.

Bug: 758630
Change-Id: I2e8434a975be2908427e5c6fd0091a5f1490f5bf
Reviewed-on: https://chromium-review.googlesource.com/629962
Commit-Queue: Ned Nguyen <nednguyen@google.com>
Reviewed-by: Emily Hanley <eyaich@chromium.org>
Reviewed-by: Ned Nguyen <nednguyen@google.com>
Reviewed-by: Marc-Antoine Ruel <maruel@chromium.org>
Cr-Commit-Position: refs/heads/master@{#502820}

tools/perf/perf_device_trigger.py

+#!/usr/bin/env python
+# Copyright 2017 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 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.
+
+This script must have roughly the same command line interface as swarming.py
+trigger. It modifies it in the following ways:
+ * Inserts bot id into swarming trigger dimensions, and swarming task arguments.
+ * Intercepts the dump-json argument, and creates its own by combining the
+   results from each trigger call.
+
+This script is normally called from the swarming recipe module in tools/build.
+"""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import tempfile
+
+from core import path_util
+
+
+def get_swarming_py_path():
+  return os.path.join(
+      path_util.GetChromiumSrcDir(), 'tools', 'swarming_client', 'swarming.py')
+
+
+def modify_args(all_args, bot_id, 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 a bot id dimension.
+    * Adds the bot id as an argument to the test, so it knows which tests to
+      run.
+
+  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('--')
+
+  return all_args[:dash_ind] + [
+      '--dump_json', temp_file, '--dimension', 'id', bot_id] + all_args[
+          dash_ind:] + ['--id', bot_id]
+
+
+def trigger_tasks(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.
+  """
+  merged_json = {'tasks': {}}
+
+  # TODO: Do these in parallel
+  for shard_num, bot_id in enumerate(args.bot_id):
+    # Holds the results of the swarming.py trigger call.
+    temp_fd, json_temp = tempfile.mkstemp(prefix='perf_device_trigger')
+    try:
+      args_to_pass = modify_args(remaining[:], bot_id, json_temp)
+
+      ret = subprocess.call(
+          sys.executable, [get_swarming_py_path()] + args_to_pass)
+      if ret:
+        sys.stderr.write('Failed to trigger a task, aborting\n')
+        return ret
+      with open(json_temp) as f:
+        result_json = json.load(f)
+
+      for k, v in result_json['tasks'].items():
+        v['shard_index'] = shard_num
+        merged_json['tasks'][k + ':%d:%d' % (shard_num, len(args.bot_id))] = v
+    finally:
+      os.close(temp_fd)
+      os.remove(json_temp)
+  with open(args.dump_json, 'w') as f:
+    json.dump(merged_json, f)
+  return 0
+
+
+def main():
+  parser = argparse.ArgumentParser(description=__doc__)
+  parser.add_argument('--bot-id', action='append', required=True,
+                      help='Which bot IDs to trigger tasks on. Number of bot'
+                      ' IDs must match the number of shards. This is because'
+                      ' the recipe code uses the --shard argument to determine'
+                      ' how many shards it expects in the output json.')
+  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',
+                      help='How many shards to trigger. Duplicated from the'
+                      ' `swarming.py trigger` command. If passed, will be'
+                      ' checked to ensure there are as many shards as bots'
+                      ' specified by the --bot-id argument. If omitted,'
+                      ' --bot-id will be used to determine the number of shards'
+                      ' to create.')
+  args, remaining = parser.parse_known_args()
+
+  if args.shards and args.shards != len(args.bot_id):
+    raise parser.error(
+        'Number of bots to use must equal number of shards (--shards)')
+  return trigger_tasks(args, remaining)
+
+
+if __name__ == '__main__':
+  sys.exit(main())

tools/perf/perf_device_trigger_unittest.py

+# Copyright 2017 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.
+
+import json
+import os
+import tempfile
+import unittest
+import sys
+
+import mock
+
+import perf_device_trigger
+
+
+class FakeArgs(object):
+  def __init__(self, shards, bot_id, dump_json):
+    self.shards = shards
+    self.bot_id = bot_id
+    self.dump_json = dump_json
+
+class PerfDeviceTriggerUnittest(unittest.TestCase):
+  def testBasic(self):
+    with mock.patch('perf_device_trigger.get_swarming_py_path') as py_path_mock:
+      py_path_mock.return_value = '/path/to/swarming.py'
+      with mock.patch('perf_device_trigger.subprocess.call') as call_mock:
+        call_mock.return_value = 0
+        m = mock.mock_open(read_data=json.dumps({
+            'tasks': {},
+        }))
+        with mock.patch('perf_device_trigger.open', m, create=True):
+          temp_fd, json_temp = tempfile.mkstemp(
+              prefix='perf_device_trigger_unittest')
+          try:
+            perf_device_trigger.trigger_tasks(
+                FakeArgs(1, ['build1'], json_temp),
+                ['trigger', '--some', '--test', '--', 'args'])
+
+            call_mock.assert_called_once()
+
+            called_args, keyword = call_mock.call_args
+            self.assertEqual(keyword, {})
+            self.assertEqual(called_args[0], sys.executable)
+            python_args = called_args[1]
+            json_ind = python_args.index('--dump_json')
+            # Remove --dump_json and its arg
+            python_args.pop(json_ind)
+            temp_json_path = python_args.pop(json_ind)
+            # We can't assert the exact name, since it's a temp file name. So
+            # just make sure it has the proper prefix. In order to do this in a
+            # platform independent way, just make sure that the path provided
+            # shares a common prefix with the path we made above (that's not
+            # just '/').
+            self.assertTrue(
+                len(os.path.commonprefix([temp_json_path, json_temp])) > 1)
+
+            self.assertEqual(
+              python_args, [
+              '/path/to/swarming.py', 'trigger',
+              '--some', '--test',
+              '--dimension', 'id', 'build1', '--',
+              'args', '--id', 'build1'])
+          finally:
+            os.close(temp_fd)
+            os.remove(json_temp)