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

sharding_supervisor.py

Index: sharding_supervisor.py
===================================================================
--- sharding_supervisor.py	(revision 130689)
+++ sharding_supervisor.py	(working copy)
@@ -42,6 +42,8 @@
 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_SUITE_TOTAL_SHARDS = 1  # run the whole suite.
+SS_DEFAULT_SUITE_SHARD_INDEX = 0  # run the first shard.
 
 
 def DetectNumCores():
@@ -194,15 +196,22 @@
     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.
+    suite_total_shards: Total number of shards the suite is split into if we

[cmp, 2012/04/04 23:00:28]

we agreed that suite_*_shards and suite_shard_* could be replaced with better
names

[nsylvain, 2012/04/04 23:24:34]

Done.

+                        want to run only a subset of the suite.
+    suite_shard_index: Shard index to run if we want to run only a subset of
+                       the suite.
   """
 
   SHARD_COMPLETED = object()
 
   def __init__(self, test, num_shards, num_runs, color, original_order,
-               prefix, retry_percent, timeout, gtest_args):
+               prefix, retry_percent, timeout, suite_total_shards,
+               suite_shard_index, gtest_args):
     """Inits ShardingSupervisor with given options and gtest arguments."""
     self.test = test
-    self.num_shards = num_shards
+    self.num_shards_to_run = num_shards
+    self.num_shards = num_shards * suite_total_shards
+    self.suite_shard_index = suite_shard_index
     self.num_runs = num_runs
     self.color = color
     self.original_order = original_order
@@ -212,8 +221,8 @@
     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):
@@ -249,7 +258,8 @@
 
     workers = []
     counter = Queue.Queue()
-    for i in range(self.num_shards):
+    start_point = self.num_shards_to_run * self.suite_shard_index
+    for i in range(start_point, start_point + self.num_shards_to_run):
       counter.put(i)
 
     for i in range(self.num_runs):
@@ -295,7 +305,7 @@
     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)
@@ -316,7 +326,7 @@
     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)
@@ -333,12 +343,14 @@
     """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.suite_shard_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
@@ -350,7 +362,9 @@
     """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.suite_shard_index)
+    self.shard_output[array_index].put(self.SHARD_COMPLETED)
 
   def RetryFailedTests(self):
     """Reruns any failed tests serially and prints another summary of the
@@ -447,6 +461,13 @@
   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(
+      "--suite-total-shards", type="int", default=SS_DEFAULT_SUITE_TOTAL_SHARDS,
+      help="number of shards when running only a subset of the suite")

[cmp, 2012/04/04 23:00:28]

change to 'if running a subset, ...'

[nsylvain, 2012/04/04 23:24:34]

rewrote that, but still not happy.

+  parser.add_option(
+      "--suite-shard-index", type="int", default=SS_DEFAULT_SUITE_SHARD_INDEX,
+      help="index of the shard to run when running only a subset of the suite")

[cmp, 2012/04/04 23:00:28]

change to 'if running a subset, ...'

+
   parser.disable_interspersed_args()
   (options, args) = parser.parse_args()
 
@@ -498,7 +519,8 @@
   # 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)
+      options.prefix, options.retry_percent, options.timeout,
+      options.suite_total_shards, options.suite_shard_index, gtest_args)
   return ss.ShardTest()