[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

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

Fixed!

[ychen]

LGTM

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

4–5 ↗

(On Diff #459169)

[jloser]

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

4–5 ↗

(On Diff #459169)

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.getenv("GTEST_TOTAL_SHARDS", total_shards),
+            'GTEST_SHARD_INDEX': os.getenv("GTEST_SHARD_INDEX", shard_idx)
         }
         test.config.environment.update(shard_env)
 
         cmd = [testPath]
         cmd = self.prepareCmd(cmd)
         if litConfig.useValgrind:
             cmd = litConfig.valgrindArgs + cmd