[SPARK-54605][K8S] Override default spark env in driver pod

resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/Constants.scala

@@ -87,6 +87,7 @@ object Constants {
   val SPARK_CONF_VOLUME_EXEC = "spark-conf-volume-exec"
   val SPARK_CONF_DIR_INTERNAL = "/opt/spark/conf"
   val SPARK_CONF_FILE_NAME = "spark.properties"
+  val SPARK_ENV_FILE_NAME = "spark-env.sh"
   val SPARK_CONF_PATH = s"$SPARK_CONF_DIR_INTERNAL/$SPARK_CONF_FILE_NAME"
   val ENV_HADOOP_TOKEN_FILE_LOCATION = "HADOOP_TOKEN_FILE_LOCATION"
 

resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/submit/KubernetesClientApplication.scala

@@ -104,8 +104,9 @@ private[spark] class Client(
   def run(): Unit = {
     val resolvedDriverSpec = builder.buildFromFeatures(conf, kubernetesClient)
     val configMapName = KubernetesClientUtils.configMapNameDriver
-    val confFilesMap = KubernetesClientUtils.buildSparkConfDirFilesMap(configMapName,
+    val rawConfFilesMap = KubernetesClientUtils.buildSparkConfDirFilesMap(configMapName,
       conf.sparkConf, resolvedDriverSpec.systemProperties)
+    val confFilesMap = KubernetesClientUtils.overrideDefaultSparkEnv(conf, rawConfFilesMap)
     val configMap = KubernetesClientUtils.buildConfigMap(configMapName, confFilesMap +
         (KUBERNETES_NAMESPACE.key -> conf.namespace))
 

resource-managers/kubernetes/core/src/main/scala/org/apache/spark/deploy/k8s/submit/KubernetesClientUtils.scala

@@ -30,9 +30,9 @@ import io.fabric8.kubernetes.api.model.{ConfigMap, ConfigMapBuilder, KeyToPath}
 
 import org.apache.spark.SparkConf
 import org.apache.spark.annotation.{DeveloperApi, Since, Unstable}
-import org.apache.spark.deploy.k8s.{Config, Constants, KubernetesUtils}
+import org.apache.spark.deploy.k8s.{Config, Constants, KubernetesDriverConf, KubernetesUtils}
 import org.apache.spark.deploy.k8s.Config.{KUBERNETES_DNS_SUBDOMAIN_NAME_MAX_LENGTH, KUBERNETES_NAMESPACE}
-import org.apache.spark.deploy.k8s.Constants.ENV_SPARK_CONF_DIR
+import org.apache.spark.deploy.k8s.Constants.{ENV_SPARK_CONF_DIR, SPARK_ENV_FILE_NAME}
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.LogKeys.{CONFIG, PATH, PATHS}
 import org.apache.spark.util.ArrayImplicits._
@@ -116,6 +116,31 @@ object KubernetesClientUtils extends Logging {
     }.toList.sortBy(x => x.getKey) // List is sorted to make mocking based tests work
   }
 
+  @Since("4.2.0")
+  def overrideDefaultSparkEnv(conf: KubernetesDriverConf,
+      confFilesMap: Map[String, String]): Map[String, String] = {
+    if (conf.environment.isEmpty) {
+      return confFilesMap
+    }
+    val customEnvs = conf.environment.map {
+      case (k, v) => s"export $k=$v"
+    }.mkString("\n", "\n", "\n")
+    val mapSize = confFilesMap.map {
+      case (k, v) => k.length + v.length
+    }.sum
+    val maxSize = conf.sparkConf.get(Config.CONFIG_MAP_MAXSIZE)
+    if (mapSize + customEnvs.length < maxSize) {
+      confFilesMap ++ Map(SPARK_ENV_FILE_NAME ->
+        s"${confFilesMap.getOrElse(SPARK_ENV_FILE_NAME, "")}$customEnvs")
+    } else {
+      logWarning(log"Skipped custom driver env, due to size constraint" +
+        log" Please see, config: `${MDC(CONFIG, Config.CONFIG_MAP_MAXSIZE.key)}`" +
+        log" for more details.")
+      confFilesMap
+    }
+  }
+
+
   /**
    * Build a ConfigMap that will hold the content for environment variable SPARK_CONF_DIR
    * on remote pods. (Java-friendly)

resource-managers/kubernetes/core/src/test/scala/org/apache/spark/deploy/k8s/submit/KubernetesClientUtilsSuite.scala

@@ -29,7 +29,7 @@ import io.fabric8.kubernetes.api.model.ConfigMapBuilder
 import org.scalatest.BeforeAndAfter
 
 import org.apache.spark.{SparkConf, SparkFunSuite}
-import org.apache.spark.deploy.k8s.Config
+import org.apache.spark.deploy.k8s.{Config, KubernetesTestConf}
 import org.apache.spark.util.Utils
 
 class KubernetesClientUtilsSuite extends SparkFunSuite with BeforeAndAfter {
@@ -128,4 +128,52 @@ class KubernetesClientUtilsSuite extends SparkFunSuite with BeforeAndAfter {
         .build()
     assert(outputConfigMap === expectedConfigMap)
   }
+
+  test("SPARK-54605: check custom driver env appended as expected when config map size is " +
+    "below threshold.") {
+    val input = Map("spark-env.sh" -> "test12345\ntest12345", "testConf.1" -> "test123456")
+    val sparkConf = testSetup(input.map(f => f._1 -> f._2.getBytes(StandardCharsets.UTF_8)))
+      .set(Config.CONFIG_MAP_MAXSIZE.key, "90")
+    val conf = KubernetesTestConf.createDriverConf(
+      sparkConf = sparkConf,
+      environment = Map("AAA" -> "value1", "BBB" -> "value2")
+    )
+    val confFileMap = KubernetesClientUtils.overrideDefaultSparkEnv(conf, input)
+    val expectedConfFileMap = Map(
+      "spark-env.sh" -> "test12345\ntest12345\nexport AAA=value1\nexport BBB=value2\n",
+      "testConf.1" -> "test123456")
+    assert(confFileMap === expectedConfFileMap)
+  }
+
+  test("SPARK-54605: check custom driver env not appended when config map size is " +
+    "larger than threshold.") {
+    val input = Map("spark-env.sh" -> "test12345\ntest12345", "testConf.1" -> "test123456")
+    val sparkConf = testSetup(input.map(f => f._1 -> f._2.getBytes(StandardCharsets.UTF_8)))
+      .set(Config.CONFIG_MAP_MAXSIZE.key, "87")
+    val conf = KubernetesTestConf.createDriverConf(
+      sparkConf = sparkConf,
+      environment = Map("AAA" -> "value1", "BBB" -> "value2")
+    )
+    val confFileMap = KubernetesClientUtils.overrideDefaultSparkEnv(conf, input)
+    val expectedConfFileMap = Map(
+      "spark-env.sh" -> "test12345\ntest12345",
+      "testConf.1" -> "test123456")
+    assert(confFileMap === expectedConfFileMap)
+  }
+
+  test("SPARK-54605: check custom driver env not appended when config map size is " +
+    "equal to threshold.") {
+    val input = Map("spark-env.sh" -> "test12345\ntest12345", "testConf.1" -> "test123456")
+    val sparkConf = testSetup(input.map(f => f._1 -> f._2.getBytes(StandardCharsets.UTF_8)))
+      .set(Config.CONFIG_MAP_MAXSIZE.key, "88")
+    val conf = KubernetesTestConf.createDriverConf(
+      sparkConf = sparkConf,
+      environment = Map("AAA" -> "value1", "BBB" -> "value2")
+    )
+    val confFileMap = KubernetesClientUtils.overrideDefaultSparkEnv(conf, input)
+    val expectedConfFileMap = Map(
+      "spark-env.sh" -> "test12345\ntest12345",
+      "testConf.1" -> "test123456")
+    assert(confFileMap === expectedConfFileMap)
+  }
 }