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

[tkaymak]

Fixes #40101. Found in review of #40090. Shape agreed with @Abacn in the thread below.

Bug

cancel() interrupted the execution thread with Future.cancel(true) and then stopped the SparkSession from the caller thread. An interrupt does not cancel a Spark job, so with useActiveSparkSession=true a batch pipeline was never cancelled. The session stop landed under a thread that was still translating, which is the NoClassDefFoundError cascade seen in the Spark Versions PreCommit on #40090.

Fix, shared code compiled for Spark 3 and Spark 4

• The execution thread runs the pipeline under a job group. cancel() stops the evaluation, cancels the job group and returns. It does not wait. waitUntilFinish() reports CANCELLED once an execution that was cancelled ends, normally or with an exception, and logs the exception at WARN. Batch EvaluationContext.stop() ends the leaf loop, StreamingEvaluationContext reuses that flag and keeps stopping its queries.
• The job group is set with interruptOnCancel=true, which is what Spark's own StreamExecution uses for streaming queries. Without it a cancel waits for every running task to finish its partition. The legacy runner stops the whole SparkContext on cancel, which is harsher.
• SparkSessionFactory.acquire and release replace getOrCreateSession. Sessions the runner creates are counted per pipeline, the last pipeline to release stops the session, on the execution thread after evaluation. A session the runner did not create is never stopped. Pipelines that run one after another therefore each get a fresh session with their own configuration, as before. A pipeline that starts while another one still holds the session shares it and logs that its configuration is not applied, which is today's behavior. Nothing waits for another pipeline. The lock covers the stop itself, so a pipeline starting during a stop creates a new session rather than adopting a stopping one. getOrCreateSession had no other callers and is removed.
• No blocking cancel, no blocking submission, no timeouts.

Tests

• SparkStructuredStreamingPipelineResultTest: the job cancel hook runs once, cancel() returns before the execution ended and waitUntilFinish() then reports CANCELLED.
• StructuredStreamingPipelineStateTest: a running batch job is cancelled and its session is gone after waitUntilFinish(), a session created outside the runner survives, and a pipeline started right after cancel() without waiting finishes DONE and leaves no session behind, the shape of the CI failure. The legacy running, cancelled and timeout cases join after cancel(), the session stop runs on the execution thread and a test that walks away leaks it into the next test class.
• Trigger files touched for the SparkStructuredStreaming and Spark4 ValidatesRunner PostCommits.

Follow up: #40120, the state after an asynchronous cancel() should be observed from the execution rather than set by cancel() itself.

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.

[tkaymak]

Thanks, that framing is clearer than mine and I will follow it. Shape of the rework:

• Job group cancellation stays as it is.
• cancel() requests, marks CANCELLED and returns. No join. waitUntilFinish() reports CANCELLED once the execution ended after a cancel request.
• SparkSessionFactory keeps a count of pipelines per session the runner created, the last one to release stops it, on the execution thread after evaluation.
  • A pipeline that starts while another still holds the session shares it and logs that its conf is not applied, which is today's behavior.
  • Sessions the runner did not create are never stopped.
  • Only the map updates are synchronized, nothing waits.
• Tests: Two unit and three live cases, one of them starts a pipeline right after a cancel without waiting

Your point on the state after an asynchronous cancel() is #40120, follow up.

[Abacn]

Thank you!