[SPARK-58513][CORE] Init UnifiedMemeoryManager should check driver config in driver , check executor config in executor

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

@@ -498,7 +498,10 @@ class SparkEnv (
     } else {
       conf.clone.set(MEMORY_OFFHEAP_ENABLED, false).set(MEMORY_OFFHEAP_SIZE, 0L)
     }
-    _memoryManager = UnifiedMemoryManager(memoryManagerConf, numUsableCores)
+    _memoryManager = UnifiedMemoryManager(
+      memoryManagerConf,
+      numUsableCores,
+      isDriver = Some(SparkContext.isDriver(executorId)))
   }
 }
 

core/src/main/scala/org/apache/spark/memory/UnifiedMemoryManager.scala

@@ -447,7 +447,14 @@ object UnifiedMemoryManager extends Logging {
   }
 
   def apply(conf: SparkConf, numCores: Int): UnifiedMemoryManager = {
-    val maxMemory = getMaxMemory(conf)
+    apply(conf, numCores, isDriver = None)
+  }
+
+  def apply(
+      conf: SparkConf,
+      numCores: Int,
+      isDriver: Option[Boolean]): UnifiedMemoryManager = {
+    val maxMemory = getMaxMemory(conf, isDriver)
     new UnifiedMemoryManager(
       conf,
       maxHeapMemory = maxMemory,
@@ -459,12 +466,14 @@ object UnifiedMemoryManager extends Logging {
   /**
    * Return the total amount of memory shared between execution and storage, in bytes.
    */
-  private def getMaxMemory(conf: SparkConf): Long = {
+  private def getMaxMemory(conf: SparkConf, isDriver: Option[Boolean]): Long = {
     val systemMemory = conf.get(TEST_MEMORY)
     val reservedMemory = conf.getLong(TEST_RESERVED_MEMORY.key,
       if (conf.contains(IS_TESTING)) 0 else RESERVED_SYSTEM_MEMORY_BYTES)
     val minSystemMemory = (reservedMemory * 1.5).ceil.toLong
-    if (systemMemory < minSystemMemory) {
+    val checkDriverMemory = isDriver.isEmpty || isDriver.contains(true)
+    val checkExecutorMemory = isDriver.isEmpty || isDriver.contains(false)
+    if (checkDriverMemory && systemMemory < minSystemMemory) {
       throw new SparkIllegalArgumentException(
         errorClass = "INVALID_DRIVER_MEMORY",
         messageParameters = Map(
@@ -473,7 +482,7 @@ object UnifiedMemoryManager extends Logging {
           "config" -> config.DRIVER_MEMORY.key))
     }
     // SPARK-12759 Check executor memory to fail fast if memory is insufficient
-    if (conf.contains(config.EXECUTOR_MEMORY)) {
+    if (checkExecutorMemory && conf.contains(config.EXECUTOR_MEMORY)) {
       val executorMemory = conf.getSizeAsBytes(config.EXECUTOR_MEMORY.key)
       if (executorMemory < minSystemMemory) {
         throw new SparkIllegalArgumentException(

core/src/test/scala/org/apache/spark/memory/UnifiedMemoryManagerSuite.scala

@@ -259,6 +259,36 @@ class UnifiedMemoryManagerSuite extends MemoryManagerSuite with PrivateMethodTes
     assert(exception.getMessage.contains("increase executor memory"))
   }
 
+  test("SPARK-58513: executor does not validate driver heap") {
+    val systemMemory = 400L * 1024
+    val reservedMemory = 300L * 1024
+    val memoryFraction = 0.8
+    val conf = new SparkConf()
+      .set(MEMORY_FRACTION, memoryFraction)
+      .set(TEST_MEMORY, systemMemory)
+      .set(TEST_RESERVED_MEMORY, reservedMemory)
+      .set(EXECUTOR_MEMORY.key, (500L * 1024).toString)
+
+    val mm = UnifiedMemoryManager(conf, numCores = 1, isDriver = Some(false))
+    val expectedMaxMemory = ((systemMemory - reservedMemory) * memoryFraction).toLong
+    assert(mm.maxHeapMemory === expectedMaxMemory)
+  }
+
+  test("SPARK-58513: driver does not validate executor memory") {
+    val systemMemory = 1024L * 1024
+    val reservedMemory = 300L * 1024
+    val memoryFraction = 0.8
+    val conf = new SparkConf()
+      .set(MEMORY_FRACTION, memoryFraction)
+      .set(TEST_MEMORY, systemMemory)
+      .set(TEST_RESERVED_MEMORY, reservedMemory)
+      .set(EXECUTOR_MEMORY.key, (reservedMemory / 2).toString)
+
+    val mm = UnifiedMemoryManager(conf, numCores = 1, isDriver = Some(true))
+    val expectedMaxMemory = ((systemMemory - reservedMemory) * memoryFraction).toLong
+    assert(mm.maxHeapMemory === expectedMaxMemory)
+  }
+
   test("execution can evict cached blocks when there are multiple active tasks (SPARK-12155)") {
     val conf = new SparkConf()
       .set(MEMORY_FRACTION, 1.0)