From 195fcd2ecbed06bab35f7b584b1701b4c4f6092d Mon Sep 17 00:00:00 2001 From: Lauri Peltonen Date: Wed, 15 Sep 2021 12:01:51 +0300 Subject: [PATCH] Implement extra resources for test actions Add support for user-specified resource types in the resource manager. This generalizes the CPU, RAM and "test count" resource support for other resource types such as the number of GPUs, available GPU memory, the number of embedded devices connected to the host, etc. The available amount of extra resources can be specified using the new --local_extra_resources== command line flag, which is analoguous to the existing --local_cpu_resources and --local_memory_resources flags. Tests can then declare the amount of extra resources they need by using a "resources::" tag. --- .../google/devtools/build/lib/actions/BUILD | 2 + .../lib/actions/ExecutionRequirements.java | 29 ++++++++ .../build/lib/actions/ResourceManager.java | 67 ++++++++++++++++++- .../build/lib/actions/ResourceSet.java | 37 ++++++++-- .../analysis/test/TestTargetProperties.java | 67 ++++++++++++++----- .../build/lib/buildtool/ExecutionTool.java | 5 ++ .../build/lib/exec/ExecutionOptions.java | 23 +++++-- .../build/lib/packages/TargetUtils.java | 3 +- .../devtools/common/options/Converters.java | 18 +++++ .../lib/actions/ResourceManagerTest.java | 51 +++++++++++++- 10 files changed, 272 insertions(+), 30 deletions(-) diff --git a/src/main/java/com/google/devtools/build/lib/actions/BUILD b/src/main/java/com/google/devtools/build/lib/actions/BUILD index b859a4ac7beaab..425b90d4bbdf40 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/BUILD +++ b/src/main/java/com/google/devtools/build/lib/actions/BUILD @@ -124,6 +124,7 @@ java_library( "//third_party:flogger", "//third_party:guava", "//third_party:jsr305", + "//third_party:rxjava3", "//third_party/protobuf:protobuf_java", ], ) @@ -305,6 +306,7 @@ java_library( "//src/main/java/com/google/devtools/common/options", "//third_party:flogger", "//third_party:guava", + "//third_party:rxjava3", ], ) diff --git a/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java b/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java index dcdd4358b497c4..e029c5ac28122f 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java +++ b/src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java @@ -152,6 +152,35 @@ public String parseIfMatches(String tag) throws ValidationException { return null; }); + /** How many extra resources an action requires for execution. */ + public static final ParseableRequirement RESOURCES = + ParseableRequirement.create( + "resources::", + Pattern.compile("resources:(.+:.+)"), + s -> { + Preconditions.checkNotNull(s); + + int splitIndex = s.indexOf(":"); + String resourceCount = s.substring(splitIndex+1); + int value; + try { + value = Integer.parseInt(resourceCount); + } catch (NumberFormatException e) { + return "can't be parsed as an integer"; + } + + // De-and-reserialize & compare to only allow canonical integer formats. + if (!Integer.toString(value).equals(resourceCount)) { + return "must be in canonical format (e.g. '4' instead of '+04')"; + } + + if (value < 1) { + return "can't be zero or negative"; + } + + return null; + }); + /** If an action supports running in persistent worker mode. */ public static final String SUPPORTS_WORKERS = "supports-workers"; diff --git a/src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java b/src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java index e8a57f6266fcce..fcb3f1b8415309 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java +++ b/src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java @@ -25,9 +25,13 @@ import com.google.devtools.build.lib.util.OS; import com.google.devtools.build.lib.util.Pair; import java.io.IOException; +import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.LinkedList; import java.util.List; +import java.util.Map; +import java.util.Set; import java.util.concurrent.CountDownLatch; /** @@ -137,6 +141,10 @@ public static ResourceManager instance() { // definition in the ResourceSet class. private double usedRam; + // Used amount of extra resources. Corresponds to the extra resource + // definition in the ResourceSet class. + private Map usedExtraResources; + // Used local test count. Corresponds to the local test count definition in the ResourceSet class. private int usedLocalTestCount; @@ -145,6 +153,7 @@ public static ResourceManager instance() { private ResourceManager() { requestList = new LinkedList<>(); + usedExtraResources = new HashMap<>(); } @VisibleForTesting public static ResourceManager instanceForTestingOnly() { @@ -158,6 +167,7 @@ private ResourceManager() { public synchronized void resetResourceUsage() { usedCpu = 0; usedRam = 0; + usedExtraResources = new HashMap<>(); usedLocalTestCount = 0; for (Pair request : requestList) { // CountDownLatch can be set only to 0 or 1. @@ -177,6 +187,7 @@ public synchronized void setAvailableResources(ResourceSet resources) { ResourceSet.create( staticResources.getMemoryMb(), staticResources.getCpuUsage(), + staticResources.getExtraResourceUsage(), staticResources.getLocalTestCount()); processWaitingThreads(); } @@ -265,6 +276,17 @@ ResourceHandle tryAcquire(ActionExecutionMetadata owner, ResourceSet resources) private void incrementResources(ResourceSet resources) { usedCpu += resources.getCpuUsage(); usedRam += resources.getMemoryMb(); + + resources.getExtraResourceUsage().entrySet().forEach( + resource -> { + String key = (String)resource.getKey(); + float value = resource.getValue(); + if (usedExtraResources.containsKey(key)) { + value += (float)usedExtraResources.get(key); + } + usedExtraResources.put(key, value); + }); + usedLocalTestCount += resources.getLocalTestCount(); } @@ -272,7 +294,7 @@ private void incrementResources(ResourceSet resources) { * Return true if any resources have been claimed through this manager. */ public synchronized boolean inUse() { - return usedCpu != 0.0 || usedRam != 0.0 || usedLocalTestCount != 0 || !requestList.isEmpty(); + return usedCpu != 0.0 || usedRam != 0.0 || !usedExtraResources.isEmpty() || usedLocalTestCount != 0 || !requestList.isEmpty(); } @@ -323,6 +345,13 @@ private synchronized CountDownLatch acquire(ResourceSet resources) { private synchronized boolean release(ResourceSet resources) { usedCpu -= resources.getCpuUsage(); usedRam -= resources.getMemoryMb(); + + for (Map.Entry resource : resources.getExtraResourceUsage().entrySet()) { + String key = (String)resource.getKey(); + float value = (float)usedExtraResources.get(key) - resource.getValue(); + usedExtraResources.put(key, value); + } + usedLocalTestCount -= resources.getLocalTestCount(); // TODO(bazel-team): (2010) rounding error can accumulate and value below can end up being @@ -334,6 +363,20 @@ private synchronized boolean release(ResourceSet resources) { if (usedRam < epsilon) { usedRam = 0; } + + Set toRemove = new HashSet<>(); + usedExtraResources.entrySet().forEach( + resource -> { + String key = (String)resource.getKey(); + float value = (float)usedExtraResources.get(key); + if (value < epsilon) { + toRemove.add(key); + } + }); + for (String key : toRemove) { + usedExtraResources.remove(key); + } + if (!requestList.isEmpty()) { processWaitingThreads(); return true; @@ -361,12 +404,29 @@ private synchronized void processWaitingThreads() { } } + /** + * Return true iff all requested extra resources are considered to be available. + */ + private boolean areExtraResourcesAvailable(ResourceSet resources) { + for (Map.Entry resource : resources.getExtraResourceUsage().entrySet()) { + String key = (String)resource.getKey(); + float used = (float)usedExtraResources.getOrDefault(key, 0f); + float requested = resource.getValue(); + float available = (float)availableResources.getExtraResourceUsage().getOrDefault(key, 0f); + float epsilon = 0.0001f; // Account for possible rounding errors. + if (requested != 0.0 && used != 0.0 && requested + used > available + epsilon) { + return false; + } + } + return true; + } + // Method will return true if all requested resources are considered to be available. private boolean areResourcesAvailable(ResourceSet resources) { Preconditions.checkNotNull(availableResources); // Comparison below is robust, since any calculation errors will be fixed // by the release() method. - if (usedCpu == 0.0 && usedRam == 0.0 && usedLocalTestCount == 0) { + if (usedCpu == 0.0 && usedRam == 0.0 && usedExtraResources.isEmpty() && usedLocalTestCount == 0) { return true; } // Use only MIN_NECESSARY_???_RATIO of the resource value to check for @@ -410,9 +470,10 @@ private boolean areResourcesAvailable(ResourceSet resources) { // 3) If used resource amount is less than total available resource amount. boolean cpuIsAvailable = cpu == 0.0 || usedCpu == 0.0 || usedCpu + cpu <= availableCpu; boolean ramIsAvailable = ram == 0.0 || usedRam == 0.0 || ram <= remainingRam; + boolean extraResourcesIsAvailable = areExtraResourcesAvailable(resources); boolean localTestCountIsAvailable = localTestCount == 0 || usedLocalTestCount == 0 || usedLocalTestCount + localTestCount <= availableLocalTestCount; - return cpuIsAvailable && ramIsAvailable && localTestCountIsAvailable; + return cpuIsAvailable && ramIsAvailable && extraResourcesIsAvailable && localTestCountIsAvailable; } @VisibleForTesting diff --git a/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java b/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java index b62f74d416907a..583aafd87e7914 100644 --- a/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java +++ b/src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java @@ -14,13 +14,16 @@ package com.google.devtools.build.lib.actions; +import com.google.common.base.Joiner; import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableMap; import com.google.common.primitives.Doubles; import com.google.devtools.build.lib.collect.nestedset.NestedSet; import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable; import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec; import com.google.devtools.common.options.Converter; import com.google.devtools.common.options.OptionsParsingException; +import io.reactivex.rxjava3.annotations.NonNull; import java.util.Iterator; import java.util.NoSuchElementException; @@ -43,12 +46,23 @@ public class ResourceSet implements ResourceSetOrBuilder { /** The number of CPUs, or fractions thereof. */ private final double cpuUsage; + /** + * Map of extra resources (for example: GPUs, embedded boards, ...) mapping + * name of the resource to a value. + */ + private final ImmutableMap extraResourceUsage; + /** The number of local tests. */ private final int localTestCount; private ResourceSet(double memoryMb, double cpuUsage, int localTestCount) { + this(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount); + } + + private ResourceSet(double memoryMb, double cpuUsage, @NonNull ImmutableMap extraResourceUsage, int localTestCount) { this.memoryMb = memoryMb; this.cpuUsage = cpuUsage; + this.extraResourceUsage = extraResourceUsage; this.localTestCount = localTestCount; } @@ -74,18 +88,28 @@ public static ResourceSet createWithLocalTestCount(int localTestCount) { } /** - * Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, ioUsage, and + * Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, and + * localTestCount. Most action resource definitions should use {@link #createWithRamCpu} or + * {@link #createWithLocalTestCount(int)}. Use this method primarily when constructing + * ResourceSets that represent available resources. + */ + public static ResourceSet create(double memoryMb, double cpuUsage, int localTestCount) { + return ResourceSet.create(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount); + } + + /** + * Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, extraResources, and * localTestCount. Most action resource definitions should use {@link #createWithRamCpu} or * {@link #createWithLocalTestCount(int)}. Use this method primarily when constructing * ResourceSets that represent available resources. */ @AutoCodec.Instantiator public static ResourceSet create( - double memoryMb, double cpuUsage, int localTestCount) { - if (memoryMb == 0 && cpuUsage == 0 && localTestCount == 0) { + double memoryMb, double cpuUsage, ImmutableMap extraResourceUsage, int localTestCount) { + if (memoryMb == 0 && cpuUsage == 0 && extraResourceUsage.size() == 0 && localTestCount == 0) { return ZERO; } - return new ResourceSet(memoryMb, cpuUsage, localTestCount); + return new ResourceSet(memoryMb, cpuUsage, extraResourceUsage, localTestCount); } /** Returns the amount of real memory (resident set size) used in MB. */ @@ -105,6 +129,10 @@ public double getCpuUsage() { return cpuUsage; } + public ImmutableMap getExtraResourceUsage() { + return extraResourceUsage; + } + /** Returns the local test count used. */ public int getLocalTestCount() { return localTestCount; @@ -115,6 +143,7 @@ public String toString() { return "Resources: \n" + "Memory: " + memoryMb + "M\n" + "CPU: " + cpuUsage + "\n" + + Joiner.on("\n").withKeyValueSeparator(": ").join(extraResourceUsage.entrySet()) + "Local tests: " + localTestCount + "\n"; } 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 0c97167cfc8390..5f9bc5023d6f3f 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 @@ -31,6 +31,7 @@ import com.google.devtools.build.lib.server.FailureDetails.FailureDetail; import com.google.devtools.build.lib.server.FailureDetails.TestAction; import com.google.devtools.build.lib.server.FailureDetails.TestAction.Code; +import java.util.HashMap; import java.util.List; import java.util.Map; @@ -162,32 +163,22 @@ public boolean isPersistentTestRunner() { return isPersistentTestRunner; } - public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs) - throws UserExecException { - if (usingLocalTestJobs) { - return LOCAL_TEST_JOBS_BASED_RESOURCES; - } - + private double getLocalCpuResourceUsage(Label label) throws UserExecException { ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size); - // Tests can override their CPU reservation with a "cpu:" tag. - ResourceSet testResourcesFromTag = null; + double cpuCount = -1.0; for (String tag : executionInfo.keySet()) { try { String cpus = ExecutionRequirements.CPU.parseIfMatches(tag); if (cpus != null) { - if (testResourcesFromTag != null) { + if (cpuCount != -1.0) { String message = String.format( "%s has more than one '%s' tag, but duplicate tags aren't allowed", label, ExecutionRequirements.CPU.userFriendlyName()); throw new UserExecException(createFailureDetail(message, Code.DUPLICATE_CPU_TAGS)); } - testResourcesFromTag = - ResourceSet.create( - testResourcesFromSize.getMemoryMb(), - Float.parseFloat(cpus), - testResourcesFromSize.getLocalTestCount()); + cpuCount = Float.parseFloat(cpus); } } catch (ValidationException e) { String message = @@ -200,8 +191,54 @@ public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs throw new UserExecException(createFailureDetail(message, Code.INVALID_CPU_TAG)); } } + return cpuCount != -1.0 ? cpuCount : testResourcesFromSize.getCpuUsage(); + } - return testResourcesFromTag != null ? testResourcesFromTag : testResourcesFromSize; + private ImmutableMap getLocalExtraResourceUsage(Label label) throws UserExecException { + // Tests can specify requirements for extra resources using "resources::" tag. + Map extraResources = new HashMap<>(); + for (String tag : executionInfo.keySet()) { + try { + String extras = ExecutionRequirements.RESOURCES.parseIfMatches(tag); + if (extras != null) { + int splitIndex = extras.indexOf(":"); + String resourceName = extras.substring(0, splitIndex); + String resourceCount = extras.substring(splitIndex+1); + if (extraResources.get(resourceName) != null) { + String message = + String.format( + "%s has more than one '%s' tag, but duplicate tags aren't allowed", + label, ExecutionRequirements.RESOURCES.userFriendlyName()); + throw new UserExecException(createFailureDetail(message, Code.DUPLICATE_CPU_TAGS)); + } + extraResources.put(resourceName, Float.parseFloat(resourceCount)); + } + } catch (ValidationException e) { + String message = + String.format( + "%s has a '%s' tag, but its value '%s' didn't pass validation: %s", + label, + ExecutionRequirements.RESOURCES.userFriendlyName(), + e.getTagValue(), + e.getMessage()); + throw new UserExecException(createFailureDetail(message, Code.INVALID_CPU_TAG)); + } + } + return ImmutableMap.copyOf(extraResources); + } + + public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs) + throws UserExecException { + if (usingLocalTestJobs) { + return LOCAL_TEST_JOBS_BASED_RESOURCES; + } + + ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size); + return ResourceSet.create( + testResourcesFromSize.getMemoryMb(), + getLocalCpuResourceUsage(label), + getLocalExtraResourceUsage(label), + testResourcesFromSize.getLocalTestCount()); } private static FailureDetail createFailureDetail(String message, Code detailedCode) { diff --git a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java index a6b182555c07f6..17d10ebc1b15ae 100644 --- a/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java +++ b/src/main/java/com/google/devtools/build/lib/buildtool/ExecutionTool.java @@ -107,6 +107,7 @@ import java.util.HashSet; import java.util.LinkedHashSet; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; @@ -883,10 +884,14 @@ public static void configureResourceManager(ResourceManager resourceMgr, BuildRe resources = ResourceSet.createWithRamCpu(options.localRamResources, options.localCpuResources); resourceMgr.setUseLocalMemoryEstimate(options.localMemoryEstimate); + ImmutableMap extraResources = options.localExtraResources.stream().collect( + ImmutableMap.toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (v1,v2) -> v1)); + resourceMgr.setAvailableResources( ResourceSet.create( resources.getMemoryMb(), resources.getCpuUsage(), + extraResources, request.getExecutionOptions().usingLocalTestJobs() ? request.getExecutionOptions().localTestJobs : Integer.MAX_VALUE)); diff --git a/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java b/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java index 387c04b9987e96..f135c00ba8ae98 100644 --- a/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java +++ b/src/main/java/com/google/devtools/build/lib/exec/ExecutionOptions.java @@ -302,8 +302,7 @@ public boolean shouldMaterializeParamFiles() { + "an integer, or \"HOST_CPUS\", optionally followed by [-|*] " + "(eg. HOST_CPUS*.5 to use half the available CPU cores)." + "By default, (\"HOST_CPUS\"), Bazel will query system configuration to estimate " - + "number of CPU cores available for the locally executed build actions. " - + "Note: This is a no-op if --local_resources is set.", + + "number of CPU cores available for the locally executed build actions.", converter = CpuResourceConverter.class) public float localCpuResources; @@ -318,11 +317,27 @@ public boolean shouldMaterializeParamFiles() { + "(eg. HOST_RAM*.5 to use half the available RAM)." + "By default, (\"HOST_RAM*.67\"), Bazel will query system configuration to estimate " + "amount of RAM available for the locally executed build actions and will use 67% " - + "of available RAM. " - + "Note: This is a no-op if --local_resources is set.", + + "of available RAM.", converter = RamResourceConverter.class) public float localRamResources; + @Option( + name = "local_extra_resources", + defaultValue = "null", + documentationCategory = OptionDocumentationCategory.UNCATEGORIZED, + effectTags = {OptionEffectTag.UNKNOWN}, + allowMultiple = true, + help = + "Set the number of extra resources available to Bazel. " + + "Takes in a string-float pair. Can be used multiple times to specify multiple " + + "types of extra resources. Bazel will limit concurrently running test actions " + + "based on the available extra resources and the extra resources required " + + "by the test actions. Tests can declare the amount of extra resources they need " + + "by using a tag of the \"resources::\" format. " + + "Available CPU, RAM and test job resources cannot be set with this flag.", + converter = Converters.StringToFloatAssignmentConverter.class) + public List> localExtraResources; + @Option( name = "experimental_local_memory_estimate", defaultValue = "false", diff --git a/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java b/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java index 0947cf5cf5e8dc..b1960e878dc85c 100644 --- a/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java +++ b/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java @@ -56,7 +56,8 @@ private static boolean legalExecInfoKeys(String tag) { || tag.startsWith("disable-") || tag.startsWith("cpu:") || tag.equals(ExecutionRequirements.LOCAL) - || tag.equals(ExecutionRequirements.WORKER_KEY_MNEMONIC); + || tag.equals(ExecutionRequirements.WORKER_KEY_MNEMONIC) + || tag.startsWith("resources:"); } private TargetUtils() {} // Uninstantiable. diff --git a/src/main/java/com/google/devtools/common/options/Converters.java b/src/main/java/com/google/devtools/common/options/Converters.java index 8563e9cff816a6..76ee1169f2077e 100644 --- a/src/main/java/com/google/devtools/common/options/Converters.java +++ b/src/main/java/com/google/devtools/common/options/Converters.java @@ -469,6 +469,24 @@ public String getTypeDescription() { } } + /** + * A converter for for assignments from a string value to a float value. + */ + public static class StringToFloatAssignmentConverter implements Converter> { + private static final AssignmentConverter baseConverter = new AssignmentConverter(); + + @Override + public Map.Entry convert(String input) throws OptionsParsingException, NumberFormatException { + Map.Entry stringEntry = baseConverter.convert(input); + return Maps.immutableEntry(stringEntry.getKey(), Float.parseFloat(stringEntry.getValue())); + } + + @Override + public String getTypeDescription() { + return "a named float, 'name=value'"; + } + } + /** * A converter for command line flag aliases. It does additional validation on the name and value * of the assignment to ensure they conform to the naming limitations. diff --git a/src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java b/src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java index bac53ea4ede5c5..813fdc5dd5ad11 100644 --- a/src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java +++ b/src/test/java/com/google/devtools/build/lib/actions/ResourceManagerTest.java @@ -51,7 +51,14 @@ public class ResourceManagerTest { @Before public final void configureResourceManager() throws Exception { rm.setAvailableResources( - ResourceSet.create(/*memoryMb=*/ 1000, /*cpuUsage=*/ 1, /* localTestCount= */ 2)); + ResourceSet.create( + /*memoryMb=*/ 1000, + /*cpuUsage=*/ 1, + /*extraResourceUsage*/ ImmutableMap.of( + "gpu", 2.0f, + "fancyresource", 1.5f + ), + /* localTestCount= */ 2)); counter = new AtomicInteger(0); sync = new CyclicBarrier(2); sync2 = new CyclicBarrier(2); @@ -63,14 +70,27 @@ private ResourceHandle acquire(double ram, double cpu, int tests) return rm.acquireResources(resourceOwner, ResourceSet.create(ram, cpu, tests)); } + private ResourceHandle acquire(double ram, double cpu, int tests, ImmutableMap extraResources) + throws InterruptedException { + return rm.acquireResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests)); + } + private ResourceHandle acquireNonblocking(double ram, double cpu, int tests) { return rm.tryAcquire(resourceOwner, ResourceSet.create(ram, cpu, tests)); } + private ResourceHandle acquireNonblocking(double ram, double cpu, int tests, ImmutableMap extraResources) { + return rm.tryAcquire(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests)); + } + private void release(double ram, double cpu, int tests) { rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests)); } + private void release(double ram, double cpu, int tests, ImmutableMap extraResources) { + rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests)); + } + private void validate(int count) { assertThat(counter.incrementAndGet()).isEqualTo(count); } @@ -104,6 +124,13 @@ public void testOverBudgetRequests() throws Exception { assertThat(rm.inUse()).isTrue(); release(0, 0, bigTests); assertThat(rm.inUse()).isFalse(); + + // Ditto, for extra resources (even if they don't exist in the available resource set): + ImmutableMap bigExtraResources = ImmutableMap.of("gpu", 10.0f, "fancyresource", 10.0f, "nonexisting", 10.0f); + acquire(0, 0, 0, bigExtraResources); + assertThat(rm.inUse()).isTrue(); + release(0, 0, 0, bigExtraResources); + assertThat(rm.inUse()).isFalse(); } @Test @@ -183,11 +210,25 @@ public void testThatTestsCannotBeOverallocated() throws Exception { thread1.joinAndAssertState(10000); } + @Test + public void testThatExtraResourcesCannotBeOverallocated() throws Exception { + assertThat(rm.inUse()).isFalse(); + + // Given a partially acquired extra resources: + acquire(0, 0, 1, ImmutableMap.of("gpu", 1.0f)); + + // When a request for extra resources is made that would overallocate, + // Then the request fails: + TestThread thread1 = new TestThread(() -> assertThat(acquireNonblocking(0, 0, 0, ImmutableMap.of("gpu", 1.1f))).isNull()); + thread1.start(); + thread1.joinAndAssertState(10000); + } + @Test public void testHasResources() throws Exception { assertThat(rm.inUse()).isFalse(); assertThat(rm.threadHasResources()).isFalse(); - acquire(1, 0.1, 1); + acquire(1, 0.1, 1, ImmutableMap.of("gpu", 1.0f)); assertThat(rm.threadHasResources()).isTrue(); // We have resources in this thread - make sure other threads @@ -208,11 +249,15 @@ public void testHasResources() throws Exception { assertThat(rm.threadHasResources()).isTrue(); release(0, 0, 1); assertThat(rm.threadHasResources()).isFalse(); + acquire(0, 0, 0, ImmutableMap.of("gpu", 1.0f)); + assertThat(rm.threadHasResources()).isTrue(); + release(0, 0, 0, ImmutableMap.of("gpu", 1.0f)); + assertThat(rm.threadHasResources()).isFalse(); }); thread1.start(); thread1.joinAndAssertState(10000); - release(1, 0.1, 1); + release(1, 0.1, 1, ImmutableMap.of("gpu", 1.0f)); assertThat(rm.threadHasResources()).isFalse(); assertThat(rm.inUse()).isFalse(); }