[SPARK-54217] Synchronize PythonRunner's MonitorThread kill decision

[st-tran]

What changes were proposed in this pull request?

This diff addresses the synchronization issue described in SPARK-54217 by respecting the existing releasedOrClosed AtomicBoolean in the PythonRunner's kill codepath, which is currently only used in the "released" codepath - not the "closed" one. In doing so, we avoid erroneously destroying a still-healthy Python worker; in the current state, it will be destroyed & a new one will be created.

Jira ticket description follows...

PythonWorkerFactory in daemon mode will allow for worker reuse, where possible, as long as the worker successfully completed its last-assigned task (via releasePythonWorker). The worker will be released into the idle queue to be picked up by the next createPythonWorker call.

However, there is a race condition that can result in a released worker in the PythonWorkerFactory idle queue getting killed. i.e. the PythonRunner lacks synchronization between:

1. the main task thread's decision to release its associated Python worker (when work is complete), and
2. the MonitorThread's decision to kill the associated Python worker (when requested by the executor, e.g. speculative execution where another attempt succeeds).

So, the following sequence of events is possible:

1. PythonRunner is running
2. The Python worker finishes its work and writes END_OF_STREAM to signal back to PythonRunner's main task thread that it is done
3. PythonRunner's main task thread receives this instruction and releases the worker for reuse
4. For a separate reason: Executor decides to kill this task (e.g. speculative execution)
5. PythonRunner's MonitorThread receives this instruction and kills the already-relinquished PythonWorker

So the next task that pulls this Python worker from the idle pool will have a dead Python worker.

Why are the changes needed?

In the latest Spark release, this change is NOT critical, however, it avoids the unnecessary killing of a still-healthy Python worker which results in another one being created.

• Prior to SPARK-47565, this would result in a crash (failing the task) as we would reuse this now-closed worker.
• With SPARK-47565, this is less of an issue, as we check that the worker is alive before we use it.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

It is not possible (or very hard) to unit test this change.

However, I've created a minimal repro of the issue. The error occurs when running without this change, and goes away with this change.

• Without this change: details in SPARK-54217; we see the following log line indicating that we are attempting to reuse killed Python workers:
  • PythonWorker(java.nio.channels.SocketChannel[closed]) process from idle queue is dead, discarding.
• With this change:
  • First, cherry-pick this PR onto this fix.
  • Then, rerun the job with this new build and we see that the errors go away:

sttran@Stevens-MacBook-Pro ~/sttran-spark/work
 % ls -lathr
total 0
drwxr-xr-x@ 18 sttran  staff   576B Nov  5 21:22 app-20251105212207-0012  # Without the fix (has errors shown in https://github.com/st-tran/spark/pull/1)
drwxr-xr-x@ 53 sttran  staff   1.7K Nov  5 22:22 ..
drwxr-xr-x@  4 sttran  staff   128B Nov  5 22:29 .
drwxr-xr-x@ 18 sttran  staff   576B Nov  5 22:29 app-20251105222956-0013  # With the fix (no errors per search below)
sttran@Stevens-MacBook-Pro ~/sttran-spark/work
 % grep -nri discard *0013
sttran@Stevens-MacBook-Pro ~/sttran-spark/work
 %

Was this patch authored or co-authored using generative AI tooling?

No.