From cf6d44d3eed5da5bbf7b163af84fe8440e558ae8 Mon Sep 17 00:00:00 2001
From: xiaojunxiang <xiaojunxiang@kingsoft.com>
Date: Mon, 2 Oct 2023 11:33:27 +0800
Subject: [PATCH] YARN-11582. Improve WebUI diagnosticMessage to show AM
 Container resource request size.

---
 .../common/fica/FiCaSchedulerApp.java         |  3 +
 .../TestApplicationLimitsByPartition.java     | 75 +++++++++++++++++++
 2 files changed, 78 insertions(+)

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/common/fica/FiCaSchedulerApp.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/common/fica/FiCaSchedulerApp.java
index 3a0fd347e5a0b3..fe5e02cc55ce80 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/common/fica/FiCaSchedulerApp.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/scheduler/common/fica/FiCaSchedulerApp.java
@@ -1026,6 +1026,9 @@ protected void getActivedAppDiagnosticMessage(
         .append(appAMNodePartitionName.isEmpty()
             ? NodeLabel.DEFAULT_NODE_LABEL_PARTITION : appAMNodePartitionName)
         .append(" ; ")
+        .append("AM Resource Request = ")
+        .append(getAMResource(appAMNodePartitionName))
+        .append(" ; ")
         .append("Partition Resource = ")
         .append(rmContext.getNodeLabelManager()
             .getResourceByLabel(appAMNodePartitionName, Resources.none()))
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/TestApplicationLimitsByPartition.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/TestApplicationLimitsByPartition.java
index ef50e5271e7188..a3a2ca3ba810e7 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/TestApplicationLimitsByPartition.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/TestApplicationLimitsByPartition.java
@@ -1027,4 +1027,79 @@ public void testAMLimitByAllResources() throws Exception {
     rm.close();
 
   }
+
+  @Test(timeout = 120000)
+  public void testDiagnosticWhenAMActivated() throws Exception {
+    /*
+     * Test Case:
+     * Verify AM resource limit per partition level and per queue level. So
+     * we use 2 queues to verify this case.
+     * Queue a1 supports labels (x,y). Configure am-resource-limit as 0.2 (x)
+     * Queue c1 supports default label. Configure am-resource-limit as 0.2
+     *
+     * Queue a1 for label X can only support 2Gb AM resource.
+     * Queue c1 (empty label) can support 2Gb AM resource.
+     *
+     * Verify at least one AM is launched, and AM resources should not go more
+     * than 2GB in each queue.
+     */
+
+    simpleNodeLabelMappingToManager();
+    CapacitySchedulerConfiguration config = (CapacitySchedulerConfiguration)
+         TestUtils.getConfigurationWithQueueLabels(conf);
+
+    // After getting queue conf, configure AM resource percent for Queue a1
+    // as 0.2 (Label X) and for Queue c1 as 0.2 (Empty Label)
+    config.setMaximumAMResourcePercentPerPartition(A1, "x", 0.2f);
+    config.setMaximumApplicationMasterResourcePerQueuePercent(C1, 0.2f);
+
+    // Now inject node label manager with this updated config.
+    MockRM rm = new MockRM(config) {
+      @Override
+      public RMNodeLabelsManager createNodeLabelManager() {
+        return mgr;
+      }
+    };
+
+    rm.getRMContext().setNodeLabelManager(mgr);
+    rm.start();
+    rm.registerNode("h1:1234", 10 * GB); // label = x
+    rm.registerNode("h2:1234", 10 * GB); // label = y
+    rm.registerNode("h3:1234", 10 * GB); // label = <empty>
+
+    // Submit app1 with 1Gb AM resource to Queue a1 for label X
+    MockRMAppSubmissionData data1 =
+         MockRMAppSubmissionData.Builder.createWithMemory(GB, rm)
+             .withAppName("app")
+             .withUser("user")
+             .withAcls(null)
+             .withQueue("a1")
+             .withAmLabel("x")
+             .build();
+    RMApp app1 = MockRMAppSubmitter.submit(rm, data1);
+
+    // Submit app2 with 1Gb AM resource to Queue a1 for label X
+    MockRMAppSubmissionData data2 =
+         MockRMAppSubmissionData.Builder.createWithMemory(GB, rm)
+             .withAppName("app")
+             .withUser("user")
+             .withAcls(null)
+             .withQueue("a1")
+             .withAmLabel("x")
+             .build();
+    RMApp app2 = MockRMAppSubmitter.submit(rm, data2);
+
+    CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler();
+    LeafQueue leafQueue = (LeafQueue) cs.getQueue("a1");
+    Assert.assertNotNull(leafQueue);
+
+    // Only one AM will be activated here, and second AM will be still pending.
+    Assert.assertEquals(2, leafQueue.getNumActiveApplications());
+    String activatedDiagnostics="AM Resource Request = ";
+    Assert.assertTrue("still doesn't show AMResource When Activated", app1.getDiagnostics()
+         .toString().contains(activatedDiagnostics));
+    Assert.assertTrue("still doesn't show AMResource When Activated", app2.getDiagnostics()
+         .toString().contains(activatedDiagnostics));
+    rm.close();
+  }
 }