[llvm][lit] Respect GTEST_TOTAL_SHARDS and GTEST_SHARD_INDEX env vars

[jloser]

Summary

There are a variety of issues with using GTest sharding by default for users of
lit using the Google Test formatter as mentioned in
https://github.com/llvm/llvm-project/issues/56492 and
https://github.com/llvm/llvm-project/issues/56491.

Currently, there is no way for users to explicitly control the sharding
behavior, even with the environment variables that GTest provides. This patch
teaches the googletest formatter to actually respect GTEST_TOTAL_SHARDS
and GTEST_SHARD_INDEX environment variables if they are set.

In practice, we could go one step further and not do any of the post-processing
of the JSON files if GTEST_TOTAL_SHARDS is 1 for example, but that it left
as a follow-up if desired. There may be preferred alternative approaches to
disabling sharding entirely through another mechanism, such as a lit config
variable.

[ychen]

Looks good. I think this is useful to discover whether tests are affecting each other or to disable sharding if the platform has trouble correctly emitting & parsing the GTest JSON result file (like the Github issues you mentioned). On a side note, if you think the LIT sharding implementation is buggy, it is better to submit a bug report.

Could you please add a test? Thanks.

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

115–116

use os.environ.get()

[jloser]

Add test and use os.environ.get

[jloser]

In D133542#3778957, @ychen wrote:

Looks good. I think this is useful to discover whether tests are affecting each other or to disable sharding if the platform has trouble correctly emitting & parsing the GTest JSON result file (like the Github issues you mentioned). On a side note, if you think the LIT sharding implementation is buggy, it is better to submit a bug report.

Could you please add a test? Thanks.

Just added a test! Let me know if you see a way to simplify it — I'm new to using FileCheck tests, so I just created a new directory so the environment variables wouldn't conflict with the other test in googletest-format.py. Thanks!

As to the use cases that this patch enables, I want to disable sharding downstream since we're seeing some issues with emitting malformed JSON files that the issues I linked mentioned. If I come across any bugs in the lit sharding itself, I will certainly file GitHub issues.

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

115–116

Fixed!

[ychen]

LGTM

llvm/utils/lit/tests/googletest-format-respect-gtest-sharding-env-vars.py

5–6

[jloser]

llvm/utils/lit/tests/googletest-format-respect-gtest-sharding-env-vars.py

5–6

Thanks, I fixed this before landing it.

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

def execute(self, test, litConfig):
 
         testName,total_shards = os.path.split(testName)
         testName,shard_idx = os.path.split(testName)
         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': total_shards,
-            'GTEST_SHARD_INDEX': shard_idx
+            '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]
         cmd = self.prepareCmd(cmd)
         if litConfig.useValgrind:
             cmd = litConfig.valgrindArgs + cmd
 

llvm/utils/lit/tests/Inputs/googletest-format-respect-gtest-sharding-env-vars/DummySubDir/OneTest.py

+#!/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("unexpected argument: %s" % (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("unexpected argument: %r" % (' '.join(sys.argv[1:])))
+
+for e in ['GTEST_TOTAL_SHARDS', 'GTEST_SHARD_INDEX', 'GTEST_OUTPUT']:
+    if e not in os.environ:
+        raise ValueError("missing environment variables: " + e)
+
+if not os.environ['GTEST_OUTPUT'].startswith('json:'):
+    raise ValueError("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:
+    if os.environ['GTEST_TOTAL_SHARDS'] == '1':
+        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
+    else:
+        f.write(dummy_output)
+        exit_code = 0
+
+sys.exit(exit_code)

llvm/utils/lit/tests/Inputs/googletest-format-respect-gtest-sharding-env-vars/lit.cfg

+import lit.formats
+config.name = 'googletest-format'
+config.test_format = lit.formats.GoogleTest('DummySubDir', 'Test')

llvm/utils/lit/tests/googletest-format-respect-gtest-sharding-env-vars.py

+# Check GTEST_TOTAL_SHARDS and GTEST_SHARD_INDEX environment variabls are
+# respected when using the googletest formatter.
+
+# RUN: env GTEST_TOTAL_SHARDS=1 GTEST_SHARD_INDEX=0 \
+# RUN: not %{lit} -v --order=random %{inputs}/googletest-format-respect-gtest-sharding-env-vars > %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-format :: [[PATH:[Dd]ummy[Ss]ub[Dd]ir/]][[FILE:OneTest\.py]]/0
+# CHECK: *** TEST 'googletest-format :: [[PATH]][[FILE]]/0{{.*}} FAILED ***
+# CHECK-NEXT: Script(shard):
+# CHECK-NEXT: --
+# CHECK-NEXT: GTEST_OUTPUT=json:{{[^[:space:]]*}} GTEST_SHUFFLE=1 GTEST_TOTAL_SHARDS=1 GTEST_SHARD_INDEX=0 GTEST_RANDOM_SEED=123 {{.*}}[[FILE]]