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

@@ -124,6 +124,7 @@ java_library(
         "//third_party:flogger",
         "//third_party:guava",
         "//third_party:jsr305",
+        "//third_party:rxjava3",
         "//third_party/protobuf:protobuf_java",
     ],
 )
@@ -304,6 +305,7 @@ java_library(
         "//src/main/java/com/google/devtools/common/options",
         "//third_party:flogger",
         "//third_party:guava",
+        "//third_party:rxjava3",
     ],
 )
 

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:<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

@@ -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<String, Float> 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<ResourceSet, CountDownLatch> 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,14 +276,25 @@ 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();
   }
 
   /**
    * 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<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
@@ -334,6 +363,20 @@ private synchronized boolean release(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);
+    }
+
     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<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.
   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

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<String, Float> 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<String, Float> 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, @NonNull ImmutableMap<String, Float> 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<String, Float> 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";
   }
 

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:<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 =
@@ -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<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);
+  }
+
+  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) {