瀏覽代碼

YARN-10198. Remove Validation of Queue Exist under Managed Parent Queue

Contributed by Peter Bacsko. Reviewed by Manikandan R.
Prabhu Joseph 5 年之前
父節點
當前提交
d312991ead

+ 79 - 43
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<QueueMapping> 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<String> 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)

+ 71 - 1
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(