different memoryOverhead

Author: witgo

yarn/alpha/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocationHandler.scala

@@ -89,8 +89,16 @@ private[yarn] class YarnAllocationHandler(
   private val pendingReleaseContainers = new ConcurrentHashMap[ContainerId, Boolean]
 
   // Additional memory overhead - in mb.
-  private val memoryOverhead = sparkConf.getInt("spark.yarn.container.memoryOverhead",
-    YarnAllocationHandler.MEMORY_OVERHEAD)
+  private def memoryOverhead: Int = {
+    var defaultMemoryOverhead = YarnAllocationHandler.MEMORY_OVERHEAD
+    sparkConf.getOption("spark.yarn.container.memoryOverhead").foreach { s =>
+      defaultMemoryOverhead = s.toInt
+    }
+    sparkConf.getOption(s"spark.yarn.executor.memoryOverhead").foreach { s =>
+      defaultMemoryOverhead = s.toInt
+    }
+    defaultMemoryOverhead
+  }
 
   private val numExecutorsRunning = new AtomicInteger()
   // Used to generate a unique id per executor

yarn/common/src/main/scala/org/apache/spark/deploy/yarn/ClientBase.scala

@@ -66,8 +66,16 @@ trait ClientBase extends Logging {
     FsPermission.createImmutable(Integer.parseInt("644", 8).toShort)
 
   // Additional memory overhead - in mb.
-  def memoryOverhead = sparkConf.getInt("spark.yarn.container.memoryOverhead",
-    YarnAllocationHandler.MEMORY_OVERHEAD)
+  protected def memoryOverhead: Int = {
+    var defaultMemoryOverhead = (args.amMemory * 0.25D).ceil.toInt
+    sparkConf.getOption("spark.yarn.container.memoryOverhead").foreach { s =>
+      defaultMemoryOverhead = s.toInt
+    }
+    sparkConf.getOption(s"spark.yarn.driver.memoryOverhead").foreach { s =>
+      defaultMemoryOverhead = s.toInt
+    }
+    defaultMemoryOverhead
+  }
 
   // TODO(harvey): This could just go in ClientArguments.
   def validateArgs() = {

yarn/stable/src/main/scala/org/apache/spark/deploy/yarn/YarnAllocationHandler.scala

@@ -90,11 +90,20 @@ private[yarn] class YarnAllocationHandler(
   // Containers to be released in next request to RM
   private val pendingReleaseContainers = new ConcurrentHashMap[ContainerId, Boolean]
 
-  // TODO: Here we can dynamically calculate the default value.
-  // eg: val memoryOverhead = (executorMemory * 0.25D).ceil.toInt
   // Additional memory overhead - in mb.
-  private val memoryOverhead = sparkConf.getInt("spark.yarn.container.memoryOverhead",
-    YarnAllocationHandler.MEMORY_OVERHEAD)
+  private def memoryOverhead: Int = {
+    // TODO: Here we can dynamically calculate the default value.
+    // eg: val defaultMemoryOverhead = (executorMemory * 0.25D).ceil.toInt
+    var defaultMemoryOverhead = YarnAllocationHandler.MEMORY_OVERHEAD
+    sparkConf.getOption("spark.yarn.container.memoryOverhead").foreach { s =>
+      defaultMemoryOverhead = s.toInt
+    }
+    sparkConf.getOption(s"spark.yarn.executor.memoryOverhead").foreach { s =>
+      defaultMemoryOverhead = s.toInt
+    }
+    defaultMemoryOverhead
+  }
+
 
   // Number of container requests that have been sent to, but not yet allocated by the
   // ApplicationMaster.