From 7b2870d9098c5af5faf7472a8f5aa0fc7da75ee2 Mon Sep 17 00:00:00 2001
From: Gabor Roczei <roczei@cloudera.com>
Date: Sat, 18 Oct 2025 10:39:46 +0200
Subject: [PATCH] yarn_noexec_tmpdir_env

---
 .../org/apache/spark/deploy/yarn/Client.scala |  3 +
 .../spark/deploy/yarn/ExecutorRunnable.scala  |  3 +
 .../spark/deploy/yarn/YarnClusterSuite.scala  | 56 +++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala b/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
index 09655007acb34..41b3555f4e38e 100644
--- a/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
+++ b/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
@@ -969,6 +969,9 @@ private[spark] class Client(
     env("SPARK_YARN_STAGING_DIR") = stagingDirPath.toString
     env("SPARK_PREFER_IPV6") = Utils.preferIPv6.toString
 
+    val tmpDir = new Path(Environment.PWD.$$(), YarnConfiguration.DEFAULT_CONTAINER_TEMP_DIR)
+    sparkConf.setIfMissing("spark.yarn.appMasterEnv.TMPDIR", tmpDir.toString)
+
     // Pick up any environment variables for the AM provided through spark.yarn.appMasterEnv.*
     val amEnvPrefix = "spark.yarn.appMasterEnv."
     sparkConf.getAll
diff --git a/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/ExecutorRunnable.scala b/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/ExecutorRunnable.scala
index 052b759e3f5c5..c79e46c31b088 100644
--- a/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/ExecutorRunnable.scala
+++ b/resource-managers/yarn/src/main/scala/org/apache/spark/deploy/yarn/ExecutorRunnable.scala
@@ -211,6 +211,9 @@ private[yarn] class ExecutorRunnable(
     System.getenv().asScala.filter { case (k, _) => k.startsWith("SPARK") }
       .foreach { case (k, v) => env(k) = v }
 
+    val tmpDir = new Path(Environment.PWD.$$(), YarnConfiguration.DEFAULT_CONTAINER_TEMP_DIR)
+    sparkConf.setIfMissing("spark.executorEnv.TMPDIR", tmpDir.toString)
+
     sparkConf.getExecutorEnv.foreach { case (key, value) =>
       if (key == Environment.CLASSPATH.name()) {
         // If the key of env variable is CLASSPATH, we assume it is a path and append it.
diff --git a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnClusterSuite.scala b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnClusterSuite.scala
index 32cf9ea66803a..5c8893b96f0f8 100644
--- a/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnClusterSuite.scala
+++ b/resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnClusterSuite.scala
@@ -26,6 +26,7 @@ import java.util.{HashMap => JHashMap}
 import scala.collection.mutable
 import scala.concurrent.duration._
 import scala.io.Source
+import scala.sys.env
 
 import org.apache.hadoop.yarn.conf.YarnConfiguration
 import org.scalatest.concurrent.Eventually._
@@ -129,6 +130,16 @@ class YarnClusterSuite extends BaseYarnClusterSuite {
     |    return 42
     """.stripMargin
 
+  test("Check the value of the TMPDIR environment " +
+    "variable within a YARN container in client mode") {
+    testTmpDirEnv(true)
+  }
+
+  test("Check the value of the TMPDIR environment" +
+    "variable within a YARN container in cluster mode") {
+    testTmpDirEnv(false)
+  }
+
   test("run Spark in yarn-client mode") {
     testBasicYarnApp(true)
   }
@@ -503,6 +514,16 @@ class YarnClusterSuite extends BaseYarnClusterSuite {
     checkResult(finalState, result, "true")
   }
 
+  private def testTmpDirEnv(clientMode: Boolean): Unit = {
+    val result = File.createTempFile("result", null, tempDir)
+    val finalState = runSpark(
+      clientMode = clientMode,
+      klass = mainClassName(QueryExecutorTmpDirEnv.getClass),
+      appArgs = Seq(result.getAbsolutePath)
+    )
+    checkResult(finalState, result)
+  }
+
   def createEmptyIvySettingsFile: File = {
     val emptyIvySettings = File.createTempFile("ivy", ".xml")
     Files.writeString(emptyIvySettings.toPath, "<ivysettings />")
@@ -903,3 +924,38 @@ private class PyConnectDepChecker(python: String, libPath: Seq[String]) {
     sys.props.getOrElse("spark.test.home", sys.env("SPARK_HOME"))
   }
 }
+
+private object QueryExecutorTmpDirEnv extends Logging {
+
+  val variableName = "TMPDIR"
+  val notFound = "Not Found"
+  val executorPrefix = s"executor-${variableName}"
+  val driverPrefix = s"driver-${variableName}"
+
+  def getExecutorEnvVariables(partitionIndex: Int, iterator: Iterator[Int]): Iterator[String] = {
+    val executorVarValue = env.getOrElse(variableName, notFound)
+    assert(!executorVarValue.equals(notFound))
+    Iterator(s"${executorPrefix}:$executorVarValue")
+  }
+
+  def main(args: Array[String]): Unit = {
+    val status = args(0)
+    val sc = new SparkContext(new SparkConf()
+      .setAppName("TMPDIR application"))
+
+    val clientMode = sc.getConf.get(SparkLauncher.DEPLOY_MODE) == "client"
+    val numPartitions = 2
+    val rdd = sc.parallelize(1 to numPartitions, numPartitions)
+    val results = rdd.mapPartitionsWithIndex(getExecutorEnvVariables).collect()
+    logInfo(results.map(_.toString).mkString)
+    val driverVarValue = env.getOrElse(variableName, notFound)
+    logInfo(s"${driverPrefix}:$driverVarValue")
+
+    if (!clientMode) {
+      assert(!driverVarValue.equals(notFound))
+    }
+    Files.writeString(new File(status).toPath, "success")
+    sc.stop()
+  }
+
+}