viirya commented on code in PR #7181: URL: https://github.com/apache/datafusion-comet/pull/7181#discussion_r4418695204
########## 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: Thanks. I'd rather keep the scope here, because a general name would sit on top of behavior that is specific to Python UDFs, and plan cancellation would need each of these to be different: - The watch runs only for plans that contain an Arrow UDF (`hasPythonUdf`). #6295 needs every plan watched. - It waits `spark.python.task.killTimeout` before interrupting, to give a call the grace period Spark's `MonitorThread` gives a worker. Cancelling a plan would act at once. - `interruptOnKill.enabled` exists because an asynchronous exception can leave a lock held in the shared interpreter, which does not apply to stopping a native stream. What does carry over is the plumbing: a token unique in the process, bound to the plan before the first `executePlan`, and a watch that keeps running after the task completes until its plans are released. I've left a pointer on #6295. ########## 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: Good catch, both fixed in 53b6a5303. The `finally` now cancels the group as well, which does nothing when the job has already ended. The `Future` clears the job group from its pooled thread once the job returns. -- 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]
