瀏覽代碼

YARN-8987. Usability improvements node-attributes CLI. Contributed by Bibin A Chundatt.

Weiwei Yang 6 年之前
父節點
當前提交
c741109522

+ 3 - 0
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeAttributesCLI.java

@@ -205,6 +205,9 @@ public class NodeAttributesCLI extends Configured implements Tool {
       // print admin command detail
       // print admin command detail
       printUsage(true, handler);
       printUsage(true, handler);
       return exitCode;
       return exitCode;
+    } catch (YarnException e) {
+      errOut.println(e.toString());
+      return exitCode;
     } catch (Exception e) {
     } catch (Exception e) {
       errOut.println(e.toString());
       errOut.println(e.toString());
       printUsage(true, handler);
       printUsage(true, handler);

+ 22 - 0
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/main/java/org/apache/hadoop/yarn/server/resourcemanager/AdminService.java

@@ -993,6 +993,7 @@ public class AdminService extends CompositeService implements
         nodeAttributesManager.addNodeAttributes(nodeAttributeMapping);
         nodeAttributesManager.addNodeAttributes(nodeAttributeMapping);
         break;
         break;
       case REMOVE:
       case REMOVE:
+        validateAttributesExists(nodesToAttributes);
         nodeAttributesManager.removeNodeAttributes(nodeAttributeMapping);
         nodeAttributesManager.removeNodeAttributes(nodeAttributeMapping);
         break;
         break;
       case REPLACE:
       case REPLACE:
@@ -1013,6 +1014,27 @@ public class AdminService extends CompositeService implements
         .newRecordInstance(NodesToAttributesMappingResponse.class);
         .newRecordInstance(NodesToAttributesMappingResponse.class);
   }
   }
 
 
+  private void validateAttributesExists(
+      List<NodeToAttributes> nodesToAttributes) throws IOException {
+    NodeAttributesManager nodeAttributesManager =
+        rm.getRMContext().getNodeAttributesManager();
+    for (NodeToAttributes nodeToAttrs : nodesToAttributes) {
+      String hostname = nodeToAttrs.getNode();
+      if (hostname == null) {
+        continue;
+      }
+      Set<NodeAttribute> attrs =
+          nodeAttributesManager.getAttributesForNode(hostname).keySet();
+      List<NodeAttribute> attributes = nodeToAttrs.getNodeAttributes();
+      for (NodeAttribute nodeAttr : attributes) {
+        if (!attrs.contains(nodeAttr)) {
+          throw new IOException("Node attribute [" + nodeAttr.getAttributeKey()
+              + "] doesn't exist on node " + nodeToAttrs.getNode());
+        }
+      }
+    }
+  }
+
   /**
   /**
    * @param nodesToAttributesMapping input to be validated
    * @param nodesToAttributesMapping input to be validated
    * @param failOnUnknownNodes indicates to fail if the nodes are not available.
    * @param failOnUnknownNodes indicates to fail if the nodes are not available.

+ 28 - 0
hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/TestRMAdminService.java

@@ -1585,6 +1585,16 @@ public class TestRMAdminService {
     Mockito.verify(spiedAttributesManager, Mockito.times(1))
     Mockito.verify(spiedAttributesManager, Mockito.times(1))
         .addNodeAttributes(Mockito.anyMap());
         .addNodeAttributes(Mockito.anyMap());
 
 
+    // add attributes for later removals
+    request =
+        NodesToAttributesMappingRequest
+            .newInstance(AttributeMappingOperationType.ADD,
+                ImmutableList.of(NodeToAttributes.newInstance("host4",
+                    ImmutableList.of(NodeAttribute.newInstance(
+                        NodeAttribute.PREFIX_CENTRALIZED, "x",
+                        NodeAttributeType.STRING, "dfasdf")))),
+                true);
+    rm.adminService.mapAttributesToNodes(request);
     request =
     request =
         NodesToAttributesMappingRequest
         NodesToAttributesMappingRequest
             .newInstance(AttributeMappingOperationType.REMOVE,
             .newInstance(AttributeMappingOperationType.REMOVE,
@@ -1601,6 +1611,24 @@ public class TestRMAdminService {
     Mockito.verify(spiedAttributesManager, Mockito.times(1))
     Mockito.verify(spiedAttributesManager, Mockito.times(1))
         .removeNodeAttributes(Mockito.anyMap());
         .removeNodeAttributes(Mockito.anyMap());
 
 
+    // Assert node to attributes mappings are empty.
+    Assert.assertTrue("Attributes of host4 should be empty",
+        rm.getRMContext().getNodeAttributesManager()
+            .getAttributesForNode("host4").isEmpty());
+    // remove non existing attributes.
+    request = NodesToAttributesMappingRequest
+        .newInstance(AttributeMappingOperationType.REMOVE, ImmutableList
+            .of(NodeToAttributes.newInstance("host4", ImmutableList
+                .of(NodeAttribute
+                    .newInstance(NodeAttribute.PREFIX_CENTRALIZED, "x",
+                        NodeAttributeType.STRING, "dfasdf")))), true);
+    try {
+      rm.adminService.mapAttributesToNodes(request);
+      fail("Should have failed for non exists attribute");
+    } catch (Exception ex) {
+      assertTrue("Exception expected if attributes does not exist", true);
+    }
+
     request =
     request =
         NodesToAttributesMappingRequest
         NodesToAttributesMappingRequest
             .newInstance(AttributeMappingOperationType.ADD,
             .newInstance(AttributeMappingOperationType.ADD,