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

[HyukjinKwon]

Previous behaviour

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

1. The execution block is wrapped with SessionHolder.withSession here.
2. SessionHolder.withContextClassLoader is then called here which in turn calls JobArtifactSet.withActive here and sets the active set to SessionHolder.connectJobArtifactSet
3. The actual JobArtifactSet that is used is built up here 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.

[zhengruifeng]

qq: If we make existing addXXX methods support destination directory, can we keep those mapping outside of spark core (in connect server) ?

[HyukjinKwon]

Yeah, let me try to decouple this.

[HyukjinKwon]

Actually, I already faced the limitation here when the test takes longer than 5 mins. I increased it as a workaround for now. We should fix them before the release.

[HyukjinKwon]

Yeah, I think with the helper instance, it makes much more sense now.

[HyukjinKwon]

Both tests were flaky. I manually checked them locally.

Merged to master.