From 7e6c5e5ad2768499d181af3a8b21874d064b9197 Mon Sep 17 00:00:00 2001
From: Szilard Nemeth <snemeth@apache.org>
Date: Tue, 26 Nov 2019 16:42:33 +0100
Subject: [PATCH] YARN-9444. YARN API ResourceUtils's
 getRequestedResourcesFromConfig doesn't recognize yarn.io/gpu as a valid
 resource. Contributed by Gergely Pollak

(cherry picked from commit 52e9ee39a12ce91b3a545603dcf1103518ad2920)
---
 .../yarn/util/resource/ResourceUtils.java     |  4 +-
 .../yarn/util/resource/TestResourceUtils.java | 45 +++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/util/resource/ResourceUtils.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/util/resource/ResourceUtils.java
index 142dafa910a..ea21d5e48fb 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/util/resource/ResourceUtils.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-api/src/main/java/org/apache/hadoop/yarn/util/resource/ResourceUtils.java
@@ -66,6 +66,8 @@ public class ResourceUtils {
       "^(((\\p{Alnum}([\\p{Alnum}-]*\\p{Alnum})?\\.)*"
           + "\\p{Alnum}([\\p{Alnum}-]*\\p{Alnum})?)/)?\\p{Alpha}([\\w.-]*)$");
 
+  public static final String YARN_IO_OPTIONAL = "(yarn\\.io/)?";
+
   private static volatile boolean initializedResources = false;
   private static final Map<String, Integer> RESOURCE_NAME_TO_INDEX =
       new ConcurrentHashMap<String, Integer>();
@@ -657,7 +659,7 @@ public class ResourceUtils {
       Configuration configuration, String prefix) {
     List<ResourceInformation> result = new ArrayList<>();
     Map<String, String> customResourcesMap = configuration
-        .getValByRegex("^" + Pattern.quote(prefix) + "[^.]+$");
+        .getValByRegex("^" + Pattern.quote(prefix) + YARN_IO_OPTIONAL + "[^.]+$");
     for (Entry<String, String> resource : customResourcesMap.entrySet()) {
       String resourceName = resource.getKey().substring(prefix.length());
       Matcher matcher =
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/util/resource/TestResourceUtils.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/util/resource/TestResourceUtils.java
index 2671de85057..1eb79911ec5 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/util/resource/TestResourceUtils.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-common/src/test/java/org/apache/hadoop/yarn/util/resource/TestResourceUtils.java
@@ -31,8 +31,12 @@ import org.junit.Before;
 import org.junit.Test;
 
 import java.io.File;
+import java.util.Arrays;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * Test class to verify all resource utility methods.
@@ -152,6 +156,47 @@ public class TestResourceUtils {
     }
   }
 
+  @Test
+  public void testGetRequestedResourcesFromConfig() {
+    Configuration conf = new Configuration();
+
+    //these resource type configurations should be recognised
+    String propertyPrefix = "mapreduce.mapper.proper.rt.";
+    String[] expectedKeys = {
+        "yarn.io/gpu",
+        "yarn.io/fpga",
+        "yarn.io/anything_without_a_dot",
+        "regular_rt",
+        "regular_rt/with_slash"};
+
+    String[] invalidKeys = {
+        propertyPrefix + "too.many_parts",
+        propertyPrefix + "yarn.notio/gpu",
+        "incorrect.prefix.yarn.io/gpu",
+        propertyPrefix + "yarn.io/",
+        propertyPrefix};
+
+    for (String s : expectedKeys) {
+      //setting the properties which are expected to be in the resource list
+      conf.set(propertyPrefix + s, "42");
+    }
+
+    for (String s : invalidKeys) {
+      //setting the properties which are expected to be in the resource list
+      conf.set(s, "24");
+    }
+
+    List<ResourceInformation> properList =
+        ResourceUtils.getRequestedResourcesFromConfig(conf, propertyPrefix);
+    Set<String> expectedSet =
+        new HashSet<>(Arrays.asList(expectedKeys));
+
+    Assert.assertEquals(properList.size(), expectedKeys.length);
+    properList.forEach(
+        item -> Assert.assertTrue(expectedSet.contains(item.getName())));
+
+  }
+
   @Test
   public void testGetResourceTypesConfigErrors() throws Exception {
     Configuration conf = new YarnConfiguration();