fix the bug produced in rebase and some improvements

Author: WangTaoTheTonic

yarn/src/main/scala/org/apache/spark/deploy/yarn/ClientArguments.scala

@@ -38,30 +38,29 @@ private[spark] class ClientArguments(args: Array[String], sparkConf: SparkConf)
   var amMemory: Int = 512 // MB
   var appName: String = "Spark"
   var priority = 0
-  var isClusterMode = false
+  def isClusterMode: Boolean = userClass != null
 
+  private var driverMemory: Int = 512 // MB
   private val driverMemKey = "spark.driver.memory"
   private val driverMemOverheadKey = "spark.yarn.driver.memoryOverhead"
   private val amMemKey = "spark.yarn.am.memory"
   private val amMemOverheadKey = "spark.yarn.am.memoryOverhead"
-  private var isDriverMemSet = false
-
-  parseArgs(args.toList)
-
   private val isDynamicAllocationEnabled =
     sparkConf.getBoolean("spark.dynamicAllocation.enabled", false)
 
+  parseArgs(args.toList)
   loadEnvironmentArgs()
   validateArgs()
 
   // Additional memory to allocate to containers
-  val amMemoryOverhead = sparkConf.getInt("spark.yarn.driver.memoryOverhead",
+  val amMemoryOverheadConf = if (isClusterMode) driverMemOverheadKey else amMemOverheadKey
+  val amMemoryOverhead = sparkConf.getInt(amMemoryOverheadConf,
     math.max((MEMORY_OVERHEAD_FACTOR * amMemory).toInt, MEMORY_OVERHEAD_MIN))
 
   val executorMemoryOverhead = sparkConf.getInt("spark.yarn.executor.memoryOverhead",
     math.max((MEMORY_OVERHEAD_FACTOR * executorMemory).toInt, MEMORY_OVERHEAD_MIN))
 
-   /** Load any default arguments provided through environment variables and Spark properties. */
+  /** Load any default arguments provided through environment variables and Spark properties. */
   private def loadEnvironmentArgs(): Unit = {
     // For backward compatibility, SPARK_YARN_DIST_{ARCHIVES/FILES} should be resolved to hdfs://,
     // while spark.yarn.dist.{archives/files} should be resolved to file:// (SPARK-2051).
@@ -99,16 +98,13 @@ private[spark] class ClientArguments(args: Array[String], sparkConf: SparkConf)
           println(s"$key is set but does not apply in cluster mode.")
         }
       }
+      amMemory = driverMemory
     } else {
       for (key <- Seq(driverMemKey, driverMemOverheadKey)) {
         if (sparkConf.getOption(key).isDefined) {
           println(s"$key is set but does not apply in client mode.")
         }
       }
-      if (isDriverMemSet) {
-        println("--driver-memory is set but does not apply in client mode.")
-      }
-      // In cluster mode, the driver and the AM live in the same JVM, so this does not apply
       amMemory = Utils.memoryStringToMb(sparkConf.get(amMemKey, "512m"))
     }
   }
@@ -125,7 +121,6 @@ private[spark] class ClientArguments(args: Array[String], sparkConf: SparkConf)
 
         case ("--class") :: value :: tail =>
           userClass = value
-          isClusterMode = true
           args = tail
 
         case ("--args" | "--arg") :: value :: tail =>
@@ -143,8 +138,7 @@ private[spark] class ClientArguments(args: Array[String], sparkConf: SparkConf)
           if (args(0) == "--master-memory") {
             println("--master-memory is deprecated. Use --driver-memory instead.")
           }
-          amMemory = value
-          isDriverMemSet = true
+          driverMemory = value
           args = tail
 
         case ("--num-workers" | "--num-executors") :: IntParam(value) :: tail =>