[SPARK-56413][SQL][UDF] gRPC implementation of the UDF worker protocol dispatcher by haiyangsun-db · Pull Request #55683 · apache/spark

haiyangsun-db · 2026-05-05T07:33:23Z

What changes were proposed in this pull request?

Adds the first transport implementation of the UDF worker protocol: gRPC bidirectional streaming over a Unix domain socket, as a new udf-worker-grpc module.

This PR builds on the engine-side core abstractions and the proto split introduced in the base PR#56397 (which is the base of this branch); those are not part of this diff. For context, the base PR provides:

the proto split into udf_message.proto (message types) and
udf_service.proto (the UdfWorker gRPC service), keeping gRPC off the
classpath of the message consumers (core, catalyst, sql-core);
the WorkerSession SessionState state machine, the Termination type,
and the WorkerHandle trait.

This PR is scoped to the gRPC implementation on top of those abstractions:

New udf-worker-grpc module:

DirectGrpcDispatcher -- a DirectWorkerDispatcher that wires UDS gRPC
into the existing worker spawn / lifecycle machinery (private 0700 socket
dir, atomic POSIX permissions with an owner-only non-POSIX fallback,
recursive cleanup, UDS path-length-safe socket naming).
validateTransportSupport requires the spec to declare UDS transport and
BIDIRECTIONAL_STREAMING.
GrpcWorkerSession -- the engine-side WorkerSession driving one
bidirectional Execute stream per UDF execution. Honors the proto ordering
invariants (Init -> DataRequest* -> Finish, Cancel-not-before-Finish on the
engine side; InitResponse -> DataResponse* -> terminator on the worker
side), serializes wire writes under a single requestLock, surfaces
transport / execution / protocol failures as GrpcWorkerSessionException,
and bounds every wait by an explicit timeout. It keeps no parallel state
machine -- it drives the single SessionState owned by the base
WorkerSession, advancing only the protocol-event edges.
GrpcWorkerChannel -- per-worker ManagedChannel plus a dedicated Netty
EventLoopGroup over the UDS, with deterministic close ordering.
UnixDomainSocketTransport -- runtime detection of Linux epoll /
macOS kqueue native UDS transports, with a loud failure when neither is
available.

Integration: sql/core gains a test-scope dependency on the grpc test-jar so
PythonUDFWorkerSpecificationSuite can reuse the grpc test dispatcher; gRPC
stays off sql/core's main classpath. The module is @Experimental and has
no production callers yet.

Why are the changes needed?

The UDF worker protocol defines a language-agnostic execution surface but did
not yet have an implementation. This PR delivers gRPC-over-UDS as the first
transport so the dispatcher can run UDFs end to end. Native UDS via Netty
avoids TCP overhead on the local socket; gRPC bidirectional streaming gives us
ordered delivery, per-message size limits, and well-tested flow control.

Does this PR introduce any user-facing change?

No. udf/worker is @Experimental and not yet wired into any Spark execution
path.

How was this patch tested?

New GrpcWorkerSessionConcurrencySuite (in-process tests) pins the
wire-ordering and fast-fail invariants: a worker-emitted ERROR / FINISH /
CANCEL / onCompleted() before InitResponse all fail init() fast;
repeated hasNext after exhaustion returns immediately; close() concurrent
with process() terminates cleanly; and the same fast-fail / data-phase
paths are re-run under cross-thread (non-directExecutor) delivery so
correctness does not silently depend on reentrant delivery.
New DirectGrpcDispatcherIntegrationSuite (subprocess tests via
EchoGrpcWorkerMain): end-to-end echo, sequential and concurrent sessions,
worker-crash mid-process, init failure, capability-mismatch rejection
(including a spec with no declared capabilities), socket-directory cleanup,
reuse-readiness.
Existing EchoProtocolSuite refactored onto a shared EchoWorkerService and
continues to cover the wire protocol end to end.
DirectWorkerDispatcherSuite ported to the gRPC test dispatcher: spawn /
SIGTERM / SIGKILL escalation / environment lifecycle / concurrent close /
timeout clamping.
All suites green; sql/Test/compile clean.

Was this patch authored or co-authored using generative AI tooling?

Yes

sven-weber-db · 2026-05-30T08:17:35Z

+ *  - [[cancel]] is thread-safe and idempotent.
+ *  - [[doClose]] half-closes the request stream after the terminator arrives.
+ *
+ * NOTE: this class does not yet implement payload chunking; the entire


nit: Add a TODO with the JIRA tracking ticket for this

haiyangsun-db · 2026-06-01T12:36:10Z

Hi @dtenedor @hvanhovell @Yicong-Huang , please feel free to review this changes. It's the Dispatcher/Session implementation of the gRPC protocol following last PR.

Spark can then use the Dispatcher and session interface to talk to udf worker, which encapsulates the protocol details or worker management.

hvanhovell · 2026-06-02T15:08:18Z

+  def detect(): UnixDomainSocketTransport = {
+    val os = System.getProperty("os.name", "").toLowerCase(Locale.ROOT)
+    if (isLinux(os) && epollAvailable) {
+      EpollTransport


I think Netty is also adding support iouring.

hvanhovell · 2026-06-02T15:20:46Z

+      throw new UnsupportedOperationException(
+        s"No Netty native UDS transport available on os=$os " +
+          s"(epoll: $epollErr, kqueue: $kqueueErr). " +
+          "UDS-backed gRPC requires netty-transport-native-epoll on Linux or " +


Is there a socket based fallback?

could be, but out-of-scoped this PR

hvanhovell · 2026-06-02T15:27:31Z

+    message: String,
+    cause: Throwable = null,
+    val executionError: ExecutionError = null)
+  extends RuntimeException(message, cause) {


Should this be an SparkRuntimeException? I understand that this requires an additional dependency.

Spark side should wrap exception thrown in this module into SparkRuntimeException, let me update the doc about that.
The intention is to make this module as isolated from spark as possible.

hvanhovell · 2026-06-02T15:28:07Z

+
+/**
+ * :: Experimental ::
+ * Exception thrown by [[GrpcWorkerSession]] when the UDF execution fails


Is this a mix of system and user errors? If we can, it would be preferable to separate these.

We shall have three layers of errors:

error from udf

error from worker / protocol (still logic at worker level)

transport error (grpc, etc)

I was trying to merge 1/2 into one big category as these are programmable errors usually and use subclasses to tell udf-only error from worker/protocol error.

hvanhovell · 2026-06-02T15:41:01Z

+  val DEFAULT_TERMINAL_TIMEOUT_MS: Long = 30000L
+
+  // Items the gRPC response observer pushes onto outputQueue. Using a sealed
+  // ADT (rather than an identity-typed `Array[Byte]` sentinel) means an


A bit of a pedantic comment. Yes, ADTs are easier to reason about that naked byte arrays.

hvanhovell · 2026-06-02T15:41:49Z

+      val u = err.getUser
+      val cls = if (u.hasErrorClass) s"[${u.getErrorClass}] " else ""
+      s"$cls${u.getMessage}"
+    case ExecutionError.KindCase.WORKER =>


Can a worker error be caused by the user?

It's possible, that's why i tend to combine worker error with user error, and distinguish using subclasses => they are definitely different from transport-level errors.

hvanhovell · 2026-06-02T18:43:50Z

+    terminalTimeoutMs: Long = DEFAULT_TERMINAL_TIMEOUT_MS)
+  extends WorkerSession(workerHandle, logger) {
+
+  require(channel != null, "channel is required")


Not checking the handle?

hvanhovell · 2026-06-02T18:44:40Z

+  // We deliberately do NOT bound this queue: blocking the gRPC receiving
+  // thread is worse than the current unbounded-but-fast behaviour. A proto
+  // change is out of scope here.
+  private val outputQueue = new LinkedBlockingQueue[QueueItem]()


We should have metrics for this.

hvanhovell · 2026-06-02T18:45:10Z

+  // Latch fired when `InitResponse` (success or error) or a transport error
+  // arrives. init() blocks on this; until it fires we have no proof the
+  // worker actually accepted the session.
+  private val initLatch = new CountDownLatch(1)


A Phaser might be more appropriate. You can combine this with other latch.

hvanhovell · 2026-06-02T20:03:35Z

+
+        // (2) Send next input batch when input has data and session is live.
+        if (!cancelRequested.get() && !finishedSending.get() && input.hasNext) {
+          val batch = try input.next() catch {


This sends a singular batch per iteration. Shouldn't we try to fill the send queue to a certain level?

It's a while loop, it will keep sending input until seeing a response.

we prefer consuming results over sending to avoid results filling up buffer

hvanhovell · 2026-06-02T20:10:00Z

+          // (4) Block for late output or the terminator. See class doc
+          //     above for the per-poll-vs-total-session timeout semantics.
+          val item = try {
+            outputQueue.poll(terminalTimeoutMs, TimeUnit.MILLISECONDS)


How does this work well with things like mapPartitions? In theory you can have a mapPartitions function that needs to consume all data first (e.g. home grown aggregation). I think that will time out if the data is partitioned into multi batches.

This is a loop, and this poll with timeout only happens when all input batches have been sent in this class.
If input are not all sent, we use poll without timeout (non-blocking) to check if there is any response.

We know at this moment, receiving responses is the only thing to do.

when we hit this branch, it means we sent all input and finish messages, and the only thing left is to wait for result response.

hvanhovell · 2026-06-08T19:04:55Z

+      // runs on the WorkerSession.close() path and an exception here would
+      // mask any in-flight failure the caller is already handling.
+      val msg = s"releaseSession called without a matching acquireSession (count=$c)"
+      assert(c >= 0, msg)


IMO this should be a harder error than an assertion error.

haiyangsun-db · 2026-06-09T11:36:02Z

Split the WorkerSession state-machine change into a separate PR for easy review: https://github.com/apache/spark/pull/56397/changes

uros-b

@haiyangsun-db FYI: please coordinate with #56741, because of possible merge order overlap.

haiyangsun-db · 2026-06-29T07:47:20Z

@uros-b thank you for the reminder, i will update

haiyangsun-db force-pushed the SPARK-56413-grpc-impl branch 2 times, most recently from 93ec09e to eef803e Compare May 6, 2026 04:03

haiyangsun-db changed the title ~~[Spark-56413] Implement bidirectional streaming for grpc protocol.~~ [SPARK-56413][CORE] Add gRPC bi-directional streaming UDF worker protocol May 6, 2026

haiyangsun-db changed the title ~~[SPARK-56413][CORE] Add gRPC bi-directional streaming UDF worker protocol~~ [SPARK-56413][UDF] Add gRPC bi-directional streaming UDF worker protocol May 6, 2026

haiyangsun-db force-pushed the SPARK-56413-grpc-impl branch from eef803e to 8d45dfe Compare May 6, 2026 04:57

haiyangsun-db force-pushed the SPARK-56413-grpc-impl branch 2 times, most recently from e45886e to 9ada1b2 Compare May 27, 2026 21:13

haiyangsun-db marked this pull request as ready for review May 28, 2026 08:18

haiyangsun-db marked this pull request as draft May 28, 2026 09:47

haiyangsun-db changed the title ~~[SPARK-56413][UDF] Add gRPC bi-directional streaming UDF worker protocol~~ [SPARK-56413][SQL][UDF] gRPC implementation of the UDF worker protocol dispatcher May 28, 2026

haiyangsun-db marked this pull request as ready for review May 28, 2026 11:13

sven-weber-db reviewed May 30, 2026

View reviewed changes

hvanhovell reviewed Jun 2, 2026

View reviewed changes

Comment thread udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala Outdated

hvanhovell reviewed Jun 2, 2026

View reviewed changes

Comment thread udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala Outdated

hvanhovell reviewed Jun 2, 2026

View reviewed changes

Comment thread udf/worker/grpc/src/main/scala/org/apache/spark/udf/worker/grpc/DirectGrpcDispatcher.scala Outdated

hvanhovell reviewed Jun 2, 2026

View reviewed changes

haiyangsun-db force-pushed the SPARK-56413-grpc-impl branch from d749c06 to c50e125 Compare June 8, 2026 08:32

haiyangsun-db marked this pull request as draft June 8, 2026 13:46

hvanhovell reviewed Jun 8, 2026

View reviewed changes

haiyangsun-db marked this pull request as ready for review June 8, 2026 21:38

adding grpc backend.

1ffaa80

haiyangsun-db force-pushed the SPARK-56413-grpc-impl branch from ca42ad4 to 1ffaa80 Compare June 10, 2026 14:16

haiyangsun-db added 2 commits June 10, 2026 16:32

nit fix.

ed087e1

merge connection into handle.

3223a97

haiyangsun-db mentioned this pull request Jun 23, 2026

[SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. #56702

Open

uros-b reviewed Jun 24, 2026

View reviewed changes

Uh oh!

Conversation

haiyangsun-db commented May 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

haiyangsun-db commented Jun 1, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

haiyangsun-db Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hvanhovell Jun 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

haiyangsun-db commented Jun 9, 2026

Uh oh!

uros-b left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

haiyangsun-db commented Jun 29, 2026

Uh oh!

haiyangsun-db commented May 5, 2026 •

edited

Loading

haiyangsun-db Jun 4, 2026 •

edited

Loading

hvanhovell Jun 2, 2026 •

edited

Loading

uros-b left a comment •

edited

Loading