YARN-11161. Support getAttributesToNodes, getClusterNodeAttributes, getNodesToAttributes API's for Federation.

Author: slfan1989

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-router/src/test/java/org/apache/hadoop/yarn/server/router/clientrm/TestRouterYarnClientUtils.java

@@ -37,6 +37,8 @@
 import org.apache.hadoop.yarn.api.protocolrecords.GetAllResourceTypeInfoResponse;
 import org.apache.hadoop.yarn.api.protocolrecords.GetResourceProfileResponse;
 import org.apache.hadoop.yarn.api.protocolrecords.GetAllResourceProfilesResponse;
+import org.apache.hadoop.yarn.api.protocolrecords.GetAttributesToNodesResponse;
+import org.apache.hadoop.yarn.api.protocolrecords.GetClusterNodeAttributesResponse;
 import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.ApplicationReport;
@@ -53,6 +55,11 @@
 import org.apache.hadoop.yarn.api.records.ReservationDefinition;
 import org.apache.hadoop.yarn.api.records.ReservationId;
 import org.apache.hadoop.yarn.api.records.ResourceTypeInfo;
+import org.apache.hadoop.yarn.api.records.NodeAttribute;
+import org.apache.hadoop.yarn.api.records.NodeAttributeType;
+import org.apache.hadoop.yarn.api.records.NodeAttributeKey;
+import org.apache.hadoop.yarn.api.records.NodeToAttributeValue;
+import org.apache.hadoop.yarn.api.records.NodeAttributeInfo;
 import org.apache.hadoop.yarn.util.Records;
 import org.apache.hadoop.yarn.server.uam.UnmanagedApplicationManager;
 import org.junit.Assert;
@@ -610,4 +617,100 @@ public void testMergeResourceProfile() {
     Assert.assertEquals(3, resource.getVirtualCores());
     Assert.assertEquals(3072, resource.getMemorySize());
   }
+
+  @Test
+  public void testMergeAttributesToNodesResponse() {
+    // normal response1
+    NodeAttribute gpu = NodeAttribute.newInstance(NodeAttribute.PREFIX_CENTRALIZED, "GPU",
+        NodeAttributeType.STRING, "nvidia");
+    Map<NodeAttributeKey, List<NodeToAttributeValue>> map1 = new HashMap<>();
+    List<NodeToAttributeValue> lists1 = new ArrayList<>();
+    NodeToAttributeValue attributeValue1 = NodeToAttributeValue.newInstance("node1", gpu.getAttributeValue());
+    lists1.add(attributeValue1);
+    map1.put(gpu.getAttributeKey(), lists1);
+    GetAttributesToNodesResponse response1 = GetAttributesToNodesResponse.newInstance(map1);
+
+    // normal response2
+    NodeAttribute docker = NodeAttribute.newInstance(NodeAttribute.PREFIX_DISTRIBUTED, "DOCKER",
+        NodeAttributeType.STRING, "docker0");
+    Map<NodeAttributeKey, List<NodeToAttributeValue>> map2 = new HashMap<>();
+    List<NodeToAttributeValue> lists2 = new ArrayList<>();
+    NodeToAttributeValue attributeValue2 = NodeToAttributeValue.newInstance("node2", docker.getAttributeValue());
+    lists2.add(attributeValue2);
+    map2.put(docker.getAttributeKey(), lists2);
+    GetAttributesToNodesResponse response2 = GetAttributesToNodesResponse.newInstance(map2);
+
+    // empty response3
+    GetAttributesToNodesResponse response3 = GetAttributesToNodesResponse.newInstance(new HashMap<>());
+
+    // null response4
+    GetAttributesToNodesResponse response4 = null;
+
+    List<GetAttributesToNodesResponse> responses = new ArrayList<>();
+    responses.add(response1);
+    responses.add(response2);
+    responses.add(response3);
+    responses.add(response4);
+
+    GetAttributesToNodesResponse response =
+        RouterYarnClientUtils.mergeAttributesToNodesResponse(responses);
+
+    Assert.assertNotNull(response);
+    Assert.assertEquals(2, response.getAttributesToNodes().size());
+
+    Map<NodeAttributeKey, List<NodeToAttributeValue>> attrs = response.getAttributesToNodes();
+    Assert.assertTrue(findHostnameAndValInMapping("node2", "docker0",
+        attrs.get(docker.getAttributeKey())));
+  }
+
+  @Test
+  public void testMergeClusterNodeAttributesResponse() {
+    // normal response1
+    NodeAttributeInfo nodeAttributeInfo1 =
+        NodeAttributeInfo.newInstance(NodeAttributeKey.newInstance("GPU"), NodeAttributeType.STRING);
+    Set<NodeAttributeInfo> attributes1 = new HashSet<>();
+    attributes1.add(nodeAttributeInfo1);
+    GetClusterNodeAttributesResponse response1 = GetClusterNodeAttributesResponse.newInstance(attributes1);
+
+    // normal response2
+    NodeAttributeInfo nodeAttributeInfo2 =
+        NodeAttributeInfo.newInstance(NodeAttributeKey.newInstance("CPU"), NodeAttributeType.STRING);
+    Set<NodeAttributeInfo> attributes2 = new HashSet<>();
+    attributes2.add(nodeAttributeInfo2);
+    GetClusterNodeAttributesResponse response2 = GetClusterNodeAttributesResponse.newInstance(attributes2);
+
+    // empty response3
+    GetClusterNodeAttributesResponse response3 = GetClusterNodeAttributesResponse.newInstance(new HashSet<>());
+
+    // null response4
+    GetClusterNodeAttributesResponse response4 = null;
+
+    List<GetClusterNodeAttributesResponse> responses = new ArrayList<>();
+    responses.add(response1);
+    responses.add(response2);
+    responses.add(response3);
+    responses.add(response4);
+
+    GetClusterNodeAttributesResponse response =
+        RouterYarnClientUtils.mergeClusterNodeAttributesResponse(responses);
+
+    Assert.assertNotNull(response);
+
+    Set<NodeAttributeInfo> nodeAttributeInfos = response.getNodeAttributes();
+    Assert.assertEquals(2, nodeAttributeInfos.size());
+
+    Object[] objectArr = nodeAttributeInfos.toArray();
+    Assert.assertEquals("rm.yarn.io/GPU(STRING)", objectArr[0].toString());
+    Assert.assertEquals("rm.yarn.io/CPU(STRING)", objectArr[1].toString());
+  }
+
+  private boolean findHostnameAndValInMapping(String hostname, String attrVal,
+      List<NodeToAttributeValue> mappingVals) {
+    for (NodeToAttributeValue value : mappingVals) {
+      if (value.getHostname().equals(hostname)) {
+        return attrVal.equals(value.getAttributeValue());
+      }
+    }
+    return false;
+  }
 }