add archives path to PYTHONPATH

jeanlyn · Apr 16, 2015 · f72987c · f72987c
1 parent de4fa6b
commit f72987c
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 0 deletions.
diff --git a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
@@ -328,6 +328,10 @@ object SparkSubmit {
       }
     }
 
+    if (args.isPython && System.getenv("PYSPARK_ARCHIVES_PATH") != null) {
+      args.files = mergeFileLists(args.files, System.getenv("PYSPARK_ARCHIVES_PATH"))
+    }
+
     // If we're running a R app, set the main class to our specific R runner
     if (args.isR && deployMode == CLIENT) {
       if (args.primaryResource == SPARKR_SHELL) {

diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/Client.scala
@@ -326,6 +326,12 @@ private[spark] class Client(
     distCacheMgr.setDistFilesEnv(env)
     distCacheMgr.setDistArchivesEnv(env)
 
+    if (System.getenv("PYSPARK_ARCHIVES_PATH") != null) {
+      val pythonPath = System.getenv("PYSPARK_ARCHIVES_PATH").split(",").map(
+        p => (new Path(p)).getName).mkString(":")
+      env("PYTHONPATH") = pythonPath
+    }
+
     // Pick up any environment variables for the AM provided through spark.yarn.appMasterEnv.*
     val amEnvPrefix = "spark.yarn.appMasterEnv."
     sparkConf.getAll

diff --git a/yarn/src/main/scala/org/apache/spark/deploy/yarn/ExecutorRunnable.scala b/yarn/src/main/scala/org/apache/spark/deploy/yarn/ExecutorRunnable.scala
@@ -285,6 +285,10 @@ class ExecutorRunnable(
       YarnSparkHadoopUtil.addPathToEnvironment(env, key, value)
     }
 
+    if (System.getenv("PYTHONPATH") != null) {
+      env("PYTHONPATH") = System.getenv("PYTHONPATH")
+    }
+
     // Keep this for backwards compatibility but users should move to the config
     sys.env.get("SPARK_YARN_USER_ENV").foreach { userEnvs =>
       YarnSparkHadoopUtil.setEnvFromInputString(env, userEnvs)