Fix testGeneratePerfSharding and makes it not depend on actual perf benchmarks

Bug: 855689
Change-Id: Ic48e91c461ba00a9c3224defd6bda1036bf4af4c
Reviewed-on: https://chromium-review.googlesource.com/1112717Reviewed-by: Ashley Enstad <ashleymarie@chromium.org>
Reviewed-by: Emily Hanley <eyaich@chromium.org>
Commit-Queue: Ned Nguyen <nednguyen@google.com>
Cr-Commit-Position: refs/heads/master@{#570219}

tools/perf/core/sharding_map_generator.py

@@ -65,41 +65,54 @@ import argparse
 from collections import OrderedDict
 import core.path_util
 import json
-import optparse
 
 core.path_util.AddTelemetryToPath()
 
-from telemetry.internal.browser import browser_options
 
+def main(options, benchmarks_data):
+  """
+    benchmarks_data is a dictionary of all benchmarks to be sharded. Its
+    structure is as follows:
+    {
+      "benchmark_1": {
+         "stories": [ "storyA", "storyB",...],
+         "repeat": <number of pageset_repeat> },
+      "benchmark_2": {
+         ...
+      },...
+    }
 
-def main(args, benchmarks):
+    The "stories" field contains a list of ordered story names. Notes that
+    this should match the actual order of how the benchmark stories are
+    executed for the sharding algorithm to be effective.
+  """
   story_timing_ordered_dict = _load_timing_data_from_file(
-      benchmarks, args.timing_data, True)
+      benchmarks_data, options.timing_data, True)
 
   all_stories = {}
-  for b in benchmarks:
-    all_stories[b.Name()] = _get_stories_for_benchmark(b)
+  for b in benchmarks_data:
+    all_stories[b] = benchmarks_data[b]['stories']
 
   sharding_map = generate_sharding_map(story_timing_ordered_dict,
-      all_stories, args.num_shards, args.debug)
+      all_stories, options.num_shards, options.debug)
 
-  with open(args.output_file, 'w') as output_file:
+  with open(options.output_file, 'w') as output_file:
     json.dump(sharding_map, output_file, indent = 4, separators=(',', ': '))
 
-  if args.test_data:
+  if options.test_data:
     story_timing_ordered_dict = _load_timing_data_from_file(
-        benchmarks, args.test_data, False)
+        benchmarks_data, options.test_data, False)
     test_results = test_sharding_map(
-        args.output_file, story_timing_ordered_dict, all_stories)
-    if args.test_data_output:
-      with open(args.test_data_output, 'w') as output_file:
+        options.output_file, story_timing_ordered_dict, all_stories)
+    if options.test_data_output:
+      with open(options.test_data_output, 'w') as output_file:
         json.dump(test_results, output_file, indent = 4, separators=(',', ': '))
 
     else:
       print test_results
 
 
-def get_args():
+def get_parser():
   parser = argparse.ArgumentParser(
       description='Generate perf test sharding map.')
   parser.add_argument(
@@ -217,9 +230,8 @@ def _add_benchmarks_to_shard(sharding_map, shard_index, stories_in_shard,
   sharding_map[str(shard_index)] = {'benchmarks': benchmarks_in_shard}
 
 
-def _load_timing_data_from_file(benchmarks, timing_data_file, repeat):
-  story_timing_ordered_dict = _init_timing_dict_for_benchmarks(benchmarks)
-  pageset_repeat_dict = _init_pageset_repeat_dict_for_benchmarks(benchmarks)
+def _load_timing_data_from_file(benchmarks_data, timing_data_file, repeat):
+  story_timing_ordered_dict = _init_timing_dict_for_benchmarks(benchmarks_data)
   with open(timing_data_file, 'r') as timing_data_file:
     story_timing = json.load(timing_data_file)
     for run in story_timing:
@@ -228,47 +240,21 @@ def _load_timing_data_from_file(benchmarks, timing_data_file, repeat):
         if run['duration']:
           if repeat:
             story_timing_ordered_dict[run['name']] = (float(run['duration'])
-                * pageset_repeat_dict[benchmark])
+                * benchmarks_data[benchmark]['repeat'])
           else:
             story_timing_ordered_dict[run['name']] += float(run['duration'])
   return story_timing_ordered_dict
 
 
-def _init_pageset_repeat_dict_for_benchmarks(benchmarks):
-  pageset_repeat = {}
-  for b in benchmarks:
-    pageset_repeat[b.Name()] = b().options.get('pageset_repeat', 1)
-  return pageset_repeat
-
-
-def _init_timing_dict_for_benchmarks(benchmarks):
+def _init_timing_dict_for_benchmarks(benchmarks_data):
   timing_data = OrderedDict()
-  for b in benchmarks:
-    story_list = _get_stories_for_benchmark(b)
+  for b in benchmarks_data:
+    story_list = benchmarks_data[b]['stories']
     for story in story_list:
-      timing_data[b.Name() + '/' + story] = 0
+      timing_data[b + '/' + story] = 0
   return timing_data
 
 
-def _get_stories_for_benchmark(b):
-    story_list = []
-    benchmark = b()
-    options = browser_options.BrowserFinderOptions()
-    # Add default values for any extra commandline options
-    # provided by the benchmark.
-    parser = optparse.OptionParser()
-    before, _ = parser.parse_args([])
-    benchmark.AddBenchmarkCommandLineArgs(parser)
-    after, _ = parser.parse_args([])
-    for extra_option in dir(after):
-      if extra_option not in dir(before):
-        setattr(options, extra_option, getattr(after, extra_option))
-    for story in benchmark.CreateStorySet(options).stories:
-      if story.name not in story_list:
-        story_list.append(story.name)
-    return story_list
-
-
 def _generate_empty_sharding_map(num_shards):
   sharding_map = OrderedDict()
   for i in range(0, num_shards):

tools/perf/core/sharding_map_generator_unittest.py

@@ -7,12 +7,11 @@ import copy
 from core import sharding_map_generator
 import json
 import os
-import subprocess
-import sys
-from telemetry import decorators
 import tempfile
 import unittest
 
+from py_utils import tempfile_ext
+
 
 class TestShardingMapGenerator(unittest.TestCase):
 
@@ -46,24 +45,26 @@ class TestShardingMapGenerator(unittest.TestCase):
     finally:
       os.remove(map_path)
 
-  # Failing everywhere (see
-  # https://chromium-review.googlesource.com/c/chromium/src/+/1112978)
-  @decorators.Disabled('all')
   def testGeneratePerfSharding(self):
-    path_output = tempfile.mkstemp(suffix='.json')[1]
-    path_results = tempfile.mkstemp(suffix='.json')[1]
-    try:
-      cmd = [sys.executable,
-          os.path.normpath('tools/perf/generate_perf_sharding')]
+    with tempfile_ext.NamedTemporaryDirectory() as temp_dir:
+      path_output = os.path.join(temp_dir, 'path_output')
+      path_results = os.path.join(temp_dir, 'path_results')
+      test_data_dir = os.path.join(os.path.dirname(__file__), 'test_data')
       args = [
           '--output-file', path_output,
-          '--timing-data', 'tools/perf/core/test_data/test_timing_data.json',
+          '--timing-data', os.path.join(test_data_dir, 'test_timing_data.json'),
           '--num-shards', '5',
           '--test-data',
-          'tools/perf/core/test_data/test_timing_data_1_build.json',
+          os.path.join(test_data_dir, 'test_timing_data_1_build.json'),
           '--test-data-output', path_results
       ]
-      subprocess.check_call(cmd + args)
+
+      with open(os.path.join(test_data_dir, 'benchmarks_data.json')) as f:
+        benchmark_data = json.load(f)
+
+      parser = sharding_map_generator.get_parser()
+      options = parser.parse_args(args)
+      sharding_map_generator.main(options, benchmark_data)
 
       with open(path_results, 'r') as test_results:
         results = json.load(test_results)
@@ -71,6 +72,3 @@ class TestShardingMapGenerator(unittest.TestCase):
         for shard in results:
           shard_total_timing.append(results[shard]['full_time'])
         self.assertTrue(max(shard_total_timing) - min(shard_total_timing) < 400)
-    finally:
-      os.remove(path_output)
-      os.remove(path_results)

tools/perf/generate_perf_sharding

@@ -12,17 +12,48 @@ devices we shard to.
 Run -h to see available commands.
 """
 
+import optparse
 import sys
 
 from core import sharding_map_generator
 from core import perf_data_generator
 
+from telemetry.internal.browser import browser_options
+
+
+def _get_stories_for_benchmark(b):
+    story_list = []
+    benchmark = b()
+    options = browser_options.BrowserFinderOptions()
+    # Add default values for any extra commandline options
+    # provided by the benchmark.
+    parser = optparse.OptionParser()
+    before, _ = parser.parse_args([])
+    benchmark.AddBenchmarkCommandLineArgs(parser)
+    after, _ = parser.parse_args([])
+    for extra_option in dir(after):
+      if extra_option not in dir(before):
+        setattr(options, extra_option, getattr(after, extra_option))
+    for story in benchmark.CreateStorySet(options).stories:
+      if story.name not in story_list:
+        story_list.append(story.name)
+    return story_list
+
+
 
 if __name__ == '__main__':
-  parser = sharding_map_generator.get_args()
-  args = parser.parse_args()
+  parser = sharding_map_generator.get_parser()
+  options = parser.parse_args()
 
   # Parse args and generate data here to avoid circular dependencies between
   # sharding_map_generator.py and perf_data_generator.py
   benchmarks = perf_data_generator.current_benchmarks()
-  sys.exit(sharding_map_generator.main(args, benchmarks))
+
+  benchmarks_data = {}
+  for b in benchmarks:
+    benchmarks_data[b.Name()] = {
+        'repeat':  b().options.get('pageset_repeat', 1),
+        'stories': _get_stories_for_benchmark(b)
+    }
+
+  sys.exit(sharding_map_generator.main(options, benchmarks_data))