[SPARK-44290][CONNECT] Session-based files and archives in Spark Connect

### Previous behaviour

Previously, we kept `JobArtifactSet` and leveraged thread local for each client.

1. The execution block is wrapped with `SessionHolder.withSession` [here](https://github.com/apache/spark/blob/master/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala#L53).
2. `SessionHolder.withContextClassLoader` is then called [here](https://github.com/apache/spark/blob/master/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala#L130) which in turn calls `JobArtifactSet.withActive` [here](https://github.com/apache/spark/blob/master/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala#L118) and sets the active set to `SessionHolder.connectJobArtifactSet`
3. The actual `JobArtifactSet` that is used is built up [here](https://github.com/apache/spark/blob/master/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/artifact/SparkConnectArtifactManager.scala#L157) in `SparkConnectArtifactManager.jobArtifactSet`

Since each client has their own `JobArtifactSet` made `active` when executing an operation, the `TaskDescription` would have artifacts specific to that client and subsequently, `IsolatedSessionState` in Executor. Therefore, we were able to separate the Spark Connect specific logic to the Spark Connect module.

### Problem

Mainly it was all good; however, the problem is that we don't call `SparkContext.addFile` or `SparkContext.addJar`, but we just pass it directly at the scheduler (to `TaskDescription`). This is fine in general but exposes several problems by not directly calling `SparkContext.addFile`:

- `SparkContext.postEnvironmentUpdate` is not invoked at `SparkContext.addFile` which matters in, for example, recording the events for History Server.
- Specifically for archives, `Utils.unpack(source, dest)` is not invoked at `SparkContext.addFile` in order to untar properly in the Driver.

Therefore, we should duplicate those logics in Spark Connect server side, which is not ideal.

In addition, we already added the isolation logic into the Executor. Driver and Executor are the symmetry (not Spark Connect Server <> Executor). Therefore, it matters about code readability, and expectation in their roles.

### Solution in this PR

This PR proposes to support session-based files and archives in Spark Connect. This PR leverages the basework for #41701 and #41625 (for jars in Spark Connect Scala client).

The changed logic is as follows:

- Keep the session UUID, and Spark Connect Server specific information such as REPL class path within a thread local.
- Add session ID when we add files or archives. `SparkContext` keeps them with a map `Map(session -> Map(file and timestamp))` in order to reuse the existing logic to address the problem mentioned

After that, on executor side,

- Executors create additional directory, named by session UUID, on the top of the default directory (that is the current working directory, see `SparkFiles.getRootDirectory`).
- When we execute Python workers, it sets the current working directory to the one created above.
- End users access to these files via using the current working directory e.g., `./blahblah.txt` in their Python UDF. Therefore, compatible with/without Spark Connect.

Note that:
- Here it creates Python workers for individual session because we set the session UUID as an environment variable, and we create new Python workers if environment variables are different, see also `SparkEnv.createPythonWorker`
- It already kills the daemon and Python workers if they are not used for a while.

### TODOs and limitations

Executor also maintains the file list but with a cache so it can evict the cache. However, it has a problem - It works as follows:
- New `IsolatedSessionState` is created.
- Task is executed once, and `IsolatedSessionState` holds the file list.
- Later `IsolatedSessionState` is evicted at https://github.com/apache/spark/pull/41625/files#diff-d7a989c491f3cb77cca02c701496a9e2a3443f70af73b0d1ab0899239f3a789dR187
- Executor will create a new `IsolatedSessionState` with empty file lists.
- Executor will attempt to redownload and overwrite the files (see https://github.com/apache/spark/blob/master/core/src/main/scala/org/apache/spark/executor/Executor.scala#L1058-L1064)
- `spark.files.overwrite` is `false` by default. So the task will suddenly fail at this point.

Possible solutions are:
- For 1., we should maintain a cache with TTL, and remove them
- For 2. we should have a dedicated directory (which this PR does) and remove the directory away when the cache is evicted. So the overwrite does not happen

### Why are the changes needed?

In order to allow session-based artifact control and multi tenancy.

### Does this PR introduce _any_ user-facing change?

Yes, this PR now allows multiple sessions to have their own space. For example, session A and session B can add a file in the same name. Previously this was not possible.

### How was this patch tested?

Unittests were added.

Closes #41495 from HyukjinKwon/session-base-exec-dir.

Authored-by: Hyukjin Kwon <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>

Author: HyukjinKwon

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/artifact/SparkConnectArtifactManager.scala

@@ -29,7 +29,7 @@ import scala.reflect.ClassTag
 import org.apache.commons.io.FileUtils
 import org.apache.hadoop.fs.{LocalFileSystem, Path => FSPath}
 
-import org.apache.spark.{JobArtifactSet, SparkContext, SparkEnv}
+import org.apache.spark.{JobArtifactState, SparkContext, SparkEnv}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.connect.artifact.util.ArtifactUtils
@@ -56,15 +56,15 @@ import org.apache.spark.util.Utils
 class SparkConnectArtifactManager(sessionHolder: SessionHolder) extends Logging {
   import SparkConnectArtifactManager._
 
-  private val sessionUUID = sessionHolder.session.sessionUUID
   // The base directory/URI where all artifacts are stored for this `sessionUUID`.
   val (artifactPath, artifactURI): (Path, String) =
     getArtifactDirectoryAndUriForSession(sessionHolder)
   // The base directory/URI where all class file artifacts are stored for this `sessionUUID`.
   val (classDir, classURI): (Path, String) = getClassfileDirectoryAndUriForSession(sessionHolder)
+  val state: JobArtifactState =
+    JobArtifactState(sessionHolder.session.sessionUUID, Option(classURI))
 
   private val jarsList = new CopyOnWriteArrayList[Path]
-  private val jarsURI = new CopyOnWriteArrayList[String]
   private val pythonIncludeList = new CopyOnWriteArrayList[String]
 
   /**
@@ -132,10 +132,16 @@ class SparkConnectArtifactManager(sessionHolder: SessionHolder) extends Logging
       }
       Files.move(serverLocalStagingPath, target)
       if (remoteRelativePath.startsWith(s"jars${File.separator}")) {
+        sessionHolder.session.sessionState.resourceLoader
+          .addJar(target.toString, state.uuid)
         jarsList.add(target)
-        jarsURI.add(artifactURI + "/" + remoteRelativePath.toString)
       } else if (remoteRelativePath.startsWith(s"pyfiles${File.separator}")) {
-        sessionHolder.session.sparkContext.addFile(target.toString)
+        sessionHolder.session.sparkContext.addFile(
+          target.toString,
+          recursive = false,
+          addedOnSubmit = false,
+          isArchive = false,
+          sessionUUID = state.uuid)
         val stringRemotePath = remoteRelativePath.toString
         if (stringRemotePath.endsWith(".zip") || stringRemotePath.endsWith(
             ".egg") || stringRemotePath.endsWith(".jar")) {
@@ -144,35 +150,28 @@ class SparkConnectArtifactManager(sessionHolder: SessionHolder) extends Logging
       } else if (remoteRelativePath.startsWith(s"archives${File.separator}")) {
         val canonicalUri =
           fragment.map(UriBuilder.fromUri(target.toUri).fragment).getOrElse(target.toUri)
-        sessionHolder.session.sparkContext.addArchive(canonicalUri.toString)
+        sessionHolder.session.sparkContext.addFile(
+          canonicalUri.toString,
+          recursive = false,
+          addedOnSubmit = false,
+          isArchive = true,
+          sessionUUID = state.uuid)
       } else if (remoteRelativePath.startsWith(s"files${File.separator}")) {
-        sessionHolder.session.sparkContext.addFile(target.toString)
+        sessionHolder.session.sparkContext.addFile(
+          target.toString,
+          recursive = false,
+          addedOnSubmit = false,
+          isArchive = false,
+          sessionUUID = state.uuid)
       }
     }
   }
 
-  /**
-   * Returns a [[JobArtifactSet]] pointing towards the session-specific jars and class files.
-   */
-  def jobArtifactSet: JobArtifactSet = {
-    val builder = Map.newBuilder[String, Long]
-    jarsURI.forEach { jar =>
-      builder += jar -> 0
-    }
-
-    new JobArtifactSet(
-      uuid = Option(sessionUUID),
-      replClassDirUri = Option(classURI),
-      jars = builder.result(),
-      files = Map.empty,
-      archives = Map.empty)
-  }
-
   /**
    * Returns a [[ClassLoader]] for session-specific jar/class file resources.
    */
   def classloader: ClassLoader = {
-    val urls = jarsList.asScala.map(_.toUri.toURL) :+ classDir.toUri.toURL
+    val urls = getSparkConnectAddedJars :+ classDir.toUri.toURL
     new URLClassLoader(urls.toArray, Utils.getContextOrSparkClassLoader)
   }
 
@@ -183,6 +182,12 @@ class SparkConnectArtifactManager(sessionHolder: SessionHolder) extends Logging
     logDebug(
       s"Cleaning up resources for session with userId: ${sessionHolder.userId} and " +
         s"sessionId: ${sessionHolder.sessionId}")
+
+    // Clean up added files
+    sessionHolder.session.sparkContext.addedFiles.remove(state.uuid)
+    sessionHolder.session.sparkContext.addedArchives.remove(state.uuid)
+    sessionHolder.session.sparkContext.addedJars.remove(state.uuid)
+
     // Clean up cached relations
     val blockManager = sessionHolder.session.sparkContext.env.blockManager
     blockManager.removeCache(sessionHolder.userId, sessionHolder.sessionId)

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala

@@ -88,11 +88,6 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio
     artifactManager.addArtifact(remoteRelativePath, serverLocalStagingPath, fragment)
   }
 
-  /**
-   * A [[JobArtifactSet]] for this SparkConnect session.
-   */
-  def connectJobArtifactSet: JobArtifactSet = artifactManager.jobArtifactSet
-
   /**
    * A [[ClassLoader]] for jar/class file resources specific to this SparkConnect session.
    */
@@ -114,8 +109,7 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio
   def withContextClassLoader[T](f: => T): T = {
     // Needed for deserializing and evaluating the UDF on the driver
     Utils.withContextClassLoader(classloader) {
-      // Needed for propagating the dependencies to the executors.
-      JobArtifactSet.withActive(connectJobArtifactSet) {
+      JobArtifactSet.withActiveJobArtifactState(artifactManager.state) {
         f
       }
     }

connector/connect/server/src/test/scala/org/apache/spark/sql/connect/artifact/ArtifactManagerSuite.scala

@@ -51,20 +51,6 @@ class ArtifactManagerSuite extends SharedSparkSession with ResourceHelper {
     super.afterEach()
   }
 
-  test("Jar artifacts are added to spark session") {
-    val copyDir = Utils.createTempDir().toPath
-    FileUtils.copyDirectory(artifactPath.toFile, copyDir.toFile)
-    val stagingPath = copyDir.resolve("smallJar.jar")
-    val remotePath = Paths.get("jars/smallJar.jar")
-    artifactManager.addArtifact(remotePath, stagingPath, None)
-
-    val expectedPath = SparkConnectArtifactManager.artifactRootPath
-      .resolve(s"$sessionUUID/jars/smallJar.jar")
-    assert(expectedPath.toFile.exists())
-    val jars = artifactManager.jobArtifactSet.jars
-    assert(jars.exists(_._1.contains(remotePath.toString)))
-  }
-
   test("Class artifacts are added to the correct directory.") {
     val copyDir = Utils.createTempDir().toPath
     FileUtils.copyDirectory(artifactPath.toFile, copyDir.toFile)

core/src/main/scala/org/apache/spark/JobArtifactSet.scala

@@ -20,105 +20,100 @@ package org.apache.spark
 import java.io.Serializable
 import java.util.Objects
 
+/**
+ * Job artifact state. For example, Spark Connect client sets the state specifically
+ * for the current client.
+ *
+ * @param uuid UUID to use in the current context of jab artifact set. Usually this is from
+ *             a Spark Connect client.
+ * @param replClassDirUri The URI for the directory that stores REPL classes.
+ */
+private[spark] case class JobArtifactState(uuid: String, replClassDirUri: Option[String])
+
 /**
  * Artifact set for a job.
  * This class is used to store session (i.e `SparkSession`) specific resources/artifacts.
  *
  * When Spark Connect is used, this job-set points towards session-specific jars and class files.
  * Note that Spark Connect is not a requirement for using this class.
  *
- * @param uuid An optional UUID for this session. If unset, a default session will be used.
- * @param replClassDirUri An optional custom URI to point towards class files.
+ * @param state Job artifact state.
  * @param jars Jars belonging to this session.
  * @param files Files belonging to this session.
  * @param archives Archives belonging to this session.
  */
-class JobArtifactSet(
-    val uuid: Option[String],
-    val replClassDirUri: Option[String],
+private[spark] class JobArtifactSet(
+    val state: Option[JobArtifactState],
     val jars: Map[String, Long],
     val files: Map[String, Long],
     val archives: Map[String, Long]) extends Serializable {
-  def withActive[T](f: => T): T = JobArtifactSet.withActive(this)(f)
-
   override def hashCode(): Int = {
-    Objects.hash(uuid, replClassDirUri, jars.toSeq, files.toSeq, archives.toSeq)
+    Objects.hash(state, jars.toSeq, files.toSeq, archives.toSeq)
   }
 
   override def equals(obj: Any): Boolean = {
     obj match {
       case that: JobArtifactSet =>
-        this.getClass == that.getClass && this.uuid == that.uuid &&
-          this.replClassDirUri == that.replClassDirUri && this.jars.toSeq == that.jars.toSeq &&
+        this.getClass == that.getClass && this.state == that.state &&
+          this.jars.toSeq == that.jars.toSeq &&
           this.files.toSeq == that.files.toSeq && this.archives.toSeq == that.archives.toSeq
     }
   }
-
 }
 
-object JobArtifactSet {
-
-  private[this] val current = new ThreadLocal[Option[JobArtifactSet]] {
-    override def initialValue(): Option[JobArtifactSet] = None
-  }
-
-  /**
-   * When Spark Connect isn't used, we default back to the shared resources.
-   * @param sc The active [[SparkContext]]
-   * @return A [[JobArtifactSet]] containing a copy of the jars/files/archives from the underlying
-   *         [[SparkContext]] `sc`.
-   */
-  def apply(sc: SparkContext): JobArtifactSet = {
-    new JobArtifactSet(
-      uuid = None,
-      replClassDirUri = sc.conf.getOption("spark.repl.class.uri"),
-      jars = sc.addedJars.toMap,
-      files = sc.addedFiles.toMap,
-      archives = sc.addedArchives.toMap)
-  }
 
-  private lazy val emptyJobArtifactSet = new JobArtifactSet(
-    None,
-    None,
-    Map.empty,
-    Map.empty,
-    Map.empty)
+private[spark] object JobArtifactSet {
+  // For testing.
+  val emptyJobArtifactSet: JobArtifactSet = new JobArtifactSet(
+    None, Map.empty, Map.empty, Map.empty)
+  // For testing.
+  def defaultJobArtifactSet: JobArtifactSet = SparkContext.getActive.map(
+    getActiveOrDefault).getOrElse(emptyJobArtifactSet)
 
-  /**
-   * Empty artifact set for use in tests.
-   */
-  private[spark] def apply(): JobArtifactSet = emptyJobArtifactSet
+  private[this] val currentClientSessionState: ThreadLocal[Option[JobArtifactState]] =
+    new ThreadLocal[Option[JobArtifactState]] {
+      override def initialValue(): Option[JobArtifactState] = None
+    }
 
-  /**
-   * Used for testing. Returns artifacts from [[SparkContext]] if one exists or otherwise, an
-   * empty set.
-   */
-  private[spark] def defaultArtifactSet(): JobArtifactSet = {
-    SparkContext.getActive.map(sc => JobArtifactSet(sc)).getOrElse(JobArtifactSet())
-  }
+  def getCurrentClientSessionState: Option[JobArtifactState] = currentClientSessionState.get()
 
   /**
-   * Execute a block of code with the currently active [[JobArtifactSet]].
-   * @param active
-   * @param block
-   * @tparam T
+   * Set the Spark Connect specific information in the active client to the underlying
+   * [[JobArtifactSet]].
+   *
+   * @param state Job artifact state.
+   * @return the result from the function applied with [[JobArtifactSet]] specific to
+   *         the active client.
    */
-  def withActive[T](active: JobArtifactSet)(block: => T): T = {
-    val old = current.get()
-    current.set(Option(active))
+  def withActiveJobArtifactState[T](state: JobArtifactState)(block: => T): T = {
+    val oldState = currentClientSessionState.get()
+    currentClientSessionState.set(Option(state))
     try block finally {
-      current.set(old)
+      currentClientSessionState.set(oldState)
     }
   }
 
   /**
-   * Optionally returns the active [[JobArtifactSet]].
-   */
-  def active: Option[JobArtifactSet] = current.get()
-
-  /**
-   * Return the active [[JobArtifactSet]] or creates the default set using the [[SparkContext]].
-   * @param sc
+   * When Spark Connect isn't used, we default back to the shared resources.
+   *
+   * @param sc The active [[SparkContext]]
+   * @return A [[JobArtifactSet]] containing a copy of the jars/files/archives.
+   *         If there is an active client, it sets the information from them.
+   *         Otherwise, it falls back to the default in the [[SparkContext]].
    */
-  def getActiveOrDefault(sc: SparkContext): JobArtifactSet = active.getOrElse(JobArtifactSet(sc))
+  def getActiveOrDefault(sc: SparkContext): JobArtifactSet = {
+    val maybeState = currentClientSessionState.get().map(s => s.copy(
+        replClassDirUri = s.replClassDirUri.orElse(sc.conf.getOption("spark.repl.class.uri"))))
+    new JobArtifactSet(
+      state = maybeState,
+      jars = maybeState
+        .map(s => sc.addedJars.getOrElse(s.uuid, sc.allAddedJars))
+        .getOrElse(sc.allAddedJars).toMap,
+      files = maybeState
+        .map(s => sc.addedFiles.getOrElse(s.uuid, sc.allAddedFiles))
+        .getOrElse(sc.allAddedFiles).toMap,
+      archives = maybeState
+        .map(s => sc.addedArchives.getOrElse(s.uuid, sc.allAddedArchives))
+        .getOrElse(sc.allAddedArchives).toMap)
+  }
 }