Extra resources (#17229)

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

Co-authored-by: kshyanashree <[email protected]>

Author: Drew Macrae

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:<str>:<float>",
+          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";
 

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;
 
@@ -171,14 +176,16 @@ 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<String, Float> 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();
@@ -192,6 +199,7 @@ public static ResourceManager instanceForTestingOnly() {
   public synchronized void resetResourceUsage() {
     usedCpu = 0;
     usedRam = 0;
+    usedExtraResources = new HashMap<>();
     usedLocalTestCount = 0;
     for (Pair<ResourceSet, LatchWithWorker> request : localRequests) {
       request.second.latch.countDown();
@@ -286,6 +294,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) {
@@ -298,6 +320,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()
@@ -357,7 +380,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);
@@ -405,6 +428,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
@@ -416,6 +440,19 @@ private synchronized void releaseResourcesOnly(ResourceSet resources) {
     if (usedRam < epsilon) {
       usedRam = 0;
     }
+
+    Set<String> toRemove = new HashSet<>();
+    for (Map.Entry<String, Float> 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 {
@@ -454,9 +491,35 @@ private synchronized void processWaitingThreads(Deque<Pair<ResourceSet, LatchWit
     }
   }
 
+  /** Throws an exception if requested extra resource isn't being tracked */
+  private void assertExtraResourcesTracked(ResourceSet resources) throws NoSuchElementException {
+    for (Map.Entry<String, Float> 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<String, Float> 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.
@@ -472,7 +535,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
@@ -503,7 +574,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

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<String, Float> 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<String, Float> 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<String, Float> 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<String, Float> 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<String, Float> 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";