remote test execution run_as_bundled attr

Summary:
This configures test bundling on the `remote_test_execution_toolchain`, which is needed to avoid overloading zippy qps. It is essentially a "part 2" for D62482846.

#buildall

Reviewed By: akrieger

Differential Revision: D62984203

fbshipit-source-id: 90a4501aaa5def15686c74e281d0410686dcc6ec

test/inject_test_run_info.bzl

@@ -5,14 +5,21 @@
 # License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 # of this source tree.
 
+load(
+    "@prelude//tests:re_utils.bzl",
+    "maybe_add_run_as_bundle_label",
+)
+
 def inject_test_run_info(ctx: AnalysisContext, test_info: ExternalRunnerTestInfo) -> list[Provider]:
     # Access this here so we get failures in CI if we forget to inject it
     # anywhere, regardless of whether an `env` is used.
     inject_test_env = ctx.attrs._inject_test_env[RunInfo]
 
-    # If forcing RE on tpx, check if the test suite should be run as a bundle
-    if read_config("tpx", "force_run_as_bundle") == "True":
-        test_info.labels.extend(["run_as_bundle"])
+    # `if test_info.labels != None` doesn't work because `None` is not of type `list[str]`,
+    # yet it is None in some cases... this hack lets us check for None without a type error.
+    if getattr(test_info, "labels", None) != None:
+        # If forcing RE on tpx, check if the test suite should be run as a bundle
+        maybe_add_run_as_bundle_label(ctx, test_info.labels)
 
     if (not test_info.env) or _exclude_test_env_from_run_info(ctx):
         return [test_info, RunInfo(args = test_info.command)]

tests/re_utils.bzl

@@ -9,24 +9,43 @@ load("@prelude//:build_mode.bzl", "BuildModeInfo")
 load("@prelude//tests:remote_test_execution_toolchain.bzl", "RemoteTestExecutionToolchainInfo")
 load("@prelude//utils:expect.bzl", "expect_non_none")
 
-def _get_re_arg(ctx: AnalysisContext):
+ReArg = record(
+    re_props = field(dict | None),
+    default_run_as_bundle = field(bool | None),
+)
+
+def _get_re_arg(ctx: AnalysisContext) -> ReArg:
     if not hasattr(ctx.attrs, "remote_execution"):
-        return None
+        return ReArg(re_props = None, default_run_as_bundle = False)
 
     if ctx.attrs.remote_execution != None:
-        # If this is a string, look up the profile on the RE toolchain.
+        # If this is a string, look up the re_props on the RE toolchain.
         if type(ctx.attrs.remote_execution) == type(""):
             expect_non_none(ctx.attrs._remote_test_execution_toolchain)
-            return ctx.attrs._remote_test_execution_toolchain[RemoteTestExecutionToolchainInfo].profiles[ctx.attrs.remote_execution]
+            return ReArg(
+                re_props =
+                    ctx.attrs._remote_test_execution_toolchain[RemoteTestExecutionToolchainInfo].profiles[ctx.attrs.remote_execution],
+                default_run_as_bundle =
+                    ctx.attrs._remote_test_execution_toolchain[RemoteTestExecutionToolchainInfo].default_run_as_bundle,
+            )
 
-        return ctx.attrs.remote_execution
+        return ReArg(re_props = ctx.attrs.remote_execution, default_run_as_bundle = False)
 
     # Check for a default RE option on the toolchain.
     re_toolchain = ctx.attrs._remote_test_execution_toolchain
     if re_toolchain != None and re_toolchain[RemoteTestExecutionToolchainInfo].default_profile != None:
-        return re_toolchain[RemoteTestExecutionToolchainInfo].default_profile
+        return ReArg(
+            re_props = re_toolchain[RemoteTestExecutionToolchainInfo].default_profile,
+            default_run_as_bundle = re_toolchain[RemoteTestExecutionToolchainInfo].default_run_as_bundle,
+        )
+
+    return ReArg(re_props = None, default_run_as_bundle = False)
 
-    return None
+def maybe_add_run_as_bundle_label(ctx: AnalysisContext, labels: list[str]) -> None:
+    re_arg = _get_re_arg(ctx)
+    run_as_bundle = re_arg.default_run_as_bundle and "re_ignore_default_run_as_bundle" not in ctx.attrs.labels
+    if run_as_bundle or read_config("tpx", "force_run_as_bundle") == "True":
+        labels.extend(["run_as_bundle"])
 
 def get_re_executors_from_props(ctx: AnalysisContext) -> ([CommandExecutorConfig, None], dict[str, CommandExecutorConfig]):
     """
@@ -35,7 +54,7 @@ def get_re_executors_from_props(ctx: AnalysisContext) -> ([CommandExecutorConfig
     Returns (default_executor, executor_overrides).
     """
 
-    re_props = _get_re_arg(ctx)
+    re_props = _get_re_arg(ctx).re_props
     if re_props == None:
         # If no RE args are set and an RE config is specified
         if bool(read_config("tpx", "force_re_props")):

tests/remote_test_execution_toolchain.bzl

@@ -12,5 +12,9 @@ RemoteTestExecutionToolchainInfo = provider(
         # A dictionary of string names to pre-registered profiles.  Rules can
         # use the profile name to references these.
         "profiles",
+        # A bool indicating whether the test suites executed by this toolchain
+        # should be run in a bundle. This makes all tests in a suite run in
+        # a single RE action as opposed to one action per test.
+        "default_run_as_bundle",
     ],
 )

toolchains/remote_test_execution.bzl

@@ -10,10 +10,17 @@ load("@prelude//tests:remote_test_execution_toolchain.bzl", "RemoteTestExecution
 load("@prelude//utils:utils.bzl", "map_val")
 
 def _impl(ctx: AnalysisContext) -> list[Provider]:
+    default_profile = map_val(ctx.attrs.profiles.get, ctx.attrs.default_profile)
+    if ctx.attrs.default_run_as_bundle != None:
+        default_run_as_bundle = ctx.attrs.default_run_as_bundle
+    else:
+        default_run_as_bundle = bool(default_profile)
+
     return [
         DefaultInfo(),
         RemoteTestExecutionToolchainInfo(
-            default_profile = map_val(ctx.attrs.profiles.get, ctx.attrs.default_profile),
+            default_profile = default_profile,
+            default_run_as_bundle = default_run_as_bundle,
             profiles = ctx.attrs.profiles,
         ),
     ]
@@ -23,6 +30,7 @@ remote_test_execution_toolchain = rule(
     is_toolchain_rule = True,
     attrs = {
         "default_profile": attrs.option(attrs.string(), default = None),
+        "default_run_as_bundle": attrs.option(attrs.bool(), default = None),
         "profiles": attrs.dict(
             key = attrs.string(),
             value = attrs.option(re_test_common.opts_for_tests_arg()),