Introduce --default_test_resources flag

Add `--default_test_resources=<resource>=<value(s)>` that allows setting the default resource utilization for tests. The flag follow a syntax simialar to `--local_resources=<resource>=<value>`, in that it allow assigning different resource types, and `--test_timeout=<value(s)>`, which accepts either 1 or 4 intergers to assign to all test sizes or to each individually.

Relates to #19679

Closes #20839.

PiperOrigin-RevId: 606233269
Change-Id: Ia53e42820ba9aa646b0600fe4e9f95f146d7b2b9

src/main/java/com/google/devtools/build/lib/analysis/BUILD

@@ -2734,7 +2734,10 @@ java_library(
 
 java_library(
     name = "test/test_configuration",
-    srcs = ["test/TestConfiguration.java"],
+    srcs = [
+        "test/TestConfiguration.java",
+        "test/TestResourcesConverter.java",
+    ],
     deps = [
         ":config/build_options",
         ":config/core_option_converters",
@@ -2747,6 +2750,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/cmdline",
         "//src/main/java/com/google/devtools/build/lib/packages",
         "//src/main/java/com/google/devtools/build/lib/util",
+        "//src/main/java/com/google/devtools/build/lib/util:resource_converter",
         "//src/main/java/com/google/devtools/common/options",
         "//third_party:guava",
     ],

src/main/java/com/google/devtools/build/lib/analysis/test/TestConfiguration.java

@@ -28,7 +28,9 @@
 import com.google.devtools.build.lib.analysis.test.CoverageConfiguration.CoverageOptions;
 import com.google.devtools.build.lib.analysis.test.TestShardingStrategy.ShardingStrategyConverter;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.devtools.build.lib.packages.TestSize;
 import com.google.devtools.build.lib.packages.TestTimeout;
+import com.google.devtools.build.lib.util.Pair;
 import com.google.devtools.build.lib.util.RegexFilter;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.OptionDefinition;
@@ -84,6 +86,26 @@ public static class TestOptions extends FragmentOptions {
                 + "-1 tells blaze to use its default timeouts for that category.")
     public Map<TestTimeout, Duration> testTimeout;
 
+    @Option(
+        name = "default_test_resources",
+        defaultValue = "null",
+        converter = TestResourcesConverter.class,
+        allowMultiple = true,
+        documentationCategory = OptionDocumentationCategory.TESTING,
+        effectTags = {OptionEffectTag.UNKNOWN},
+        help =
+            "Override the default resources amount for tests. The expected format is"
+                + " <resource>=<value>. If a single positive number is specified as <value>"
+                + " it will override the default resources for all test sizes. If 4"
+                + " comma-separated numbers are specified, they will override the resource"
+                + " amount for respectively the small, medium, large, enormous test sizes."
+                + " Values can also be HOST_RAM/HOST_CPU, optionally followed"
+                + " by [-|*]<float> (eg. memory=HOST_RAM*.1,HOST_RAM*.2,HOST_RAM*.3,HOST_RAM*.4)."
+                + " The default test resources specified by this flag are overridden by explicit"
+                + " resources specified in tags.")
+    // We need to store these as Pair(s) instead of Map.Entry(s) so that they are serializable.
+    public List<Pair<String, Map<TestSize, Double>>> testResources;
+
     @Option(
       name = "test_filter",
       allowMultiple = false,
@@ -320,16 +342,28 @@ public FragmentOptions getExec() {
   private final TestOptions options;
   private final ImmutableMap<TestTimeout, Duration> testTimeout;
   private final boolean shouldInclude;
+  private final ImmutableMap<TestSize, ImmutableMap<String, Double>> testResources;
 
   public TestConfiguration(BuildOptions buildOptions) {
     this.shouldInclude = buildOptions.contains(TestOptions.class);
     if (shouldInclude) {
       TestOptions options = buildOptions.get(TestOptions.class);
       this.options = options;
       this.testTimeout = ImmutableMap.copyOf(options.testTimeout);
+      ImmutableMap.Builder<TestSize, ImmutableMap<String, Double>> testResources =
+          ImmutableMap.builderWithExpectedSize(TestSize.values().length);
+      for (TestSize size : TestSize.values()) {
+        ImmutableMap.Builder<String, Double> resources = ImmutableMap.builder();
+        for (Pair<String, Map<TestSize, Double>> resource : options.testResources) {
+          resources.put(resource.getFirst(), resource.getSecond().get(size));
+        }
+        testResources.put(size, resources.buildKeepingLast());
+      }
+      this.testResources = testResources.buildOrThrow();
     } else {
       this.options = null;
       this.testTimeout = null;
+      this.testResources = ImmutableMap.of();
     }
   }
 
@@ -343,6 +377,11 @@ public ImmutableMap<TestTimeout, Duration> getTestTimeout() {
     return testTimeout;
   }
 
+  /** Returns test resource mapping as set by --default_test_resources options. */
+  public ImmutableMap<String, Double> getTestResources(TestSize size) {
+    return testResources.getOrDefault(size, ImmutableMap.of());
+  }
+
   public String getTestFilter() {
     return options.testFilter;
   }

src/main/java/com/google/devtools/build/lib/analysis/test/TestResourcesConverter.java

@@ -0,0 +1,67 @@
+// Copyright 2024 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.analysis.test;
+
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Maps;
+import com.google.devtools.build.lib.packages.TestSize;
+import com.google.devtools.build.lib.util.Pair;
+import com.google.devtools.build.lib.util.ResourceConverter;
+import com.google.devtools.common.options.Converter;
+import com.google.devtools.common.options.Converters;
+import com.google.devtools.common.options.OptionsParsingException;
+import java.util.ArrayList;
+import java.util.EnumMap;
+import java.util.Map;
+
+public class TestResourcesConverter
+    extends Converter.Contextless<Pair<String, Map<TestSize, Double>>> {
+  private static final Converters.AssignmentConverter assignmentConverter =
+      new Converters.AssignmentConverter();
+  private static final ResourceConverter.DoubleConverter resourceConverter =
+      new ResourceConverter.DoubleConverter(
+          /* keywords= */ ImmutableMap.of(
+              ResourceConverter.HOST_CPUS_KEYWORD,
+                  () -> (double) ResourceConverter.HOST_CPUS_SUPPLIER.get(),
+              ResourceConverter.HOST_RAM_KEYWORD,
+                  () -> (double) ResourceConverter.HOST_RAM_SUPPLIER.get()),
+          /* minValue= */ 0.0,
+          /* maxValue= */ Double.MAX_VALUE);
+
+  @Override
+  public String getTypeDescription() {
+    return "a resource name followed by equal and 1 float or 4 float, e.g memory=10,30,60,100";
+  }
+
+  @Override
+  public Pair<String, Map<TestSize, Double>> convert(String input) throws OptionsParsingException {
+    Map.Entry<String, String> assignment = assignmentConverter.convert(input);
+    ArrayList<Double> values = new ArrayList<>(TestSize.values().length);
+    for (String s : Splitter.on(",").splitToList(assignment.getValue())) {
+      values.add(resourceConverter.convert(s));
+    }
+
+    if (values.size() != 1 && values.size() != TestSize.values().length) {
+      throw new OptionsParsingException("Invalid number of comma-separated entries in " + input);
+    }
+
+    EnumMap<TestSize, Double> amounts = Maps.newEnumMap(TestSize.class);
+    for (TestSize size : TestSize.values()) {
+      amounts.put(size, values.get(Math.min(values.size() - 1, size.ordinal())));
+    }
+    return Pair.of(assignment.getKey(), amounts);
+  }
+}

src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java

@@ -71,6 +71,7 @@ private static ResourceSet getResourceSetFromSize(TestSize size) {
   private final boolean isExternal;
   private final String language;
   private final ImmutableMap<String, String> executionInfo;
+  private final TestConfiguration testConfiguration;
 
   /**
    * Creates test target properties instance. Constructor expects that it will be called only for
@@ -93,7 +94,7 @@ private static ResourceSet getResourceSetFromSize(TestSize size) {
 
     boolean incompatibleExclusiveTestSandboxed = false;
 
-    TestConfiguration testConfiguration = ruleContext.getFragment(TestConfiguration.class);
+    testConfiguration = ruleContext.getFragment(TestConfiguration.class);
     if (testConfiguration != null) {
       incompatibleExclusiveTestSandboxed = testConfiguration.incompatibleExclusiveTestSandboxed();
     }
@@ -201,11 +202,21 @@ public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs
       return LOCAL_TEST_JOBS_BASED_RESOURCES;
     }
 
-    Map<String, Double> resources = parseTags(label, executionInfo);
-    ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size);
-    testResourcesFromSize.getResources().forEach(resources::putIfAbsent);
+    ResourceSet defaultResources = getResourceSetFromSize(size);
+    Map<String, Double> resourcesFromTags = parseTags(label, executionInfo);
+    Map<String, Double> configResources =
+        testConfiguration == null ? ImmutableMap.of() : testConfiguration.getTestResources(size);
+    if (resourcesFromTags.isEmpty() && configResources.isEmpty()) {
+      return defaultResources;
+    }
+
     return ResourceSet.create(
-        ImmutableMap.copyOf(resources), testResourcesFromSize.getLocalTestCount());
+        ImmutableMap.<String, Double>builder()
+            .putAll(defaultResources.getResources())
+            .putAll(configResources)
+            .putAll(resourcesFromTags)
+            .buildKeepingLast(),
+        defaultResources.getLocalTestCount());
   }
 
   private static FailureDetail createFailureDetail(String message, Code detailedCode) {

src/test/java/com/google/devtools/build/lib/rules/test/TestTargetPropertiesTest.java

@@ -52,6 +52,37 @@ public void testTestWithCpusTagHasCorrectLocalResourcesEstimate() throws Excepti
     assertThat(localResourceUsage.getCpuUsage()).isEqualTo(4.0);
   }
 
+  @Test
+  public void testTestResourcesFlag() throws Exception {
+    scratch.file("tests/test.sh", "#!/bin/bash", "exit 0");
+    scratch.file(
+        "tests/BUILD",
+        "sh_test(",
+        "  name = 'test',",
+        "  size = 'medium',",
+        "  srcs = ['test.sh'],",
+        "  tags = ['resources:gpu:4'],",
+        ")");
+    useConfiguration(
+        "--default_test_resources=memory=10,20,30,40",
+        "--default_test_resources=cpu=1,2,3,4",
+        "--default_test_resources=gpu=1",
+        "--default_test_resources=cpu=5");
+    ConfiguredTarget testTarget = getConfiguredTarget("//tests:test");
+    TestRunnerAction testAction =
+        (TestRunnerAction)
+            getGeneratingAction(TestProvider.getTestStatusArtifacts(testTarget).get(0));
+    ResourceSet localResourceUsage =
+        testAction
+            .getTestProperties()
+            .getLocalResourceUsage(testAction.getOwner().getLabel(), false);
+    // Tags-specified resources overrides --default_test_resources=gpu.
+    assertThat(localResourceUsage.getResources()).containsEntry("gpu", 4.0);
+    // The last-specified value of --default_test_resources=cpu is used.
+    assertThat(localResourceUsage.getCpuUsage()).isEqualTo(5.0);
+    assertThat(localResourceUsage.getMemoryMb()).isEqualTo(20);
+  }
+
   @Test
   public void testTestWithExclusiveRunLocallyByDefault() throws Exception {
     useConfiguration("--noincompatible_exclusive_test_sandboxed");