diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacitySchedulerConfiguration.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacitySchedulerConfiguration.java index 90979dc94dd..c345a552075 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacitySchedulerConfiguration.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/CapacitySchedulerConfiguration.java @@ -413,7 +413,7 @@ * Different resource types supported. */ public enum AbsoluteResourceType { - MEMORY, VCORES; + MEMORY, VCORES, GPUS; } AppPriorityACLConfigurationParser priorityACLConfig = new AppPriorityACLConfigurationParser(); @@ -2483,12 +2483,24 @@ private void updateMinMaxResourceToConf(String label, String queue, "Cannot set resource, root queue will take 100% of cluster capacity"); } + Integer gpuIndex = ResourceUtils.getResourceTypeIndex() + .get(ResourceInformation.GPU_URI); StringBuilder resourceString = new StringBuilder(); - resourceString - .append("[" + AbsoluteResourceType.MEMORY.toString().toLowerCase() + "=" - + resource.getMemorySize() + "," - + AbsoluteResourceType.VCORES.toString().toLowerCase() + "=" - + resource.getVirtualCores() + "]"); + if (gpuIndex == null) { + resourceString + .append("[" + AbsoluteResourceType.MEMORY.toString().toLowerCase() + "=" + + resource.getMemorySize() + "," + + AbsoluteResourceType.VCORES.toString().toLowerCase() + "=" + + resource.getVirtualCores() + "]"); + } else { + resourceString + .append("[" + AbsoluteResourceType.MEMORY.toString().toLowerCase() + "=" + + resource.getMemorySize() + "," + + AbsoluteResourceType.VCORES.toString().toLowerCase() + "=" + + resource.getVirtualCores() + "," + + AbsoluteResourceType.GPUS.toString().toLowerCase() + "=" + + resource.getResourceValue(ResourceInformation.GPU_URI) + "]"); + } String prefix = getQueuePrefix(queue) + type; if (!label.isEmpty()) { @@ -2586,6 +2598,16 @@ private void updateResourceValuesFromConfig(Set resourceTypes, case VCORES : resource.setVirtualCores(resourceValue.intValue()); break; + case GPUS : + Integer gpuIndex = ResourceUtils.getResourceTypeIndex() + .get(ResourceInformation.GPU_URI); + if (gpuIndex != null) { + resource.setResourceValue(ResourceInformation.GPU_URI, + resourceValue.intValue()); + } else { + LOG.error("GPU is not supported in conf."); + } + break; default : resource.setResourceInformation(splits[0].trim(), ResourceInformation .newInstance(splits[0].trim(), units, resourceValue)); diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCSAllocateCustomResource.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCSAllocateCustomResource.java index 65473b9eea7..2459f949749 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCSAllocateCustomResource.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/capacity/TestCSAllocateCustomResource.java @@ -21,6 +21,7 @@ import org.apache.commons.io.FileUtils; import org.apache.hadoop.yarn.api.records.ContainerId; import org.apache.hadoop.yarn.api.records.Resource; +import org.apache.hadoop.yarn.api.records.ResourceInformation; import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.server.resourcemanager.MockAM; import org.apache.hadoop.yarn.server.resourcemanager.MockNM; @@ -46,7 +47,7 @@ import java.io.File; import java.io.IOException; -import java.util.ArrayList; +import java.util.*; import static org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacitySchedulerConfiguration.MAXIMUM_ALLOCATION_MB; @@ -182,4 +183,65 @@ public void testCapacitySchedulerInitWithCustomResourceType() .getResourceValue("yarn.io/gpu")); rm.close(); } + + /** + * Test CS absolute conf with GPU resource type. + * */ + @Test + public void testCapacitySchedulerAbsoluteConfWithGPUResourceType() + throws IOException { + Resource QUEUE_A_MINRES = Resource.newInstance(1000, + 10, Collections.singletonMap(ResourceInformation.GPU_URI, 10L)); + Resource QUEUE_A_MAXRES = Resource.newInstance(10000, + 100, Collections.singletonMap(ResourceInformation.GPU_URI, 100L)); + + // Add corresponding gpu supported in conf. + Set resourceTypes = new HashSet<>( + Arrays.asList("memory", "vcores", "gpus")); + + // reset resource types + ResourceUtils.resetResourceTypes(); + String resourceTypesFileName = "resource-types-test.xml"; + File source = new File( + conf.getClassLoader().getResource(resourceTypesFileName).getFile()); + resourceTypesFile = new File(source.getParent(), "resource-types.xml"); + FileUtils.copyFile(source, resourceTypesFile); + + CapacitySchedulerConfiguration newConf = + new CapacitySchedulerConfiguration(conf); + + // Define top-level queues + newConf.setQueues(CapacitySchedulerConfiguration.ROOT, + new String[] { "a", "b", "c" }); + newConf.setMinimumResourceRequirement("", "root.a", + QUEUE_A_MINRES); + newConf.setMaximumResourceRequirement("", "root.a", + QUEUE_A_MAXRES); + + newConf.setClass(CapacitySchedulerConfiguration.RESOURCE_CALCULATOR_CLASS, + DominantResourceCalculator.class, ResourceCalculator.class); + + //start RM + MockRM rm = new MockRM(newConf); + rm.start(); + + // Check the gpu resource conf is right. + CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler(); + Assert.assertEquals(QUEUE_A_MINRES, + cs.getConfiguration(). + getMinimumResourceRequirement("", "root.a", resourceTypes)); + Assert.assertEquals(QUEUE_A_MAXRES, + cs.getConfiguration(). + getMaximumResourceRequirement("", "root.a", resourceTypes)); + + // Check the gpu resource of queue is right. + Assert.assertEquals(QUEUE_A_MINRES, cs.getQueue("root.a"). + getQueueResourceQuotas().getConfiguredMinResource()); + Assert.assertEquals(QUEUE_A_MAXRES, cs.getQueue("root.a"). + getQueueResourceQuotas().getConfiguredMaxResource()); + + rm.close(); + + } + }