[unittest] Add option to allow disabling sharding in unittest

llvm/utils/lit/lit/LitConfig.py

@@ -37,6 +37,7 @@ def __init__(
         maxIndividualTestTime=0,
         parallelism_groups={},
         per_test_coverage=False,
+        disableGTestSharding=False,
     ):
         # The name of the test runner.
         self.progname = progname
@@ -87,6 +88,7 @@ def __init__(
         self.maxIndividualTestTime = maxIndividualTestTime
         self.parallelism_groups = parallelism_groups
         self.per_test_coverage = per_test_coverage
+        self.disableGTestSharding = bool(disableGTestSharding)
 
     @property
     def maxIndividualTestTime(self):

llvm/utils/lit/lit/cl_arguments.py

@@ -118,6 +118,12 @@ def parse_args():
         )
 
     execution_group = parser.add_argument_group("Test Execution")
+    execution_group.add_argument(
+        "--disable-gtest-sharding",
+        dest="disableGTestSharding",
+        help="Disable sharding for GoogleTest format",
+        action="store_true",
+    )
     execution_group.add_argument(
         "--path",
         help="Additional paths to add to testing environment",

llvm/utils/lit/lit/formats/googletest.py

@@ -68,24 +68,49 @@ def getTestsInDirectory(self, testSuite, path_in_suite, litConfig, localConfig):
                     self.seen_executables.add(execpath)
                 num_tests = self.get_num_tests(execpath, litConfig, localConfig)
                 if num_tests is not None:
-                    # Compute the number of shards.
-                    shard_size = init_shard_size
-                    nshard = int(math.ceil(num_tests / shard_size))
-                    while nshard < core_count and shard_size > 1:
-                        shard_size = shard_size // 2
+                    if not litConfig.disableGTestSharding:
+                        # Compute the number of shards.
+                        shard_size = init_shard_size
                         nshard = int(math.ceil(num_tests / shard_size))
-
-                    # Create one lit test for each shard.
-                    for idx in range(nshard):
-                        testPath = path_in_suite + (subdir, fn, str(idx), str(nshard))
+                        while nshard < core_count and shard_size > 1:
+                            shard_size = shard_size // 2
+                            nshard = int(math.ceil(num_tests / shard_size))
+
+                        # Create one lit test for each shard.
+                        for idx in range(nshard):
+                            testPath = path_in_suite + (
+                                subdir,
+                                fn,
+                                str(idx),
+                                str(nshard),
+                            )
+                            json_file = (
+                                "-".join(
+                                    [
+                                        execpath,
+                                        testSuite.config.name,
+                                        str(os.getpid()),
+                                        str(idx),
+                                        str(nshard),
+                                    ]
+                                )
+                                + ".json"
+                            )
+                            yield lit.Test.Test(
+                                testSuite,
+                                testPath,
+                                localConfig,
+                                file_path=execpath,
+                                gtest_json_file=json_file,
+                            )
+                    else:
+                        testPath = path_in_suite + (subdir, fn)
                         json_file = (
                             "-".join(
                                 [
                                     execpath,
                                     testSuite.config.name,
                                     str(os.getpid()),
-                                    str(idx),
-                                    str(nshard),
                                 ]
                             )
                             + ".json"
@@ -118,24 +143,32 @@ def execute(self, test, litConfig):
         if test.gtest_json_file is None:
             return lit.Test.FAIL, ""
 
-        testPath, testName = os.path.split(test.getSourcePath())
-        while not os.path.exists(testPath):
-            # Handle GTest parametrized and typed tests, whose name includes
-            # some '/'s.
-            testPath, namePrefix = os.path.split(testPath)
-            testName = namePrefix + "/" + testName
-
-        testName, total_shards = os.path.split(testName)
-        testName, shard_idx = os.path.split(testName)
+        testPath = test.getSourcePath()
         from lit.cl_arguments import TestOrder
 
         use_shuffle = TestOrder(litConfig.order) == TestOrder.RANDOM
         shard_env = {
             "GTEST_OUTPUT": "json:" + test.gtest_json_file,
             "GTEST_SHUFFLE": "1" if use_shuffle else "0",
-            "GTEST_TOTAL_SHARDS": os.environ.get("GTEST_TOTAL_SHARDS", total_shards),
-            "GTEST_SHARD_INDEX": os.environ.get("GTEST_SHARD_INDEX", shard_idx),
         }
+        if not litConfig.disableGTestSharding:
+            testPath, testName = os.path.split(test.getSourcePath())
+            while not os.path.exists(testPath):
+                # Handle GTest parameterized and typed tests, whose name includes
+                # some '/'s.
+                testPath, namePrefix = os.path.split(testPath)
+                testName = namePrefix + "/" + testName
+
+            testName, total_shards = os.path.split(testName)
+            testName, shard_idx = os.path.split(testName)
+            shard_env.update(
+                {
+                    "GTEST_TOTAL_SHARDS": os.environ.get(
+                        "GTEST_TOTAL_SHARDS", total_shards
+                    ),
+                    "GTEST_SHARD_INDEX": os.environ.get("GTEST_SHARD_INDEX", shard_idx),
+                }
+            )
         test.config.environment.update(shard_env)
 
         cmd = [testPath]

llvm/utils/lit/lit/main.py

@@ -41,6 +41,7 @@ def main(builtin_params={}):
         params=params,
         config_prefix=opts.configPrefix,
         per_test_coverage=opts.per_test_coverage,
+        disableGTestSharding=opts.disableGTestSharding,
     )
 
     discovered_tests = lit.discovery.find_tests_for_inputs(

llvm/utils/lit/tests/Inputs/googletest-no-sharding/DummySubDir/OneTest.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+
+import os
+import sys
+
+if len(sys.argv) == 3 and sys.argv[1] == "--gtest_list_tests":
+    if sys.argv[2] != "--gtest_filter=-*DISABLED_*":
+        raise ValueError(f"unexpected argument: {sys.argv[2]}")
+    print(
+        """\
+FirstTest.
+  subTestA
+  subTestB
+  subTestC
+  subTestD
+ParameterizedTest/0.
+  subTest
+ParameterizedTest/1.
+  subTest"""
+    )
+    sys.exit(0)
+elif len(sys.argv) != 1:
+    # sharding and json output are specified using environment variables
+    raise ValueError(f"unexpected argument: {' '.join(sys.argv[1:])!r}")
+
+for e in ["GTEST_OUTPUT"]:
+    if e not in os.environ:
+        raise ValueError(f"missing environment variables: {e}")
+
+if not os.environ["GTEST_OUTPUT"].startswith("json:"):
+    raise ValueError(f"must emit json output: {os.environ['GTEST_OUTPUT']}")
+
+output = """\
+{
+"random_seed": 123,
+"testsuites": [
+    {
+        "name": "FirstTest",
+        "testsuite": [
+            {
+                "name": "subTestA",
+                "result": "COMPLETED",
+                "time": "0.001s"
+            },
+            {
+                "name": "subTestB",
+                "result": "COMPLETED",
+                "time": "0.001s",
+                "failures": [
+                    {
+                        "failure": "I am subTest B, I FAIL\\nAnd I have two lines of output",
+                        "type": ""
+                    }
+                ]
+            },
+            {
+                "name": "subTestC",
+                "result": "SKIPPED",
+                "time": "0.001s"
+            },
+            {
+                "name": "subTestD",
+                "result": "UNRESOLVED",
+                "time": "0.001s"
+            }
+        ]
+    },
+    {
+        "name": "ParameterizedTest/0",
+        "testsuite": [
+            {
+                "name": "subTest",
+                "result": "COMPLETED",
+                "time": "0.001s"
+            }
+        ]
+    },
+    {
+        "name": "ParameterizedTest/1",
+        "testsuite": [
+            {
+                "name": "subTest",
+                "result": "COMPLETED",
+                "time": "0.001s"
+            }
+        ]
+    }
+]
+}"""
+
+dummy_output = """\
+{
+"testsuites": [
+]
+}"""
+
+json_filename = os.environ["GTEST_OUTPUT"].split(":", 1)[1]
+with open(json_filename, "w", encoding="utf-8") as f:
+    print("[ RUN      ] FirstTest.subTestB", flush=True)
+    print("I am subTest B output", file=sys.stderr, flush=True)
+    print("[  FAILED  ] FirstTest.subTestB (8 ms)", flush=True)
+    f.write(output)
+    exit_code = 1
+
+sys.exit(exit_code)

llvm/utils/lit/tests/Inputs/googletest-no-sharding/lit.cfg

@@ -0,0 +1,4 @@
+import lit.formats
+
+config.name = "googletest-no-sharding"
+config.test_format = lit.formats.GoogleTest("DummySubDir", "Test")

llvm/utils/lit/tests/googletest-no-sharding.py

@@ -0,0 +1,43 @@
+# Check the various features of the GoogleTest format.
+
+# RUN: not %{lit} -v --disable-gtest-sharding --order=random %{inputs}/googletest-no-sharding > %t.out
+# FIXME: Temporarily dump test output so we can debug failing tests on
+# buildbots.
+# RUN: cat %t.out
+# RUN: FileCheck < %t.out %s
+#
+# END.
+
+# CHECK: -- Testing:
+# CHECK: FAIL: googletest-no-sharding :: [[PATH:[Dd]ummy[Ss]ub[Dd]ir/]][[FILE:OneTest\.py]]
+# CHECK: *** TEST 'googletest-no-sharding :: [[PATH]][[FILE]]' FAILED ***
+# CHECK-NEXT: Script(shard):
+# CHECK-NEXT: --
+# CHECK-NEXT: GTEST_OUTPUT=json:{{[^[:space:]]*}} GTEST_SHUFFLE=1 GTEST_RANDOM_SEED=123 {{.*}}[[FILE]]
+# CHECK-NEXT: --
+# CHECK-EMPTY:
+# CHECK-NEXT: Script:
+# CHECK-NEXT: --
+# CHECK-NEXT: [[FILE]] --gtest_filter=FirstTest.subTestB
+# CHECK-NEXT: --
+# CHECK-NEXT: I am subTest B output
+# CHECK-EMPTY:
+# CHECK-NEXT: I am subTest B, I FAIL
+# CHECK-NEXT: And I have two lines of output
+# CHECK-EMPTY:
+# CHECK: Script:
+# CHECK-NEXT: --
+# CHECK-NEXT: [[FILE]] --gtest_filter=FirstTest.subTestD
+# CHECK-NEXT: --
+# CHECK-NEXT: unresolved test result
+# CHECK: ***
+# CHECK: ***
+# CHECK: Unresolved Tests (1):
+# CHECK-NEXT:   googletest-no-sharding :: FirstTest/subTestD
+# CHECK: ***
+# CHECK-NEXT: Failed Tests (1):
+# CHECK-NEXT:   googletest-no-sharding :: FirstTest/subTestB
+# CHECK: Skipped{{ *}}: 1
+# CHECK: Passed{{ *}}: 3
+# CHECK: Unresolved{{ *}}: 1
+# CHECK: Failed{{ *}}: 1