[lit] Move sharding logic into separate function

4 files changed, 36 insertions, 23 deletions

diff --git a/llvm/utils/lit/lit/cl_arguments.py b/llvm/utils/lit/lit/cl_arguments.py
index 8d4bf06939f..6c553586480 100644
--- a/llvm/utils/lit/lit/cl_arguments.py
+++ b/llvm/utils/lit/lit/cl_arguments.py
@@ -198,6 +198,9 @@ def parse_args():
             parser.error("--num-shards and --run-shard must be used together")
         if opts.runShard > opts.numShards:
             parser.error("--run-shard must be between 1 and --num-shards (inclusive)")
+        opts.shard = (opts.runShard, opts.numShards)
+    else:
+      opts.shard = None
 
     return opts
 
diff --git a/llvm/utils/lit/lit/main.py b/llvm/utils/lit/lit/main.py
index e50c706c706..e56cd44dd40 100755
--- a/llvm/utils/lit/lit/main.py
+++ b/llvm/utils/lit/lit/main.py
@@ -68,24 +68,12 @@ def main(builtinParameters = {}):
     if opts.filter:
         tests = [t for t in tests if opts.filter.search(t.getFullName())]
 
-    order_tests(tests, opts)
+    determine_order(tests, opts)
 
     # Then optionally restrict our attention to a shard of the tests.
-    if (opts.numShards is not None) or (opts.runShard is not None):
-        num_tests = len(tests)
-        # Note: user views tests and shard numbers counting from 1.
-        test_ixs = range(opts.runShard - 1, num_tests, opts.numShards)
-        tests = [tests[i] for i in test_ixs]
-        # Generate a preview of the first few test indices in the shard
-        # to accompany the arithmetic expression, for clarity.
-        preview_len = 3
-        ix_preview = ", ".join([str(i+1) for i in test_ixs[:preview_len]])
-        if len(test_ixs) > preview_len:
-            ix_preview += ", ..."
-        litConfig.note('Selecting shard %d/%d = size %d/%d = tests #(%d*k)+%d = [%s]' %
-                       (opts.runShard, opts.numShards,
-                        len(tests), num_tests,
-                        opts.numShards, opts.runShard, ix_preview))
+    if opts.shard:
+        (run, shards) = opts.shard
+        tests = filter_by_shard(tests, run, shards, litConfig)
 
     # Finally limit the number of tests, if desired.
     if opts.maxTests is not None:
@@ -96,6 +84,7 @@ def main(builtinParameters = {}):
 
     testing_time = run_tests(tests, litConfig, opts, numTotalTests)
 
+    # move into print_summary
     if not opts.quiet:
         print('Testing Time: %.2fs' % (testing_time,))
 
@@ -160,21 +149,37 @@ def print_suites_or_tests(tests, opts):
             for test in ts_tests:
                 print('  %s' % (test.getFullName(),))
 
-def order_tests(tests, opts):
+def determine_order(tests, opts):
     if opts.shuffle:
         import random
         random.shuffle(tests)
     elif opts.incremental:
+        def by_mtime(test):
+            try:
+                return os.path.getmtime(test.getFilePath())
+            except:
+                return 0
         tests.sort(key=by_mtime, reverse=True)
     else:
         tests.sort(key=lambda t: (not t.isEarlyTest(), t.getFullName()))
 
-def by_mtime(test):
-    fname = test.getFilePath()
-    try:
-        return os.path.getmtime(fname)
-    except:
-        return 0
+def filter_by_shard(tests, run, shards, litConfig):
+    test_ixs = range(run - 1, len(tests), shards)
+    selected_tests = [tests[i] for i in test_ixs]
+
+    # For clarity, generate a preview of the first few test indices in the shard
+    # to accompany the arithmetic expression.
+    preview_len = 3
+    preview = ", ".join([str(i + 1) for i in test_ixs[:preview_len]])
+    if len(test_ixs) > preview_len:
+        preview += ", ..."
+    # TODO(python3): string interpolation
+    msg = 'Selecting shard {run}/{shards} = size {sel_tests}/{total_tests} = ' \
+          'tests #({shards}*k)+{run} = [{preview}]'.format(
+              run=run, shards=shards, sel_tests=len(selected_tests),
+              total_tests=len(tests), preview=preview)
+    litConfig.note(msg)
+    return selected_tests
 
 def update_incremental_cache(test):
     if not test.result.code.isFailure:
diff --git a/llvm/utils/lit/lit/run.py b/llvm/utils/lit/lit/run.py
index 88ce445bc42..d24cfd47e5a 100644
--- a/llvm/utils/lit/lit/run.py
+++ b/llvm/utils/lit/lit/run.py
@@ -66,6 +66,8 @@ class Run(object):
 
         return end - start
 
+    # TODO(yln): as the comment says.. this is racing with the main thread waiting
+    # for results
     def _process_result(self, test, result):
         # Don't add any more test results after we've hit the maximum failure
         # count.  Otherwise we're racing with the main thread, which is going
@@ -147,6 +149,7 @@ class ParallelRun(Run):
                 pool.terminate()
                 break
 
+    # TODO(yln): interferes with progress bar
     # Some tests use threads internally, and at least on Linux each of these
     # threads counts toward the current process limit. Try to raise the (soft)
     # process limit so that tests don't fail due to resource exhaustion.
diff --git a/llvm/utils/lit/lit/util.py b/llvm/utils/lit/lit/util.py
index 7eb5eaaa2ff..c662be3c636 100644
--- a/llvm/utils/lit/lit/util.py
+++ b/llvm/utils/lit/lit/util.py
@@ -116,6 +116,8 @@ def to_unicode(s):
     return s
 
 
+# TODO(yln): multiprocessing.cpu_count()
+# TODO(python3): len(os.sched_getaffinity(0)) and os.cpu_count()
 def detectCPUs():
     """Detects the number of CPUs on a system.