From ec16b1d3b9294e7047f9d22310e0ac272cee4fce Mon Sep 17 00:00:00 2001
From: zhuqi-lucas <821684824@qq.com>
Date: Tue, 8 Jun 2021 21:03:43 +0800
Subject: [PATCH] YARN-10807. Parents node labels are incorrectly added to
 child queues in weight mode. Contributed by Benjamin Teke.

---
 .../scheduler/capacity/ParentQueue.java       | 17 +++++++----
 .../TestCapacitySchedulerWeightMode.java      | 30 +++++++++++++++++++
 2 files changed, 42 insertions(+), 5 deletions(-)

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/ParentQueue.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/ParentQueue.java
index 6b6dd5afd3..6f850187bf 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/ParentQueue.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/ParentQueue.java
@@ -1250,16 +1250,23 @@ public void updateClusterResource(Resource clusterResource,
           float sumOfWeight = 0;
 
           for (CSQueue queue : childQueues) {
-            float weight = Math.max(0,
-                queue.getQueueCapacities().getWeight(nodeLabel));
-            sumOfWeight += weight;
+            if (queue.getQueueCapacities().getExistingNodeLabels()
+                .contains(nodeLabel)) {
+              float weight = Math.max(0,
+                  queue.getQueueCapacities().getWeight(nodeLabel));
+              sumOfWeight += weight;
+            }
           }
           // When sum of weight == 0, skip setting normalized_weight (so
           // normalized weight will be 0).
           if (Math.abs(sumOfWeight) > 1e-6) {
             for (CSQueue queue : childQueues) {
-              queue.getQueueCapacities().setNormalizedWeight(nodeLabel,
-                  queue.getQueueCapacities().getWeight(nodeLabel) / sumOfWeight);
+              if (queue.getQueueCapacities().getExistingNodeLabels()
+                  .contains(nodeLabel)) {
+                queue.getQueueCapacities().setNormalizedWeight(nodeLabel,
+                    queue.getQueueCapacities().getWeight(nodeLabel) /
+                        sumOfWeight);
+              }
             }
           }
         }
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/TestCapacitySchedulerWeightMode.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/TestCapacitySchedulerWeightMode.java
index 300e4cd097..171123a6ce 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/TestCapacitySchedulerWeightMode.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/TestCapacitySchedulerWeightMode.java
@@ -338,6 +338,36 @@ public void testGetCapacityOrWeightStringParentPctLeafWeights()
     }
   }
 
+  /**
+   * This test ensures that while iterating through a parent's Node Labels
+   * (when calculating the normalized weights) the parent's Node Labels won't
+   * be added to the children with weight -1. If the parent
+   * has a node label that a specific child doesn't the normalized calling the
+   * normalized weight setter will be skipped. The queue root.b has access to
+   * the labels "x" and "y", but root.b.b1 won't. For more information see
+   * YARN-10807.
+   * @throws Exception
+   */
+  @Test
+  public void testChildAccessibleNodeLabelsWeightMode() throws Exception {
+    MockRM rm = new MockRM(getCSConfWithQueueLabelsWeightOnly(conf));
+    rm.start();
+
+    CapacityScheduler cs =
+        (CapacityScheduler) rm.getRMContext().getScheduler();
+    LeafQueue b1 = (LeafQueue) cs.getQueue(B1);
+
+    Assert.assertNotNull(b1);
+    Assert.assertTrue(b1.getAccessibleNodeLabels().isEmpty());
+
+    Set<String> b1ExistingNodeLabels = ((CSQueue) b1).getQueueCapacities()
+        .getExistingNodeLabels();
+    Assert.assertEquals(1, b1ExistingNodeLabels.size());
+    Assert.assertEquals("", b1ExistingNodeLabels.iterator().next());
+
+    rm.close();
+  }
+
   @Test
   public void testQueueInfoWeight() throws Exception {
     MockRM rm = new MockRM(conf);