Add AM GPU param

Author: hungj

hadoop-plugin-test/expectedJobs/jobs1/jobs1_job24.job

@@ -1,6 +1,7 @@
 # This file generated from the Hadoop DSL. Do not edit by hand.
 type=TensorFlowJob
 dependencies=jobs1_job23
+am_gpus=1
 am_memory=2048
 am_vcores=1
 archive=tensorflow-starter-kit.zip

hadoop-plugin-test/src/main/gradle/positive/jobs1.gradle

@@ -385,6 +385,7 @@ hadoop {
       executes 'path/to/python/script.py'
       amMemoryMB 2048
       amCores 1
+      amGpus 1
       psMemoryMB 2048
       psCores 1
       workerMemoryMB 8192

hadoop-plugin/src/main/groovy/com/linkedin/gradle/hadoopdsl/job/TensorFlowJob.groovy

@@ -57,6 +57,14 @@ interface TensorFlowJob {
   @HadoopDslMethod
   void amCores(int amCores);
 
+  /**
+   * Sets number of GPUs for TensorFlow AM.
+   *
+   * @param amGpus Number of GPUs for TensorFlow AM
+   */
+  @HadoopDslMethod
+  void amGpus(int amGpus);
+
   /**
    * Sets memory in MB for each parameter server's YARN container.
    * Leaving this unset will default to the Hadoop application's default value.

hadoop-plugin/src/main/groovy/com/linkedin/gradle/hadoopdsl/job/TensorFlowSparkJob.groovy

@@ -90,6 +90,11 @@ class TensorFlowSparkJob extends SparkJob implements TensorFlowJob {
     super.setConf("spark.driver.cores", amCores);
   }
 
+  @Override
+  void amGpus(int amGpus) {
+    throw new Exception("Requesting AM GPUs via tensorflow on spark is not supported currently.");
+  }
+
   @Override
   void psMemoryMB(int psMemoryMB) {
     // If worker memory has already been set, ignore
@@ -120,7 +125,7 @@ class TensorFlowSparkJob extends SparkJob implements TensorFlowJob {
 
   @Override
   void workerGpus(int workerGpus) {
-    throw new Exception("Requesting GPUs via tensorflow on spark is not supported currently.");
+    throw new Exception("Requesting worker GPUs via tensorflow on spark is not supported currently.");
   }
 
   @Override

hadoop-plugin/src/main/groovy/com/linkedin/gradle/hadoopdsl/job/TensorFlowTonyJob.groovy

@@ -38,6 +38,7 @@ import com.linkedin.gradle.hadoopdsl.HadoopDslMethod;
  *     executes path/to/python/script.py
  *     amMemoryMB 2048
  *     amCores 1
+ *     amGpus 1
  *     psMemoryMB 2048
  *     psCores 1
  *     workerMemoryMB 8192
@@ -57,6 +58,7 @@ class TensorFlowTonyJob extends HadoopJavaProcessJob implements TensorFlowJob {
   String executePath;
   int amMemory;
   int amCores;
+  int amGpus;
   int psMemory;
   int psCores;
   int workerMemory;
@@ -93,6 +95,7 @@ class TensorFlowTonyJob extends HadoopJavaProcessJob implements TensorFlowJob {
     cloneJob.executePath = executePath;
     cloneJob.amMemory = amMemory;
     cloneJob.amCores = amCores;
+    cloneJob.amGpus = amGpus;
     cloneJob.psMemory = psMemory;
     cloneJob.psCores = psCores;
     cloneJob.workerMemory = workerMemory;
@@ -123,6 +126,12 @@ class TensorFlowTonyJob extends HadoopJavaProcessJob implements TensorFlowJob {
     setJobProperty("am_vcores", this.amCores);
   }
 
+  @Override
+  void amGpus(int amGpus) {
+    this.amGpus = amGpus;
+    setJobProperty("am_gpus", this.amGpus);
+  }
+
   @Override
   void psMemoryMB(int psMemoryMB) {
     this.psMemory = psMemoryMB;