viirya commented on code in PR #7181:
URL: https://github.com/apache/datafusion-comet/pull/7181#discussion_r4365458903
##########
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:
Good catch, thanks. I reproduced it on Python 3.10 and 3.13 with a lock
released in a `finally` after a 0.4 s sleep: a second injection during the
cleanup leaves the lock held. Fixed in 89da483c1. Each recorded call now
carries an `interrupted` flag, and `interrupt_task` raises only in calls it has
not interrupted yet, so a running call gets `TaskKilled` exactly once. The
killed flag is unchanged, so later calls of the task still fail without
entering Python. The monitor keeps polling, because a task killed before its
plan registers is not recorded, and a plan the task registers later still needs
the flag set.
The trade-off is that a call that catches `TaskKilled` (e.g. `except
BaseException`) and keeps running is no longer interrupted again. The user
guide and the `PythonUdfMonitor` comment now say so. The new
`interrupts_a_call_once_so_its_cleanup_runs` test interrupts a call, waits
until its `finally` cleanup has started, checks that a second `interrupt_task`
raises in nothing, and checks that the cleanup released its lock. It fails with
the previous repeated delivery.
##########
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:
Fixed in 89da483c1 with `val _ = watchedTasks.remove(taskAttemptId)`,
matching `CometExecIterator`. `test-compile -Pspark-3.5,delta
-Pstrict-warnings` now passes. The four `@nowarn` errors in the shuffle
managers came from the same failed compile and are gone with it.
--
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]
