Skip to content

Commit 0c5b6e8

Browse files
AlessandroPatticopybara-github
AlessandroPatti
authored and
copybara-github
committed
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
1 parent af28cae commit 0c5b6e8

File tree

5 files changed

+158
-6
lines changed

5 files changed

+158
-6
lines changed

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2761,7 +2761,10 @@ java_library(
27612761

27622762
java_library(
27632763
name = "test/test_configuration",
2764-
srcs = ["test/TestConfiguration.java"],
2764+
srcs = [
2765+
"test/TestConfiguration.java",
2766+
"test/TestResourcesConverter.java",
2767+
],
27652768
deps = [
27662769
":config/build_options",
27672770
":config/core_option_converters",
@@ -2774,6 +2777,7 @@ java_library(
27742777
"//src/main/java/com/google/devtools/build/lib/cmdline",
27752778
"//src/main/java/com/google/devtools/build/lib/packages",
27762779
"//src/main/java/com/google/devtools/build/lib/util",
2780+
"//src/main/java/com/google/devtools/build/lib/util:resource_converter",
27772781
"//src/main/java/com/google/devtools/common/options",
27782782
"//third_party:guava",
27792783
],

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,9 @@
2828
import com.google.devtools.build.lib.analysis.test.CoverageConfiguration.CoverageOptions;
2929
import com.google.devtools.build.lib.analysis.test.TestShardingStrategy.ShardingStrategyConverter;
3030
import com.google.devtools.build.lib.cmdline.Label;
31+
import com.google.devtools.build.lib.packages.TestSize;
3132
import com.google.devtools.build.lib.packages.TestTimeout;
33+
import com.google.devtools.build.lib.util.Pair;
3234
import com.google.devtools.build.lib.util.RegexFilter;
3335
import com.google.devtools.common.options.Option;
3436
import com.google.devtools.common.options.OptionDefinition;
@@ -84,6 +86,26 @@ public static class TestOptions extends FragmentOptions {
8486
+ "-1 tells blaze to use its default timeouts for that category.")
8587
public Map<TestTimeout, Duration> testTimeout;
8688

89+
@Option(
90+
name = "default_test_resources",
91+
defaultValue = "null",
92+
converter = TestResourcesConverter.class,
93+
allowMultiple = true,
94+
documentationCategory = OptionDocumentationCategory.TESTING,
95+
effectTags = {OptionEffectTag.UNKNOWN},
96+
help =
97+
"Override the default resources amount for tests. The expected format is"
98+
+ " <resource>=<value>. If a single positive number is specified as <value>"
99+
+ " it will override the default resources for all test sizes. If 4"
100+
+ " comma-separated numbers are specified, they will override the resource"
101+
+ " amount for respectively the small, medium, large, enormous test sizes."
102+
+ " Values can also be HOST_RAM/HOST_CPU, optionally followed"
103+
+ " by [-|*]<float> (eg. memory=HOST_RAM*.1,HOST_RAM*.2,HOST_RAM*.3,HOST_RAM*.4)."
104+
+ " The default test resources specified by this flag are overridden by explicit"
105+
+ " resources specified in tags.")
106+
// We need to store these as Pair(s) instead of Map.Entry(s) so that they are serializable.
107+
public List<Pair<String, Map<TestSize, Double>>> testResources;
108+
87109
@Option(
88110
name = "test_filter",
89111
allowMultiple = false,
@@ -308,16 +330,28 @@ public static class TestOptions extends FragmentOptions {
308330
private final TestOptions options;
309331
private final ImmutableMap<TestTimeout, Duration> testTimeout;
310332
private final boolean shouldInclude;
333+
private final ImmutableMap<TestSize, ImmutableMap<String, Double>> testResources;
311334

312335
public TestConfiguration(BuildOptions buildOptions) {
313336
this.shouldInclude = buildOptions.contains(TestOptions.class);
314337
if (shouldInclude) {
315338
TestOptions options = buildOptions.get(TestOptions.class);
316339
this.options = options;
317340
this.testTimeout = ImmutableMap.copyOf(options.testTimeout);
341+
ImmutableMap.Builder<TestSize, ImmutableMap<String, Double>> testResources =
342+
ImmutableMap.builderWithExpectedSize(TestSize.values().length);
343+
for (TestSize size : TestSize.values()) {
344+
ImmutableMap.Builder<String, Double> resources = ImmutableMap.builder();
345+
for (Pair<String, Map<TestSize, Double>> resource : options.testResources) {
346+
resources.put(resource.getFirst(), resource.getSecond().get(size));
347+
}
348+
testResources.put(size, resources.buildKeepingLast());
349+
}
350+
this.testResources = testResources.buildOrThrow();
318351
} else {
319352
this.options = null;
320353
this.testTimeout = null;
354+
this.testResources = ImmutableMap.of();
321355
}
322356
}
323357

@@ -331,6 +365,11 @@ public ImmutableMap<TestTimeout, Duration> getTestTimeout() {
331365
return testTimeout;
332366
}
333367

368+
/** Returns test resource mapping as set by --default_test_resources options. */
369+
public ImmutableMap<String, Double> getTestResources(TestSize size) {
370+
return testResources.getOrDefault(size, ImmutableMap.of());
371+
}
372+
334373
public String getTestFilter() {
335374
return options.testFilter;
336375
}

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
// Copyright 2024 The Bazel Authors. All rights reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package com.google.devtools.build.lib.analysis.test;
16+
17+
import com.google.common.base.Splitter;
18+
import com.google.common.collect.ImmutableMap;
19+
import com.google.common.collect.Maps;
20+
import com.google.devtools.build.lib.packages.TestSize;
21+
import com.google.devtools.build.lib.util.Pair;
22+
import com.google.devtools.build.lib.util.ResourceConverter;
23+
import com.google.devtools.common.options.Converter;
24+
import com.google.devtools.common.options.Converters;
25+
import com.google.devtools.common.options.OptionsParsingException;
26+
import java.util.ArrayList;
27+
import java.util.EnumMap;
28+
import java.util.Map;
29+
30+
public class TestResourcesConverter
31+
extends Converter.Contextless<Pair<String, Map<TestSize, Double>>> {
32+
private static final Converters.AssignmentConverter assignmentConverter =
33+
new Converters.AssignmentConverter();
34+
private static final ResourceConverter.DoubleConverter resourceConverter =
35+
new ResourceConverter.DoubleConverter(
36+
/* keywords= */ ImmutableMap.of(
37+
ResourceConverter.HOST_CPUS_KEYWORD,
38+
() -> (double) ResourceConverter.HOST_CPUS_SUPPLIER.get(),
39+
ResourceConverter.HOST_RAM_KEYWORD,
40+
() -> (double) ResourceConverter.HOST_RAM_SUPPLIER.get()),
41+
/* minValue= */ 0.0,
42+
/* maxValue= */ Double.MAX_VALUE);
43+
44+
@Override
45+
public String getTypeDescription() {
46+
return "a resource name followed by equal and 1 float or 4 float, e.g memory=10,30,60,100";
47+
}
48+
49+
@Override
50+
public Pair<String, Map<TestSize, Double>> convert(String input) throws OptionsParsingException {
51+
Map.Entry<String, String> assignment = assignmentConverter.convert(input);
52+
ArrayList<Double> values = new ArrayList<>(TestSize.values().length);
53+
for (String s : Splitter.on(",").splitToList(assignment.getValue())) {
54+
values.add(resourceConverter.convert(s));
55+
}
56+
57+
if (values.size() != 1 && values.size() != TestSize.values().length) {
58+
throw new OptionsParsingException("Invalid number of comma-separated entries in " + input);
59+
}
60+
61+
EnumMap<TestSize, Double> amounts = Maps.newEnumMap(TestSize.class);
62+
for (TestSize size : TestSize.values()) {
63+
amounts.put(size, values.get(Math.min(values.size() - 1, size.ordinal())));
64+
}
65+
return Pair.of(assignment.getKey(), amounts);
66+
}
67+
}

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

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,7 @@ private static ResourceSet getResourceSetFromSize(TestSize size) {
7171
private final boolean isExternal;
7272
private final String language;
7373
private final ImmutableMap<String, String> executionInfo;
74+
private final TestConfiguration testConfiguration;
7475

7576
/**
7677
* Creates test target properties instance. Constructor expects that it will be called only for
@@ -93,7 +94,7 @@ private static ResourceSet getResourceSetFromSize(TestSize size) {
9394

9495
boolean incompatibleExclusiveTestSandboxed = false;
9596

96-
TestConfiguration testConfiguration = ruleContext.getFragment(TestConfiguration.class);
97+
testConfiguration = ruleContext.getFragment(TestConfiguration.class);
9798
if (testConfiguration != null) {
9899
incompatibleExclusiveTestSandboxed = testConfiguration.incompatibleExclusiveTestSandboxed();
99100
}
@@ -201,11 +202,21 @@ public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs
201202
return LOCAL_TEST_JOBS_BASED_RESOURCES;
202203
}
203204

204-
Map<String, Double> resources = parseTags(label, executionInfo);
205-
ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size);
206-
testResourcesFromSize.getResources().forEach(resources::putIfAbsent);
205+
ResourceSet defaultResources = getResourceSetFromSize(size);
206+
Map<String, Double> resourcesFromTags = parseTags(label, executionInfo);
207+
Map<String, Double> configResources =
208+
testConfiguration == null ? ImmutableMap.of() : testConfiguration.getTestResources(size);
209+
if (resourcesFromTags.isEmpty() && configResources.isEmpty()) {
210+
return defaultResources;
211+
}
212+
207213
return ResourceSet.create(
208-
ImmutableMap.copyOf(resources), testResourcesFromSize.getLocalTestCount());
214+
ImmutableMap.<String, Double>builder()
215+
.putAll(defaultResources.getResources())
216+
.putAll(configResources)
217+
.putAll(resourcesFromTags)
218+
.buildKeepingLast(),
219+
defaultResources.getLocalTestCount());
209220
}
210221

211222
private static FailureDetail createFailureDetail(String message, Code detailedCode) {

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,37 @@ public void testTestWithCpusTagHasCorrectLocalResourcesEstimate() throws Excepti
5252
assertThat(localResourceUsage.getCpuUsage()).isEqualTo(4.0);
5353
}
5454

55+
@Test
56+
public void testTestResourcesFlag() throws Exception {
57+
scratch.file("tests/test.sh", "#!/bin/bash", "exit 0");
58+
scratch.file(
59+
"tests/BUILD",
60+
"sh_test(",
61+
" name = 'test',",
62+
" size = 'medium',",
63+
" srcs = ['test.sh'],",
64+
" tags = ['resources:gpu:4'],",
65+
")");
66+
useConfiguration(
67+
"--default_test_resources=memory=10,20,30,40",
68+
"--default_test_resources=cpu=1,2,3,4",
69+
"--default_test_resources=gpu=1",
70+
"--default_test_resources=cpu=5");
71+
ConfiguredTarget testTarget = getConfiguredTarget("//tests:test");
72+
TestRunnerAction testAction =
73+
(TestRunnerAction)
74+
getGeneratingAction(TestProvider.getTestStatusArtifacts(testTarget).get(0));
75+
ResourceSet localResourceUsage =
76+
testAction
77+
.getTestProperties()
78+
.getLocalResourceUsage(testAction.getOwner().getLabel(), false);
79+
// Tags-specified resources overrides --default_test_resources=gpu.
80+
assertThat(localResourceUsage.getResources()).containsEntry("gpu", 4.0);
81+
// The last-specified value of --default_test_resources=cpu is used.
82+
assertThat(localResourceUsage.getCpuUsage()).isEqualTo(5.0);
83+
assertThat(localResourceUsage.getMemoryMb()).isEqualTo(20);
84+
}
85+
5586
@Test
5687
public void testTestWithExclusiveRunLocallyByDefault() throws Exception {
5788
useConfiguration("--noincompatible_exclusive_test_sandboxed");

0 commit comments

Comments
 (0)