Skip to content

[SQL][SPARK-57318] Refactor WorkerSession into a state-machine interface #56397

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
haiyangsun-db wants to merge 6 commits into from
Closed

[SQL][SPARK-57318] Refactor WorkerSession into a state-machine interface #56397

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
3e279bc
[SQL] Refactor WorkerSession into a state-machine interface with Ter…
haiyangsun-db Jun 9, 2026
856db47
unit testing state machine.
haiyangsun-db Jun 9, 2026
d4e6759
updates and improvements
haiyangsun-db Jun 9, 2026
abcdd32
fix ci, address comments.
haiyangsun-db Jun 9, 2026
e4bb07f
last style fix.
haiyangsun-db Jun 9, 2026
0d50e03
improvements.
haiyangsun-db Jun 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 30 additions & 11 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExternalUDFExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ trait ExternalUDFExec extends UnaryExecNode {
// ---------------------------------------------------------------------------

protected def externalUdfMetrics: Map[String, SQLMetric] = Map(
// TODO [SPARK-55278]: Emit the correct metrics here
// TODO [SPARK-57324]: Emit the correct metrics here
)

override lazy val metrics: Map[String, SQLMetric] = externalUdfMetrics
Expand All @@ -60,10 +60,12 @@ trait ExternalUDFExec extends UnaryExecNode {

/**
* Creates a [[WorkerSession]] via [[SparkEnv#getExternalUDFDispatcher]].
* Registers session cancellation on task failure and session termination
* on task completion. The provided function receives the session
* and must return the result iterator. The function may use the
* session but MUST NOT cancel or close it.
* Finalizes the session on task completion (which fires on both success and
* failure). [[WorkerSession#close]] is the single finalizer: it fetches the
* `FinishResponse` if processing completed, or cancels anything still in
* flight and waits for the `CancelResponse`. The provided function receives
* the session and must return the result iterator. It may use the session
* but MUST NOT close it.
*/
protected def withUDFWorkerSession(
taskContext: TaskContext,
Expand All @@ -74,12 +76,29 @@ trait ExternalUDFExec extends UnaryExecNode {
workerSpec)
val session = dispatcher.createSession(securityScope)

// Make sure to cancel the session, if the task fails
taskContext.addTaskFailureListener { (_, _) =>
session.cancel()
}

// Make sure to close the session once we are done
// Finalize the session when the task ends. The completion listener fires on
// both success and failure, and close() is the single finalizer that
// resolves to whichever terminator the stream reached:
//
// - Task completed and the result iterator was fully consumed: process()
// sent Finish (input exhausted) and the data drained. close() then
// typically returns the Finished termination -- but a failure during the
// finish/cleanup phase (raised after the data drained, so it reached no one
// through the iterator) is surfaced here too, as Failed / TransportFailed.
// - Task failed, was killed, or stopped before draining (e.g. a downstream
// LIMIT or exception): the stream has not finished, so close() sends a
// Cancel, the worker runs its cleanup, and its CancelResponse is returned
// as a Cancelled termination. An empty Cancel is enough here -- there is
// no extra information to convey to the worker on cancellation -- so we
// rely on close()'s default and pass no cancel thunk.
// - The stream died without a terminator (transport failure / timeout):
// close() returns a best-effort TransportFailed termination rather than
// raising it (a thread interrupt may still propagate); the underlying
// failure has already surfaced through the result iterator.
//
// The returned Termination (per-execution metrics, finish/cancel callback
// result) is not consumed yet -- TODO [SPARK-57324] surface it once metrics
// wiring lands.
taskContext.addTaskCompletionListener[Unit] { _ =>
session.close()
}
Expand Down
84 changes: 36 additions & 48 deletions .../scala/org/apache/spark/sql/execution/externalUDF/PythonUDFWorkerSpecificationSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
package org.apache.spark.sql.execution.externalUDF

import java.io.File
import java.net.{StandardProtocolFamily, UnixDomainSocketAddress}
import java.nio.channels.SocketChannel
import java.nio.charset.StandardCharsets
import java.nio.file.Files

Expand All @@ -26,72 +28,57 @@ import scala.jdk.CollectionConverters._
import org.apache.spark.api.python.SimplePythonFunction
import org.apache.spark.sql.IntegratedUDFTestUtils
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.udf.worker.UDFWorkerSpecification
import org.apache.spark.udf.worker.core.{TestDirectWorkerDispatcher,
UnixSocketWorkerConnection}
import org.apache.spark.udf.worker.{Cancel, UDFWorkerSpecification}
import org.apache.spark.udf.worker.core.{TestDirectWorkerDispatcher, WorkerConnection}

/**
* A test [[UnixSocketWorkerConnection]] that opens a real Unix
* domain socket channel to the worker.
* A [[WorkerConnection]] that opens a real Unix-domain-socket channel to the
* worker at construction and considers itself active while the channel is
* open. Verifies the worker's socket is actually connectable, not merely
* that the file exists.
*/
private class RealSocketConnection(
socketPath: String,
private val channel: java.nio.channels.SocketChannel)
extends UnixSocketWorkerConnection(socketPath) {

private class ConnectingSocketConnection(
private val channel: SocketChannel) extends WorkerConnection {
override def isActive: Boolean = channel.isOpen

override def close(): Unit = {
channel.close()
super.close()
}
override def close(): Unit = channel.close()
}

private object RealSocketConnection {
def connect(socketPath: String): RealSocketConnection = {
val address =
java.net.UnixDomainSocketAddress.of(socketPath)
val channel = java.nio.channels.SocketChannel.open(
java.net.StandardProtocolFamily.UNIX)
channel.connect(address)
new RealSocketConnection(socketPath, channel)
private object ConnectingSocketConnection {
def connect(socketPath: String): ConnectingSocketConnection = {
val channel = SocketChannel.open(StandardProtocolFamily.UNIX)
channel.connect(UnixDomainSocketAddress.of(socketPath))
new ConnectingSocketConnection(channel)
}
}

/**
* Extends [[TestDirectWorkerDispatcher]] to use a real Unix
* domain socket connection instead of just checking file
* existence.
* A [[TestDirectWorkerDispatcher]] whose connection actually connects to the
* worker socket (rather than just checking file existence), so the test
* verifies the spawned Python worker is reachable.
*/
private class RealSocketTestDispatcher(
spec: UDFWorkerSpecification)
private class ConnectingTestDispatcher(spec: UDFWorkerSpecification)
extends TestDirectWorkerDispatcher(spec) {

// Use /tmp to avoid UDS path length limits in deep
// worktree paths.
override protected def newEndpointAddress(
workerId: String): String = {
val socketDir = java.nio.file.Files.createTempDirectory(
java.nio.file.Paths.get("/tmp"), "udf-test-")
socketDir.toFile.deleteOnExit()
socketDir.resolve(s"w-$workerId.sock").toString
}

override protected def createConnection(
socketPath: String): UnixSocketWorkerConnection =
RealSocketConnection.connect(socketPath)
override protected def newConnection(address: String): WorkerConnection =
ConnectingSocketConnection.connect(address)
}

/**
* Tests that [[PythonUDFWorkerSpecification#fromPythonFunction]]
* produces a valid [[UDFWorkerSpecification]] that can be used by
* a [[org.apache.spark.udf.worker.core.WorkerDispatcher]]
* produces a valid [[org.apache.spark.udf.worker.UDFWorkerSpecification]]
* that can be used by a
* [[org.apache.spark.udf.worker.core.WorkerDispatcher]]
* to spawn a real Python worker process.
*
* The test overrides `spark.python.worker.module` to point to a
* test-only Python module that creates the UDS socket and waits
* for SIGTERM, verifying that pythonExec, PYTHONPATH, env vars,
* and command construction all work end-to-end.
*
* Reuses [[TestDirectWorkerDispatcher]]'s spawn / wait-for-ready machinery
* (shared from `udf-worker-core` test sources) but overrides the connection
* to open a real UDS channel -- the Python test worker only binds a raw
* socket, it does not host a gRPC server, so a full gRPC session is not
* exercised here.
*/
class PythonUDFWorkerSpecificationSuite
extends SharedSparkSession {
Expand Down Expand Up @@ -185,15 +172,16 @@ class PythonUDFWorkerSpecificationSuite
val workerSpec =
PythonUDFWorkerSpecification.fromPythonFunction(func, conf)

// Verify the spec works end-to-end with a real
// socket connection
val dispatcher = new RealSocketTestDispatcher(workerSpec)
// Verify the spec works end-to-end: the dispatcher spawns the Python
// worker, waits for the socket, and opens a real UDS connection to it
// (proving the worker is reachable, not just that the file exists).
val dispatcher = new ConnectingTestDispatcher(workerSpec)
try {
val session = dispatcher.createSession(
securityScope = None)
assert(session != null,
"Expected a non-null session from the dispatcher")
session.close()
session.close(() => Cancel.getDefaultInstance)
} finally {
dispatcher.close()
}
Expand Down
66 changes: 66 additions & 0 deletions udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/Termination.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
/*
* 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.udf.worker.core

import org.apache.spark.annotation.Experimental
import org.apache.spark.udf.worker.{CancelResponse, ExecutionError, FinishResponse}

/**
* :: Experimental ::
* The terminal outcome a [[WorkerSession]] settles on, returned by
* [[WorkerSession#close]]. Mirrors the four terminal `WorkerSession.SessionState`s,
* so close() reports the outcome faithfully rather than collapsing failures into
* a clean cancel.
*
* '''Clean outcomes''' ([[Finished]] / [[Cancelled]]) wrap the worker's
* `FinishResponse` / `CancelResponse` -- per-execution metrics, an optional
* inline `data` payload, and an optional finish/cancel '''callback error''' --
* which the engine surfaces to the UDF's client-side finish/cancel callbacks.
* The proto allows `Cancel` to follow `Finish` on the same stream, so close()
* may observe either depending on arrival timing (see `udf_message.proto`): if
* `FinishResponse` was already produced when a `Cancel` arrives the engine still
* receives [[Finished]], otherwise [[Cancelled]].
*
* '''Failure outcomes''' ([[Failed]] / [[TransportFailed]]) carry the cause
* instead of a proto terminator (none arrived, so they have no metrics). They
* exist so a failure is not reported as a benign cancel -- in particular an
* error raised during finish/close, '''after all data has been drained''',
* reaches the caller only through this value, never through the result iterator.
*/
@Experimental
sealed trait Termination

object Termination {
/** The stream ended with a `FinishResponse` (normal completion). */
final case class Finished(response: FinishResponse) extends Termination

/** The stream ended with a `CancelResponse` (cooperative cancellation). */
final case class Cancelled(response: CancelResponse) extends Termination

/**
* An execution error settled the session without a proto terminator (e.g. an
* `ErrorResponse`, or an error before init completed). Carries the structured
* [[ExecutionError]].
*/
final case class Failed(error: ExecutionError) extends Termination

/**
* A transport failure, timeout, or interrupt tore the stream down before any
* terminator arrived. Carries the underlying cause.
*/
final case class TransportFailed(cause: Throwable) extends Termination
}
4 changes: 2 additions & 2 deletions udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerConnection.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ import org.apache.spark.annotation.Experimental
* multiplexed streams over this channel.
*
* Implementations expose only lifecycle. Data transmission happens at
* the [[WorkerSession]] level -- this class is solely about whether the
* the [[WorkerSession]] level -- this trait is solely about whether the
* channel is open.
*
* '''Relationship to other classes (direct creation mode):'''
Expand All @@ -43,7 +43,7 @@ import org.apache.spark.annotation.Experimental
* }}}
*/
@Experimental
abstract class WorkerConnection extends AutoCloseable {
trait WorkerConnection extends AutoCloseable {
/** Returns true if the underlying transport channel is still usable. */
def isActive: Boolean
}
66 changes: 66 additions & 0 deletions udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerHandle.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
/*
* 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.udf.worker.core

import org.apache.spark.annotation.Experimental

/**
* :: Experimental ::
* The lifecycle-side handle the engine uses to release a worker back to
* its dispatcher.
*
* Sessions hold a [[WorkerHandle]] -- not the concrete worker type --
* so the engine-facing layer (the `core` package) is decoupled from any
* specific worker provisioning model (direct local spawn, indirect
* daemon-provided, etc.). Each [[WorkerDispatcher]] supplies its own
* implementation; the session calls back into it on close.
*
* Contract:
* - [[releaseSession]] decrements the dispatcher-side reference count.
* The caller ([[WorkerSession.close]]) invokes it exactly once per
* session, so implementations need not be idempotent and MAY fail fast
* on an unbalanced (negative) count. Sessions sharing the same worker
* call it independently (each exactly once).
* - [[markInvalid]] is sticky: once set, the dispatcher MUST NOT return
* this worker to any reuse pool. Sessions set this when the worker
* is observed in an unsafe state (transport error, hung worker,
* protocol violation).
* - [[id]] is a stable string for diagnostics. The session uses it in
* error messages; the dispatcher uses it to key the worker in its
* own tracking map. Implementations must keep it constant for the
* lifetime of the handle.
*/
@Experimental
trait WorkerHandle {

/** Stable identifier for logs and diagnostic messages. */
def id: String

/**
* Marks the worker as unsafe to recycle. Idempotent and sticky.
* Sessions call this in their close path when the protocol layer
* reports an unsalvageable termination.
*/
def markInvalid(): Unit

/**
* Decrements the dispatcher-side session ref count. The caller
* ([[WorkerSession.close]]) invokes this exactly once per session, so
* implementations need not be idempotent.
*/
def releaseSession(): Unit
}
Loading