Support partitions for other tests.

PiperOrigin-RevId: 345399936

3 files changed, 8 insertions, 6 deletions

diff --git a/tools/bazel.mk b/tools/bazel.mk
index 9063eebaf..ca5621a9c 100644
--- a/tools/bazel.mk
+++ b/tools/bazel.mk
@@ -41,6 +41,9 @@ DOCKER_CONFIG := /etc/docker
 # Bazel flags.
 BAZEL := bazel $(STARTUP_OPTIONS)
 BASE_OPTIONS := --color=no --curses=no
+ifneq (,$(BAZEL_CONFIG))
+BASE_OPTIONS += --config=$(BAZEL_CONFIG)
+endif
 
 # Basic options.
 UID := $(shell id -u ${USER})
@@ -103,11 +106,6 @@ FULL_DOCKER_RUN_OPTIONS += --group-add $(KVM_GROUP)
 endif
 endif
 
-# Load the appropriate config.
-ifneq (,$(BAZEL_CONFIG))
-OPTIONS += --config=$(BAZEL_CONFIG)
-endif
-
 bazel-image: load-default
 	@if docker ps --all | grep $(BUILDER_NAME); then docker rm -f $(BUILDER_NAME); fi
 	docker run --user 0:0 --entrypoint "" --name $(BUILDER_NAME) \
diff --git a/tools/bazeldefs/defs.bzl b/tools/bazeldefs/defs.bzl
index c2f94bb9c..51e17a79a 100644
--- a/tools/bazeldefs/defs.bzl
+++ b/tools/bazeldefs/defs.bzl
@@ -7,6 +7,8 @@ build_test = _build_test
 bzl_library = _bzl_library
 rbe_platform = native.platform
 rbe_toolchain = native.toolchain
+more_shards = 4
+most_shards = 8
 
 def short_path(path):
     return path
diff --git a/tools/defs.bzl b/tools/defs.bzl
index 2c8129e7e..b6f188aeb 100644
--- a/tools/defs.bzl
+++ b/tools/defs.bzl
@@ -8,7 +8,7 @@ change for Google-internal and bazel-compatible rules.
 load("//tools/go_stateify:defs.bzl", "go_stateify")
 load("//tools/go_marshal:defs.bzl", "go_marshal", "marshal_deps", "marshal_test_deps")
 load("//tools/nogo:defs.bzl", "nogo_test")
-load("//tools/bazeldefs:defs.bzl", _build_test = "build_test", _bzl_library = "bzl_library", _coreutil = "coreutil", _default_installer = "default_installer", _default_net_util = "default_net_util", _proto_library = "proto_library", _rbe_platform = "rbe_platform", _rbe_toolchain = "rbe_toolchain", _select_arch = "select_arch", _select_system = "select_system", _short_path = "short_path")
+load("//tools/bazeldefs:defs.bzl", _build_test = "build_test", _bzl_library = "bzl_library", _coreutil = "coreutil", _default_installer = "default_installer", _default_net_util = "default_net_util", _more_shards = "more_shards", _most_shards = "most_shards", _proto_library = "proto_library", _rbe_platform = "rbe_platform", _rbe_toolchain = "rbe_toolchain", _select_arch = "select_arch", _select_system = "select_system", _short_path = "short_path")
 load("//tools/bazeldefs:cc.bzl", _cc_binary = "cc_binary", _cc_flags_supplier = "cc_flags_supplier", _cc_grpc_library = "cc_grpc_library", _cc_library = "cc_library", _cc_proto_library = "cc_proto_library", _cc_test = "cc_test", _cc_toolchain = "cc_toolchain", _gbenchmark = "gbenchmark", _grpcpp = "grpcpp", _gtest = "gtest", _vdso_linker_option = "vdso_linker_option")
 load("//tools/bazeldefs:go.bzl", _gazelle = "gazelle", _go_binary = "go_binary", _go_embed_data = "go_embed_data", _go_grpc_and_proto_libraries = "go_grpc_and_proto_libraries", _go_library = "go_library", _go_path = "go_path", _go_proto_library = "go_proto_library", _go_test = "go_test", _select_goarch = "select_goarch", _select_goos = "select_goos")
 load("//tools/bazeldefs:pkg.bzl", _pkg_deb = "pkg_deb", _pkg_tar = "pkg_tar")
@@ -26,6 +26,8 @@ short_path = _short_path
 rbe_platform = _rbe_platform
 rbe_toolchain = _rbe_toolchain
 coreutil = _coreutil
+more_shards = _more_shards
+most_shards = _most_shards
 
 # C++ rules.
 cc_binary = _cc_binary