diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/placement/UserGroupMappingPlacementRule.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/placement/UserGroupMappingPlacementRule.java index 6dce9c76c91..d401dd7dcaa 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/placement/UserGroupMappingPlacementRule.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/placement/UserGroupMappingPlacementRule.java @@ -29,9 +29,9 @@ import org.apache.hadoop.yarn.api.records.ApplicationId; import org.apache.hadoop.yarn.api.records.ApplicationSubmissionContext; import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.exceptions.YarnException; -import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler; import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.MappingType; import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.QueueMappingBuilder; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.ResourceScheduler; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.AutoCreatedLeafQueue; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CSQueue; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacityScheduler; @@ -65,13 +65,18 @@ public class UserGroupMappingPlacementRule extends PlacementRule { this(false, null, null); } - public UserGroupMappingPlacementRule(boolean overrideWithQueueMappings, + @VisibleForTesting + UserGroupMappingPlacementRule(boolean overrideWithQueueMappings, List newMappings, Groups groups) { this.mappings = newMappings; this.overrideWithQueueMappings = overrideWithQueueMappings; this.groups = groups; } + private String getPrimaryGroup(String user) throws IOException { + return groups.getGroups(user).get(0); + } + private String getSecondaryGroup(String user) throws IOException { List groupsList = groups.getGroups(user); String secondaryGroup = null; @@ -100,60 +105,27 @@ public class UserGroupMappingPlacementRule extends PlacementRule { if (mapping.getParentQueue() != null && mapping.getParentQueue().equals(PRIMARY_GROUP_MAPPING) && mapping.getQueue().equals(CURRENT_USER_MAPPING)) { - if (this.queueManager - .getQueue(groups.getGroups(user).get(0)) != null) { - QueueMapping queueMapping = - QueueMappingBuilder.create() - .type(mapping.getType()) - .source(mapping.getSource()).queue(user) - .parentQueue(groups.getGroups(user).get(0)) - .build(); - validateQueueMapping(queueMapping); - return getPlacementContext(queueMapping, user); - } else { - return null; - } + return getContextForGroupParent(user, mapping, + getPrimaryGroup(user)); } else if (mapping.getParentQueue() != null && mapping.getParentQueue().equals(SECONDARY_GROUP_MAPPING) && mapping.getQueue().equals(CURRENT_USER_MAPPING)) { - String secondaryGroup = getSecondaryGroup(user); - if (secondaryGroup != null) { - QueueMapping queueMapping = - QueueMappingBuilder.create() - .type(mapping.getType()) - .source(mapping.getSource()) - .queue(user) - .parentQueue(secondaryGroup) - .build(); - validateQueueMapping(queueMapping); - return getPlacementContext(queueMapping, user); - } else { - return null; - } + return getContextForGroupParent(user, mapping, + getSecondaryGroup(user)); } else if (mapping.getQueue().equals(CURRENT_USER_MAPPING)) { return getPlacementContext(mapping, user); } else if (mapping.getQueue().equals(PRIMARY_GROUP_MAPPING)) { - if (this.queueManager - .getQueue(groups.getGroups(user).get(0)) != null) { - return getPlacementContext(mapping, - groups.getGroups(user).get(0)); - } else { - return null; - } + return getContextForPrimaryGroup(user, mapping); } else if (mapping.getQueue().equals(SECONDARY_GROUP_MAPPING)) { - String secondaryGroup = getSecondaryGroup(user); - if (secondaryGroup != null) { - return getPlacementContext(mapping, secondaryGroup); - } else { - return null; - } + return getContextForSecondaryGroup(user, mapping); } else { return getPlacementContext(mapping); } } + if (user.equals(mapping.getSource())) { if (mapping.getQueue().equals(PRIMARY_GROUP_MAPPING)) { - return getPlacementContext(mapping, groups.getGroups(user).get(0)); + return getPlacementContext(mapping, getPrimaryGroup(user)); } else if (mapping.getQueue().equals(SECONDARY_GROUP_MAPPING)) { String secondaryGroup = getSecondaryGroup(user); if (secondaryGroup != null) { @@ -180,6 +152,70 @@ public class UserGroupMappingPlacementRule extends PlacementRule { return null; } + // invoked for mappings: + // u:%user:[parent].%primary_group + // u:%user:%primary_group + private ApplicationPlacementContext getContextForPrimaryGroup( + String user, + QueueMapping mapping) throws IOException { + String group = getPrimaryGroup(user); + + CSQueue parent = queueManager.getQueue(mapping.getParentQueue()); + + if (parent instanceof ManagedParentQueue) { + return getPlacementContext(mapping, group); + } else { + if (this.queueManager.getQueue(group) != null) { + return getPlacementContext(mapping, group); + } else { + return null; + } + } + } + + // invoked for mappings + // u:%user:%secondary_group + // u:%user:[parent].%secondary_group + private ApplicationPlacementContext getContextForSecondaryGroup( + String user, + QueueMapping mapping) throws IOException { + String secondaryGroup = getSecondaryGroup(user); + + if (secondaryGroup != null) { + if (this.queueManager.getQueue(secondaryGroup) != null) { + return getPlacementContext(mapping, secondaryGroup); + } else { + return null; + } + } else { + return null; + } + } + + // invoked for mappings: + // u:%user:%primary_group.%user + // u:%user:%secondary_group.%user + private ApplicationPlacementContext getContextForGroupParent( + String user, + QueueMapping mapping, + String group) throws IOException { + + if (this.queueManager.getQueue(group) != null) { + // replace the group string + QueueMapping resolvedGroupMapping = + QueueMappingBuilder.create() + .type(mapping.getType()) + .source(mapping.getSource()) + .queue(user) + .parentQueue(group) + .build(); + validateQueueMapping(resolvedGroupMapping); + return getPlacementContext(resolvedGroupMapping, user); + } else { + return null; + } + } + @Override public ApplicationPlacementContext getPlacementForApp( ApplicationSubmissionContext asc, String user) diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestUserGroupMappingPlacementRule.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestUserGroupMappingPlacementRule.java index e663a14c462..35bd9494a17 100644 --- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestUserGroupMappingPlacementRule.java +++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/placement/TestUserGroupMappingPlacementRule.java @@ -36,6 +36,7 @@ import org.apache.hadoop.yarn.server.resourcemanager.placement.QueueMapping.Queu import org.apache.hadoop.yarn.server.resourcemanager.placement.TestUserGroupMappingPlacementRule.QueueMappingTestData.QueueMappingTestDataBuilder; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.CapacitySchedulerQueueManager; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.LeafQueue; +import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.ManagedParentQueue; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.capacity.ParentQueue; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.PrimaryGroupMapping; import org.apache.hadoop.yarn.server.resourcemanager.scheduler.fair.SimpleGroupsMapping; @@ -74,6 +75,9 @@ public class TestUserGroupMappingPlacementRule { ParentQueue bsubgroup2 = mock(ParentQueue.class); when(bsubgroup2.getQueueName()).thenReturn("bsubgroup2"); + ManagedParentQueue managedParent = mock(ManagedParentQueue.class); + when(managedParent.getQueueName()).thenReturn("managedParent"); + LeafQueue a = mock(LeafQueue.class); when(a.getQueueName()).thenReturn("a"); when(a.getParent()).thenReturn(agroup); @@ -88,6 +92,7 @@ public class TestUserGroupMappingPlacementRule { when(queueManager.getQueue("agroup")).thenReturn(agroup); when(queueManager.getQueue("bsubgroup2")).thenReturn(bsubgroup2); when(queueManager.getQueue("asubgroup2")).thenReturn(asubgroup2); + when(queueManager.getQueue("managedParent")).thenReturn(managedParent); rule.setQueueManager(queueManager); ApplicationSubmissionContext asc = Records.newRecord( @@ -236,7 +241,72 @@ public class TestUserGroupMappingPlacementRule { .inputUser("a") .expectedQueue("q1") .build()); - + + // "agroup" queue exists + verifyQueueMapping( + QueueMappingTestDataBuilder.create() + .queueMapping(QueueMappingBuilder.create() + .type(MappingType.USER) + .source("%user") + .queue("%primary_group") + .parentQueue("bsubgroup2") + .build()) + .inputUser("a") + .expectedQueue("agroup") + .build()); + + // "abcgroup" queue doesn't exist, %primary_group queue, not managed parent + verifyQueueMapping( + QueueMappingTestDataBuilder.create() + .queueMapping(QueueMappingBuilder.create() + .type(MappingType.USER) + .source("%user") + .queue("%primary_group") + .parentQueue("bsubgroup2") + .build()) + .inputUser("abc") + .expectedQueue("default") + .build()); + + // "abcgroup" queue doesn't exist, %primary_group queue, managed parent + verifyQueueMapping( + QueueMappingTestDataBuilder.create() + .queueMapping(QueueMappingBuilder.create() + .type(MappingType.USER) + .source("%user") + .queue("%primary_group") + .parentQueue("managedParent") + .build()) + .inputUser("abc") + .expectedQueue("abcgroup") + .build()); + + // "abcgroup" queue doesn't exist, %secondary_group queue + verifyQueueMapping( + QueueMappingTestDataBuilder.create() + .queueMapping(QueueMappingBuilder.create() + .type(MappingType.USER) + .source("%user") + .queue("%secondary_group") + .parentQueue("bsubgroup2") + .build()) + .inputUser("abc") + .expectedQueue("default") + .build()); + + // "asubgroup2" queue exists, %secondary_group queue + verifyQueueMapping( + QueueMappingTestDataBuilder.create() + .queueMapping(QueueMappingBuilder.create() + .type(MappingType.USER) + .source("%user") + .queue("%secondary_group") + .parentQueue("bsubgroup2") + .build()) + .inputUser("a") + .expectedQueue("asubgroup2") + .build()); + // specify overwritten, and see if user specified a queue, and it will be // overridden verifyQueueMapping(