Android: adds device affinity for perf tests.

build/android/pylib/perf/setup.py

Index: build/android/pylib/perf/setup.py
diff --git a/build/android/pylib/perf/setup.py b/build/android/pylib/perf/setup.py
index cf844560f00c944a178ff4d0858e8afc436003ab..5314d36fa8a874ffd5d447c250524b6384567c48 100644
--- a/build/android/pylib/perf/setup.py
+++ b/build/android/pylib/perf/setup.py
@@ -6,15 +6,73 @@
 
 import json
 import fnmatch
+import logging
 import os
 import shutil
 
+from pylib import android_commands
 from pylib import constants
 from pylib import forwarder
+from pylib.device import device_list
 from pylib.perf import test_runner
 from pylib.utils import test_environment
 
 
+def _GetAllDevices():
+  devices_path = os.path.join(os.environ.get('CHROMIUM_OUT_DIR', 'out'),
+                              device_list.LAST_DEVICES_FILENAME)
+  try:
+    devices = device_list.GetPersistentDeviceList(devices_path)
+  except IOError as e:
+    logging.error('Unable to find %s [%s]', devices_path, e)
+    devices = android_commands.GetAttachedDevices()
+  return sorted(devices)
+
+
+def _GetStepsDictFromSingleStep(test_options):
+  # Running a single command, build the tests structure.
+  steps_dict = {
+    'version': 1,
+    'steps': {
+        'single_step': {
+          'device_affinity': 0,
+          'cmd': test_options.single_step
+        },
+    }
+  }
+  return steps_dict
+
+# TODO(bulach): remove once it rolls downstream, crbug.com/378862.
+def _GetStepsDictFromV0(steps_v0):
+  steps_dict = {
+    'version': 1,
+    'steps': {},
+  }
+  affinity = 0
+  for step in steps_v0:
+    steps_dict['steps'][step[0]] = {
+        'device_affinity': affinity,
+        'cmd': step[1],
+    }
+    affinity += 1
+  return steps_dict
+
+
+def _GetStepsDict(test_options):
+  if test_options.single_step:
+    return _GetStepsDictFromSingleStep(test_options)
+  if test_options.steps:
+    with file(test_options.steps, 'r') as f:
+      steps = json.load(f)
+      # TODO(bulach): remove once it rolls downstream, crbug.com/378862.
+      if isinstance(steps, list):
+        return _GetStepsDictFromV0(steps)
+
+      # Already using the new format.
+      assert steps['version'] == 1
+      return steps
+
+
 def Setup(test_options):
   """Create and return the test runner factory and tests.
 
@@ -22,7 +80,7 @@ def Setup(test_options):
     test_options: A PerformanceOptions object.
 
   Returns:
-    A tuple of (TestRunnerFactory, tests).
+    A tuple of (TestRunnerFactory, tests, devices).
   """
   # TODO(bulach): remove this once the bot side lands. BUG=318369
   constants.SetBuildType('Release')
@@ -34,33 +92,24 @@ def Setup(test_options):
   test_environment.CleanupLeftoverProcesses()
   forwarder.Forwarder.UseMultiprocessing()
 
-  if test_options.single_step:
-    # Running a single command, build the tests structure.
-    tests = [['single_step', test_options.single_step]]
-
-  if test_options.steps:
-    with file(test_options.steps, 'r') as f:
-      tests = json.load(f)
+  # We want to keep device affinity, so return all devices ever seen.
+  all_devices = _GetAllDevices()
 
-  # The list is necessary to keep the steps order, but internally
-  # the format is squashed from a list of lists into a single dict:
-  # [["A", "cmd"], ["B", "cmd"]] into {"A": "cmd", "B": "cmd"}
-  sorted_test_names = [i[0] for i in tests]
-  tests_dict = dict(tests)
+  steps_dict = _GetStepsDict(test_options)
+  sorted_step_names = sorted(steps_dict['steps'].keys())
 
   if test_options.test_filter:
-    sorted_test_names = fnmatch.filter(sorted_test_names,
+    sorted_step_names = fnmatch.filter(sorted_step_names,
                                        test_options.test_filter)
-    tests_dict = dict((k, v) for k, v in tests_dict.iteritems()
-                      if k in sorted_test_names)
 
   flaky_steps = []
   if test_options.flaky_steps:
     with file(test_options.flaky_steps, 'r') as f:
       flaky_steps = json.load(f)
 
-  def TestRunnerFactory(device, _shard_index):
+  def TestRunnerFactory(device, shard_index):
     return test_runner.TestRunner(
-        test_options, device, tests_dict, flaky_steps)
+        test_options, device, shard_index, len(all_devices),
+        steps_dict, flaky_steps)
 
-  return (TestRunnerFactory, sorted_test_names)
+  return (TestRunnerFactory, sorted_step_names, all_devices)