Skip to content

Commit 2b2bea4

Browse files
Drew Macraecopybara-github
Drew Macrae
authored and
copybara-github
committed
Extra resources
This recreates a [closed PR](#13996) to implement extra resources which we're hoping to use in lowRISC/opentitan#16436 Fixes:#16817 Closes #16785. PiperOrigin-RevId: 498557024 Change-Id: I60d8f8f4a4a02748147cabb4cd60a2a9b95a2c68
1 parent d42aedd commit 2b2bea4

File tree

10 files changed

+338
-24
lines changed

10 files changed

+338
-24
lines changed

‎src/main/java/com/google/devtools/build/lib/actions/ExecutionRequirements.java‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,30 @@ public String parseIfMatches(String tag) throws ValidationException {
154154
return null;
155155
});
156156

157+
/** How many extra resources an action requires for execution. */
158+
public static final ParseableRequirement RESOURCES =
159+
ParseableRequirement.create(
160+
"resources:<str>:<float>",
161+
Pattern.compile("resources:(.+:.+)"),
162+
s -> {
163+
Preconditions.checkNotNull(s);
164+
165+
int splitIndex = s.indexOf(":");
166+
String resourceCount = s.substring(splitIndex + 1);
167+
float value;
168+
try {
169+
value = Float.parseFloat(resourceCount);
170+
} catch (NumberFormatException e) {
171+
return "can't be parsed as a float";
172+
}
173+
174+
if (value < 0) {
175+
return "can't be negative";
176+
}
177+
178+
return null;
179+
});
180+
157181
/** If an action supports running in persistent worker mode. */
158182
public static final String SUPPORTS_WORKERS = "supports-workers";
159183

‎src/main/java/com/google/devtools/build/lib/actions/ResourceManager.java‎

Lines changed: 82 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,13 @@
2727
import com.google.devtools.build.lib.worker.WorkerPool;
2828
import java.io.IOException;
2929
import java.util.Deque;
30+
import java.util.HashMap;
31+
import java.util.HashSet;
3032
import java.util.Iterator;
3133
import java.util.LinkedList;
34+
import java.util.Map;
35+
import java.util.NoSuchElementException;
36+
import java.util.Set;
3237
import java.util.concurrent.CountDownLatch;
3338
import javax.annotation.Nullable;
3439

@@ -183,14 +188,16 @@ public double getUsedCPU() {
183188
// definition in the ResourceSet class.
184189
private double usedRam;
185190

191+
// Used amount of extra resources. Corresponds to the extra resource
192+
// definition in the ResourceSet class.
193+
private Map<String, Float> usedExtraResources;
194+
186195
// Used local test count. Corresponds to the local test count definition in the ResourceSet class.
187196
private int usedLocalTestCount;
188197

189198
/** If set, local-only actions are given priority over dynamically run actions. */
190199
private boolean prioritizeLocalActions;
191200

192-
private ResourceManager() {}
193-
194201
@VisibleForTesting
195202
public static ResourceManager instanceForTestingOnly() {
196203
return new ResourceManager();
@@ -204,6 +211,7 @@ public static ResourceManager instanceForTestingOnly() {
204211
public synchronized void resetResourceUsage() {
205212
usedCpu = 0;
206213
usedRam = 0;
214+
usedExtraResources = new HashMap<>();
207215
usedLocalTestCount = 0;
208216
for (Pair<ResourceSet, LatchWithWorker> request : localRequests) {
209217
request.second.latch.countDown();
@@ -298,6 +306,20 @@ private Worker incrementResources(ResourceSet resources)
298306
throws IOException, InterruptedException {
299307
usedCpu += resources.getCpuUsage();
300308
usedRam += resources.getMemoryMb();
309+
310+
resources
311+
.getExtraResourceUsage()
312+
.entrySet()
313+
.forEach(
314+
resource -> {
315+
String key = (String) resource.getKey();
316+
float value = resource.getValue();
317+
if (usedExtraResources.containsKey(key)) {
318+
value += (float) usedExtraResources.get(key);
319+
}
320+
usedExtraResources.put(key, value);
321+
});
322+
301323
usedLocalTestCount += resources.getLocalTestCount();
302324

303325
if (resources.getWorkerKey() != null) {
@@ -310,6 +332,7 @@ private Worker incrementResources(ResourceSet resources)
310332
public synchronized boolean inUse() {
311333
return usedCpu != 0.0
312334
|| usedRam != 0.0
335+
|| !usedExtraResources.isEmpty()
313336
|| usedLocalTestCount != 0
314337
|| !localRequests.isEmpty()
315338
|| !dynamicWorkerRequests.isEmpty()
@@ -369,7 +392,7 @@ public void acquireResourceOwnership() {
369392
* wait.
370393
*/
371394
private synchronized LatchWithWorker acquire(ResourceSet resources, ResourcePriority priority)
372-
throws IOException, InterruptedException {
395+
throws IOException, InterruptedException, NoSuchElementException {
373396
if (areResourcesAvailable(resources)) {
374397
Worker worker = incrementResources(resources);
375398
return new LatchWithWorker(/* latch= */ null, worker);
@@ -417,6 +440,7 @@ private boolean release(ResourceSet resources, @Nullable Worker worker)
417440
private synchronized void releaseResourcesOnly(ResourceSet resources) {
418441
usedCpu -= resources.getCpuUsage();
419442
usedRam -= resources.getMemoryMb();
443+
420444
usedLocalTestCount -= resources.getLocalTestCount();
421445

422446
// 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) {
428452
if (usedRam < epsilon) {
429453
usedRam = 0;
430454
}
455+
456+
Set<String> toRemove = new HashSet<>();
457+
for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
458+
String key = (String) resource.getKey();
459+
float value = (float) usedExtraResources.get(key) - resource.getValue();
460+
usedExtraResources.put(key, value);
461+
if (value < epsilon) {
462+
toRemove.add(key);
463+
}
464+
}
465+
for (String key : toRemove) {
466+
usedExtraResources.remove(key);
467+
}
431468
}
432469

433470
private synchronized boolean processAllWaitingThreads() throws IOException, InterruptedException {
@@ -466,9 +503,35 @@ private synchronized void processWaitingThreads(Deque<Pair<ResourceSet, LatchWit
466503
}
467504
}
468505

506+
/** Throws an exception if requested extra resource isn't being tracked */
507+
private void assertExtraResourcesTracked(ResourceSet resources) throws NoSuchElementException {
508+
for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
509+
String key = (String) resource.getKey();
510+
if (!availableResources.getExtraResourceUsage().containsKey(key)) {
511+
throw new NoSuchElementException(
512+
"Resource " + key + " is not tracked in this resource set.");
513+
}
514+
}
515+
}
516+
517+
/** Return true iff all requested extra resources are considered to be available. */
518+
private boolean areExtraResourcesAvailable(ResourceSet resources) throws NoSuchElementException {
519+
for (Map.Entry<String, Float> resource : resources.getExtraResourceUsage().entrySet()) {
520+
String key = (String) resource.getKey();
521+
float used = (float) usedExtraResources.getOrDefault(key, 0f);
522+
float requested = resource.getValue();
523+
float available = availableResources.getExtraResourceUsage().get(key);
524+
float epsilon = 0.0001f; // Account for possible rounding errors.
525+
if (requested != 0.0 && used != 0.0 && requested + used > available + epsilon) {
526+
return false;
527+
}
528+
}
529+
return true;
530+
}
531+
469532
// Method will return true if all requested resources are considered to be available.
470533
@VisibleForTesting
471-
boolean areResourcesAvailable(ResourceSet resources) {
534+
boolean areResourcesAvailable(ResourceSet resources) throws NoSuchElementException {
472535
Preconditions.checkNotNull(availableResources);
473536
// Comparison below is robust, since any calculation errors will be fixed
474537
// by the release() method.
@@ -484,7 +547,15 @@ boolean areResourcesAvailable(ResourceSet resources) {
484547
workerKey == null
485548
|| (activeWorkers < availableWorkers && workerPool.couldBeBorrowed(workerKey));
486549

487-
if (usedCpu == 0.0 && usedRam == 0.0 && usedLocalTestCount == 0 && workerIsAvailable) {
550+
// We test for tracking of extra resources whenever acquired and throw an
551+
// exception before acquiring any untracked resource.
552+
assertExtraResourcesTracked(resources);
553+
554+
if (usedCpu == 0.0
555+
&& usedRam == 0.0
556+
&& usedExtraResources.isEmpty()
557+
&& usedLocalTestCount == 0
558+
&& workerIsAvailable) {
488559
return true;
489560
}
490561
// Use only MIN_NECESSARY_???_RATIO of the resource value to check for
@@ -515,7 +586,12 @@ boolean areResourcesAvailable(ResourceSet resources) {
515586
localTestCount == 0
516587
|| usedLocalTestCount == 0
517588
|| usedLocalTestCount + localTestCount <= availableLocalTestCount;
518-
return cpuIsAvailable && ramIsAvailable && localTestCountIsAvailable && workerIsAvailable;
589+
boolean extraResourcesIsAvailable = areExtraResourcesAvailable(resources);
590+
return cpuIsAvailable
591+
&& ramIsAvailable
592+
&& extraResourcesIsAvailable
593+
&& localTestCountIsAvailable
594+
&& workerIsAvailable;
519595
}
520596

521597
@VisibleForTesting

‎src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java‎

Lines changed: 58 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@
1414

1515
package com.google.devtools.build.lib.actions;
1616

17+
import com.google.common.base.Joiner;
1718
import com.google.common.base.Splitter;
19+
import com.google.common.collect.ImmutableMap;
1820
import com.google.common.primitives.Doubles;
1921
import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
2022
import com.google.devtools.build.lib.util.OS;
@@ -43,6 +45,12 @@ public class ResourceSet implements ResourceSetOrBuilder {
4345
/** The number of CPUs, or fractions thereof. */
4446
private final double cpuUsage;
4547

48+
/**
49+
* Map of extra resources (for example: GPUs, embedded boards, ...) mapping name of the resource
50+
* to a value.
51+
*/
52+
private final ImmutableMap<String, Float> extraResourceUsage;
53+
4654
/** The number of local tests. */
4755
private final int localTestCount;
4856

@@ -51,8 +59,18 @@ public class ResourceSet implements ResourceSetOrBuilder {
5159

5260
private ResourceSet(
5361
double memoryMb, double cpuUsage, int localTestCount, @Nullable WorkerKey workerKey) {
62+
this(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, workerKey);
63+
}
64+
65+
private ResourceSet(
66+
double memoryMb,
67+
double cpuUsage,
68+
@Nullable ImmutableMap<String, Float> extraResourceUsage,
69+
int localTestCount,
70+
@Nullable WorkerKey workerKey) {
5471
this.memoryMb = memoryMb;
5572
this.cpuUsage = cpuUsage;
73+
this.extraResourceUsage = extraResourceUsage;
5674
this.localTestCount = localTestCount;
5775
this.workerKey = workerKey;
5876
}
@@ -83,21 +101,51 @@ public static ResourceSet createWithLocalTestCount(int localTestCount) {
83101
}
84102

85103
/**
86-
* Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, ioUsage, and
87-
* localTestCount. Most action resource definitions should use {@link #createWithRamCpu} or {@link
104+
* Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, and localTestCount.
105+
* Most action resource definitions should use {@link #createWithRamCpu} or {@link
88106
* #createWithLocalTestCount(int)}. Use this method primarily when constructing ResourceSets that
89107
* represent available resources.
90108
*/
91109
public static ResourceSet create(double memoryMb, double cpuUsage, int localTestCount) {
92-
return createWithWorkerKey(memoryMb, cpuUsage, localTestCount, /* workerKey= */ null);
110+
return ResourceSet.createWithWorkerKey(
111+
memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, /* workerKey= */ null);
112+
}
113+
114+
/**
115+
* Returns a new ResourceSet with the provided values for memoryMb, cpuUsage, extraResources, and
116+
* localTestCount. Most action resource definitions should use {@link #createWithRamCpu} or {@link
117+
* #createWithLocalTestCount(int)}. Use this method primarily when constructing ResourceSets that
118+
* represent available resources.
119+
*/
120+
public static ResourceSet create(
121+
double memoryMb,
122+
double cpuUsage,
123+
ImmutableMap<String, Float> extraResourceUsage,
124+
int localTestCount) {
125+
return createWithWorkerKey(
126+
memoryMb, cpuUsage, extraResourceUsage, localTestCount, /* workerKey= */ null);
93127
}
94128

95129
public static ResourceSet createWithWorkerKey(
96130
double memoryMb, double cpuUsage, int localTestCount, WorkerKey workerKey) {
97-
if (memoryMb == 0 && cpuUsage == 0 && localTestCount == 0 && workerKey == null) {
131+
return ResourceSet.createWithWorkerKey(
132+
memoryMb, cpuUsage, /* extraResourceUsage= */ ImmutableMap.of(), localTestCount, workerKey);
133+
}
134+
135+
public static ResourceSet createWithWorkerKey(
136+
double memoryMb,
137+
double cpuUsage,
138+
ImmutableMap<String, Float> extraResourceUsage,
139+
int localTestCount,
140+
WorkerKey workerKey) {
141+
if (memoryMb == 0
142+
&& cpuUsage == 0
143+
&& extraResourceUsage.size() == 0
144+
&& localTestCount == 0
145+
&& workerKey == null) {
98146
return ZERO;
99147
}
100-
return new ResourceSet(memoryMb, cpuUsage, localTestCount, workerKey);
148+
return new ResourceSet(memoryMb, cpuUsage, extraResourceUsage, localTestCount, workerKey);
101149
}
102150

103151
/** Returns the amount of real memory (resident set size) used in MB. */
@@ -124,6 +172,10 @@ public double getCpuUsage() {
124172
return cpuUsage;
125173
}
126174

175+
public ImmutableMap<String, Float> getExtraResourceUsage() {
176+
return extraResourceUsage;
177+
}
178+
127179
/** Returns the local test count used. */
128180
public int getLocalTestCount() {
129181
return localTestCount;
@@ -138,6 +190,7 @@ public String toString() {
138190
+ "CPU: "
139191
+ cpuUsage
140192
+ "\n"
193+
+ Joiner.on("\n").withKeyValueSeparator(": ").join(extraResourceUsage.entrySet())
141194
+ "Local tests: "
142195
+ localTestCount
143196
+ "\n";

0 commit comments

Comments
 (0)