viirya opened a new pull request, #7181:
URL: https://github.com/apache/datafusion-comet/pull/7181
## Which issue does this PR close?
Closes #6811.
## Rationale for this change
The native scalar Arrow UDF path from #6130 runs the Python callable on the
Spark task thread, inside `tokio::task::block_in_place`. A Spark kill does not
reach it: the Java interrupt is ignored while the thread is in Python, and
Spark's `PythonRunner` `MonitorThread`, which destroys the Python worker of a
killed task, does not apply. A cancelled job or the losing copy of a
speculative task therefore kept its executor core, and kept competing for the
embedded interpreter's GIL, until the UDF returned on its own. A stuck UDF
never released them.
## What changes are included in this PR?
Native side (`python_udf.rs`, `arrow_python_udf.rs`):
- A global registry maps a task attempt id to the Python thread ids
currently running one of its UDF calls. `ArrowPythonUdfExec::execute` registers
the plan under its task attempt (now passed from `createPlan` through
`ExecutionContext` and the planner). Each call records its thread right before
entering the callable and removes it right after, on success or error. The
entry is removed when the task's last plan drops its registration, i.e. when
the plan is released.
- `interrupt_task` marks the task as killed and raises `comet.TaskKilled` in
its recorded threads with `PyThreadState_SetAsyncExc`. `TaskKilled` subclasses
`BaseException`, so a user `except Exception` block does not swallow it. Once a
task is killed, later calls fail without entering Python, and a call that
returned just as the interrupt arrived clears the pending exception on its
thread and fails with the same error.
- Lock ordering: holders of the registry lock never wait for the GIL.
`interrupt_task` sets the flag without the GIL and attaches only when a call is
running, so it does not touch Python for a task that is not in a UDF call.
- New JNI functions: `NativeBase.interruptPythonUdfs(taskAttemptId)` (a
no-op that returns 0 in builds without `python-udf`) and
`Native.hasPythonUdf(plan)`.
JVM side:
- `CometExecIterator` asks `hasPythonUdf` only when
`NativeBase.supportsPythonUdf()`. For a plan that contains an Arrow UDF, it
starts `PythonUdfMonitor`, one daemon thread per task like Spark's
`MonitorThread`, which polls `TaskContext.isInterrupted` every 100 ms until the
task completes and calls the interrupt while the task is killed. The interrupt
is repeated, so it also reaches calls that start after an earlier attempt and
calls that caught the exception. Unlike Spark's monitor, it does not wait
`spark.python.task.killTimeout`, because the interrupt is cooperative and lets
the call unwind instead of killing a process. Plans without a native Arrow UDF,
and builds without the feature, start no thread.
- The UDF then fails with a non-fatal error. Spark's `Executor` reports a
non-fatal failure of a task whose `reasonIfKilled` is set as `TaskKilled`; this
is the same code in Spark 4.1.3 and 4.2.0 (`Executor.scala`, `case _:
InterruptedException | NonFatal(_) if task != null &&
task.reasonIfKilled.isDefined`).
Known limitation, now documented in the PyArrow UDF guide in place of the
old "killing a task does not interrupt" note: an asynchronous exception takes
effect only when the thread next runs Python bytecode. A call blocked in native
code, such as a long PyArrow kernel that released the GIL, keeps the task's
core until that code returns.
## How are these changes tested?
- Rust: `interrupts_running_call_of_killed_task` runs a `while True` UDF
that wraps its body in `try`/`except Exception`, interrupts it from another
thread, and checks that the call fails with `TaskKilled`, that its thread is
unregistered, that a later call of the killed task fails without entering
Python, and that the registry entry is removed when the registration is
dropped. With the `PyThreadState_SetAsyncExc` call stubbed out, the test fails
("the interrupted call did not return").
- Scala: `CometArrowPythonUdfSuite` "killing a task interrupts its running
native Arrow UDF" cancels the job group of a query whose native UDF loops
forever, and checks that the task ends as `TaskKilled` and the job fails within
30 seconds. With the monitor disabled, the task is never killed and the test
fails.
- I ran the Rust tests and the full `CometArrowPythonUdfSuite` with
`-Pspark-4.1` on Python 3.10 and 3.13.
This touches the planner and a native operator, so it would be good to run
it with the `run-pyarrow-udf-tests` and `run-spark-4.1-tests` labels.
This pull request and its description were written by Isaac.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
