yarn_noexec_tmpdir_env

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

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.

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()
+  }
+
+}