Improve sharing of conda environments across Python and R

By storing previous environments that are the same within a single
Java executor process, we are able to reuse environments that we
know are exactly the same.

This cache needs to be protected by a lock to prevent concurrent
modifications since we do not want extra copies of the environment
on disk.

Author: mnazbro

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

@@ -20,9 +20,11 @@ import java.nio.file.Files
 import java.nio.file.Path
 import java.nio.file.Paths
 import java.nio.file.attribute.PosixFilePermission
+import java.util.concurrent.locks.{Lock, ReentrantLock}
 import java.util.regex.Pattern
 
 import scala.collection.JavaConverters._
+import scala.collection.mutable
 import scala.sys.process.BasicIO
 import scala.sys.process.Process
 import scala.sys.process.ProcessBuilder
@@ -211,6 +213,9 @@ object CondaEnvironmentManager extends Logging {
   private[this] val httpUrlToken =
     Pattern.compile("(\\b\\w+://[^:/@]*:)([^/@]+)(?=@([\\w-.]+)(:\\d+)?\\b)")
 
+  private val initializedEnvironmentsLock = new ReentrantLock()
+  private val initializedEnvironments = mutable.HashMap[CondaSetupInstructions, CondaEnvironment]()
+
   private[conda] def redactCredentials(line: String): String = {
     httpUrlToken.matcher(line).replaceAll("$1<password>")
   }
@@ -241,4 +246,22 @@ object CondaEnvironmentManager extends Logging {
     condaEnvManager.create(envDir, condaPackages, instructions.channels, instructions.extraArgs)
   }
 
+  /**
+   * Helper method that will cache precreated conda environments.
+   */
+  def createOrGetCondaEnvironment(instructions: CondaSetupInstructions): CondaEnvironment = {
+    initializedEnvironmentsLock.lock()
+    try {
+      initializedEnvironments.get(instructions) match {
+        case Some(condaEnv) => condaEnv
+        case None =>
+          val condaEnv = createCondaEnvironment(instructions)
+          initializedEnvironments.put(instructions, condaEnv)
+          condaEnv
+      }
+    } finally {
+      initializedEnvironmentsLock.unlock()
+    }
+  }
+
 }

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

@@ -94,7 +94,7 @@ private[spark] class PythonWorkerFactory(requestedPythonExec: Option[String],
 
   private[this] val condaEnv = {
     // Set up conda environment if there are any conda packages requested
-    condaInstructions.map(CondaEnvironmentManager.createCondaEnvironment)
+    condaInstructions.map(CondaEnvironmentManager.createOrGetCondaEnvironment)
   }
 
   private[this] val envVars: Map[String, String] = {

core/src/main/scala/org/apache/spark/api/r/RRunner.scala

@@ -348,7 +348,7 @@ private[r] object RRunner {
     val sparkConf = SparkEnv.get.conf
     val requestedRCommand = Provenance.fromConf("spark.r.command")
       .orElse(Provenance.fromConf("spark.sparkr.r.command"))
-    val condaEnv = condaSetupInstructions.map(CondaEnvironmentManager.createCondaEnvironment)
+    val condaEnv = condaSetupInstructions.map(CondaEnvironmentManager.createOrGetCondaEnvironment)
     val rCommand = condaEnv.map { conda =>
       requestedRCommand.foreach(exec => sys.error(s"It's forbidden to set the r executable " +
         s"when using conda, but found: $exec"))