From 6cb90005a7d0651474883ac4e1b6961ef74fe513 Mon Sep 17 00:00:00 2001
From: Eric Badger <ebadger@verizonmedia.com>
Date: Tue, 20 Apr 2021 17:34:49 +0000
Subject: [PATCH] YARN-10723. Change CS nodes page in UI to support custom
 resource. Contributed by Qi Zhu

---
 .../resourcemanager/webapp/NodesPage.java     | 65 +++++++++++++------
 .../resourcemanager/webapp/TestNodesPage.java | 26 +++++++-
 2 files changed, 69 insertions(+), 22 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/webapp/NodesPage.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodesPage.java
index 8e8a2610f0..545ae88180 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodesPage.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/NodesPage.java
@@ -39,6 +39,7 @@
 import org.apache.hadoop.yarn.webapp.view.HtmlBlock;
 
 import java.util.Collection;
+import java.util.Map;
 
 import static org.apache.hadoop.yarn.webapp.YarnWebParams.NODE_LABEL;
 import static org.apache.hadoop.yarn.webapp.YarnWebParams.NODE_STATE;
@@ -90,9 +91,7 @@ protected void render(Block html) {
             .th(".mem", "Phys Mem Used %")
             .th(".vcores", "VCores Used")
             .th(".vcores", "VCores Avail")
-            .th(".vcores", "Phys VCores Used %")
-            .th(".gpus", "GPUs Used")
-            .th(".gpus", "GPUs Avail");
+            .th(".vcores", "Phys VCores Used %");
       } else {
         trbody.th(".containers", "Running Containers (G)")
             .th(".allocationTags", "Allocation Tags")
@@ -102,14 +101,26 @@ protected void render(Block html) {
             .th(".vcores", "VCores Used (G)")
             .th(".vcores", "VCores Avail (G)")
             .th(".vcores", "Phys VCores Used %")
-            .th(".gpus", "GPUs Used (G)")
-            .th(".gpus", "GPUs Avail (G)")
             .th(".containers", "Running Containers (O)")
             .th(".mem", "Mem Used (O)")
             .th(".vcores", "VCores Used (O)")
             .th(".containers", "Queued Containers");
       }
 
+      for (Map.Entry<String, Integer> integerEntry :
+          ResourceUtils.getResourceTypeIndex().entrySet()) {
+        if (integerEntry.getKey().equals(ResourceInformation.MEMORY_URI)
+            || integerEntry.getKey().equals(ResourceInformation.VCORES_URI)) {
+          continue;
+        }
+
+        trbody.th("." + integerEntry.getKey(),
+            integerEntry.getKey() + " " + "Used");
+
+        trbody.th("." + integerEntry.getKey(),
+            integerEntry.getKey() + " " + "Avail");
+      }
+
       TBODY<TABLE<Hamlet>> tbody =
           trbody.th(".nodeManagerVersion", "Version").__().__().tbody();
 
@@ -175,17 +186,7 @@ protected void render(Block html) {
           nodeTableData.append("\",\"<a ").append("href='" + "//" + httpAddress)
               .append("'>").append(httpAddress).append("</a>\",").append("\"");
         }
-        Integer gpuIndex = ResourceUtils.getResourceTypeIndex()
-            .get(ResourceInformation.GPU_URI);
-        long usedGPUs = 0;
-        long availableGPUs = 0;
-        if (gpuIndex != null && info.getUsedResource() != null
-            && info.getAvailableResource() != null) {
-          usedGPUs = info.getUsedResource().getResource()
-              .getResourceValue(ResourceInformation.GPU_URI);
-          availableGPUs = info.getAvailableResource().getResource()
-              .getResourceValue(ResourceInformation.GPU_URI);
-        }
+
         nodeTableData.append("<br title='")
             .append(String.valueOf(info.getLastHealthUpdate())).append("'>")
             .append(Times.format(info.getLastHealthUpdate())).append("\",\"")
@@ -205,10 +206,6 @@ protected void render(Block html) {
             .append(String.valueOf(info.getAvailableVirtualCores()))
             .append("\",\"")
             .append(String.valueOf((int) info.getVcoreUtilization()))
-            .append("\",\"")
-            .append(String.valueOf(usedGPUs))
-            .append("\",\"")
-            .append(String.valueOf(availableGPUs))
             .append("\",\"");
 
         // If opportunistic containers are enabled, add extra fields.
@@ -226,6 +223,34 @@ protected void render(Block html) {
               .append("\",\"");
         }
 
+        for (Map.Entry<String, Integer> integerEntry :
+            ResourceUtils.getResourceTypeIndex().entrySet()) {
+          if (integerEntry.getKey().equals(ResourceInformation.MEMORY_URI)
+              || integerEntry.getKey().equals(ResourceInformation.VCORES_URI)) {
+            continue;
+          }
+
+          long usedCustomResource = 0;
+          long availableCustomResource = 0;
+
+          String resourceName = integerEntry.getKey();
+          Integer index = integerEntry.getValue();
+
+          if (index != null && info.getUsedResource() != null
+              && info.getAvailableResource() != null) {
+            usedCustomResource = info.getUsedResource().getResource()
+                .getResourceValue(resourceName);
+            availableCustomResource = info.getAvailableResource().getResource()
+                .getResourceValue(resourceName);
+
+            nodeTableData
+                .append(usedCustomResource)
+                .append("\",\"")
+                .append(availableCustomResource)
+                .append("\",\"");
+          }
+        }
+
         nodeTableData.append(ni.getNodeManagerVersion())
             .append("\"],\n");
       }
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestNodesPage.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestNodesPage.java
index 9ab6583b06..dd271fd34d 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestNodesPage.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/webapp/TestNodesPage.java
@@ -53,7 +53,7 @@ public class TestNodesPage {
   // Number of Actual Table Headers for NodesPage.NodesBlock might change in
   // future. In that case this value should be adjusted to the new value.
   private final int numberOfThInMetricsTable = 23;
-  private final int numberOfActualTableHeaders = 18;
+  private final int numberOfActualTableHeaders = 16;
   private final int numberOfThForOpportunisticContainers = 4;
 
   private Injector injector;
@@ -119,12 +119,34 @@ public void testNodesBlockRenderForLostNodesWithGPUResources()
         initResourceTypes(ResourceInformation.GPU_URI);
     this.setUpInternal(true);
     try {
-      this.testNodesBlockRenderForLostNodes();
+      // Test gpu as a custom resource.
+      //<th class="yarn io/gpu">
+      //  yarn.io/gpu Used
+      //</th>
+      //<th class="yarn io/gpu">
+      //   yarn.io/gpu Avail
+      //</th>
+      this.testNodesBlockRenderForLostNodesWithGPU();
     } finally {
       ResourceUtils.initializeResourcesFromResourceInformationMap(oldRtMap);
     }
   }
 
+  public void testNodesBlockRenderForLostNodesWithGPU() {
+    NodesBlock nodesBlock = injector.getInstance(NodesBlock.class);
+    nodesBlock.set("node.state", "lost");
+    nodesBlock.render();
+    PrintWriter writer = injector.getInstance(PrintWriter.class);
+    WebAppTests.flushOutput(injector);
+
+    Mockito.verify(writer,
+        Mockito.times(numberOfActualTableHeaders
+            + numberOfThInMetricsTable + 2))
+        .print("<th");
+    Mockito.verify(writer, Mockito.times(numberOfThInMetricsTable))
+        .print("<td");
+  }
+
   @Test
   public void testNodesBlockRenderForNodeLabelFilterWithNonEmptyLabel() {
     NodesBlock nodesBlock = injector.getInstance(NodesBlock.class);