comphead commented on code in PR #7181: URL: https://github.com/apache/datafusion-comet/pull/7181#discussion_r4407318798
########## spark/src/main/scala/org/apache/comet/PythonUdfMonitor.scala: ########## @@ -0,0 +1,111 @@ +/* + * 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, TimeUnit} +import java.util.concurrent.atomic.{AtomicInteger, AtomicLong} + +import org.apache.spark.{SparkEnv, TaskContext} +import org.apache.spark.internal.Logging + +/** + * Stops 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. Once the task is killed, its later calls fail + * without entering Python. A call still running `spark.python.task.killTimeout` after the kill + * gets an exception raised in it, once so that its cleanup can run, unless that is disabled. + */ +object PythonUdfMonitor extends Logging { Review Comment: #6295 is the same problem for every native plan: a killed task stays in `executePlan` until the plan yields its next batch, and rich7420 reported a prototype there with a native cancellation token and a JVM watcher. This PR adds most of what that needs, a per-task kill poller, a process-unique task token and a JNI call that binds the token to a plan before the first `executePlan`. All of it is named and gated on Python UDFs, though (`hasPythonUdf`, `setPythonUdfTask`, `killPythonUdfs`, `python_udf_task`). `CelebornNativeShuffleDestination.watchForCancellation` already polls the same signal for another consumer. Could the token and the watcher be named after the task, so #6295 can add plan cancellation as a second consumer instead of a second watcher and a second per-plan binding? If you would rather keep the scope, a pointer on #6295 to this mechanism would help whoever picks it up. ########## spark/src/test/spark-4.1+/org/apache/spark/sql/comet/CometArrowPythonUdfSuite.scala: ########## @@ -583,6 +590,101 @@ class CometArrowPythonUdfSuite extends CometTestBase { } } + /** + * Runs `job` in a job group, waits until `started` exists, cancels the group, and checks that + * the task ends as killed. + */ + private def assertCancelKillsTask(started: File)(job: => Any): Unit = { + val killed = new CountDownLatch(1) + val listener = new SparkListener { + override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = taskEnd.reason match { + case _: TaskKilled => killed.countDown() + case _ => + } + } + spark.sparkContext.addSparkListener(listener) + try { + val group = "comet-arrow-udf-kill" + val running = Future { + spark.sparkContext.setJobGroup(group, "native Arrow UDF loop", interruptOnCancel = true) + job + }(ExecutionContext.global) + val deadline = 60.seconds.fromNow + while (!started.exists()) { + assert(deadline.hasTimeLeft(), "the UDF did not start") Review Comment: If the UDF never creates `started`, this assertion throws before `cancelJobGroup` runs, and the looping UDF keeps its task alive for the rest of the suite, holding one of the `local[5]` cores. Could the `finally` cancel the group as well? The `Future` also leaves the job group and `interruptOnCancel` properties on an `ExecutionContext.global` thread for whatever runs there next, which `CometTaskMetricsSuite` avoids with `clearJobGroup()`. -- 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]
