sunchao commented on code in PR #7181:
URL: https://github.com/apache/datafusion-comet/pull/7181#discussion_r4363245616
##########
native/core/src/execution/python_udf.rs:
##########
@@ -29,14 +29,164 @@ use pyo3::ffi::Py_uintptr_t;
use pyo3::prelude::*;
use pyo3::sync::MutexExt;
use pyo3::types::{PyBytes, PyList, PyTuple};
-use std::sync::{Mutex, PoisonError};
+use pyo3::PyTypeInfo;
+use std::collections::HashMap;
+use std::os::raw::{c_long, c_ulong};
+use std::sync::{Arc, LazyLock, Mutex, MutexGuard, PoisonError};
// Tasks create their callables concurrently. Concurrent first imports of a
package
// can hand a thread a partially initialized module (observed with PyArrow on
Python
// 3.10), so callables are unpickled one at a time. Take it with
`lock_py_attached`,
// which waits for the lock without holding the GIL.
static IMPORT_LOCK: Mutex<()> = Mutex::new(());
+pyo3::create_exception!(
+ comet,
+ TaskKilled,
+ pyo3::exceptions::PyBaseException,
+ "Raised in a running Arrow UDF when Spark kills its task."
+);
+
+unsafe extern "C" {
+ // Part of the stable Python C API, but not declared by pyo3-ffi.
+ fn PyThread_get_thread_ident() -> c_ulong;
+}
+
+/// The native plans of one Spark task attempt that hold Arrow UDFs.
+#[derive(Default)]
+struct TaskCalls {
+ plans: usize,
+ killed: bool,
+ /// Python thread ids of the calls running now, one entry per call.
+ threads: Vec<c_ulong>,
+}
+
+// Holders of this lock never wait for the GIL, so a thread attached to Python
may take it.
+static TASKS: LazyLock<Mutex<HashMap<i64, TaskCalls>>> =
LazyLock::new(Default::default);
+
+fn tasks() -> MutexGuard<'static, HashMap<i64, TaskCalls>> {
+ TASKS.lock().unwrap_or_else(PoisonError::into_inner)
+}
+
+/// Registers a native plan's Arrow UDFs under their Spark task attempt, so
that
+/// [`interrupt_task`] reaches their calls. The task's entry is removed when
its last
+/// registration is dropped, which happens when the plan is released.
+pub struct TaskRegistration {
+ task_attempt_id: i64,
+}
+
+impl TaskRegistration {
+ pub fn new(task_attempt_id: i64) -> Self {
+ tasks().entry(task_attempt_id).or_default().plans += 1;
+ Self { task_attempt_id }
+ }
+
+ fn is_killed(&self) -> bool {
+ tasks()
+ .get(&self.task_attempt_id)
+ .is_some_and(|task| task.killed)
+ }
+
+ fn killed_error(&self) -> ArrowError {
+ ArrowError::ComputeError(format!(
+ "Arrow UDF stopped because task {} was killed",
+ self.task_attempt_id
+ ))
+ }
+}
+
+impl Drop for TaskRegistration {
+ fn drop(&mut self) {
+ let mut tasks = tasks();
+ if let Some(task) = tasks.get_mut(&self.task_attempt_id) {
+ task.plans -= 1;
+ if task.plans == 0 {
+ tasks.remove(&self.task_attempt_id);
+ }
+ }
+ }
+}
+
+/// A call to the user's callable, recorded so that [`interrupt_task`] can
raise in its thread.
+struct RunningCall<'a> {
+ task: &'a TaskRegistration,
+ thread: c_ulong,
+}
+
+impl<'a> RunningCall<'a> {
+ /// Fails without recording the call when the task has already been killed.
+ fn enter(_py: Python<'_>, task: &'a TaskRegistration) -> Result<Self> {
+ // SAFETY: the caller is attached to the interpreter.
+ let thread = unsafe { PyThread_get_thread_ident() };
+ let mut tasks = tasks();
+ let calls = tasks.entry(task.task_attempt_id).or_default();
+ if calls.killed {
+ return Err(task.killed_error());
+ }
+ calls.threads.push(thread);
+ Ok(Self { task, thread })
+ }
+}
+
+impl Drop for RunningCall<'_> {
+ fn drop(&mut self) {
+ let mut tasks = tasks();
+ let Some(calls) = tasks.get_mut(&self.task.task_attempt_id) else {
+ return;
+ };
+ if let Some(index) = calls.threads.iter().position(|t| *t ==
self.thread) {
+ calls.threads.swap_remove(index);
+ }
+ if calls.killed {
+ // An interrupt that arrived as the callable returned is still
pending on this
+ // thread. Nothing can set another one now that the call is
unrecorded, so clear
+ // it rather than let it fire in later Python code on this thread.
+ // SAFETY: the call is made while attached, on the thread it names.
+ unsafe {
+ pyo3::ffi::PyThreadState_SetAsyncExc(self.thread as c_long,
std::ptr::null_mut())
+ };
+ }
+ }
+}
+
+/// Marks a Spark task attempt as killed and raises [`TaskKilled`] in its
running Arrow UDF
+/// calls. Later calls of the task fail without entering Python. Returns the
number of calls
+/// interrupted. A task with no registered plan is not recorded, so the caller
repeats this
+/// until the task completes.
+///
+/// The exception takes effect when the thread next runs Python bytecode: a
call blocked in
+/// native code that released the GIL is interrupted only when that code
returns.
+pub fn interrupt_task(task_attempt_id: i64) -> usize {
+ let running = match tasks().get_mut(&task_attempt_id) {
+ Some(task) => {
+ task.killed = true;
+ !task.threads.is_empty()
+ }
+ None => false,
+ };
+ if !running {
+ return 0;
+ }
+ // Calls cannot start after the flag is set, so only calls recorded before
it remain.
+ Python::attach(|py| {
+ let exception = TaskKilled::type_object(py);
+ let tasks = tasks();
+ let Some(task) = tasks.get(&task_attempt_id) else {
+ return 0;
+ };
+ task.threads
+ .iter()
+ .map(|thread| {
+ // SAFETY: attached to the interpreter, and `exception` is a
live type object.
+ let set = unsafe {
+ pyo3::ffi::PyThreadState_SetAsyncExc(*thread as c_long,
exception.as_ptr())
Review Comment:
[P2] [P2] Avoid reinjecting cancellation while a call is unwinding
A call remains in `task.threads` until its entire Python invocation returns,
including its `finally` blocks. Because the monitor calls this method every 100
ms, subsequent `TaskKilled` exceptions also interrupt cleanup that is already
handling the first one. For example, a UDF holding a module-level lock with
`finally: flush(); lock.release()` leaves that lock permanently held when the
cleanup takes longer than a polling interval. Other tasks share that module and
can then block indefinitely. A single interrupt lets the same cleanup finish.
Could cancellation delivery be recorded per active call so subsequent polls do
not reinject into it, while retaining the killed flag that rejects later calls?
Evidence: A bounded Rust harness copied this head's TaskRegistration,
RunningCall, and interrupt_task implementation unchanged. With PyO3 0.28.3 and
CPython 3.12.13, the callable acquired a shared threading.Lock, entered a loop,
and used `finally: cleaning.set(); time.sleep(0.4); lock.release()`. One
interrupt produced TaskKilled and left the lock available. A second interrupt
100 ms later, after confirming cleanup had begun, also produced TaskKilled but
left the lock unavailable after the call ended. Both calls were correctly
unregistered, so the retained lock outlived cancellation. Command: `cargo test
--offline -- --nocapture --test-threads=1` in
`/tmp/comet7181-908ae-root-probe-xk1ssxaq`. Output: `repeat=false, shared lock
available after cancellation=true` and `repeat=true, shared lock available
after cancellation=false`. The production monitor repeats delivery at this same
interval.
##########
spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala:
##########
@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet
+
+import java.util.concurrent.ConcurrentHashMap
+
+import org.apache.spark.TaskContext
+import org.apache.spark.internal.Logging
+
+/**
+ * Interrupts the native Python UDF calls of killed tasks. Spark's Python
runner has a
+ * `MonitorThread` that destroys the worker of a killed task. A native UDF
runs in the embedded
+ * interpreter on the task thread instead, where neither a Java interrupt nor
a stopped input
+ * reaches it. So one daemon thread per task polls its kill state and, while
the task is killed
+ * but still running, raises an exception in the task's running Python calls.
Calls that start
+ * after the first interrupt fail without entering Python.
+ */
+object PythonUdfMonitor extends Logging {
+
+ // Raising the exception is cheap and lets the call unwind on its own, so
unlike Spark's
+ // monitor this one does not wait `spark.python.task.killTimeout` before
acting. Repeating it
+ // reaches calls that start after an earlier attempt and calls that caught
the exception.
+ private val PollIntervalMs = 100L
+
+ private val watchedTasks = ConcurrentHashMap.newKeySet[java.lang.Long]()
+
+ /** Whether the native library can run Python UDFs, and so may need its
calls interrupted. */
+ lazy val supported: Boolean = NativeBase.supportsPythonUdf()
+
+ /** Starts watching the task unless it is watched already. The watch ends
with the task. */
+ def watch(context: TaskContext): Unit = {
+ val taskAttemptId = context.taskAttemptId()
+ if (watchedTasks.add(taskAttemptId)) {
+ val thread = new Thread(s"Comet Python UDF monitor for task
$taskAttemptId") {
+ override def run(): Unit = {
+ try {
+ while (!context.isCompleted()) {
+ if (context.isInterrupted() &&
NativeBase.interruptPythonUdfs(taskAttemptId) > 0) {
+ logInfo(s"Interrupted Python UDF calls of killed task
$taskAttemptId")
+ }
+ Thread.sleep(PollIntervalMs)
+ }
+ } finally {
+ watchedTasks.remove(taskAttemptId)
Review Comment:
[P2] [P2] Explicitly discard the removal result for strict Scala compilation
`watchedTasks.remove(taskAttemptId)` returns a Boolean, which is discarded
by this `finally` block. The supported Spark 3.5 strict-warnings profile treats
that discard as a compilation error, so the new shared source fails the PR's
strict Scala job even though native Arrow UDF execution is disabled on that
Spark version. Explicitly discard the result, for example with `val _ =
watchedTasks.remove(taskAttemptId)`.
Evidence: Exact-head run 40003235556, job 120591854928, executes `./mvnw -B
test-compile -Pspark-3.5,delta -Pstrict-warnings -DskipTests`. Its compiler
output reports `PythonUdfMonitor.scala:61: discarded non-Unit value`, then
exits with compilation failure. This diagnostic points to the newly added file.
Log:
https://github.com/apache/datafusion-comet/actions/runs/40003235556/job/120591854928.
--
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]
