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=<resourcename>=<amount> 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:<resourcename>:<amount>" tag.

Author: scele

src/main/java/com/google/devtools/build/lib/actions/BUILD

@@ -134,6 +134,7 @@ java_library(
         "//third_party:flogger",
         "//third_party:guava",
         "//third_party:jsr305",
+        "//third_party:rxjava3",
         "//third_party/protobuf:protobuf_java",
     ],
 )
@@ -318,6 +319,7 @@ java_library(
         "//third_party:flogger",
         "//third_party:guava",
         "//third_party:jsr305",
+        "//third_party:rxjava3",
     ],
 )
 

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

@@ -154,6 +154,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:<str>:<int>",
+          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";
 

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

@@ -27,8 +27,12 @@
 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.Set;
 import java.util.concurrent.CountDownLatch;
 import javax.annotation.Nullable;
 
@@ -183,13 +187,19 @@ 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<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() {}
+  private ResourceManager() {
+    usedExtraResources = new HashMap<>();
+  }
 
   @VisibleForTesting
   public static ResourceManager instanceForTestingOnly() {
@@ -204,6 +214,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();
@@ -298,6 +309,17 @@ 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()
@@ -417,6 +440,13 @@ private boolean release(ResourceSet resources, @Nullable Worker worker)
   private synchronized void releaseResourcesOnly(ResourceSet resources) {
     usedCpu -= resources.getCpuUsage();
     usedRam -= resources.getMemoryMb();
+
+    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);
+    }
+
     usedLocalTestCount -= resources.getLocalTestCount();
 
     // TODO(bazel-team): (2010) rounding error can accumulate and value below can end up being
@@ -428,6 +458,19 @@ private synchronized void releaseResourcesOnly(ResourceSet resources) {
     if (usedRam < epsilon) {
       usedRam = 0;
     }
+
+    Set<String> 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);
+    }
   }
 
   private synchronized boolean processAllWaitingThreads() throws IOException, InterruptedException {
@@ -466,6 +509,23 @@ private synchronized void processWaitingThreads(Deque<Pair<ResourceSet, LatchWit
     }
   }
 
+  /**
+   * Return true iff all requested extra resources are considered to be available.
+   */
+  private boolean areExtraResourcesAvailable(ResourceSet resources) {
+    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 = (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.
   @VisibleForTesting
   boolean areResourcesAvailable(ResourceSet resources) {
@@ -484,7 +544,7 @@ boolean areResourcesAvailable(ResourceSet resources) {
         workerKey == null
             || (activeWorkers < availableWorkers && workerPool.couldBeBorrowed(workerKey));
 
-    if (usedCpu == 0.0 && usedRam == 0.0 && usedLocalTestCount == 0 && workerIsAvailable) {
+    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 +575,8 @@ 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,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.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.util.OS;
 import com.google.devtools.build.lib.worker.WorkerKey;
 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;
 import javax.annotation.Nullable;
@@ -43,16 +46,26 @@ 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;
 
   /** The workerKey of used worker. Null if no worker is used. */
   @Nullable private final WorkerKey workerKey;
 
-  private ResourceSet(
-      double memoryMb, double cpuUsage, int localTestCount, @Nullable WorkerKey workerKey) {
+  private ResourceSet(double memoryMb, double cpuUsage, int localTestCount, @Nullable WorkerKey workerKey) {
+    this(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, workerKey);
+  }
+
+  private ResourceSet(double memoryMb, double cpuUsage, @NonNull 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 +96,31 @@ 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 createWithWorkerKey(memoryMb, cpuUsage, localTestCount, /* workerKey= */ null);
+    return ResourceSet.create(memoryMb, cpuUsage, ImmutableMap.of(), localTestCount, /* wolkerKey= */ 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, extraResourceUseage, localTestCount, /* workerKey= */ null);
   }
 
   public static ResourceSet createWithWorkerKey(
-      double memoryMb, double cpuUsage, int localTestCount, WorkerKey workerKey) {
-    if (memoryMb == 0 && cpuUsage == 0 && localTestCount == 0 && workerKey == null) {
+      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 +147,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 +165,7 @@ public String toString() {
         + "CPU: "
         + cpuUsage
         + "\n"
+        + Joiner.on("\n").withKeyValueSeparator(": ").join(extraResourceUsage.entrySet())
         + "Local tests: "
         + localTestCount
         + "\n";

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:<n>" 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,44 @@ 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<String, Float> getLocalExtraResourceUsage(Label label) throws UserExecException {
+    // Tests can specify requirements for extra resources using "resources:<resourcename>:<amount>" tag.
+    Map<String, Float> 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)