Modify the sharding supervisor to support sharding only a sub-shard of a test suite.

sharding_supervisor.py

 #!/usr/bin/env python
 # Copyright (c) 2011 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.
 
 """Shards a given test suite and runs the shards in parallel.
 
 ShardingSupervisor is called to process the command line options and creates
 the specified number of worker threads. These threads then run each shard of
 the test in a separate process and report on the results. When all the shards
@@ skipping 24 matching lines @@
 except ImportError:
   # Unable to find depot_tools, so just use standard subprocess
   import subprocess
 
 SS_USAGE = "python %prog [options] path/to/test [gtest_args]"
 SS_DEFAULT_NUM_CORES = 4
 SS_DEFAULT_SHARDS_PER_CORE = 5  # num_shards = cores * SHARDS_PER_CORE
 SS_DEFAULT_RUNS_PER_CORE = 1  # num_workers = cores * RUNS_PER_CORE
 SS_DEFAULT_RETRY_PERCENT = 5  # --retry-failed ignored if more than 5% fail
 SS_DEFAULT_TIMEOUT = 530  # Slightly less than buildbot's default 600 seconds
+SS_DEFAULT_TOTAL_SLAVES = 1  # run the whole suite.

[cmp, 2012/04/05 19:07:49]

yes!

+SS_DEFAULT_SLAVE_INDEX = 0  # run the tests for the first slave.
 
 
 def DetectNumCores():
   """Detects the number of cores on the machine.
 
   Returns:
     The number of cores on the machine or DEFAULT_NUM_CORES if it could not
     be found.
   """
   try:
     # Linux, Unix, MacOS
     if hasattr(os, "sysconf"):
       if "SC_NPROCESSORS_ONLN" in os.sysconf_names:
         # Linux, Unix
         return int(os.sysconf("SC_NPROCESSORS_ONLN"))
       else:
         # OSX
         return int(os.popen2("sysctl -n hw.ncpu")[1].read())
     # Windows
     return int(os.environ["NUMBER_OF_PROCESSORS"])
   except ValueError:
     return SS_DEFAULT_NUM_CORES
 
 
-def RunShard(test, num_shards, index, gtest_args, stdout, stderr):
+def RunShard(test, total_shards, index, gtest_args, stdout, stderr):

[cmp, 2012/04/05 19:07:49]

thank you!

   """Runs a single test shard in a subprocess.
 
   Returns:
     The Popen object representing the subprocess handle.
   """
   args = [test]
   args.extend(gtest_args)
   env = os.environ.copy()
-  env["GTEST_TOTAL_SHARDS"] = str(num_shards)
+  env["GTEST_TOTAL_SHARDS"] = str(total_shards)
   env["GTEST_SHARD_INDEX"] = str(index)
 
   # Use a unique log file for each shard
   # Allows ui_tests to be run in parallel on the same machine
   env["CHROME_LOG_FILE"] = "chrome_log_%d" % index
 
   return subprocess.Popen(
       args, stdout=stdout,
       stderr=stderr,
       env=env,
@@ skipping 59 matching lines @@
     as it is finished.
     """
     while True:
       try:
         index = self.counter.get_nowait()
       except Queue.Empty:
         break
       chars = cStringIO.StringIO()
       shard_running = True
       shard = RunShard(
-          self.supervisor.test, self.supervisor.num_shards, index,
+          self.supervisor.test, self.supervisor.total_shards, index,
           self.supervisor.gtest_args, subprocess.PIPE, subprocess.STDOUT)
       while shard_running:
         char = shard.stdout.read(1)
         if not char and shard.poll() is not None:
           shard_running = False
         chars.write(char)
         if char == "\n" or not shard_running:
           line = chars.getvalue()
           if not line and not shard_running:
             break
           self.ProcessLine(index, line)
           self.supervisor.LogOutputLine(index, line)
           chars.close()
           chars = cStringIO.StringIO()
       if self.current_test:
         self.ReportFailure("INCOMPLETE", index, self.current_test)
       self.supervisor.ShardIndexCompleted(index)
       if shard.returncode != 0:
         self.supervisor.LogShardFailure(index)
 
 
 class ShardingSupervisor(object):
   """Supervisor object that handles the worker threads.
 
   Attributes:
     test: Name of the test to shard.
-    num_shards: Total number of shards to split the test into.
+    num_shards_to_run: Total number of shards to split the test into.
     num_runs: Total number of worker threads to create for running shards.
     color: Indicates which coloring mode to use in the output.
     original_order: True if shard output should be printed as it comes.
     prefix: True if each line should indicate the shard index.
     retry_percent: Integer specifying the max percent of tests to retry.
     gtest_args: The options to pass to gtest.
     failed_tests: List of statements from shard output indicating a failure.
     failed_shards: List of shards that contained failing tests.
     shards_completed: List of flags indicating which shards have finished.
     shard_output: Buffer that stores the output from each shard.
     test_counter: Stores the total number of tests run.
+    total_slaves: Total number of slaves running this test.
+    slave_index: Current slave to run tests for.
+
+  If total_slaves is set, we run only a subset of the tests. This is meant to be
+  used when we want to shard across machines as well as across cpus. In that
+  case the number of shards to execute will be the same, but they will be
+  smaller, as the total number of shards in the test suite will be multiplied
+  by 'total_slaves'.
+ 
+  For example, if you are on a quad core machine, the sharding supervisor by
+  default will use 20 shards for the whole suite. However, if you set
+  total_slaves to 2, it will split the suite in 40 shards and will only run
+  shards [0-19] or shards [20-39] depending if you set slave_index to 0 or 1.

[cmp, 2012/04/05 19:07:49]

will help a lot, thanks

   """
 
   SHARD_COMPLETED = object()
 
-  def __init__(self, test, num_shards, num_runs, color, original_order,
-               prefix, retry_percent, timeout, gtest_args):
+  def __init__(self, test, num_shards_to_run, num_runs, color, original_order,
+               prefix, retry_percent, timeout, total_slaves, slave_index,
+               gtest_args):
     """Inits ShardingSupervisor with given options and gtest arguments."""
     self.test = test
-    self.num_shards = num_shards
+    # Number of shards to run locally.
+    self.num_shards_to_run = num_shards_to_run
+    # Total shards in the test suite running across all slaves.
+    self.total_shards = num_shards_to_run * total_slaves
+    self.slave_index = slave_index
     self.num_runs = num_runs
     self.color = color
     self.original_order = original_order
     self.prefix = prefix
     self.retry_percent = retry_percent
     self.timeout = timeout
     self.gtest_args = gtest_args
     self.failed_tests = []
     self.failed_shards = []
-    self.shards_completed = [False] * num_shards
-    self.shard_output = [Queue.Queue() for _ in range(num_shards)]
+    self.shards_completed = [False] * self.num_shards_to_run
+    self.shard_output = [Queue.Queue() for _ in range(self.num_shards_to_run)]
     self.test_counter = itertools.count()
 
   def ShardTest(self):
     """Runs the test and manages the worker threads.
 
     Runs the test and outputs a summary at the end. All the tests in the
     suite are run by creating (cores * runs_per_core) threads and
     (cores * shards_per_core) shards. When all the worker threads have
     finished, the lines saved in failed_tests are printed again. If enabled,
     and failed tests that do not have FLAKY or FAILS in their names are run
@@ skipping 15 matching lines @@
 
     test_start = re.compile(
         ansi_regex + r"\[\s+RUN\s+\] " + ansi_regex + test_name_regex)
     test_ok = re.compile(
         ansi_regex + r"\[\s+OK\s+\] " + ansi_regex + test_name_regex)
     test_fail = re.compile(
         ansi_regex + r"\[\s+FAILED\s+\] " + ansi_regex + test_name_regex)
 
     workers = []
     counter = Queue.Queue()
-    for i in range(self.num_shards):
+    start_point = self.num_shards_to_run * self.slave_index
+    for i in range(start_point, start_point + self.num_shards_to_run):
       counter.put(i)
 
     for i in range(self.num_runs):
       worker = ShardRunner(
           self, counter, test_start, test_ok, test_fail)
       worker.start()
       workers.append(worker)
     if self.original_order:
       for worker in workers:
         worker.join()
@@ skipping 25 matching lines @@
 
   def LogShardFailure(self, index):
     """Records that a test in the given shard has failed."""
     self.failed_shards.append(index)
 
   def WaitForShards(self):
     """Prints the output from each shard in consecutive order, waiting for
     the current shard to finish before starting on the next shard.
     """
     try:
-      for shard_index in range(self.num_shards):
+      for shard_index in range(self.num_shards_to_run):
         while True:
           try:
             line = self.shard_output[shard_index].get(True, self.timeout)
           except Queue.Empty:
             # Shard timed out, notice failure and move on.
             self.LogShardFailure(shard_index)
             # TODO(maruel): Print last test. It'd be simpler to have the
             # processing in the main thread.
             # TODO(maruel): Make sure the worker thread terminates.
             sys.stdout.write('TIMED OUT\n\n')
             LogTestFailure(
                 'FAILURE: SHARD %d TIMED OUT; %d seconds' % (
                     shard_index, self.timeout))
             break
           if line is self.SHARD_COMPLETED:
             break
           sys.stdout.write(line)
     except:
       sys.stdout.flush()
       print >> sys.stderr, 'CAUGHT EXCEPTION: dumping remaining data:'
-      for shard_index in range(self.num_shards):
+      for shard_index in range(self.num_shards_to_run):
         while True:
           try:
             line = self.shard_output[shard_index].get(False)
           except Queue.Empty:
             # Shard timed out, notice failure and move on.
             self.LogShardFailure(shard_index)
             break
           if line is self.SHARD_COMPLETED:
             break
           sys.stderr.write(line)
       raise
 
   def LogOutputLine(self, index, line):
     """Either prints the shard output line immediately or saves it in the
     output buffer, depending on the settings. Also optionally adds a  prefix.
     """
+    # Fix up the index.
+    array_index = index - (self.num_shards_to_run * self.slave_index)
     if self.prefix:
       line = "%i>%s" % (index, line)
     if self.original_order:
       sys.stdout.write(line)
     else:
-      self.shard_output[index].put(line)
+      self.shard_output[array_index].put(line)
 
   def IncrementTestCount(self):
     """Increments the number of tests run. This is relevant to the
     --retry-percent option.
     """
     self.test_counter.next()
 
   def ShardIndexCompleted(self, index):
     """Records that a shard has finished so the output from the next shard
     can now be printed.
     """
-    self.shard_output[index].put(self.SHARD_COMPLETED)
+    # Fix up the index.
+    array_index = index - (self.num_shards_to_run * self.slave_index)
+    self.shard_output[array_index].put(self.SHARD_COMPLETED)
 
   def RetryFailedTests(self):
     """Reruns any failed tests serially and prints another summary of the
     results if no more than retry_percent failed.
     """
     num_tests_run = self.test_counter.next()
     if len(self.failed_tests) > self.retry_percent * num_tests_run:
       sys.stderr.write("\nNOT RETRYING FAILED TESTS (too many failed)\n")
       return 1
     self.WriteText(sys.stderr, "\nRETRYING FAILED TESTS:\n", "\x1b[1;5;33m")
@@ skipping 76 matching lines @@
       "--retry-failed", action="store_true",
       help="retry tests that did not pass serially")
   parser.add_option(
       "--retry-percent", type="int",
       default=SS_DEFAULT_RETRY_PERCENT,
       help="ignore --retry-failed if more than this percent fail [0, 100]"
       " (default = %i)" % SS_DEFAULT_RETRY_PERCENT)
   parser.add_option(
       "-t", "--timeout", type="int", default=SS_DEFAULT_TIMEOUT,
       help="timeout in seconds to wait for a shard (default=%default s)")
+  parser.add_option(
+      "--total-slaves", type="int", default=SS_DEFAULT_TOTAL_SLAVES,
+      help="if running a subset, number of slaves sharing the test")
+  parser.add_option(
+      "--slave-index", type="int", default=SS_DEFAULT_SLAVE_INDEX,
+      help="if running a subset, index of the slave to run tests for")
+
   parser.disable_interspersed_args()
   (options, args) = parser.parse_args()
 
   if not args:
     parser.error("You must specify a path to test!")
   if not os.path.exists(args[0]):
     parser.error("%s does not exist!" % args[0])
 
   num_cores = DetectNumCores()
 
   if options.shards_per_core < 1:
     parser.error("You must have at least 1 shard per core!")
-  num_shards = num_cores * options.shards_per_core
+  num_shards_to_run = num_cores * options.shards_per_core
 
   if options.runs_per_core < 1:
     parser.error("You must have at least 1 run per core!")
   num_runs = num_cores * options.runs_per_core
 
   gtest_args = ["--gtest_color=%s" % {
       True: "yes", False: "no"}[options.color]] + args[1:]
 
   if options.original_order:
     options.prefix = True
 
   # TODO(charleslee): for backwards compatibility with buildbot's log_parser
   if options.reorder:
     options.original_order = False
     options.prefix = True
 
   if options.random_seed:
     seed = random.randint(1, 99999)
     gtest_args.extend(["--gtest_shuffle", "--gtest_random_seed=%i" % seed])
 
   if options.retry_failed:
     if options.retry_percent < 0 or options.retry_percent > 100:
       parser.error("Retry percent must be an integer [0, 100]!")
   else:
     options.retry_percent = -1
 
   if options.runshard != None:
     # run a single shard and exit
-    if (options.runshard < 0 or options.runshard >= num_shards):
+    if (options.runshard < 0 or options.runshard >= num_shards_to_run):
       parser.error("Invalid shard number given parameters!")
     shard = RunShard(
-        args[0], num_shards, options.runshard, gtest_args, None, None)
+        args[0], num_shards_to_run, options.runshard, gtest_args, None, None)
     shard.communicate()
     return shard.poll()
 
   # shard and run the whole test
   ss = ShardingSupervisor(
-      args[0], num_shards, num_runs, options.color, options.original_order,
-      options.prefix, options.retry_percent, options.timeout, gtest_args)
+      args[0], num_shards_to_run, num_runs, options.color,
+      options.original_order, options.prefix, options.retry_percent,
+      options.timeout, options.total_slaves, options.slave_index, gtest_args)
   return ss.ShardTest()
 
 
 if __name__ == "__main__":
   sys.exit(main())