From 2b2bea4ee9d51ad875a8d273aa344ac9c5b6742c Mon Sep 17 00:00:00 2001 From: Drew Macrae Date: Fri, 30 Dec 2022 03:41:12 -0800 Subject: [PATCH] Extra resources This recreates a [closed PR](https://github.com/bazelbuild/bazel/pull/13996) to implement extra resources which we're hoping to use in https://github.com/lowRISC/opentitan/pull/16436 Fixes:#16817 Closes #16785. PiperOrigin-RevId: 498557024 Change-Id: I60d8f8f4a4a02748147cabb4cd60a2a9b95a2c68 --- .../lib/actions/ExecutionRequirements.java | 24 +++++ .../build/lib/actions/ResourceManager.java | 88 +++++++++++++++++-- .../build/lib/actions/ResourceSet.java | 63 +++++++++++-- .../analysis/test/TestTargetProperties.java | 57 ++++++++++-- .../build/lib/buildtool/ExecutionTool.java | 8 ++ .../build/lib/exec/ExecutionOptions.java | 17 ++++ .../build/lib/packages/TargetUtils.java | 3 +- .../build/lib/worker/WorkerSpawnRunner.java | 1 + .../devtools/common/options/Converters.java | 18 ++++ .../lib/actions/ResourceManagerTest.java | 83 ++++++++++++++++- 10 files changed, 338 insertions(+), 24 deletions(-) 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 84336e6a55f267..3a24ecfdb0b3ce 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 @@ -154,6 +154,30 @@ 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); + float value; + try { + value = Float.parseFloat(resourceCount); + } catch (NumberFormatException e) { + return "can't be parsed as a float"; + } + + if (value < 0) { + return "can't be 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 1c603e17955fd6..ccb8b4e292fc6d 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 @@ -27,8 +27,13 @@ import com.google.devtools.build.lib.worker.WorkerPool; import java.io.IOException; import java.util.Deque; +import java.util.HashMap; +import java.util.HashSet; import java.util.Iterator; import java.util.LinkedList; +import java.util.Map; +import java.util.NoSuchElementException; +import java.util.Set; import java.util.concurrent.CountDownLatch; import javax.annotation.Nullable; @@ -183,14 +188,16 @@ public double getUsedCPU() { // 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; /** If set, local-only actions are given priority over dynamically run actions. */ private boolean prioritizeLocalActions; - private ResourceManager() {} - @VisibleForTesting public static ResourceManager instanceForTestingOnly() { return new ResourceManager(); @@ -204,6 +211,7 @@ public static ResourceManager instanceForTestingOnly() { public synchronized void resetResourceUsage() { usedCpu = 0; usedRam = 0; + usedExtraResources = new HashMap<>(); usedLocalTestCount = 0; for (Pair request : localRequests) { request.second.latch.countDown(); @@ -298,6 +306,20 @@ private Worker incrementResources(ResourceSet resources) throws IOException, InterruptedException { 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(); if (resources.getWorkerKey() != null) { @@ -310,6 +332,7 @@ private Worker incrementResources(ResourceSet resources) public synchronized boolean inUse() { return usedCpu != 0.0 || usedRam != 0.0 + || !usedExtraResources.isEmpty() || usedLocalTestCount != 0 || !localRequests.isEmpty() || !dynamicWorkerRequests.isEmpty() @@ -369,7 +392,7 @@ public void acquireResourceOwnership() { * wait. */ private synchronized LatchWithWorker acquire(ResourceSet resources, ResourcePriority priority) - throws IOException, InterruptedException { + throws IOException, InterruptedException, NoSuchElementException { if (areResourcesAvailable(resources)) { Worker worker = incrementResources(resources); return new LatchWithWorker(/* latch= */ null, worker); @@ -417,6 +440,7 @@ private boolean release(ResourceSet resources, @Nullable Worker worker) private synchronized void releaseResourcesOnly(ResourceSet resources) { usedCpu -= resources.getCpuUsage(); usedRam -= resources.getMemoryMb(); + usedLocalTestCount -= resources.getLocalTestCount(); // TODO(bazel-team): (2010) rounding error can accumulate and value below can end up being @@ -428,6 +452,19 @@ private synchronized void releaseResourcesOnly(ResourceSet resources) { if (usedRam < epsilon) { usedRam = 0; } + + Set toRemove = new HashSet<>(); + for (Map.Entry resource : resources.getExtraResourceUsage().entrySet()) { + String key = (String) resource.getKey(); + float value = (float) usedExtraResources.get(key) - resource.getValue(); + usedExtraResources.put(key, value); + if (value < epsilon) { + toRemove.add(key); + } + } + for (String key : toRemove) { + usedExtraResources.remove(key); + } } private synchronized boolean processAllWaitingThreads() throws IOException, InterruptedException { @@ -466,9 +503,35 @@ private synchronized void processWaitingThreads(Deque resource : resources.getExtraResourceUsage().entrySet()) { + String key = (String) resource.getKey(); + if (!availableResources.getExtraResourceUsage().containsKey(key)) { + throw new NoSuchElementException( + "Resource " + key + " is not tracked in this resource set."); + } + } + } + + /** Return true iff all requested extra resources are considered to be available. */ + private boolean areExtraResourcesAvailable(ResourceSet resources) throws NoSuchElementException { + 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 = availableResources.getExtraResourceUsage().get(key); + 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. @VisibleForTesting - boolean areResourcesAvailable(ResourceSet resources) { + boolean areResourcesAvailable(ResourceSet resources) throws NoSuchElementException { Preconditions.checkNotNull(availableResources); // Comparison below is robust, since any calculation errors will be fixed // by the release() method. @@ -484,7 +547,15 @@ boolean areResourcesAvailable(ResourceSet resources) { workerKey == null || (activeWorkers < availableWorkers && workerPool.couldBeBorrowed(workerKey)); - if (usedCpu == 0.0 && usedRam == 0.0 && usedLocalTestCount == 0 && workerIsAvailable) { + // We test for tracking of extra resources whenever acquired and throw an + // exception before acquiring any untracked resource. + assertExtraResourcesTracked(resources); + + if (usedCpu == 0.0 + && usedRam == 0.0 + && usedExtraResources.isEmpty() + && usedLocalTestCount == 0 + && workerIsAvailable) { return true; } // Use only MIN_NECESSARY_???_RATIO of the resource value to check for @@ -515,7 +586,12 @@ boolean areResourcesAvailable(ResourceSet resources) { localTestCount == 0 || usedLocalTestCount == 0 || usedLocalTestCount + localTestCount <= availableLocalTestCount; - return cpuIsAvailable && ramIsAvailable && localTestCountIsAvailable && workerIsAvailable; + boolean extraResourcesIsAvailable = areExtraResourcesAvailable(resources); + return cpuIsAvailable + && ramIsAvailable + && extraResourcesIsAvailable + && localTestCountIsAvailable + && workerIsAvailable; } @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 a0b60461e9541f..d2fd23b8bed5ce 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,7 +14,9 @@ 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.concurrent.ThreadSafety.Immutable; import com.google.devtools.build.lib.util.OS; @@ -43,6 +45,12 @@ 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; @@ -51,8 +59,18 @@ public class ResourceSet implements ResourceSetOrBuilder { private ResourceSet( double memoryMb, double cpuUsage, int localTestCount, @Nullable WorkerKey workerKey) { + this(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, workerKey); + } + + private ResourceSet( + double memoryMb, + double cpuUsage, + @Nullable ImmutableMap extraResourceUsage, + int localTestCount, + @Nullable WorkerKey workerKey) { this.memoryMb = memoryMb; this.cpuUsage = cpuUsage; + this.extraResourceUsage = extraResourceUsage; this.localTestCount = localTestCount; this.workerKey = workerKey; } @@ -83,21 +101,51 @@ public static ResourceSet createWithLocalTestCount(int localTestCount) { } /** - * Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, ioUsage, and - * localTestCount. Most action resource definitions should use {@link #createWithRamCpu} or {@link + * 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 createWithWorkerKey(memoryMb, cpuUsage, localTestCount, /* workerKey= */ null); + return ResourceSet.createWithWorkerKey( + memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, /* workerKey= */ null); + } + + /** + * 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. + */ + public static ResourceSet create( + double memoryMb, + double cpuUsage, + ImmutableMap extraResourceUsage, + int localTestCount) { + return createWithWorkerKey( + memoryMb, cpuUsage, extraResourceUsage, localTestCount, /* workerKey= */ null); } public static ResourceSet createWithWorkerKey( double memoryMb, double cpuUsage, int localTestCount, WorkerKey workerKey) { - if (memoryMb == 0 && cpuUsage == 0 && localTestCount == 0 && workerKey == null) { + return ResourceSet.createWithWorkerKey( + memoryMb, cpuUsage, /* extraResourceUsage= */ ImmutableMap.of(), localTestCount, workerKey); + } + + public static ResourceSet createWithWorkerKey( + double memoryMb, + double cpuUsage, + ImmutableMap extraResourceUsage, + int localTestCount, + WorkerKey workerKey) { + if (memoryMb == 0 + && cpuUsage == 0 + && extraResourceUsage.size() == 0 + && localTestCount == 0 + && workerKey == null) { return ZERO; } - return new ResourceSet(memoryMb, cpuUsage, localTestCount, workerKey); + return new ResourceSet(memoryMb, cpuUsage, extraResourceUsage, localTestCount, workerKey); } /** Returns the amount of real memory (resident set size) used in MB. */ @@ -124,6 +172,10 @@ public double getCpuUsage() { return cpuUsage; } + public ImmutableMap getExtraResourceUsage() { + return extraResourceUsage; + } + /** Returns the local test count used. */ public int getLocalTestCount() { return localTestCount; @@ -138,6 +190,7 @@ public String toString() { + "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 82bc1a94ebfedb..90b7a2c4af8f2a 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; @@ -160,25 +161,29 @@ public ResourceSet getLocalResourceUsage(Label label, boolean usingLocalTestJobs } ResourceSet testResourcesFromSize = TestTargetProperties.getResourceSetFromSize(size); + return ResourceSet.create( + testResourcesFromSize.getMemoryMb(), + getLocalCpuResourceUsage(label), + getLocalExtraResourceUsage(label), + testResourcesFromSize.getLocalTestCount()); + } + 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 = @@ -191,10 +196,46 @@ 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); } + + private static FailureDetail createFailureDetail(String message, Code detailedCode) { return FailureDetail.newBuilder() .setMessage(message) 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 3bdfe0db1faf62..a4345c039138eb 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 @@ -110,6 +110,7 @@ import java.time.Duration; import java.util.Collection; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; @@ -925,10 +926,17 @@ private Builder createBuilder( public static void configureResourceManager(ResourceManager resourceMgr, BuildRequest request) { ExecutionOptions options = request.getOptions(ExecutionOptions.class); resourceMgr.setPrioritizeLocalActions(options.prioritizeLocalActions); + ImmutableMap extraResources = + options.localExtraResources.stream() + .collect( + ImmutableMap.toImmutableMap( + Map.Entry::getKey, Map.Entry::getValue, (v1, v2) -> v2)); + resourceMgr.setAvailableResources( ResourceSet.create( options.localRamResources, options.localCpuResources, + extraResources, options.usingLocalTestJobs() ? options.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 4caab45cdf522e..9515288e95316d 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 @@ -325,6 +325,23 @@ public boolean shouldMaterializeParamFiles() { 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 actions " + + "based on the available extra resources and the extra resources required. " + + "Tests can declare the amount of extra resources they need " + + "by using a tag of the \"resources::\" format. " + + "Available CPU, RAM and resources cannot be set with this flag.", + converter = Converters.StringToFloatAssignmentConverter.class) + public List> localExtraResources; + @Option( name = "local_test_jobs", defaultValue = "auto", 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 744cefce6363bb..6f148c1e75f790 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 @@ -57,7 +57,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/build/lib/worker/WorkerSpawnRunner.java b/src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnRunner.java index d5beddb20cae05..0d1c14ee808b48 100644 --- a/src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnRunner.java +++ b/src/main/java/com/google/devtools/build/lib/worker/WorkerSpawnRunner.java @@ -416,6 +416,7 @@ WorkResponse execInWorkerWorkerAsResource( ResourceSet.createWithWorkerKey( spawn.getLocalResources().getMemoryMb(), spawn.getLocalResources().getCpuUsage(), + spawn.getLocalResources().getExtraResourceUsage(), spawn.getLocalResources().getLocalTestCount(), key); 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 0a8a9499dcc5f7..da70766f504060 100644 --- a/src/main/java/com/google/devtools/common/options/Converters.java +++ b/src/main/java/com/google/devtools/common/options/Converters.java @@ -472,6 +472,24 @@ public String getTypeDescription() { } } + /** A converter for for assignments from a string value to a float value. */ + public static class StringToFloatAssignmentConverter + extends Converter.Contextless> { + 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 05c817f085cf8b..56a3cef2a4b51e 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 @@ -39,9 +39,11 @@ import com.google.devtools.build.lib.worker.WorkerFactory; import com.google.devtools.build.lib.worker.WorkerKey; import com.google.devtools.build.lib.worker.WorkerPool; +import com.google.errorprone.annotations.CanIgnoreReturnValue; import java.io.IOException; import java.time.Duration; import java.util.Collection; +import java.util.NoSuchElementException; import java.util.concurrent.CountDownLatch; import java.util.concurrent.CyclicBarrier; import java.util.concurrent.TimeUnit; @@ -68,7 +70,13 @@ public final class ResourceManagerTest { @Before public 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); @@ -116,8 +124,34 @@ private ResourceHandle acquire(double ram, double cpu, int tests, String mnemoni ResourcePriority.LOCAL); } + @CanIgnoreReturnValue + private ResourceHandle acquire( + double ram, + double cpu, + ImmutableMap extraResources, + int tests, + ResourcePriority priority) + throws InterruptedException, IOException, NoSuchElementException { + return rm.acquireResources( + resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests), priority); + } + + @CanIgnoreReturnValue + private ResourceHandle acquire( + double ram, double cpu, ImmutableMap extraResources, int tests) + throws InterruptedException, IOException, NoSuchElementException { + return acquire(ram, cpu, extraResources, tests, ResourcePriority.LOCAL); + } + private void release(double ram, double cpu, int tests) throws IOException, InterruptedException { - rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests), /* worker=*/ null); + rm.releaseResources(resourceOwner, ResourceSet.create(ram, cpu, tests), /* worker= */ null); + } + + private void release( + double ram, double cpu, ImmutableMap extraResources, int tests) + throws InterruptedException, IOException { + rm.releaseResources( + resourceOwner, ResourceSet.create(ram, cpu, extraResources, tests), /* worker= */ null); } private void validate(int count) { @@ -167,6 +201,14 @@ public void testOverBudgetRequests() throws Exception { assertThat(rm.inUse()).isTrue(); release(0, 0, bigTests); assertThat(rm.inUse()).isFalse(); + + // Ditto, for extra resources: + ImmutableMap bigExtraResources = + ImmutableMap.of("gpu", 10.0f, "fancyresource", 10.0f); + acquire(0, 0, bigExtraResources, 0); + assertThat(rm.inUse()).isTrue(); + release(0, 0, bigExtraResources, 0); + assertThat(rm.inUse()).isFalse(); } @Test @@ -248,11 +290,26 @@ public void testThatTestsCannotBeOverallocated() throws Exception { assertThat(e).hasCauseThat().hasMessageThat().contains("is still alive"); } + @Test + public void testThatExtraResourcesCannotBeOverallocated() throws Exception { + assertThat(rm.inUse()).isFalse(); + + // Given a partially acquired extra resources: + acquire(0, 0, ImmutableMap.of("gpu", 1.0f), 1); + + // When a request for extra resources is made that would overallocate, + // Then the request fails: + TestThread thread1 = new TestThread(() -> acquire(0, 0, ImmutableMap.of("gpu", 1.1f), 0)); + thread1.start(); + AssertionError e = assertThrows(AssertionError.class, () -> thread1.joinAndAssertState(1000)); + assertThat(e).hasCauseThat().hasMessageThat().contains("is still alive"); + } + @Test public void testHasResources() throws Exception { assertThat(rm.inUse()).isFalse(); assertThat(rm.threadHasResources()).isFalse(); - acquire(1, 0.1, 1); + acquire(1, 0.1, ImmutableMap.of("gpu", 1.0f), 1); assertThat(rm.threadHasResources()).isTrue(); // We have resources in this thread - make sure other threads @@ -273,11 +330,15 @@ public void testHasResources() throws Exception { assertThat(rm.threadHasResources()).isTrue(); release(0, 0, 1); assertThat(rm.threadHasResources()).isFalse(); + acquire(0, 0, ImmutableMap.of("gpu", 1.0f), 0); + assertThat(rm.threadHasResources()).isTrue(); + release(0, 0, ImmutableMap.of("gpu", 1.0f), 0); + assertThat(rm.threadHasResources()).isFalse(); }); thread1.start(); thread1.joinAndAssertState(10000); - release(1, 0.1, 1); + release(1, 0.1, ImmutableMap.of("gpu", 1.0f), 1); assertThat(rm.threadHasResources()).isFalse(); assertThat(rm.inUse()).isFalse(); } @@ -668,6 +729,20 @@ private CyclicBarrier startAcquireReleaseThread(ResourcePriority priority) { return sync; } + @Test + public void testNonexistingResource() throws Exception { + // If we try to use nonexisting resource we should return an error + TestThread thread1 = + new TestThread( + () -> { + assertThrows( + NoSuchElementException.class, + () -> acquire(0, 0, ImmutableMap.of("nonexisting", 1.0f), 0)); + }); + thread1.start(); + thread1.joinAndAssertState(1000); + } + @Test public void testAcquireWithWorker_acquireAndRelease() throws Exception { int memory = 100;