YARN-8836. Add tags and attributes in resource definition. Contributed by Weiwei Yang.

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/api/records/ResourceInformation.java

@@ -19,11 +19,15 @@
 package org.apache.hadoop.yarn.api.records;
 
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import org.apache.hadoop.classification.InterfaceAudience;
 import org.apache.hadoop.yarn.api.protocolrecords.ResourceTypes;
 import org.apache.hadoop.yarn.util.UnitsConversionUtil;
 
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * Class to encapsulate information about a Resource - the name of the resource,
@@ -37,6 +41,8 @@ public class ResourceInformation implements Comparable<ResourceInformation> {
   private long value;
   private long minimumAllocation;
   private long maximumAllocation;
+  private Set<String> tags = new HashSet<>();
+  private Map<String, String> attributes = new HashMap<>();
 
   // Known resource types
   public static final String MEMORY_URI = "memory-mb";
@@ -184,6 +190,42 @@ public class ResourceInformation implements Comparable<ResourceInformation> {
     this.maximumAllocation = maximumAllocation;
   }
 
+  /**
+   * Get the attributes of the resource.
+   * @return resource attributes
+   */
+  public Map<String, String> getAttributes() {
+    return attributes;
+  }
+
+  /**
+   * Set a map of attributes to the resource.
+   * @param attributes resource attributes
+   */
+  public void setAttributes(Map<String, String> attributes) {
+    if (attributes != null) {
+      this.attributes = attributes;
+    }
+  }
+
+  /**
+   * Get resource tags.
+   * @return resource tags
+   */
+  public Set<String> getTags() {
+    return this.tags;
+  }
+
+  /**
+   * Add tags to the resource.
+   * @param tags resource tags
+   */
+  public void setTags(Set<String> tags) {
+    if (tags != null) {
+      this.tags = tags;
+    }
+  }
+
   /**
    * Create a new instance of ResourceInformation from another object.
    *
@@ -199,6 +241,15 @@ public class ResourceInformation implements Comparable<ResourceInformation> {
   public static ResourceInformation newInstance(String name, String units,
       long value, ResourceTypes type, long minimumAllocation,
       long maximumAllocation) {
+    return ResourceInformation.newInstance(name, units, value, type,
+        minimumAllocation, maximumAllocation,
+        ImmutableSet.of(), ImmutableMap.of());
+  }
+
+  public static ResourceInformation newInstance(String name, String units,
+      long value, ResourceTypes type, long minimumAllocation,
+      long maximumAllocation,
+      Set<String> tags, Map<String, String> attributes) {
     ResourceInformation ret = new ResourceInformation();
     ret.setName(name);
     ret.setResourceType(type);
@@ -206,6 +257,8 @@ public class ResourceInformation implements Comparable<ResourceInformation> {
     ret.setValue(value);
     ret.setMinimumAllocation(minimumAllocation);
     ret.setMaximumAllocation(maximumAllocation);
+    ret.setTags(tags);
+    ret.setAttributes(attributes);
     return ret;
   }
 
@@ -258,13 +311,16 @@ public class ResourceInformation implements Comparable<ResourceInformation> {
     dst.setValue(src.getValue());
     dst.setMinimumAllocation(src.getMinimumAllocation());
     dst.setMaximumAllocation(src.getMaximumAllocation());
+    dst.setTags(src.getTags());
+    dst.setAttributes(src.getAttributes());
   }
 
   @Override
   public String toString() {
     return "name: " + this.name + ", units: " + this.units + ", type: "
         + resourceType + ", value: " + value + ", minimum allocation: "
-        + minimumAllocation + ", maximum allocation: " + maximumAllocation;
+        + minimumAllocation + ", maximum allocation: " + maximumAllocation
+        + ", tags: " + tags + ", attributes " + attributes;
   }
 
   public String getShorthandRepresentation() {
@@ -284,7 +340,9 @@ public class ResourceInformation implements Comparable<ResourceInformation> {
     }
     ResourceInformation r = (ResourceInformation) obj;
     if (!this.name.equals(r.getName())
-        || !this.resourceType.equals(r.getResourceType())) {
+        || !this.resourceType.equals(r.getResourceType())
+        || !this.tags.equals(r.getTags())
+        || !this.attributes.equals(r.getAttributes())) {
       return false;
     }
     if (this.units.equals(r.units)) {
@@ -302,6 +360,12 @@ public class ResourceInformation implements Comparable<ResourceInformation> {
     result = prime * result + resourceType.hashCode();
     result = prime * result + units.hashCode();
     result = prime * result + Long.hashCode(value);
+    if (tags != null && !tags.isEmpty()) {
+      result = prime * result + tags.hashCode();
+    }
+    if (attributes != null && !attributes.isEmpty()) {
+      result = prime * result + attributes.hashCode();
+    }
     return result;
   }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/util/resource/ResourceUtils.java

@@ -41,9 +41,11 @@ import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -55,6 +57,7 @@ public class ResourceUtils {
 
   public static final String UNITS = ".units";
   public static final String TYPE = ".type";
+  public static final String TAGS = ".tags";
   public static final String MINIMUM_ALLOCATION = ".minimum-allocation";
   public static final String MAXIMUM_ALLOCATION = ".maximum-allocation";
 
@@ -242,6 +245,10 @@ public class ResourceUtils {
                   + "'. One of name, units or type is configured incorrectly.");
         }
         ResourceTypes resourceType = ResourceTypes.valueOf(resourceTypeName);
+        String[] resourceTags = conf.getTrimmedStrings(
+            YarnConfiguration.RESOURCE_TYPES + "." + resourceName + TAGS);
+        Set<String> resourceTagSet = new HashSet<>();
+        Collections.addAll(resourceTagSet, resourceTags);
         LOG.info("Adding resource type - name = " + resourceName + ", units = "
             + resourceUnits + ", type = " + resourceTypeName);
         if (resourceInformationMap.containsKey(resourceName)) {
@@ -250,7 +257,7 @@ public class ResourceUtils {
         }
         resourceInformationMap.put(resourceName, ResourceInformation
             .newInstance(resourceName, resourceUnits, 0L, resourceType,
-                minimumAllocation, maximumAllocation));
+                minimumAllocation, maximumAllocation, resourceTagSet, null));
       }
     }
 

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/proto/yarn_protos.proto

@@ -62,6 +62,8 @@ message ResourceInformationProto {
   optional int64 value = 2;
   optional string units = 3;
   optional ResourceTypesProto type = 4;
+  repeated string tags = 5;
+  repeated StringStringMapProto attributes = 6;
 }
 
 message ResourceTypeInfoProto {

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/test/java/org/apache/hadoop/yarn/conf/TestResourceInformation.java

@@ -18,6 +18,9 @@
 
 package org.apache.hadoop.yarn.conf;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import org.apache.hadoop.yarn.api.protocolrecords.ResourceTypes;
 import org.apache.hadoop.yarn.api.records.ResourceInformation;
 import org.junit.Assert;
 import org.junit.Test;
@@ -70,4 +73,45 @@ public class TestResourceInformation {
     Assert.assertEquals("Resource value incorrect", value, ri.getValue());
     Assert.assertEquals("Resource units incorrect", units, ri.getUnits());
   }
+
+  @Test
+  public void testEqualsWithTagsAndAttributes() {
+    // Same tags but different order
+    ResourceInformation ri01 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100,
+        ImmutableSet.of("A", "B"), null);
+    ResourceInformation ri02 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100, ImmutableSet.of("B", "A"), null);
+    Assert.assertEquals(ri01, ri02);
+
+    // Different tags
+    ResourceInformation ri11 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100, null, null);
+    ResourceInformation ri12 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100, ImmutableSet.of("B", "A"), null);
+    Assert.assertNotEquals(ri11, ri12);
+
+    // Different attributes
+    ResourceInformation ri21 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100, null,
+        ImmutableMap.of("A", "A1", "B", "B1"));
+    ResourceInformation ri22 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100, null,
+        ImmutableMap.of("A", "A1", "B", "B2"));
+    Assert.assertNotEquals(ri21, ri22);
+
+    // No tags or attributes
+    ResourceInformation ri31 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100, null, null);
+    ResourceInformation ri32 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100, null, null);
+    Assert.assertEquals(ri31, ri32);
+
+    // Null tags/attributes same as empty ones
+    ResourceInformation ri41 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100, ImmutableSet.of(), null);
+    ResourceInformation ri42 = ResourceInformation.newInstance("r1", "M", 100,
+        ResourceTypes.COUNTABLE, 0, 100, null, ImmutableMap.of());
+    Assert.assertEquals(ri41, ri42);
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/ProtoUtils.java

@@ -78,6 +78,7 @@ import org.apache.hadoop.yarn.proto.YarnProtos.QueueACLProto;
 import org.apache.hadoop.yarn.proto.YarnProtos.QueueStateProto;
 import org.apache.hadoop.yarn.proto.YarnProtos.ReservationRequestInterpreterProto;
 import org.apache.hadoop.yarn.proto.YarnProtos.ResourceProto;
+import org.apache.hadoop.yarn.proto.YarnProtos.StringStringMapProto;
 import org.apache.hadoop.yarn.proto.YarnProtos.TimedPlacementConstraintProto;
 import org.apache.hadoop.yarn.proto.YarnProtos.YarnApplicationAttemptStateProto;
 import org.apache.hadoop.yarn.proto.YarnProtos.YarnApplicationStateProto;
@@ -528,6 +529,33 @@ public class ProtoUtils {
     return ret;
   }
 
+  public static Map<String, String> convertStringStringMapProtoListToMap(
+      List<StringStringMapProto> pList) {
+    Map<String, String> ret = new HashMap<>();
+    if (pList != null) {
+      for (StringStringMapProto p : pList) {
+        if (p.hasKey()) {
+          ret.put(p.getKey(), p.getValue());
+        }
+      }
+    }
+    return ret;
+  }
+
+  public static List<YarnProtos.StringStringMapProto> convertToProtoFormat(
+      Map<String, String> stringMap) {
+    List<YarnProtos.StringStringMapProto> pList = new ArrayList<>();
+    if (stringMap != null && !stringMap.isEmpty()) {
+      StringStringMapProto.Builder pBuilder = StringStringMapProto.newBuilder();
+      for (Map.Entry<String, String> entry : stringMap.entrySet()) {
+        pBuilder.setKey(entry.getKey());
+        pBuilder.setValue(entry.getValue());
+        pList.add(pBuilder.build());
+      }
+    }
+    return pList;
+  }
+
   public static PlacementConstraintTargetProto.TargetType convertToProtoFormat(
           TargetExpression.TargetType t) {
     return PlacementConstraintTargetProto.TargetType.valueOf(t.name());

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/main/java/org/apache/hadoop/yarn/api/records/impl/pb/ResourcePBImpl.java

@@ -18,6 +18,8 @@
 
 package org.apache.hadoop.yarn.api.records.impl.pb;
 
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.classification.InterfaceAudience.Private;
@@ -32,6 +34,7 @@ import org.apache.hadoop.yarn.proto.YarnProtos.ResourceInformationProto;
 import org.apache.hadoop.yarn.util.UnitsConversionUtil;
 import org.apache.hadoop.yarn.util.resource.ResourceUtils;
 
+import java.util.HashSet;
 import java.util.Map;
 
 @Private
@@ -184,6 +187,17 @@ public class ResourcePBImpl extends Resource {
       ri.setUnits(units);
       ri.setValue(value);
     }
+    if (entry.getTagsCount() > 0) {
+      ri.setTags(new HashSet<>(entry.getTagsList()));
+    } else {
+      ri.setTags(ImmutableSet.of());
+    }
+    if (entry.getAttributesCount() > 0) {
+      ri.setAttributes(ProtoUtils
+          .convertStringStringMapProtoListToMap(entry.getAttributesList()));
+    } else {
+      ri.setAttributes(ImmutableMap.of());
+    }
     return ri;
   }
 
@@ -230,6 +244,15 @@ public class ResourcePBImpl extends Resource {
         e.setUnits(resInfo.getUnits());
         e.setType(ProtoUtils.converToProtoFormat(resInfo.getResourceType()));
         e.setValue(resInfo.getValue());
+        if (resInfo.getAttributes() != null
+            && !resInfo.getAttributes().isEmpty()) {
+          e.addAllAttributes(ProtoUtils.convertToProtoFormat(
+              resInfo.getAttributes()));
+        }
+        if (resInfo.getTags() != null
+            && !resInfo.getTags().isEmpty()) {
+          e.addAllTags(resInfo.getTags());
+        }
         builder.addResourceValueMap(e);
       }
     }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/api/TestResourcePBImpl.java

@@ -23,6 +23,7 @@ import java.io.File;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.ResourceInformation;
+import org.apache.hadoop.yarn.api.records.impl.pb.ProtoUtils;
 import org.apache.hadoop.yarn.api.records.impl.pb.ResourcePBImpl;
 import org.apache.hadoop.yarn.conf.YarnConfiguration;
 import org.apache.hadoop.yarn.proto.YarnProtos;
@@ -44,7 +45,7 @@ public class TestResourcePBImpl {
   public void setup() throws Exception {
     ResourceUtils.resetResourceTypes();
 
-    String resourceTypesFile = "resource-types-4.xml";
+    String resourceTypesFile = "resource-types-5.xml";
     Configuration conf = new YarnConfiguration();
     TestResourceUtils.setupResourceTypes(conf, resourceTypesFile);
   }
@@ -53,7 +54,7 @@ public class TestResourcePBImpl {
   public void teardown() {
     Configuration conf = new YarnConfiguration();
     File source = new File(
-        conf.getClassLoader().getResource("resource-types-4.xml").getFile());
+        conf.getClassLoader().getResource("resource-types-5.xml").getFile());
     File dest = new File(source.getParent(), "resource-types.xml");
     if (dest.exists()) {
       dest.delete();
@@ -175,4 +176,127 @@ public class TestResourcePBImpl {
     Assert.assertEquals("G",
         res2.getResourceInformation("resource1").getUnits());
   }
+
+  @Test
+  public void testResourceTags() {
+    YarnProtos.ResourceInformationProto riProto =
+        YarnProtos.ResourceInformationProto.newBuilder()
+            .setType(
+                YarnProtos.ResourceTypeInfoProto.newBuilder()
+                    .setName("yarn.io/test-volume")
+                    .setType(YarnProtos.ResourceTypesProto.COUNTABLE).getType())
+            .setValue(10)
+            .setUnits("G")
+            .setKey("yarn.io/test-volume")
+            .addTags("tag_A")
+            .addTags("tag_B")
+            .addTags("tag_C")
+            .build();
+    YarnProtos.ResourceProto proto =
+        YarnProtos.ResourceProto.newBuilder()
+            .setMemory(1024)
+            .setVirtualCores(3)
+            .addResourceValueMap(riProto)
+            .build();
+    Resource res = new ResourcePBImpl(proto);
+
+    Assert.assertNotNull(res.getResourceInformation("yarn.io/test-volume"));
+    Assert.assertEquals(10,
+        res.getResourceInformation("yarn.io/test-volume")
+            .getValue());
+    Assert.assertEquals("G",
+        res.getResourceInformation("yarn.io/test-volume")
+            .getUnits());
+    Assert.assertEquals(3,
+        res.getResourceInformation("yarn.io/test-volume")
+            .getTags().size());
+    Assert.assertFalse(res.getResourceInformation("yarn.io/test-volume")
+        .getTags().isEmpty());
+    Assert.assertTrue(res.getResourceInformation("yarn.io/test-volume")
+        .getAttributes().isEmpty());
+
+    boolean protoConvertExpected = false;
+    YarnProtos.ResourceProto protoFormat = ProtoUtils.convertToProtoFormat(res);
+    for (YarnProtos.ResourceInformationProto pf :
+        protoFormat.getResourceValueMapList()) {
+      if (pf.getKey().equals("yarn.io/test-volume")) {
+        protoConvertExpected = pf.getAttributesCount() == 0
+            && pf.getTagsCount() == 3;
+      }
+    }
+    Assert.assertTrue("Expecting resource's protobuf message"
+        + " contains 0 attributes and 3 tags",
+        protoConvertExpected);
+  }
+
+  @Test
+  public void testResourceAttributes() {
+    YarnProtos.ResourceInformationProto riProto =
+        YarnProtos.ResourceInformationProto.newBuilder()
+            .setType(
+                YarnProtos.ResourceTypeInfoProto.newBuilder()
+                    .setName("yarn.io/test-volume")
+                    .setType(YarnProtos.ResourceTypesProto.COUNTABLE).getType())
+            .setValue(10)
+            .setUnits("G")
+            .setKey("yarn.io/test-volume")
+            .addAttributes(
+                YarnProtos.StringStringMapProto
+                    .newBuilder()
+                    .setKey("driver").setValue("test-driver")
+                    .build())
+            .addAttributes(
+                YarnProtos.StringStringMapProto
+                    .newBuilder()
+                    .setKey("mount").setValue("/mnt/data")
+                    .build())
+            .build();
+    YarnProtos.ResourceProto proto =
+        YarnProtos.ResourceProto.newBuilder()
+            .setMemory(1024)
+            .setVirtualCores(3)
+            .addResourceValueMap(riProto)
+            .build();
+    Resource res = new ResourcePBImpl(proto);
+
+    Assert.assertNotNull(res.getResourceInformation("yarn.io/test-volume"));
+    Assert.assertEquals(10,
+        res.getResourceInformation("yarn.io/test-volume")
+            .getValue());
+    Assert.assertEquals("G",
+        res.getResourceInformation("yarn.io/test-volume")
+            .getUnits());
+    Assert.assertEquals(2,
+        res.getResourceInformation("yarn.io/test-volume")
+            .getAttributes().size());
+    Assert.assertTrue(res.getResourceInformation("yarn.io/test-volume")
+        .getTags().isEmpty());
+    Assert.assertFalse(res.getResourceInformation("yarn.io/test-volume")
+        .getAttributes().isEmpty());
+
+    boolean protoConvertExpected = false;
+    YarnProtos.ResourceProto protoFormat = ProtoUtils.convertToProtoFormat(res);
+    for (YarnProtos.ResourceInformationProto pf :
+        protoFormat.getResourceValueMapList()) {
+      if (pf.getKey().equals("yarn.io/test-volume")) {
+        protoConvertExpected = pf.getAttributesCount() == 2
+            && pf.getTagsCount() == 0;
+      }
+    }
+    Assert.assertTrue("Expecting resource's protobuf message"
+            + " contains 2 attributes and 0 tags",
+        protoConvertExpected);
+  }
+
+  @Test
+  public void testParsingResourceTags() {
+    ResourceInformation info =
+        ResourceUtils.getResourceTypes().get("resource3");
+    Assert.assertTrue(info.getAttributes().isEmpty());
+    Assert.assertFalse(info.getTags().isEmpty());
+    Assert.assertEquals(info.getTags().size(), 2);
+    info.getTags().remove("resource3_tag_1");
+    info.getTags().remove("resource3_tag_2");
+    Assert.assertTrue(info.getTags().isEmpty());
+  }
 }

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/resources/resource-types/resource-types-5.xml

@@ -0,0 +1,53 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<?xml-stylesheet type="text/xsl" href="configuration.xsl"?>
+<!--
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. See accompanying LICENSE file.
+-->
+
+<configuration>
+
+    <property>
+        <name>yarn.resource-types</name>
+        <value>resource1,resource2,resource3,yarn.io/gpu,yarn.io/test-volume</value>
+    </property>
+
+    <property>
+        <name>yarn.resource-types.resource1.units</name>
+        <value>G</value>
+    </property>
+
+    <property>
+        <name>yarn.resource-types.resource2.units</name>
+        <value>m</value>
+    </property>
+
+    <property>
+        <name>yarn.resource-types.resource3.units</name>
+        <value>G</value>
+    </property>
+
+    <property>
+        <name>yarn.resource-types.resource3.tags</name>
+        <value>resource3_tag_1,resource3_tag_2</value>
+    </property>
+
+    <property>
+        <name>yarn.resource-types.yarn.io/gpu.units</name>
+        <value></value>
+    </property>
+
+    <property>
+        <name>yarn.resource-types.yarn.io/test-volume.units</name>
+        <value>G</value>
+    </property>
+</configuration>