Support setting Application Master node label (#624)

tony-framework · Dec 11, 2021 · feaab9a · feaab9a
1 parent ce3f207
commit feaab9a
Showing 1 changed file with 20 additions and 8 deletions.
diff --git a/tony-core/src/main/java/com/linkedin/tony/TonyClient.java b/tony-core/src/main/java/com/linkedin/tony/TonyClient.java
@@ -81,7 +81,9 @@
 import org.apache.hadoop.yarn.api.records.LocalResource;
 import org.apache.hadoop.yarn.api.records.LocalResourceType;
 import org.apache.hadoop.yarn.api.records.LocalResourceVisibility;
+import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.Resource;
+import org.apache.hadoop.yarn.api.records.ResourceRequest;
 import org.apache.hadoop.yarn.api.records.YarnApplicationState;
 import org.apache.hadoop.yarn.client.api.YarnClient;
 import org.apache.hadoop.yarn.client.api.YarnClientApplication;
@@ -327,11 +329,6 @@ public void submitApplication(ApplicationSubmissionContext appContext)
       appContext.setApplicationTags(applicationTags);
     }
 
-    // Set up resource type requirements
-    Resource capability = Resource.newInstance((int) amMemory, amVCores);
-    Utils.setCapabilityGPU(capability, amGpus);
-    appContext.setResource(capability);
-
     // Set the queue to which this application is to be submitted in the RM
     String yarnQueue = tonyConf.get(TonyConfigurationKeys.YARN_QUEUE_NAME,
         TonyConfigurationKeys.DEFAULT_YARN_QUEUE_NAME);
@@ -341,10 +338,25 @@ public void submitApplication(ApplicationSubmissionContext appContext)
     ContainerLaunchContext amSpec =
         createAMContainerSpec(this.amMemory, getTokens());
     appContext.setAMContainerSpec(amSpec);
-    String nodeLabel = tonyConf.get(TonyConfigurationKeys.APPLICATION_NODE_LABEL);
-    if (nodeLabel != null) {
-      appContext.setNodeLabelExpression(nodeLabel);
+    String appNodeLabel = tonyConf.get(TonyConfigurationKeys.APPLICATION_NODE_LABEL);
+    if (appNodeLabel != null) {
+      appContext.setNodeLabelExpression(appNodeLabel);
+    }
+
+    // Set up resource type requirements
+    Resource capability = Resource.newInstance((int) amMemory, amVCores);
+    Utils.setCapabilityGPU(capability, amGpus);
+    ResourceRequest amRequest = Records.newRecord(ResourceRequest.class);
+    amRequest.setResourceName(ResourceRequest.ANY);
+    amRequest.setPriority(Priority.newInstance(0));
+    amRequest.setCapability(capability);
+    amRequest.setNumContainers(1);
+    String amNodeLabel = tonyConf.get(TonyConfigurationKeys.getNodeLabelKey(Constants.AM_NAME));
+    if (amNodeLabel != null) {
+      amRequest.setNodeLabelExpression(amNodeLabel);
     }
+    appContext.setAMContainerResourceRequest(amRequest);
+
     LOG.info("Submitting YARN application");
     yarnClient.submitApplication(appContext);
     ApplicationReport report = yarnClient.getApplicationReport(appId);