Fix tests

Author: Keqiu Hu

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

@@ -2,16 +2,16 @@
 type=TensorFlowJob
 dependencies=jobs1_job23
 am_gpus=1
-am_memory=2048
+am_memory=2g
 am_vcores=1
 archive=tensorflow-starter-kit.zip
 executes=path/to/python/script.py
 num_ps=2
 num_workers=4
-ps_memory=2048
+ps_memory=2g
 ps_vcores=1
 worker_env.ENV1=val1
 worker_env.ENV2=val2
 worker_gpus=2
-worker_memory=8192
+worker_memory=8g
 worker_vcores=1

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

@@ -36,12 +36,12 @@ import com.linkedin.gradle.hadoopdsl.HadoopDslMethod;
  *       'task_params': taskParams,
  *     ]
  *     executes path/to/python/script.py
- *     amMemory 2g
+ *     amMemory '2g'
  *     amCores 1
  *     amGpus 1
- *     psMemory 2g
+ *     psMemory '2g'
  *     psCores 1
- *     workerMemory 8g
+ *     workerMemory '8g'
  *     workerCores 1
  *     workerGpus 2
  *     numPs 2
@@ -116,6 +116,7 @@ class TensorFlowTonyJob extends HadoopJavaProcessJob implements TensorFlowJob {
 
   @Override
   void amMemory(String amMemory) {
+    this.amMemory = amMemory;
     setJobProperty("am_memory", this.amMemory);
   }
 
@@ -133,6 +134,7 @@ class TensorFlowTonyJob extends HadoopJavaProcessJob implements TensorFlowJob {
 
   @Override
   void psMemory(String psMemory) {
+    this.psMemory = psMemory;
     setJobProperty("ps_memory", this.psMemory);
   }
 
@@ -144,6 +146,7 @@ class TensorFlowTonyJob extends HadoopJavaProcessJob implements TensorFlowJob {
 
   @Override
   void workerMemory(String workerMemory) {
+    this.workerMemory = workerMemory;
     setJobProperty("worker_memory", this.workerMemory);
   }