From e8b368272ce14a350580ce5d6ee388cbb94f342f Mon Sep 17 00:00:00 2001 From: Alessandro Patti Date: Mon, 12 Feb 2024 06:31:58 -0800 Subject: [PATCH] Introduce --default_test_resources flag Add `--default_test_resources==` that allows setting the default resource utilization for tests. The flag follow a syntax simialar to `--local_resources==`, in that it allow assigning different resource types, and `--test_timeout=`, 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 --- .../google/devtools/build/lib/analysis/BUILD | 6 +- .../lib/analysis/test/TestConfiguration.java | 39 +++++++++++ .../analysis/test/TestResourcesConverter.java | 67 +++++++++++++++++++ .../analysis/test/TestTargetProperties.java | 21 ++++-- .../rules/test/TestTargetPropertiesTest.java | 31 +++++++++ 5 files changed, 158 insertions(+), 6 deletions(-) create mode 100644 src/main/java/com/google/devtools/build/lib/analysis/test/TestResourcesConverter.java diff --git a/src/main/java/com/google/devtools/build/lib/analysis/BUILD b/src/main/java/com/google/devtools/build/lib/analysis/BUILD index 189a962464aa20..a8c437f0b0eb42 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/BUILD +++ b/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", ], diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestConfiguration.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestConfiguration.java index 940d7ab142384a..562ea3bca15675 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/test/TestConfiguration.java +++ b/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; + @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" + + " =. If a single positive number is specified as " + + " 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 [-|*] (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>> testResources; + @Option( name = "test_filter", allowMultiple = false, @@ -320,6 +342,7 @@ public FragmentOptions getExec() { private final TestOptions options; private final ImmutableMap testTimeout; private final boolean shouldInclude; + private final ImmutableMap> testResources; public TestConfiguration(BuildOptions buildOptions) { this.shouldInclude = buildOptions.contains(TestOptions.class); @@ -327,9 +350,20 @@ public TestConfiguration(BuildOptions buildOptions) { TestOptions options = buildOptions.get(TestOptions.class); this.options = options; this.testTimeout = ImmutableMap.copyOf(options.testTimeout); + ImmutableMap.Builder> testResources = + ImmutableMap.builderWithExpectedSize(TestSize.values().length); + for (TestSize size : TestSize.values()) { + ImmutableMap.Builder resources = ImmutableMap.builder(); + for (Pair> 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 getTestTimeout() { return testTimeout; } + /** Returns test resource mapping as set by --default_test_resources options. */ + public ImmutableMap getTestResources(TestSize size) { + return testResources.getOrDefault(size, ImmutableMap.of()); + } + public String getTestFilter() { return options.testFilter; } diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestResourcesConverter.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestResourcesConverter.java new file mode 100644 index 00000000000000..5f058251dda775 --- /dev/null +++ b/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>> { + 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> convert(String input) throws OptionsParsingException { + Map.Entry assignment = assignmentConverter.convert(input); + ArrayList 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 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); + } +} diff --git a/src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java b/src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java index bcc1cc9b2f0766..3a9e9a348bc4dd 100644 --- a/src/main/java/com/google/devtools/build/lib/analysis/test/TestTargetProperties.java +++ b/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 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 resources = parseTags(label, executionInfo); - ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size); - testResourcesFromSize.getResources().forEach(resources::putIfAbsent); + ResourceSet defaultResources = getResourceSetFromSize(size); + Map resourcesFromTags = parseTags(label, executionInfo); + Map configResources = + testConfiguration == null ? ImmutableMap.of() : testConfiguration.getTestResources(size); + if (resourcesFromTags.isEmpty() && configResources.isEmpty()) { + return defaultResources; + } + return ResourceSet.create( - ImmutableMap.copyOf(resources), testResourcesFromSize.getLocalTestCount()); + ImmutableMap.builder() + .putAll(defaultResources.getResources()) + .putAll(configResources) + .putAll(resourcesFromTags) + .buildKeepingLast(), + defaultResources.getLocalTestCount()); } private static FailureDetail createFailureDetail(String message, Code detailedCode) { diff --git a/src/test/java/com/google/devtools/build/lib/rules/test/TestTargetPropertiesTest.java b/src/test/java/com/google/devtools/build/lib/rules/test/TestTargetPropertiesTest.java index 4ef6077cfab3fa..ac06118ea8a549 100644 --- a/src/test/java/com/google/devtools/build/lib/rules/test/TestTargetPropertiesTest.java +++ b/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");