Allow configuring conda env variables (apache-spark-on-k8s#468)

Allow passing env variables for conda so that we can enable instrumentation/other flags when required.

Author: geremih

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

@@ -38,7 +38,8 @@ final class CondaEnvironment(val manager: CondaEnvironmentManager,
                              val envName: String,
                              bootstrapPackages: Seq[String],
                              bootstrapChannels: Seq[String],
-                             extraArgs: Seq[String] = Nil) extends Logging {
+                             extraArgs: Seq[String] = Nil,
+                             envVars: Map[String, String] = Map.empty) extends Logging {
 
   import CondaEnvironment._
 
@@ -75,7 +76,8 @@ final class CondaEnvironment(val manager: CondaEnvironmentManager,
         ::: extraArgs.toList
         ::: "--" :: packages.toList,
       description = s"install dependencies in conda env $condaEnvDir",
-      channels = channels.iterator.map(_.url).toList
+      channels = channels.iterator.map(_.url).toList,
+      envVars = envVars
     )
 
     this.packages ++= packages
@@ -96,7 +98,7 @@ final class CondaEnvironment(val manager: CondaEnvironmentManager,
    * This is for sending the instructions to the executors so they can replicate the same steps.
    */
   def buildSetupInstructions: CondaSetupInstructions = {
-    CondaSetupInstructions(packages.toList, channels.toList, extraArgs)
+    CondaSetupInstructions(packages.toList, channels.toList, extraArgs, envVars)
   }
 }
 
@@ -155,7 +157,8 @@ object CondaEnvironment {
   case class CondaSetupInstructions(
          packages: Seq[String],
          unauthenticatedChannels: Seq[UnauthenticatedChannel],
-         extraArgs: Seq[String])
+         extraArgs: Seq[String],
+         envVars: Map[String, String])
         (userInfos: Map[UnauthenticatedChannel, String]) {
     require(unauthenticatedChannels.nonEmpty)
     require(packages.nonEmpty)
@@ -167,10 +170,11 @@ object CondaEnvironment {
   }
 
   object CondaSetupInstructions {
-    def apply(packages: Seq[String], channels: Seq[AuthenticatedChannel], extraArgs: Seq[String])
+    def apply(packages: Seq[String], channels: Seq[AuthenticatedChannel], extraArgs: Seq[String],
+              envVars: Map[String, String])
         : CondaSetupInstructions = {
       val ChannelsWithCreds(unauthed, userInfos) = unauthenticateChannels(channels)
-      CondaSetupInstructions(packages, unauthed, extraArgs)(userInfos)
+      CondaSetupInstructions(packages, unauthed, extraArgs, envVars)(userInfos)
     }
   }
 }

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

@@ -73,7 +73,8 @@ final class CondaEnvironmentManager(condaBinaryPath: String,
               baseDir: String,
               condaPackages: Seq[String],
               condaChannelUrls: Seq[String],
-              condaExtraArgs: Seq[String] = Nil): CondaEnvironment = {
+              condaExtraArgs: Seq[String] = Nil,
+              condaEnvVars: Map[String, String] = Map.empty): 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"
@@ -95,7 +96,8 @@ final class CondaEnvironmentManager(condaBinaryPath: String,
         ::: verbosityFlags
         ::: "--" :: condaPackages.toList,
       description = "create conda env",
-      channels = condaChannelUrls.toList
+      channels = condaChannelUrls.toList,
+      envVars = condaEnvVars
     )
 
     new CondaEnvironment(this, linkedBaseDir, name, condaPackages, condaChannelUrls, condaExtraArgs)
@@ -139,7 +141,8 @@ final class CondaEnvironmentManager(condaBinaryPath: String,
   private[conda] def runCondaProcess(baseRoot: Path,
                                      args: List[String],
                                      channels: List[String],
-                                     description: String): Unit = {
+                                     description: String,
+                                     envVars: Map[String, String]): Unit = {
     val condarc = generateCondarc(baseRoot, channels)
     val fakeHomeDir = baseRoot.resolve("home")
     // Attempt to create fake home dir
@@ -148,7 +151,7 @@ final class CondaEnvironmentManager(condaBinaryPath: String,
     val extraEnv = List(
       "CONDARC" -> condarc.toString,
       "HOME" -> fakeHomeDir.toString
-    )
+    ) ++ envVars
 
     val command = Process(
       condaBinaryPath :: args,

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

@@ -54,17 +54,35 @@ object CondaRunner {
       val condaBootstrapDeps = sparkConf.get(CONDA_BOOTSTRAP_PACKAGES)
       val condaChannelUrls = sparkConf.get(CONDA_CHANNEL_URLS)
       val condaExtraArgs = sparkConf.get(CONDA_EXTRA_ARGUMENTS)
+      val condaEnvVariables = extractEnvVariables(sparkConf)
       val condaBaseDir = Utils.createTempDir(Utils.getLocalDir(sparkConf), "conda").getAbsolutePath
       val condaEnvironmentManager = CondaEnvironmentManager.fromConf(sparkConf)
-      val environment = condaEnvironmentManager
-                        .create(condaBaseDir, condaBootstrapDeps, condaChannelUrls, condaExtraArgs)
+      val environment =
+        condaEnvironmentManager.create(
+          condaBaseDir,
+          condaBootstrapDeps,
+          condaChannelUrls,
+          condaExtraArgs,
+          condaEnvVariables)
       setCondaEnvironment(environment)
       Some(environment)
     } else {
       None
     }
   }
 
+  /**
+   * Extracts environment variables specified in the form
+   * "spark.conda.env.[EnvironmentVariableName]" from the sparkConf.
+   */
+  def extractEnvVariables(sparkConf: SparkConf): Map[String, String] = {
+    val condaEnvPrefix = "spark.conda.env."
+    sparkConf.getAll
+      .filter { case (k, v) => k.startsWith(condaEnvPrefix) }
+      .map { case (k, v) => (k.substring(condaEnvPrefix.length), v) }
+      .toMap
+  }
+
   /**
    * Sets the given environment as the global environment, which will be accessible by calling
    * [[SparkContext.condaEnvironment]]. This method can only be called once! If an environment

core/src/test/scala/org/apache/spark/deploy/CondaRunnerSuite.scala

@@ -0,0 +1,32 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.deploy
+
+import org.apache.spark.{SparkConf, SparkFunSuite}
+
+class CondaRunnerSuite extends SparkFunSuite {
+
+  test("correctly reads conda env vars") {
+    val conf = new SparkConf
+    conf.set("spark.conda.env.key1", "value1")
+    conf.set("spark.conda.env.key2", "value2")
+    val expected = Map("key1" -> "value1", "key2" -> "value2")
+    assert(CondaRunner.extractEnvVariables(conf) == expected)
+  }
+
+}