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

@@ -2733,7 +2733,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",
@@ -2746,6 +2749,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");