This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 1d9f7989802d [SPARK-59822][UDF] Make the external UDF worker
dispatcher factory configurable
1d9f7989802d is described below
commit 1d9f7989802d644810f95741d66cd3995147531d
Author: Yossi Eliaz <[email protected]>
AuthorDate: Wed Sep 30 07:30:47 2026 +0900
[SPARK-59822][UDF] Make the external UDF worker dispatcher factory
configurable
### What changes were proposed in this pull request?
Add the optional `spark.udf.worker.dispatcherFactory` setting, tracked by
[SPARK-59822](https://issues.apache.org/jira/browse/SPARK-59822). When an
external UDF first needs a dispatcher, `SparkEnv` loads the configured
`UDFDispatcherFactory` using Spark's class loader and supports constructors
taking `(SparkConf, Boolean)`, `SparkConf`, or no arguments. Factory validation
does not initialize an incompatible class. Lazy manager creation and
`SparkEnv.stop()` share a lifecycle gate, so [...]
A custom factory is a deployment-mode hook, not another client protocol. It
must consume the standard `UDFWorkerSpecification` and use the unified UDF
protocol. Spark should also provide a generic default dispatcher so most
clients only describe their worker; that default and the shading/runtime wiring
for the gRPC transport are separate work. Until then, an unset factory still
fails on first external UDF dispatcher use, with an error naming the setting.
### Why are the changes needed?
The factory currently installed by `SparkEnv` throws for every external
UDF. This hook lets deployments supply an implementation without adding a
compile-time gRPC dependency to `core`. The lifecycle coordination prevents an
in-flight task from publishing a manager after shutdown has passed its close
check.
### Does this PR introduce _any_ user-facing change?
Yes. It adds an experimental operator-level setting for selecting an
external UDF worker factory. Existing deployments behave the same when the
setting is unset; the first attempt to create an external UDF dispatcher still
fails, now with a message naming the setting.
### How was this patch tested?
`SparkEnvUDFDispatcherSuite` covers the unset setting, no-argument and
SparkConf/driver-aware constructors, an incompatible class (including its
initializer), an unknown class, and manager creation racing with shutdown.
`PlanExternalUDFsSuite` verifies the updated error at the SQL execution
boundary. CI for the latest follow-up commit is pending; the local sparse
checkout lacks the Spark build wrapper and Scala build tools.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5 (original patch); Codex GPT-6 (version
correction, lifecycle fix, tests, and PR description).
Closes #58859 from zozo123/SPARK-udf-dispatcher-factory.
Authored-by: Yossi Eliaz <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
.../src/main/scala/org/apache/spark/SparkEnv.scala | 92 +++++++---
.../org/apache/spark/internal/config/UDF.scala | 34 ++++
.../apache/spark/SparkEnvUDFDispatcherSuite.scala | 196 +++++++++++++++++++++
.../externalUDF/PlanExternalUDFsSuite.scala | 5 +-
4 files changed, 300 insertions(+), 27 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/SparkEnv.scala
b/core/src/main/scala/org/apache/spark/SparkEnv.scala
index f8e2ae3abcca..04b75c4bcee3 100644
--- a/core/src/main/scala/org/apache/spark/SparkEnv.scala
+++ b/core/src/main/scala/org/apache/spark/SparkEnv.scala
@@ -51,7 +51,8 @@ import org.apache.spark.shuffle.{ShuffleBlockResolver,
ShuffleManager}
import org.apache.spark.shuffle.streaming.MultiShuffleManager
import org.apache.spark.storage._
import org.apache.spark.udf.worker.UDFWorkerSpecification
-import org.apache.spark.udf.worker.core.{UDFDispatcherFactory,
UDFDispatcherManager, WorkerDispatcher}
+import org.apache.spark.udf.worker.core.{UDFDispatcherFactory,
UDFDispatcherManager,
+ WorkerDispatcher, WorkerLogger}
import org.apache.spark.util.{RpcUtils, Utils}
import org.apache.spark.util.ArrayImplicits._
import org.apache.spark.util.SparkUDFWorkerLogger
@@ -222,20 +223,12 @@ class SparkEnv (
* using the new UDF framework proposed in SPARK-55278.
* Initialized on first use via [[getExternalUDFDispatcher]].
*/
- @volatile private var udfDispatcherManager: Option[UDFDispatcherManager] =
None
+ // Guarded by this SparkEnv's monitor, together with isStopped.
+ private var udfDispatcherManager: Option[UDFDispatcherManager] = None
private def createUDFDispatcherManager(): UDFDispatcherManager = {
- val factory = new UDFDispatcherFactory {
- override def createDispatcher(
- workerSpec: UDFWorkerSpecification,
- logger: org.apache.spark.udf.worker.core.WorkerLogger
- ): WorkerDispatcher = {
- // TODO [SPARK-55278]: Wire in the correct dispatcher factory
- throw new UnsupportedOperationException(
- "No UDF dispatcher factory configured. " +
- "Set up a concrete factory for SPARK-55278.")
- }
- }
+ val factory = SparkEnv.resolveUDFDispatcherFactory(
+ conf, SparkContext.isDriver(executorId))
new UDFDispatcherManager(factory, new SparkUDFWorkerLogger())
}
@@ -246,15 +239,14 @@ class SparkEnv (
*/
private[spark] def getExternalUDFDispatcher(
workerSpec: UDFWorkerSpecification): WorkerDispatcher = {
- val manager : UDFDispatcherManager = udfDispatcherManager.getOrElse {
- synchronized {
- // Get or Else synchronized to protect
- // against concurrent creation requests.
- udfDispatcherManager.getOrElse {
- val created = createUDFDispatcherManager()
- udfDispatcherManager = Some(created)
- created
- }
+ val manager = synchronized {
+ if (isStopped) {
+ throw new IllegalStateException("SparkEnv is stopped")
+ }
+ udfDispatcherManager.getOrElse {
+ val created = createUDFDispatcherManager()
+ udfDispatcherManager = Some(created)
+ created
}
}
manager.getDispatcher(workerSpec)
@@ -283,9 +275,15 @@ class SparkEnv (
new AtomicReference[VersionedCredentials]()
private[spark] def stop(): Unit = {
-
- if (!isStopped) {
- isStopped = true
+ val shouldStop = synchronized {
+ if (isStopped) {
+ false
+ } else {
+ isStopped = true
+ true
+ }
+ }
+ if (shouldStop) {
pythonWorkers.values.foreach(_.stop())
udfDispatcherManager.foreach(_.close())
mapOutputTracker.stop()
@@ -527,6 +525,50 @@ object SparkEnv extends Logging {
private[spark] val driverSystemName = "sparkDriver"
private[spark] val executorSystemName = "sparkExecutor"
+ /**
+ * :: Experimental ::
+ * Resolves the [[UDFDispatcherFactory]] used to create dispatchers for
external UDF
+ * workers (SPARK-55278), from `spark.udf.worker.dispatcherFactory`.
+ *
+ * The factory is loaded reflectively rather than named here: `core` depends
on
+ * `udf-worker-proto` and `udf-worker-core` only, so that neither it nor its
consumers
+ * carry a gRPC dependency. Every concrete dispatcher therefore lives in a
module that
+ * core cannot reference at compile time, and the implementation has to be
selected by
+ * configuration.
+ *
+ * When the config is unset the returned factory throws on first use rather
than at
+ * SparkEnv creation, so applications that never run an external UDF are
unaffected.
+ */
+ private[spark] def resolveUDFDispatcherFactory(
+ conf: SparkConf,
+ isDriver: Boolean): UDFDispatcherFactory = {
+ conf.get(UDF.DISPATCHER_FACTORY) match {
+ case Some(className) =>
+ // Validate before constructing so that a misconfigured class is
reported against
+ // the config key, rather than as a ClassCastException from the call
site.
+ val cls = Utils.classForName[AnyRef](className, initialize = false)
+ if (!classOf[UDFDispatcherFactory].isAssignableFrom(cls)) {
+ throw new SparkException(
+ s"${UDF.DISPATCHER_FACTORY.key} is set to $className, which does
not implement " +
+ s"${classOf[UDFDispatcherFactory].getName}.")
+ }
+ Utils.instantiateSerializerOrShuffleManager[UDFDispatcherFactory](
+ className, conf, isDriver)
+
+ case None =>
+ new UDFDispatcherFactory {
+ override def createDispatcher(
+ workerSpec: UDFWorkerSpecification,
+ logger: WorkerLogger): WorkerDispatcher = {
+ throw new UnsupportedOperationException(
+ "No UDF dispatcher factory is configured, so external UDF
workers cannot be " +
+ s"created. Set ${UDF.DISPATCHER_FACTORY.key} to the name of a
" +
+ s"${classOf[UDFDispatcherFactory].getName} implementation.")
+ }
+ }
+ }
+ }
+
def set(e: SparkEnv): Unit = {
env = e
}
diff --git a/core/src/main/scala/org/apache/spark/internal/config/UDF.scala
b/core/src/main/scala/org/apache/spark/internal/config/UDF.scala
new file mode 100644
index 000000000000..af536cd5fb0b
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/internal/config/UDF.scala
@@ -0,0 +1,34 @@
+/*
+ * 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.spark.internal.config
+
+private[spark] object UDF {
+
+ val DISPATCHER_FACTORY = ConfigBuilder("spark.udf.worker.dispatcherFactory")
+ .doc("Name of the class implementing " +
+ "org.apache.spark.udf.worker.core.UDFDispatcherFactory, used to create
the " +
+ "dispatchers that provision and manage external UDF workers
(SPARK-55278). The " +
+ "class must have a no-arg constructor, a constructor taking a SparkConf,
or a " +
+ "constructor taking a SparkConf and a Boolean indicating whether it is
being " +
+ "created on the driver. When unset, no external UDF dispatcher is
available and " +
+ "executing an external UDF fails. This is an operator-level setting: it
selects " +
+ "the runtime that user code is executed by, so it is not intended to be
changed " +
+ "by individual queries.")
+ .version("4.4.0")
+ .stringConf
+ .createOptional
+}
diff --git
a/core/src/test/scala/org/apache/spark/SparkEnvUDFDispatcherSuite.scala
b/core/src/test/scala/org/apache/spark/SparkEnvUDFDispatcherSuite.scala
new file mode 100644
index 000000000000..886bbafbbdc8
--- /dev/null
+++ b/core/src/test/scala/org/apache/spark/SparkEnvUDFDispatcherSuite.scala
@@ -0,0 +1,196 @@
+/*
+ * 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.spark
+
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+
+import scala.concurrent.{ExecutionContext, Future}
+import scala.concurrent.duration._
+import scala.util.Try
+
+import org.mockito.Answers.RETURNS_DEEP_STUBS
+import org.mockito.Mockito.{mock, verify}
+
+import org.apache.spark.broadcast.BroadcastManager
+import org.apache.spark.internal.config.UDF
+import org.apache.spark.metrics.MetricsSystem
+import org.apache.spark.rpc.RpcEnv
+import org.apache.spark.scheduler.OutputCommitCoordinator
+import org.apache.spark.storage.BlockManager
+import org.apache.spark.udf.worker.UDFWorkerSpecification
+import org.apache.spark.udf.worker.core.{UDFDispatcherFactory,
WorkerDispatcher, WorkerLogger}
+import org.apache.spark.util.ThreadUtils
+
+class SparkEnvUDFDispatcherSuite extends SparkFunSuite {
+
+ private def spec = UDFWorkerSpecification.getDefaultInstance
+
+ private def newEnv(conf: SparkConf): SparkEnv = {
+ new SparkEnv(
+ SparkContext.DRIVER_IDENTIFIER,
+ mock(classOf[RpcEnv]),
+ null,
+ null,
+ null,
+ mock(classOf[MapOutputTracker]),
+ mock(classOf[BroadcastManager]),
+ mock(classOf[BlockManager], RETURNS_DEEP_STUBS),
+ null,
+ mock(classOf[MetricsSystem]),
+ mock(classOf[OutputCommitCoordinator]),
+ conf)
+ }
+
+ test("no factory configured: dispatcher creation fails with an actionable
message") {
+ val factory = SparkEnv.resolveUDFDispatcherFactory(new SparkConf(false),
isDriver = true)
+ // Resolution itself must succeed: applications that never run an external
UDF should not
+ // fail at SparkEnv creation just because no dispatcher is configured.
+ val e = intercept[UnsupportedOperationException] {
+ factory.createDispatcher(spec, WorkerLogger.NoOp)
+ }
+ assert(e.getMessage.contains(UDF.DISPATCHER_FACTORY.key))
+ }
+
+ test("configured factory is instantiated and used") {
+ val conf = new SparkConf(false)
+ .set(UDF.DISPATCHER_FACTORY, classOf[TestNoArgDispatcherFactory].getName)
+ val factory = SparkEnv.resolveUDFDispatcherFactory(conf, isDriver = true)
+ assert(factory.isInstanceOf[TestNoArgDispatcherFactory])
+ assert(factory.createDispatcher(spec, WorkerLogger.NoOp) === null)
+ }
+
+ test("factory receives the SparkConf and isDriver when it declares them") {
+ val conf = new SparkConf(false)
+ .set(UDF.DISPATCHER_FACTORY, classOf[TestConfDispatcherFactory].getName)
+ .set("spark.test.marker", "set-by-conf")
+ Seq(true, false).foreach { isDriver =>
+ val factory = SparkEnv.resolveUDFDispatcherFactory(conf, isDriver)
+ .asInstanceOf[TestConfDispatcherFactory]
+ assert(factory.conf.get("spark.test.marker") === "set-by-conf")
+ assert(factory.isDriver === isDriver)
+ }
+ }
+
+ test("a class that is not a UDFDispatcherFactory is rejected against the
config key") {
+ val conf = new SparkConf(false)
+ .set(UDF.DISPATCHER_FACTORY, classOf[NotAFactory].getName)
+ val e = intercept[SparkException] {
+ SparkEnv.resolveUDFDispatcherFactory(conf, isDriver = true)
+ }
+ assert(e.getMessage.contains(UDF.DISPATCHER_FACTORY.key))
+ assert(e.getMessage.contains(classOf[NotAFactory].getName))
+ }
+
+ test("an incompatible factory class is rejected before its initializer
runs") {
+ val key = "spark.test.incompatibleFactoryInitialized"
+ System.clearProperty(key)
+ try {
+ val conf = new SparkConf(false)
+ .set(UDF.DISPATCHER_FACTORY,
"org.apache.spark.NotAFactoryWithInitializer$")
+ intercept[SparkException] {
+ SparkEnv.resolveUDFDispatcherFactory(conf, isDriver = true)
+ }
+ assert(System.getProperty(key) == null)
+ } finally {
+ System.clearProperty(key)
+ }
+ }
+
+ test("stop waits for lazy dispatcher manager creation and closes created
dispatchers") {
+ val factory = TestBlockingDispatcherFactory
+ val conf = new SparkConf(false)
+ .set(UDF.DISPATCHER_FACTORY,
classOf[TestBlockingDispatcherFactory].getName)
+ val env = newEnv(conf)
+ val pool = ThreadUtils.newDaemonFixedThreadPool(2, "udf-dispatcher-test")
+ implicit val executionContext: ExecutionContext =
ExecutionContext.fromExecutor(pool)
+ val creator = Future { Try(env.getExternalUDFDispatcher(spec)) }
+ try {
+ assert(factory.entered.await(10, TimeUnit.SECONDS))
+ val stopStarted = new CountDownLatch(1)
+ val stopFinished = new CountDownLatch(1)
+ val stopper = Future {
+ stopStarted.countDown()
+ env.stop()
+ stopFinished.countDown()
+ }
+ assert(stopStarted.await(10, TimeUnit.SECONDS))
+ assert(!stopFinished.await(100, TimeUnit.MILLISECONDS))
+ factory.release.countDown()
+ val result = ThreadUtils.awaitResult(creator, 30.seconds)
+ ThreadUtils.awaitResult(stopper, 30.seconds)
+ assert(result.isSuccess ||
result.failed.get.isInstanceOf[IllegalStateException])
+ if (factory.dispatcherCreated) {
+ verify(factory.dispatcher).close()
+ }
+ intercept[IllegalStateException] {
+ env.getExternalUDFDispatcher(spec)
+ }
+ } finally {
+ factory.release.countDown()
+ pool.shutdownNow()
+ }
+ }
+
+ test("an unknown class name fails to resolve") {
+ val conf = new SparkConf(false).set(UDF.DISPATCHER_FACTORY,
"not.a.real.Factory")
+ intercept[ClassNotFoundException] {
+ SparkEnv.resolveUDFDispatcherFactory(conf, isDriver = true)
+ }
+ }
+}
+
+// Declared as top-level classes so that they have constructors Spark can find
reflectively.
+
+class TestNoArgDispatcherFactory extends UDFDispatcherFactory {
+ override def createDispatcher(
+ workerSpec: UDFWorkerSpecification,
+ logger: WorkerLogger): WorkerDispatcher = null
+}
+
+class TestConfDispatcherFactory(val conf: SparkConf, val isDriver: Boolean)
+ extends UDFDispatcherFactory {
+ override def createDispatcher(
+ workerSpec: UDFWorkerSpecification,
+ logger: WorkerLogger): WorkerDispatcher = null
+}
+
+class NotAFactory
+
+object NotAFactoryWithInitializer {
+ System.setProperty("spark.test.incompatibleFactoryInitialized", "true")
+}
+
+object TestBlockingDispatcherFactory {
+ val entered = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ val dispatcher: WorkerDispatcher = mock(classOf[WorkerDispatcher])
+ @volatile var dispatcherCreated = false
+}
+
+class TestBlockingDispatcherFactory extends UDFDispatcherFactory {
+ TestBlockingDispatcherFactory.entered.countDown()
+ if (!TestBlockingDispatcherFactory.release.await(30, TimeUnit.SECONDS)) {
+ throw new IllegalStateException("Timed out waiting to create the test
dispatcher factory")
+ }
+
+ override def createDispatcher(
+ workerSpec: UDFWorkerSpecification,
+ logger: WorkerLogger): WorkerDispatcher = {
+ TestBlockingDispatcherFactory.dispatcherCreated = true
+ TestBlockingDispatcherFactory.dispatcher
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFsSuite.scala
index 5161128f1103..39a70100f82a 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFsSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PlanExternalUDFsSuite.scala
@@ -731,8 +731,9 @@ class PlanExternalUDFsSuite extends QueryTest with
SharedSparkSession {
val exception = intercept[SparkException] {
execution.execute().collect()
}
- assert(Utils.exceptionString(exception).contains(
- "No UDF dispatcher factory configured. Set up a concrete factory for
SPARK-55278."))
+ val message = Utils.exceptionString(exception)
+ assert(message.contains("No UDF dispatcher factory is configured"))
+ assert(message.contains("spark.udf.worker.dispatcherFactory"))
}
test("map partitions external UDF execution reports unimplemented before
worker startup") {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
