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 ab40223..011bd97 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 @@ private static final Pattern RESOURCE_REQUEST_VALUE_PATTERN = Pattern.compile("^([0-9]+) ?([a-zA-Z]*)$"); + public static final String YARN_IO_OPTIONAL = "(yarn\\.io/)?"; + private static volatile boolean initializedResources = false; private static final Map RESOURCE_NAME_TO_INDEX = new ConcurrentHashMap(); @@ -596,7 +598,7 @@ public static String getDefaultUnit(String resourceType) { Configuration configuration, String prefix) { List result = new ArrayList<>(); Map customResourcesMap = configuration - .getValByRegex("^" + Pattern.quote(prefix) + "[^.]+$"); + .getValByRegex("^" + Pattern.quote(prefix) + YARN_IO_OPTIONAL + "[^.]+$"); for (Entry 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 b511705..c42e037 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 @@ -30,8 +30,13 @@ import org.junit.Test; import java.io.File; +import java.util.Arrays; import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; import java.util.Map; +import java.util.Set; /** * Test class to verify all resource utility methods. @@ -154,6 +159,49 @@ public void testGetResourceTypesConfigs() throws Exception { } @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 properList = + ResourceUtils.getRequestedResourcesFromConfig(conf, propertyPrefix); + Set expectedSet = + new HashSet<>(Arrays.asList(expectedKeys)); + + Assert.assertEquals(properList.size(), expectedKeys.length); + for(Iterator propIter = properList.iterator(); + propIter.hasNext();) { + ResourceInformation item = propIter.next(); + Assert.assertTrue(expectedSet.contains(item.getName())); + } + } + + @Test public void testGetResourceTypesConfigErrors() throws Exception { Configuration conf = new YarnConfiguration();