[bazelified tests] Shard bazelified bazel distribtests to make them faster (#34102)

Basically run each of the subtests (buildtest, distribtest_cpp,
distribtest_python) as a separate bazel target.

- currently the bazel distribtest are the slowest targets in
grpc_bazel_rbe_nonbazel
- the shards are basically independent tests anyway
- when split into multiple targets, they each get a separate target log
so it's easier debug issues since there isn't multiple bazel invocations
in each log.

test/distrib/bazel/test_single_bazel_version.sh

@@ -15,12 +15,36 @@
 
 set -ex
 
-if [ "$#" != "1" ] ; then
+if [ "$#" == "0" ] ; then
     echo "Must supply bazel version to be tested." >/dev/stderr
     exit 1
 fi
 
 VERSION="$1"
+shift 1
+
+# directories under test/distrib/bazel/ to test.
+TEST_DIRECTORIES=(
+  "cpp"
+  "python"
+)
+# construct list of all supported test shards
+ALL_TEST_SHARDS=("buildtest")
+for TEST_DIRECTORY in "${TEST_DIRECTORIES[@]}"
+do
+  ALL_TEST_SHARDS+=("distribtest_${TEST_DIRECTORY}")
+done
+
+# Read list of shards to run from the commandline args.
+# If ther are no args, run all the shards.
+if [ "$#" != "0" ]
+then
+  # Use remaining commandline args as test shard names.
+  TEST_SHARDS=("$@")
+else
+  # Run all supported shards.
+  TEST_SHARDS=("${ALL_TEST_SHARDS[@]}")
+fi
 
 cd "$(dirname "$0")"/../../..
 
@@ -49,11 +73,6 @@ EXCLUDED_TARGETS=(
   "-//tools/bazelify_tests/..."
 )
 
-TEST_DIRECTORIES=(
-  "cpp"
-  "python"
-)
-
 FAILED_TESTS=""
 
 export OVERRIDE_BAZEL_VERSION="$VERSION"
@@ -62,17 +81,31 @@ export OVERRIDE_BAZEL_WRAPPER_DOWNLOAD_DIR=/tmp
 
 ACTION_ENV_FLAG="--action_env=bazel_cache_invalidate=version_${VERSION}"
 
-tools/bazel version | grep "$VERSION" || { echo "Detected bazel version did not match expected value of $VERSION" >/dev/stderr; exit 1; }
-tools/bazel build "${ACTION_ENV_FLAG}" -- //... "${EXCLUDED_TARGETS[@]}" || FAILED_TESTS="${FAILED_TESTS}buildtest "
-tools/bazel build "${ACTION_ENV_FLAG}" --config fuzztest -- //fuzztest/... || FAILED_TESTS="${FAILED_TESTS}fuzztest_buildtest "
-
-for TEST_DIRECTORY in "${TEST_DIRECTORIES[@]}"; do
-  pushd "test/distrib/bazel/$TEST_DIRECTORY/"
-
-  tools/bazel version | grep "$VERSION" || { echo "Detected bazel version did not match expected value of $VERSION" >/dev/stderr; exit 1; }
-  tools/bazel test "${ACTION_ENV_FLAG}" --test_output=all //:all || FAILED_TESTS="${FAILED_TESTS}distribtest_${TEST_DIRECTORY} "
-
-  popd
+for TEST_SHARD in "${TEST_SHARDS[@]}"
+do
+  SHARD_RAN=""
+  if [ "${TEST_SHARD}" == "buildtest" ] ; then
+    tools/bazel version | grep "$VERSION" || { echo "Detected bazel version did not match expected value of $VERSION" >/dev/stderr; exit 1; }
+    tools/bazel build "${ACTION_ENV_FLAG}" -- //... "${EXCLUDED_TARGETS[@]}" || FAILED_TESTS="${FAILED_TESTS}buildtest "
+    tools/bazel build "${ACTION_ENV_FLAG}" --config fuzztest -- //fuzztest/... || FAILED_TESTS="${FAILED_TESTS}fuzztest_buildtest "
+    SHARD_RAN="true"
+  fi
+
+  for TEST_DIRECTORY in "${TEST_DIRECTORIES[@]}"
+  do
+    pushd "test/distrib/bazel/${TEST_DIRECTORY}/"
+    if [ "${TEST_SHARD}" == "distribtest_${TEST_DIRECTORY}" ] ; then
+      tools/bazel version | grep "$VERSION" || { echo "Detected bazel version did not match expected value of $VERSION" >/dev/stderr; exit 1; }
+      tools/bazel test "${ACTION_ENV_FLAG}" --test_output=all //:all || FAILED_TESTS="${FAILED_TESTS}distribtest_${TEST_DIRECTORY} "
+      SHARD_RAN="true"
+    fi
+    popd
+  done
+
+  if [ "${SHARD_RAN}" == "" ]; then
+    echo "Unknown shard '${TEST_SHARD}'"
+    exit 1
+  fi
 done
 
 if [ "$FAILED_TESTS" != "" ]

tools/bazelify_tests/test/bazel_distribtests.bzl

@@ -19,6 +19,12 @@ Generates portability tests.
 load("supported_bazel_versions.bzl", "SUPPORTED_BAZEL_VERSIONS")
 load("//tools/bazelify_tests:build_defs.bzl", "grpc_run_bazel_distribtest_test")
 
+_TEST_SHARDS = [
+    "buildtest",
+    "distribtest_cpp",
+    "distribtest_python",
+]
+
 def generate_bazel_distribtests(name):
     """Generates the bazel distribtests.
 
@@ -28,14 +34,15 @@ def generate_bazel_distribtests(name):
     test_names = []
 
     for bazel_version in SUPPORTED_BAZEL_VERSIONS:
-        test_name = "bazel_distribtest_%s" % bazel_version
-        grpc_run_bazel_distribtest_test(
-            name = test_name,
-            size = "enormous",
-            args = [bazel_version],
-            docker_image_version = "tools/dockerfile/test/bazel.current_version",
-        )
-        test_names.append(test_name)
+        for shard_name in _TEST_SHARDS:
+            test_name = "bazel_distribtest_%s_%s" % (bazel_version, shard_name)
+            grpc_run_bazel_distribtest_test(
+                name = test_name,
+                size = "enormous",
+                args = [bazel_version, shard_name],
+                docker_image_version = "tools/dockerfile/test/bazel.current_version",
+            )
+            test_names.append(test_name)
 
     # Generate test suite that allows easily running all bazel distribtests.
     native.test_suite(