[SPARK-4666] Improve YarnAllocator's parsing of "memory overhead" param

core/src/main/scala/org/apache/spark/SparkConf.scala

@@ -204,6 +204,28 @@ class SparkConf(loadDefaults: Boolean) extends Cloneable with Logging {
     getOption(key).map(_.toBoolean).getOrElse(defaultValue)
   }
 
+  // Limit of bytes for total size of results (default is 1GB)
+  def getMaxResultSize: Long = {
+    getMemory("spark.driver.maxResultSize", "1g", outputScale = 'b')
+  }
+
+  private def getMemory(
+      key: String,
+      defaultValue: String,
+      defaultInputScale: Char = 'b',
+      outputScale: Char = 'm'): Long = {
+    Utils.parseMemoryString(getOption(key).getOrElse(defaultValue), defaultInputScale, outputScale)
+  }
+
+  def getMB(
+      key: String,
+      defaultValue: Int): Int = {
+    getOption(key)
+        .map(Utils.memoryStringToMb(_, defaultInputScale = 'm'))
+        .map(_.toInt)
+        .getOrElse(defaultValue)
+  }
+
   /** Get all executor environment variables set on this SparkConf */
   def getExecutorEnv: Seq[(String, String)] = {
     val prefix = "spark.executorEnv."

core/src/main/scala/org/apache/spark/deploy/Client.scala

@@ -74,7 +74,7 @@ private class ClientActor(driverArgs: ClientArguments, conf: SparkConf)
 
         val driverDescription = new DriverDescription(
           driverArgs.jarUrl,
-          driverArgs.memory,
+          driverArgs.memoryMB,
           driverArgs.cores,
           driverArgs.supervise,
           command)

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

@@ -39,7 +39,7 @@ private[spark] class ClientArguments(args: Array[String]) {
   var jarUrl: String = ""
   var mainClass: String = ""
   var supervise: Boolean = DEFAULT_SUPERVISE
-  var memory: Int = DEFAULT_MEMORY
+  var memoryMB: Int = DEFAULT_MEMORY_MB
   var cores: Int = DEFAULT_CORES
   private var _driverOptions = ListBuffer[String]()
   def driverOptions = _driverOptions.toSeq
@@ -55,7 +55,7 @@ private[spark] class ClientArguments(args: Array[String]) {
       parse(tail)
 
     case ("--memory" | "-m") :: MemoryParam(value) :: tail =>
-      memory = value
+      memoryMB = value
       parse(tail)
 
     case ("--supervise" | "-s") :: tail =>
@@ -106,7 +106,7 @@ private[spark] class ClientArguments(args: Array[String]) {
       |
       |Options:
       |   -c CORES, --cores CORES        Number of cores to request (default: $DEFAULT_CORES)
-      |   -m MEMORY, --memory MEMORY     Megabytes of memory to request (default: $DEFAULT_MEMORY)
+      |   -m MEMORY, --memory MEMORY     Megabytes of memory to request (default: $DEFAULT_MEMORY_MB)
       |   -s, --supervise                Whether to restart the driver on failure
       |                                  (default: $DEFAULT_SUPERVISE)
       |   -v, --verbose                  Print more debugging output
@@ -118,7 +118,7 @@ private[spark] class ClientArguments(args: Array[String]) {
 
 object ClientArguments {
   private[spark] val DEFAULT_CORES = 1
-  private[spark] val DEFAULT_MEMORY = 512 // MB
+  private[spark] val DEFAULT_MEMORY_MB = 512
   private[spark] val DEFAULT_SUPERVISE = false
 
   def isValidJarUrl(s: String): Boolean = {

core/src/main/scala/org/apache/spark/deploy/DriverDescription.scala

@@ -19,19 +19,19 @@ package org.apache.spark.deploy
 
 private[spark] class DriverDescription(
     val jarUrl: String,
-    val mem: Int,
+    val memMB: Int,
     val cores: Int,
     val supervise: Boolean,
     val command: Command)
   extends Serializable {
 
   def copy(
       jarUrl: String = jarUrl,
-      mem: Int = mem,
+      memMB: Int = memMB,
       cores: Int = cores,
       supervise: Boolean = supervise,
       command: Command = command): DriverDescription =
-    new DriverDescription(jarUrl, mem, cores, supervise, command)
+    new DriverDescription(jarUrl, memMB, cores, supervise, command)
 
   override def toString: String = s"DriverDescription (${command.mainClass})"
 }

core/src/main/scala/org/apache/spark/deploy/JsonProtocol.scala

@@ -71,7 +71,7 @@ private[spark] object JsonProtocol {
     ("starttime" -> obj.startTime.toString) ~
     ("state" -> obj.state.toString) ~
     ("cores" -> obj.desc.cores) ~
-    ("memory" -> obj.desc.mem)
+    ("memory" -> obj.desc.memMB)
   }
 
   def writeMasterState(obj: MasterStateResponse) = {

core/src/main/scala/org/apache/spark/deploy/master/Master.scala

@@ -543,7 +543,7 @@ private[spark] class Master(
       while (numWorkersVisited < numWorkersAlive && !launched) {
         val worker = shuffledAliveWorkers(curPos)
         numWorkersVisited += 1
-        if (worker.memoryFree >= driver.desc.mem && worker.coresFree >= driver.desc.cores) {
+        if (worker.memoryFree >= driver.desc.memMB && worker.coresFree >= driver.desc.cores) {
           launchDriver(worker, driver)
           waitingDrivers -= driver
           launched = true

core/src/main/scala/org/apache/spark/deploy/master/WorkerInfo.scala

@@ -90,13 +90,13 @@ private[spark] class WorkerInfo(
 
   def addDriver(driver: DriverInfo) {
     drivers(driver.id) = driver
-    memoryUsed += driver.desc.mem
+    memoryUsed += driver.desc.memMB
     coresUsed += driver.desc.cores
   }
 
   def removeDriver(driver: DriverInfo) {
     drivers -= driver.id
-    memoryUsed -= driver.desc.mem
+    memoryUsed -= driver.desc.memMB
     coresUsed -= driver.desc.cores
   }
 

core/src/main/scala/org/apache/spark/deploy/master/ui/MasterPage.scala

@@ -193,8 +193,8 @@ private[spark] class MasterPage(parent: MasterWebUI) extends WebUIPage("") {
       <td sorttable_customkey={driver.desc.cores.toString}>
         {driver.desc.cores}
       </td>
-      <td sorttable_customkey={driver.desc.mem.toString}>
-        {Utils.megabytesToString(driver.desc.mem.toLong)}
+      <td sorttable_customkey={driver.desc.memMB.toString}>
+        {Utils.megabytesToString(driver.desc.memMB.toLong)}
       </td>
       <td>{driver.desc.command.arguments(2)}</td>
     </tr>

core/src/main/scala/org/apache/spark/deploy/rest/StandaloneRestServer.scala

@@ -394,7 +394,7 @@ private class SubmitRequestServlet(
       "org.apache.spark.deploy.worker.DriverWrapper",
       Seq("{{WORKER_URL}}", "{{USER_JAR}}", mainClass) ++ appArgs, // args to the DriverWrapper
       environmentVariables, extraClassPath, extraLibraryPath, javaOpts)
-    val actualDriverMemory = driverMemory.map(Utils.memoryStringToMb).getOrElse(DEFAULT_MEMORY)
+    val actualDriverMemory = driverMemory.map(Utils.memoryStringToMb).getOrElse(DEFAULT_MEMORY_MB)
     val actualDriverCores = driverCores.map(_.toInt).getOrElse(DEFAULT_CORES)
     val actualSuperviseDriver = superviseDriver.map(_.toBoolean).getOrElse(DEFAULT_SUPERVISE)
     new DriverDescription(

core/src/main/scala/org/apache/spark/deploy/worker/DriverRunner.scala

@@ -81,7 +81,7 @@ private[spark] class DriverRunner(
           }
 
           // TODO: If we add ability to submit multiple jars they should also be added here
-          val builder = CommandUtils.buildProcessBuilder(driverDesc.command, driverDesc.mem,
+          val builder = CommandUtils.buildProcessBuilder(driverDesc.command, driverDesc.memMB,
             sparkHome.getAbsolutePath, substituteVariables)
           launchDriver(builder, driverDir, driverDesc.supervise)
         }

core/src/main/scala/org/apache/spark/deploy/worker/Worker.scala

@@ -435,7 +435,7 @@ private[spark] class Worker(
       driver.start()
 
       coresUsed += driverDesc.cores
-      memoryUsed += driverDesc.mem
+      memoryUsed += driverDesc.memMB
     }
 
     case KillDriver(driverId) => {
@@ -464,7 +464,7 @@ private[spark] class Worker(
       master ! DriverStateChanged(driverId, state, exception)
       val driver = drivers.remove(driverId).get
       finishedDrivers(driverId) = driver
-      memoryUsed -= driver.driverDesc.mem
+      memoryUsed -= driver.driverDesc.memMB
       coresUsed -= driver.driverDesc.cores
     }
 

core/src/main/scala/org/apache/spark/deploy/worker/ui/WorkerPage.scala

@@ -139,8 +139,8 @@ private[spark] class WorkerPage(parent: WorkerWebUI) extends WebUIPage("") {
       <td sorttable_customkey={driver.driverDesc.cores.toString}>
         {driver.driverDesc.cores.toString}
       </td>
-      <td sorttable_customkey={driver.driverDesc.mem.toString}>
-        {Utils.megabytesToString(driver.driverDesc.mem)}
+      <td sorttable_customkey={driver.driverDesc.memMB.toString}>
+        {Utils.megabytesToString(driver.driverDesc.memMB)}
       </td>
       <td>
         <a href={s"logPage?driverId=${driver.driverId}&logType=stdout"}>stdout</a>

core/src/main/scala/org/apache/spark/executor/Executor.scala

@@ -105,7 +105,7 @@ private[spark] class Executor(
   private val akkaFrameSize = AkkaUtils.maxFrameSizeBytes(conf)
 
   // Limit of bytes for total size of results (default is 1GB)
-  private val maxResultSize = Utils.getMaxResultSize(conf)
+  private val maxResultSize = conf.getMaxResultSize
 
   // Maintains the list of running tasks.
   private val runningTasks = new ConcurrentHashMap[Long, TaskRunner]

core/src/main/scala/org/apache/spark/scheduler/TaskSetManager.scala

@@ -70,7 +70,7 @@ private[spark] class TaskSetManager(
   val SPECULATION_MULTIPLIER = conf.getDouble("spark.speculation.multiplier", 1.5)
 
   // Limit of bytes for total size of results (default is 1GB)
-  val maxResultSize = Utils.getMaxResultSize(conf)
+  val maxResultSize = conf.getMaxResultSize
 
   // Serializer for closures and tasks.
   val env = SparkEnv.get

core/src/main/scala/org/apache/spark/scheduler/local/LocalBackend.scala

@@ -20,6 +20,7 @@ package org.apache.spark.scheduler.local
 import java.nio.ByteBuffer
 
 import scala.concurrent.duration._
+import scala.language.postfixOps
 
 import akka.actor.{Actor, ActorRef, Props}
 

core/src/main/scala/org/apache/spark/util/Utils.scala

@@ -980,33 +980,76 @@ private[spark] object Utils extends Logging {
     )
   }
 
+
+  private val TB = 1L << 40
+  private val GB = 1L << 30
+  private val MB = 1L << 20
+  private val KB = 1L << 10
+
+  private val scaleCharToFactor: Map[Char, Long] = Map(
+    'b' -> 1L,
+    'k' -> KB,
+    'm' -> MB,
+    'g' -> GB,
+    't' -> TB
+  )
+
   /**
-   * Convert a Java memory parameter passed to -Xmx (such as 300m or 1g) to a number of megabytes.
-   */
-  def memoryStringToMb(str: String): Int = {
+   * Convert a Java memory parameter passed to -Xmx (such as "300m" or "1g") to a number of
+   * megabytes (or other byte-scale denominations as specified by @outputScaleChar).
+   *
+   * For @defaultInputScaleChar and @outputScaleChar, valid values are: 'b' (bytes), 'k'
+   * (kilobytes), 'm' (megabytes), 'g' (gigabytes), and 't' (terabytes).
+   *
+   * @param str String to parse an amount of memory out of
+   * @param defaultInputScaleChar if no "scale" is provided on the end of @str (i.e. @str is a
+   *                              plain numeric value), assume this scale (default: 'b' for
+   *                              'bytes')
+   * @param outputScaleChar express the output in this scale, i.e. number of bytes, kilobytes,
+   *                        megabytes, or gigabytes.
+   */
+  def parseMemoryString(
+      str: String,
+      defaultInputScaleChar: Char = 'b',
+      outputScaleChar: Char = 'm'): Long = {
+
     val lower = str.toLowerCase
-    if (lower.endsWith("k")) {
-      (lower.substring(0, lower.length-1).toLong / 1024).toInt
-    } else if (lower.endsWith("m")) {
-      lower.substring(0, lower.length-1).toInt
-    } else if (lower.endsWith("g")) {
-      lower.substring(0, lower.length-1).toInt * 1024
-    } else if (lower.endsWith("t")) {
-      lower.substring(0, lower.length-1).toInt * 1024 * 1024
-    } else {// no suffix, so it's just a number in bytes
-      (lower.toLong / 1024 / 1024).toInt
-    }
+    val lastChar = lower(lower.length - 1)
+    val (num, inputScaleChar) =
+      if (lastChar.isDigit) {
+        (lower.toLong, defaultInputScaleChar)
+      } else {
+        (lower.substring(0, lower.length - 1).toLong, lastChar)
+      }
+
+    (for {
+      inputScale <- scaleCharToFactor.get(inputScaleChar)
+      outputScale <- scaleCharToFactor.get(outputScaleChar)
+    } yield {
+      inputScale * num / outputScale
+    }).getOrElse(
+        throw new IllegalArgumentException(
+          "Invalid memory string or scale: %s, %s, %s".format(
+            str,
+            defaultInputScaleChar,
+            outputScaleChar
+          )
+        )
+      )
   }
 
+  /**
+   * Wrapper for @parseMemoryString taking default arguments and returning an int, which is safe
+   * since we are converting to a number of megabytes.
+   */
+  def memoryStringToMb(str: String): Int = memoryStringToMb(str, defaultInputScale = 'b')
+  def memoryStringToMb(str: String, defaultInputScale: Char = 'b'): Int =
+    parseMemoryString(str, defaultInputScale, 'm').toInt
+
   /**
    * Convert a quantity in bytes to a human-readable string such as "4.0 MB".
    */
   def bytesToString(size: Long): String = {
-    val TB = 1L << 40
-    val GB = 1L << 30
-    val MB = 1L << 20
-    val KB = 1L << 10
-
     val (value, unit) = {
       if (size >= 2*TB) {
         (size.asInstanceOf[Double] / TB, "TB")
@@ -1047,7 +1090,7 @@ private[spark] object Utils extends Logging {
    * Convert a quantity in megabytes to a human-readable string such as "4.0 MB".
    */
   def megabytesToString(megabytes: Long): String = {
-    bytesToString(megabytes * 1024L * 1024L)
+    bytesToString(megabytes * MB)
   }
 
   /**
@@ -1905,11 +1948,6 @@ private[spark] object Utils extends Logging {
     method.invoke(obj, values.toSeq: _*)
   }
 
-  // Limit of bytes for total size of results (default is 1GB)
-  def getMaxResultSize(conf: SparkConf): Long = {
-    memoryStringToMb(conf.get("spark.driver.maxResultSize", "1g")).toLong << 20
-  }
-
   /**
    * Return the current system LD_LIBRARY_PATH name
    */

docs/running-on-yarn.md

@@ -113,23 +113,23 @@ Most of the configs are the same for Spark on YARN as for other deployment modes
 </tr>
 <tr>
  <td><code>spark.yarn.executor.memoryOverhead</code></td>
-  <td>executorMemory * 0.07, with minimum of 384 </td>
+  <td>executorMemory * 0.07, with a minimum of 384 megabytes </td>
   <td>
-    The amount of off heap memory (in megabytes) to be allocated per executor. This is memory that accounts for things like VM overheads, interned strings, other native overheads, etc. This tends to grow with the executor size (typically 6-10%).
+    The amount of off heap memory to be allocated per executor. This is memory that accounts for things like VM overheads, interned strings, other native overheads, etc. This tends to grow with the executor size (typically 6-10%), but is specified here as an absolute amount of memory, e.g. "1g" or "384m".
   </td>
 </tr>
 <tr>
   <td><code>spark.yarn.driver.memoryOverhead</code></td>
-  <td>driverMemory * 0.07, with minimum of 384 </td>
+  <td>driverMemory * 0.07, with a minimum of 384 megabytes </td>
   <td>
-    The amount of off heap memory (in megabytes) to be allocated per driver in cluster mode. This is memory that accounts for things like VM overheads, interned strings, other native overheads, etc. This tends to grow with the container size (typically 6-10%).
+    Same as <code>spark.yarn.executor.memoryOverhead</code>, but for the driver in cluster mode.
   </td>
 </tr>
 <tr>
   <td><code>spark.yarn.am.memoryOverhead</code></td>
-  <td>AM memory * 0.07, with minimum of 384 </td>
+  <td>AM memory * 0.07, with a minimum of 384 megabytes </td>
   <td>
-    Same as <code>spark.yarn.driver.memoryOverhead</code>, but for the Application Master in client mode.
+    Same as <code>spark.yarn.executor.memoryOverhead</code>, but for the Application Master in client mode.
   </td>
 </tr>
 <tr>

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

@@ -27,7 +27,7 @@ class ApplicationMasterArguments(val args: Array[String]) {
   var primaryPyFile: String = null
   var pyFiles: String = null
   var userArgs: Seq[String] = Seq[String]()
-  var executorMemory = 1024
+  var executorMemoryMB = 1024
   var executorCores = 1
   var numExecutors = DEFAULT_NUMBER_EXECUTORS
 
@@ -67,7 +67,7 @@ class ApplicationMasterArguments(val args: Array[String]) {
           args = tail
 
         case ("--worker-memory" | "--executor-memory") :: MemoryParam(value) :: tail =>
-          executorMemory = value
+          executorMemoryMB = value
           args = tail
 
         case ("--worker-cores" | "--executor-cores") :: IntParam(value) :: tail =>