SAMZA-1014: Add property to set YARN AM cpu cores

Author: logarithm

samza-yarn/src/main/java/org/apache/samza/config/YarnConfig.java

@@ -85,6 +85,12 @@ public class YarnConfig extends MapConfig {
    */
   public static final String AM_CONTAINER_LABEL = "yarn.am.container.label";
 
+  /**
+   * Number of CPU cores to request from YARN for running the AM
+   */
+  public static final String AM_CONTAINER_MAX_CPU_CORES = "yarn.am.container.cpu.cores";
+  private static final int DEFAULT_AM_CPU_CORES = 1;
+
   /**
    * Determines the interval for the Heartbeat between the AM and the Yarn RM
    */
@@ -188,6 +194,10 @@ public String getAMContainerLabel() {
     return get(AM_CONTAINER_LABEL, null);
   }
 
+  public int getAMContainerMaxCpuCores() {
+    return getInt(AM_CONTAINER_MAX_CPU_CORES, DEFAULT_AM_CPU_CORES);
+  }
+
   public String getAmOpts() {
     return get(AM_JVM_OPTIONS, "");
   }

samza-yarn/src/main/scala/org/apache/samza/job/yarn/ClientHelper.scala

@@ -92,7 +92,7 @@ class ClientHelper(conf: Configuration) extends Logging {
 
     val packagePath = new Path(yarnConfig.getPackagePath)
     val mem = yarnConfig.getAMContainerMaxMemoryMb
-    val cpu = 1
+    val cpu = yarnConfig.getAMContainerMaxCpuCores
     val queueName = Option(yarnConfig.getQueueName)
     val appMasterLabel = Option(yarnConfig.getAMContainerLabel)