viirya commented on code in PR #7181:
URL: https://github.com/apache/datafusion-comet/pull/7181#discussion_r4369216134
##########
spark/src/main/scala/org/apache/comet/CometExecIterator.scala:
##########
@@ -235,6 +235,11 @@ class CometExecIterator(
CometExecIterator.startMemoryUsageLog()
+ // A native Python UDF ignores the task's kill, so watch the task to
interrupt its calls.
+ if (PythonUdfMonitor.supported && nativeLib.hasPythonUdf(plan)) {
Review Comment:
Right: `addTaskCompletionListener` runs the listener at once on a completed
task, and `close()` releases the plan before the probe read it. Fixed in
3d9fe9ec9. The probe now runs inside the cleanup-protected block that creates
the plan, before the listener is registered. That block also binds the task's
token (see the token thread) and counts the plan for the monitor, so a failure
there releases the plan like any other setup failure.
##########
native/core/src/execution/python_udf.rs:
##########
@@ -29,14 +29,177 @@ 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,
+ /// The calls running now.
+ calls: Vec<Call>,
+}
+
+struct Call {
+ /// Python thread id of the thread running the call.
+ thread: c_ulong,
+ /// Whether [`interrupt_task`] has raised in the call. It raises once, so
that a second
+ /// exception does not abort the cleanup the first one started.
+ interrupted: bool,
+}
+
+// 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.calls.push(Call {
+ thread,
+ interrupted: false,
+ });
+ 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;
+ };
+ let Some(index) = calls.calls.iter().position(|c| c.thread ==
self.thread) else {
+ return;
+ };
+ if calls.calls.swap_remove(index).interrupted {
+ // 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 that it has not interrupted yet. 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.calls.iter().any(|call| !call.interrupted)
+ }
+ 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 mut tasks = tasks();
+ let Some(task) = tasks.get_mut(&task_attempt_id) else {
+ return 0;
+ };
+ let mut interrupted = 0;
+ for call in task.calls.iter_mut().filter(|call| !call.interrupted) {
+ // SAFETY: attached to the interpreter, and `exception` is a live
type object.
+ let set = unsafe {
+ pyo3::ffi::PyThreadState_SetAsyncExc(call.thread as c_long,
exception.as_ptr())
Review Comment:
This one holds and the code cannot rule it out: the exception can fire right
after a C `acquire()` returns and before the `try` whose `finally` releases it,
and a kill that lands while the call waits for the lock makes that the likely
outcome rather than a narrow race. A Python `__enter__` that acquires a lock
has the same gap. Moving `acquire()` into the `try` only trades it for
releasing a lock that was never taken. Spark avoids the problem by destroying
the worker process, which a shared interpreter cannot do.
In 3d9fe9ec9 the monitor reduces and exposes the risk instead:
- A killed task's later calls fail at once without entering Python, which is
always safe.
- A running call gets `TaskKilled` only after it has run for
`spark.python.task.killTimeout` since the kill, like Spark's monitor, so calls
that finish within the timeout are never interrupted.
- `spark.comet.exec.nativeArrowPythonUDF.interruptOnKill.enabled=false`
disables raising in running calls; a killed task then waits for its call to
return. `kill_without_interrupt_lets_the_running_call_finish` covers that mode.
- The user guide describes the lock hazard, including `logging` before
Python 3.13, next to the setting.
Explicit cooperative cancellation would not stop an arbitrary loop, which is
the case #6811 is about, so I kept the forced interrupt as the default with the
opt-out.
##########
spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala:
##########
@@ -0,0 +1,70 @@
+/*
+ * 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,
once per call so
+ * that its cleanup can run. 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 of plans that the task registers after an earlier attempt.
+ 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()) {
Review Comment:
Agreed. `BatchProducer.stop` takes the stream's mutex, which the producer
holds for the whole poll, including the Python call, so `releasePlan` waits for
that call after `isCompleted` is already true. Fixed in 3d9fe9ec9. Each watch
counts the native UDF plans bound to it, `CometExecIterator.close` uncounts its
plan after `releasePlan` returns, and the monitor runs until the task has
completed and no counted plan remains. The new "killing a completed task
interrupts the UDF call its plan release waits for" test reads a two-row
Parquet file one row per batch, takes the first row while the second loops in
the UDF, completes the task, and cancels the job; it fails with the previous
lifetime.
##########
native/core/src/execution/python_udf.rs:
##########
@@ -29,14 +29,177 @@ 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,
+ /// The calls running now.
+ calls: Vec<Call>,
+}
+
+struct Call {
+ /// Python thread id of the thread running the call.
+ thread: c_ulong,
+ /// Whether [`interrupt_task`] has raised in the call. It raises once, so
that a second
+ /// exception does not abort the cleanup the first one started.
+ interrupted: bool,
+}
+
+// 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;
Review Comment:
Agreed, `taskAttemptId` is only unique within one SparkContext. Fixed in
3d9fe9ec9. `PythonUdfMonitor` assigns each watched `TaskContext` a token from a
process-wide counter. `CometExecIterator` binds it to the plan with a new
`Native.setPythonUdfTask` before the first `executePlan`, and the native
registry, `ArrowPythonUdfExec`, and `NativeBase.killPythonUdfs` all use the
token instead of the attempt id.
##########
native/core/src/execution/python_udf.rs:
##########
@@ -297,14 +468,18 @@ impl ArrowPythonUdf {
kwargs.set_item(name, arg).map_err(python_error)?;
}
}
- let result = self
- .callable
- .bind(py)
- .call(
- PyTuple::new(py, positional).map_err(python_error)?,
- Some(&kwargs),
- )
- .map_err(|error| python_traceback_error(py, error))?;
+ let positional = PyTuple::new(py,
positional).map_err(python_error)?;
+ let call = self
+ .task
+ .as_deref()
+ .map(|task| RunningCall::enter(py, task))
+ .transpose()?;
+ let result = self.callable.bind(py).call(positional,
Some(&kwargs));
+ drop(call);
Review Comment:
Fixed in 3d9fe9ec9. The call now stays recorded through the Spark 4.2
`RecordBatch.from_arrays` conversion, and the killed check runs after it, so a
cancelled conversion fails instead of returning a batch.
`interrupts_a_user_defined_array_conversion` interrupts a result whose
`__arrow_array__` loops, and fails if the call is unrecorded before the
conversion.
##########
native/core/src/execution/python_udf.rs:
##########
@@ -29,14 +29,177 @@ 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,
+ /// The calls running now.
+ calls: Vec<Call>,
+}
+
+struct Call {
+ /// Python thread id of the thread running the call.
+ thread: c_ulong,
+ /// Whether [`interrupt_task`] has raised in the call. It raises once, so
that a second
+ /// exception does not abort the cleanup the first one started.
+ interrupted: bool,
+}
+
+// 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.calls.push(Call {
+ thread,
+ interrupted: false,
+ });
+ 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;
+ };
+ let Some(index) = calls.calls.iter().position(|c| c.thread ==
self.thread) else {
+ return;
+ };
+ if calls.calls.swap_remove(index).interrupted {
+ // 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())
Review Comment:
Agreed. `PyThreadState_SetAsyncExc` signals the interpreter even for a NULL
exception, and only a thread that holds an exception clears the flag. Fixed in
3d9fe9ec9. After unregistering an interrupted call and releasing the registry
lock, the thread runs a short loop of Python bytecode, so a still-pending
`TaskKilled` fires and is discarded there, and nothing is ever cleared with
NULL. `consumes_an_interrupt_pending_when_the_call_returns` interrupts a C
callable (`functools.partial(time.sleep, 0.3)`) that returns straight to Rust,
then runs more Python on the same thread state and checks that the interrupt
does not fire there.
##########
native/core/src/execution/python_udf.rs:
##########
@@ -297,14 +468,18 @@ impl ArrowPythonUdf {
kwargs.set_item(name, arg).map_err(python_error)?;
}
}
- let result = self
- .callable
- .bind(py)
- .call(
- PyTuple::new(py, positional).map_err(python_error)?,
- Some(&kwargs),
- )
- .map_err(|error| python_traceback_error(py, error))?;
+ let positional = PyTuple::new(py,
positional).map_err(python_error)?;
+ let call = self
+ .task
+ .as_deref()
+ .map(|task| RunningCall::enter(py, task))
Review Comment:
Fixed in 3d9fe9ec9. `evaluate_named` checks the killed flag before
`Python::attach`, and `RunningCall::enter` keeps its check for a kill that
lands while the batch's arguments are imported.
--
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]
