Make conda channels dynamic on executors (apache-spark-on-k8s#160)

Author: justinuang

core/src/main/scala/org/apache/spark/api/conda/CondaEnvironmentManager.scala

@@ -34,17 +34,16 @@ import org.apache.spark.internal.config.CONDA_CHANNEL_URLS
 import org.apache.spark.internal.config.CONDA_VERBOSITY
 import org.apache.spark.util.Utils
 
-final class CondaEnvironmentManager(condaBinaryPath: String, condaChannelUrls: Seq[String],
-                                    verbosity: Int = 0)
-  extends Logging {
+final class CondaEnvironmentManager(condaBinaryPath: String, verbosity: Int = 0) extends Logging {
 
-  require(condaChannelUrls.nonEmpty, "Can't have an empty list of conda channel URLs")
   require(verbosity >= 0 && verbosity <= 3, "Verbosity must be between 0 and 3 inclusively")
 
   def create(
-      baseDir: String,
-      bootstrapPackages: Seq[String]): CondaEnvironment = {
-    require(bootstrapPackages.nonEmpty, "Expected at least one bootstrap package.")
+              baseDir: String,
+              condaPackages: Seq[String],
+              condaChannelUrls: Seq[String]): CondaEnvironment = {
+    require(condaPackages.nonEmpty, "Expected at least one conda package.")
+    require(condaChannelUrls.nonEmpty, "Can't have an empty list of conda channel URLs")
     val name = "conda-env"
 
     // must link in /tmp to reduce path length in case baseDir is very long...
@@ -62,11 +61,11 @@ final class CondaEnvironmentManager(condaBinaryPath: String, condaChannelUrls: S
       List("create", "-n", name, "-y", "--override-channels", "--no-default-packages")
         ::: verbosityFlags
         ::: condaChannelUrls.flatMap(Iterator("--channel", _)).toList
-        ::: "--" :: bootstrapPackages.toList,
+        ::: "--" :: condaPackages.toList,
       description = "create conda env"
     )
 
-    new CondaEnvironment(this, linkedBaseDir, name, bootstrapPackages, condaChannelUrls)
+    new CondaEnvironment(this, linkedBaseDir, name, condaPackages, condaChannelUrls)
   }
 
   /**
@@ -142,10 +141,7 @@ object CondaEnvironmentManager {
   def fromConf(sparkConf: SparkConf): CondaEnvironmentManager = {
     val condaBinaryPath = sparkConf.get(CONDA_BINARY_PATH).getOrElse(
       sys.error(s"Expected config ${CONDA_BINARY_PATH.key} to be set"))
-    val condaChannelUrls = sparkConf.get(CONDA_CHANNEL_URLS)
-    require(condaChannelUrls.nonEmpty,
-      s"Must define at least one conda channel in config ${CONDA_CHANNEL_URLS.key}")
     val verbosity = sparkConf.get(CONDA_VERBOSITY)
-    new CondaEnvironmentManager(condaBinaryPath, condaChannelUrls, verbosity)
+    new CondaEnvironmentManager(condaBinaryPath, verbosity)
   }
 }

core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala

@@ -69,7 +69,7 @@ private[spark] class PythonWorkerFactory(requestedPythonExec: Option[String],
         val dirId = hash % localDirs.length
         Utils.createTempDir(localDirs(dirId).getAbsolutePath, "conda").getAbsolutePath
       }
-      condaEnvManager.create(envDir, condaPackages)
+      condaEnvManager.create(envDir, condaPackages, instructions.channels)
     }
   }
 

core/src/main/scala/org/apache/spark/deploy/CondaRunner.scala

@@ -34,9 +34,11 @@ abstract class CondaRunner extends Logging {
 
     if (CondaEnvironmentManager.isConfigured(sparkConf)) {
       val condaBootstrapDeps = sparkConf.get(CONDA_BOOTSTRAP_PACKAGES)
+      val condaChannelUrls = sparkConf.get(CONDA_CHANNEL_URLS)
       val condaBaseDir = Utils.createTempDir(Utils.getLocalDir(sparkConf), "conda").getAbsolutePath
       val condaEnvironmentManager = CondaEnvironmentManager.fromConf(sparkConf)
-      val environment = condaEnvironmentManager.create(condaBaseDir, condaBootstrapDeps)
+      val environment = condaEnvironmentManager
+        .create(condaBaseDir, condaBootstrapDeps, condaChannelUrls)
 
       // Save this as a global in order for SparkContext to be able to access it later, in case we
       // are shelling out, but providing a bridge back into this JVM.

resource-managers/yarn/src/test/scala/org/apache/spark/deploy/yarn/YarnClusterSuite.scala

@@ -91,7 +91,14 @@ class YarnClusterSuite extends BaseYarnClusterSuite {
     |
     |    status = open(sys.argv[1],'w')
     |
-    |    numpy_multiply = lambda x: numpy.multiply(x, mod1.func() * mod2.func())
+    |    # Addict exists only in external-conda-forge, not anaconda
+    |    sc.addCondaChannel("https://conda.anaconda.org/conda-forge")
+    |    sc.addCondaPackages('addict=1.0.0')
+    |
+    |    def numpy_multiply(x):
+    |        # Ensure package from non-base channel is installed
+    |        import addict
+    |        numpy.multiply(x, mod1.func() * mod2.func())
     |
     |    rdd = sc.parallelize(range(10)).map(numpy_multiply)
     |    cnt = rdd.count()