[Spark] Make cancel() cancel the Spark jobs and stop only a session the runner created

[tkaymak]

Fixes #40101. Found in review of #40090.

Bug

SparkStructuredStreamingPipelineResult.cancel() interrupted the execution thread with Future.cancel(true) and then stopped the SparkSession from the caller thread. Neither is a cancellation. An interrupt does not cancel a Spark job, DAGScheduler.runJob waits on the JobWaiter until it is cancelled explicitly, so with useActiveSparkSession=true a batch pipeline was never cancelled and cancel() still reported CANCELLED. SparkSessionFactory.getOrCreateSession used Builder.getOrCreate, which adopts a usable default session, so a pipeline could stop a session shared with other pipelines in the JVM. Batch EvaluationContext.stop() was a no-op. The execution thread kept translating on a stopped SparkContext, which is how the first Spark Versions run of #40090 failed.

Fix, shared code compiled for Spark 3 and Spark 4

• SparkSessionFactory.acquire records whether the runner created the session (no usable active or default session before getOrCreate) and counts users per created session. release stops a created session when its last user leaves. useActiveSparkSession=true never owns. Spark 3.5 returns stopped sessions from getActiveSession and getDefaultSession, so the usable check filters them the way getOrCreate does.
• The execution thread sets a job group before translation, skips evaluate() when cancel arrived during translation, and releases the session in finally. The session stop therefore always happens after the last Spark call of the pipeline, on the thread that made it.
• Batch EvaluationContext.stop() sets a flag the leaf loop checks. Running jobs are cancelled through cancelJobGroup with interruptOnCancel. StreamingEvaluationContext keeps stopping its queries.
• cancel() stops the context, cancels the job group, joins the execution thread without a timeout, and reports CANCELLED. The join is bounded by Spark itself, task interruption for batch and spark.sql.streaming.stopTimeout for queries. A pipeline that already ended reports DONE or FAILED. No Future.cancel(true), no terminal state callback. waitUntilFinish reports CANCELLED for a pipeline that ended after a cancel request.

Tests, no sleeps

• SparkStructuredStreamingPipelineResultTest: the join returns only after the execution ended, a second cancel() is a no-op, a cancelled execution that ends with an exception is CANCELLED, not FAILED.
• StructuredStreamingPipelineStateTest: a batch cancel of a running job stops the job and the owned session; a pipeline on a session it did not create leaves that session running after cancel.

Behavior changes to note, Spark 3 and Spark 4 structured streaming

• cancel() blocks until the execution thread has ended. Tasks that ignore interruption delay it until Spark ends them, spark.task.reaper.enabled bounds that. Translation is not interrupted, a cancel during translation skips the evaluation.
• waitUntilFinish() after a cancel returns CANCELLED instead of throwing.
• A session is stopped only when created by the runner, on the execution thread, after its last pipeline.
• The Spark UI shows the pipeline's jobs under the group Beam <jobName>.
• SparkSessionFactory.acquire decides ownership under its own lock. A session created by other code on another thread between its check and its getOrCreate call would be attributed to the runner, no Beam code path does that.
• SparkSessionFactory.getOrCreateSession is replaced by acquire and release. The class is not annotated @Internal, the only caller in the repository was the runner.

Gates: ErrorProne on native JDK 17, checkstyle, spotbugs, spotless, :runners:spark:4:test full, :runners:spark:3:test for the touched classes, and StructuredStreamingPipelineStateTest ten times in the CI fork mode.

R: @Abacn

[github-actions]

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

[github-actions]

Assigning reviewers:

R: @kennknowles added as fallback since no labels match configuration

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[Abacn]

I understand this change tries to fix a few important gaps/bugs, however the behavior change part is likely undesiable. mainly cancel() becomes synchronous, and the reference counting may still not eliminate races of parallel jobs.

The original diagnosis is spot-on:

1. cancel() merely interrupting the execution thread (Future.cancel(true)) does not cancel running Spark jobs in the DAGScheduler.

2. Calling session.stop() synchronously from the caller thread in offerNewState tears down the SparkContext while the execution thread may still be actively translating or evaluating.

In the Beam model, PipelineResult.cancel() is designed as an asynchronous cancellation request. Making it synchrounous can hang callers. In fact, synchrounous cancel() is what introduced most of the state-machine complexity in this PR (handling interrupted cancels, unobserved completions/failures, re-checking pipelineExecution.isDone(), etc.).

Reference counting: the choice of synchrnous cancel() may be correlated with reference counting introduced. Previously we force cancel session on each job cancellation. Subsequent job then starts in a fresh new session thus its conf is honored. Now, if cancel remains async then subsequent job could run on same session and not honoring conf.

Reference counting itself does not solve the race. It is the synchrnous cancel() force job submitting sequentially thus mitigating race, and parallel jobs happen to run on same session does not crash, but just have conf messed up, as a result less likely (but still possible) test failure.

Take a step back: it is a Spark limitation limiting us running jobs in parallel that would need different pipeline options won't have all configurations honored. This stems from the fact that only one active SparkContext is allowed throughout JVM.

Need to think more about this. It's likely hard to find a proper fix. Would it possible to have a fix of miinimum behavior change (most notable the synchronous cancel) that could largely reduce the likelihood of race?

[Abacn]

These tests may be useful, however, a generic idea is to keep test and main source balanced and test concise, as nowadays it's much easier to write boilerplate tests than before

[Abacn]

Please add a few comments noting why we need to manage active sessions ourselves now.

[tkaymak]

Will do: getOrCreate adopts an existing session, a pipeline must not stop one it did not create, and the next pipeline needs the previous one's session gone to get its own config.

[tkaymak]

Agreed on both, the synchronous cancel and the reference count go. Before I push the rework, the shape I have in mind:

1. cancel() asynchronous again: stop the evaluation context, cancel the pipeline's job group, return CANCELLED. No join. waitUntilFinish() after a cancel reports CANCELLED once the execution thread ends.
2. The session stop moves onto the execution thread, in its finally after evaluation, so it never runs under a live thread, and it applies only to sessions the runner created. A session it adopted through getOrCreate (a notebook for example) is left alone. No more counting.
3. The race you describe is between the next pipeline's getOrCreate and the previous thread's stop. Spark's getOrCreate adopts a context that is mid stop, and creating a new one while the old is not fully stopped throws. So acquire waits, bounded, for a session the runner created for a still running pipeline to be stopped and then creates a fresh one, which keeps the previous behavior of a fresh session with the pipeline's own conf. Parallel pipelines in one JVM stay a Spark limitation and will be documented.

Tests down to two unit plus two live cases, plus one that starts a pipeline right after a cancel and asserts it got a new SparkContext, which is the CI failure case. A short rationale block goes into SparkSessionFactory as you asked.

Does that match what you have in mind?

[Abacn]

So acquire waits, bounded ... to be stopped and then creates a fresh one ... Parallel pipelines in one JVM stay a Spark limitation and will be documented.

Disallow more than one pipeline submitting at the same time in JVM would be another regression

I feel the agent (if used) tend to add locks is a tendency, as it is a generic (lazy) fix for concurrent issues. We should define what would be the expectation. Considering current behavior

1. what part, currently working, need to be preserved:
   a. non-blocking job cancellation
   b. non-blocking job submission
2. what part currently working, good to be preserved:

• when jobs submitted and run in sequence, every job configuration is honored (because previous one closed session so that this one starts a new)

3. what part causing issues, and, given the limitation of Spark API,

• what part need to be fixed
  a. pipelineResult.cancel() doesn't cancel the job. For batch it's less severe as long as job stops itself eventually. As we are adding streaming capability, this needs to be fixed.
  b. Session get cancelled crashing another job sharing session and is translating
  This caused flaky tests and fix it is found to be challenging under the constraint of need to be preserved
• what part is tolerable
  c. latter started job not honoring Conf due to reusing session (race exists in current master as well as the fix).

After isolating each issue and requirements, here is my proposal:

1. Make cancellation actually cancel the job (fix 3a). This part's fix is already well formed in this PR
2. parallel job submission does not crash (fixing 3b) under constraint of 1 and ideally 2.
   To do this we may still need to track active jobs for sessions in some way, but instead of introducing blocking waits, acknowledging possible conflict in the case of parallel job submission (it's current behavior so no regression), but make best effort to keep "when jobs submitted and run in sequence, every job configuration is honored" (also current behavior)

Based on the observation in #40101, it's translating on a stopped session (or stopped midway) causing crash. If there are two jobs both running, and the first one ends and stopped its session, the second one can still run and ends. It's only job management get affected (cancel request won't reach job?). If this is true, we can make use of this fact to separate the effort of fixing crash and a proper long term fix, thus make life easier

To fix crash, here is some idea

• We only need to track if there is a pipeline started translating and not yet submitted to Spark. If so the session is not safe to stop
• We lazily stop previous session on subsequent job run when acquiring for a session. If previous session is safe to stop, stop it before call getOrCreate; otherwise just return the same session. This wouldn't eliminate race bug, but it's still in align with current (incorrect but non-crashing) behavior that overlapping job submission get same session with individual Conf not honored.

By doing this we still need only one synchronized acquire, and mostly keep and simplified the current PR's structure

[Abacn]

Also async job.cancel() shouldn't set the job status to CANCELLED immediately per spec. It is subsequent job.status() will query for status and set it to cancelled if so; or subsequent job.waitUntilFinish() blocks until job actually cancelled on (mini)cluster. This can be follow ups.