[SPARK-57640] Implement gRPC WorkerSession and test it against an echo server. - #56702
[SPARK-57640] Implement gRPC WorkerSession and test it against an echo server.#56702haiyangsun-db wants to merge 10 commits into
Conversation
haiyangsun-db
changed the title
Yicong-Huang
left a comment
Yicong-Huang
left a comment
There was a problem hiding this comment.
Thanks for the effort, left some comments!
Correctness: - Fast-fail init on a DataResponse received before InitResponse instead of enqueuing the stray batch and hanging until initResponseTimeoutMs. - Guard input.hasNext (not just input.next) against exceptions, cancelling the stream before rethrowing so the worker is not stranded awaiting input. - Guard the caller-supplied finish() thunk the same way. - Carry the init error onto the terminal: the InitResponse-with-error path now settles Failed(err) before attempting Cancel, so close() returns Termination.Failed rather than a bare Cancelled. Readability: - Introduce OneShotValue to bundle the init latch and captured InitResponse so the publish-then-release ordering lives in one place. - Centralize the protocol transition edges behind a Transitions object over the single SessionState (no behavior change). - Correct the driving-model scaladoc: input is pull-driven, not strictly one-input-per-output. Tests: - Assert the WorkerHandle lifecycle (released once; invalid only on an unsalvageable TransportFailed). - Add regression tests for the four correctness fixes above.
The comment claimed settledTermination returns a best-effort empty Cancelled for the failure terminals, which was the pre-fix behavior. close() now returns the settled terminal as-is: Failed / TransportFailed carry their cause rather than collapsing to a bare Cancelled (asserted by the pre-init-error and close-timeout tests). Also complete the verbless "an init error from init()" fragment.
Yicong-Huang
left a comment
Yicong-Huang
left a comment
There was a problem hiding this comment.
Thanks, the new code makes the state transitions much easier to follow. I have two high-level comments:
-
The init-handshake logic still feels too complex. The complexity seems to come from using the same responseObserver for both init and the data phase, which spreads initResolved checks and init-waiter signaling across many branches. Would it be cleaner to isolate init as a one-shot completion and switch to data-phase handling only after it resolves?
-
For the missing post-init-error Cancel, I think Initializing should transition through Cancelling rather than directly to Failed. sendCancelInternal already supports Initializing -> Cancelling, and the timeout and interrupt paths use it. The init-error path should similarly send Cancel first, then call Transitions.failed, instead of settling Failed early and suppressing the wire-level Cancel.
Please also see inline some detailed comments.
| // the session as cancelled-before-start. The base WorkerSession still | ||
| // releases the worker handle, so the worker is torn down. | ||
| Transitions.cancelled(CancelResponse.getDefaultInstance) | ||
| return Termination.Cancelled(CancelResponse.getDefaultInstance) |
There was a problem hiding this comment.
Returns a fresh Cancelled even when a terminal already settled (e.g. the channel-shutdown TransportFailed in doInit also leaves requestObserver null), so the return disagrees with the state. Return settledTermination instead. Harmless today since base close() discards this value, but misleading.
There was a problem hiding this comment.
L529 only settles cancelled if (!currentState.isTerminal), then return settledTermination (L531).
| /** Failure terminal carrying a transport-level cause. */ | ||
| def transportFailed(cause: Throwable): Boolean = | ||
| completeTerminal(Termination.TransportFailed(cause)) |
There was a problem hiding this comment.
can this be merged to failed?
There was a problem hiding this comment.
There's no Transitions.failed anymore - transportFailed and the terminal outcomes are now distinct by design
| case _: InterruptedException => | ||
| Thread.currentThread().interrupt() | ||
| sendCancelInternal(() => cancelWithReason("interrupted during init")) | ||
| // Make sure close() does not block waiting for a terminator that | ||
| // will never arrive on this thread's behalf. | ||
| Transitions.transportFailed( | ||
| new InterruptedException("interrupted while waiting for InitResponse")) | ||
| throw new InterruptedException("interrupted while waiting for InitResponse") |
There was a problem hiding this comment.
who can interrupt this thread? if it is user, then this should be cancel? if this is the engine/system, then this should be failed?
There was a problem hiding this comment.
@Yicong-Huang In spark's case, this usually means query cancellation, that can originate from a end-user clicking cancel or interrupt the query, or AQE/Spark engine cancel a stage.
I'm trying to introduce a better and consistent way of handling interrupts here.
On the other hand, I would propose us to move faster with this first version, as we can evolve over time later.
There was a problem hiding this comment.
it's a first-class Interrupted outcome that cooperatively cancels (bounded), keeps the worker reusable, and no longer masquerades as TransportFailed
The previous init-error handling settled a Failed terminal in the gRPC response callback before attempting Cancel, so close() would report the init error. But that made the wire Cancel unreachable: sendCancelInternal bailed on both the directExecutor path (requestObserver still null) and the cross-thread path (state already terminal). udf_message.proto requires the engine to send Cancel after an init error and the worker to reply with CancelResponse, so this was a protocol violation. Root cause: the INIT-error / pre-init ERROR handling ran in the response callback (handleControl), which under directExecutor reentrancy executes inside doInit's stream.onNext -- before requestObserver is published on the next line -- so a Cancel written there could hit a null observer. Move the Cancel into doInit, which runs only after requestObserver is published. The callbacks now just record the error and wake init, leaving the session in Initializing; doInit's new failInitWithError helper sends Cancel, drains the CancelResponse (a real awaitTerminal, no longer a no-op), and throws. The reentrancy hazard is gone. As a result the settled terminal on an init error is the proto terminator Cancelled, not Failed -- the structured ExecutionError is surfaced through init()'s exception and kept on executionError for the iterator / close()'s salvage decision. Transitions.failed no longer has any producer and is removed. Also: - doClose returns settledTermination when requestObserver is null rather than a fresh Cancelled that could disagree with an already-settled TransportFailed. - OneShotValue.await throws TimeoutException instead of returning a Boolean the caller must inspect; the doInit call site handles it alongside the interrupt path. - Expand the init-interrupt comment: an interrupt leaves the worker in an unknown state after a best-effort Cancel, so TransportFailed (not cancel) is the correct, unsalvageable outcome. Update the two tests that asserted close() == Failed to the new contract: init throws the structured error and close() returns Cancelled. The cross-thread test now also asserts a Cancel was actually written to the wire, pinning the protocol requirement.
onTerminalSettled now signals initValue, so every terminal-settling path wakes a thread still blocked in doInit -- including close(), which settles a terminal (cancelled when requestObserver is null, or transportFailed on terminal timeout) without going through the response callback that would otherwise signal initValue. Previously a close() concurrent with a blocked init() left init() parked for the full initResponseTimeoutMs after the terminal had already settled, contrary to the initValue latch invariant. Because onTerminalSettled runs once, on the completeTerminal winner, after the terminal CAS, settle-before-release still holds; the per-callback terminal signalWithoutValue() calls become redundant but stay (idempotent), while the non-terminal init signals remain required. Tests: add a close-concurrent-with-blocked-init regression (terminal settled by close() itself), plus coverage for the terminator-carried callback-error paths (FinishResponse.error surfaced; data-phase ErrorResponse taking precedence over a CancelResponse callback error) and the data-phase output timeout. Full suite: 23 tests, 0 failures.
Doc-only consistency pass over the udf/worker module; no behavior change. - Termination: the scaladoc said it "mirrors the four terminal SessionStates", contradicting WorkerSession, which states the four outcomes live once on Termination and the state has a single Terminal case. Reworded to describe the single-Terminal design. - WorkerDispatcher: invalidation is driven by WorkerSession.close (via isWorkerSalvageable), not by doProcess as the doc claimed; doProcess never touches markInvalid. Corrected the described mechanism. - UDFDispatcherManager: "read lock (with upgrade)" is a misnomer -- ReentrantReadWriteLock cannot upgrade a held read lock, and the code releases the read lock before acquiring the write lock. Reworded. - GrpcWorkerSession: note that the header wire diagram omits PayloadChunk* (chunking is not yet implemented; see the TODO) so it doesn't overstate a match with udf_message.proto. - GrpcWorkerSessionConcurrencySuite: drop a stale reference to a "Closed" state (no such SessionState exists; close() settles a terminal).
An engine-side interrupt (cancelled query / killed task) is not a worker fault, so it should not be reported as TransportFailed nor make the worker unsalvageable. Introduce a distinct Termination.Interrupted and handle interrupts cooperatively at every blocking point. core: - Termination.Interrupted(cause): a new terminal, documented as an engine-side interrupt distinct from a cooperative Cancelled (which drains a CancelResponse) and a TransportFailed (a genuine transport fault). - isWorkerSalvageable leaves Interrupted salvageable (only TransportFailed and a never-settled session stay unsalvageable). Docs updated across Termination / WorkerSession / SessionState to enumerate five terminals and to stop attributing "interrupt" to TransportFailed. grpc (GrpcWorkerSession): - New handleInterrupt(waitContext) shared by every site that blocks on a worker event -- doInit's InitResponse wait, awaitTerminal, doClose's terminator wait, and ProcessIterator's output poll. It sends a best-effort Cancel, then waits up to interruptCancelTimeoutMs (new constructor param, default DEFAULT_INTERRUPT_CANCEL_TIMEOUT_MS = 2000ms, configurable) for the CancelResponse: a responsive worker settles a clean Cancelled (salvageable, response drained); on timeout it falls back to Interrupted (also salvageable). The bounded drain runs before the thread interrupt flag is restored, since InterruptedException clears it on throw. - Explicit Interrupted arms added to doInit's terminal inspection and throwIfTerminalError so an interrupt surfaces a clear message. TODO [SPARK-57640]: back the optimistic salvageable assumption with a worker liveness/heartbeat check so an interrupt that leaves a wedged (but not transport-dead) worker is still caught rather than recycled. Tests: core gains an Interrupted salvageable-mapping case; grpc adds responsive (-> Cancelled) and unresponsive (-> Interrupted, bounded by interruptCancelTimeoutMs) interrupt-during-init tests. Suites: WorkerSessionSuite 18, GrpcWorkerSessionConcurrencySuite 25, EchoProtocolSuite 18 -- all pass.
84242cb to
1244ff7
Compare
Since onTerminalSettled now signals initValue (settling any terminal wakes a blocked doInit), the explicit initValue.signalWithoutValue() calls in the branches that also settle a terminal were redundant. Drop the 7 such calls (responseObserver's DATA-before-init / malformed / onError / onCompleted, and handleControl's FINISH / CANCEL / CONTROL_NOT_SET): those branches now just settle the terminal, matching how the data phase already behaves, and onTerminalSettled is the single place a settled terminal wakes doInit. The only direct initValue signals left are the three genuinely non-terminal init paths (handleControl's INIT-ok / INIT-error / pre-init-ERROR), which must wake doInit without settling a terminal so the session stays Initializing and process() can still run. This addresses review feedback that init-waiter signaling was spread across too many branches. No behavior change: GrpcWorkerSessionConcurrencySuite (25) and EchoProtocolSuite (18) pass unchanged.
- Termination.Interrupted / Transitions.interrupted: the docs still said the post-interrupt Cancel's CancelResponse "is not awaited", which predates the bounded-drain design. It IS awaited briefly; Interrupted settles only when that short wait expires without an ack (otherwise a cooperative Cancelled settles). Correct both docs to match handleInterrupt. - doClose: add a TODO noting that an interrupted close() blocks for another interruptCancelTimeoutMs to chase a clean, salvageable Cancelled, with the current tradeoff (keep the worker recyclable) and the revisit condition (killed-task teardown latency). - awaitTerminal: drop a spurious s"" interpolator on a literal with no substitution. Doc/comment only; no behavior change.
|
Thanks for the thorough review -- this was really helpful. I've pushed a series of commits addressing everything; recap below, with the design discussion on high-level point 1 at the end. High-level point 2 (Cancel-after-init-error should go through Cancelling, not settle Failed early): done. The init-error path no longer settles Inline comments: all addressed in the threads (recap so it's in one place):
High-level point 1 (init-handshake complexity): You're right that the signaling was the main source of noise, and it turned out to be mostly accidental. On fully isolating init behind a one-shot handler and switching to a data-phase handler once it resolves: I looked at this and decided against it. gRPC gives us a single |
2 participants
What changes were proposed in this pull request?
This PR adds
GrpcWorkerSession, the engine-sideWorkerSessionimplementation that drives one bidirectional gRPCExecutestream per UDF execution. It is the first concrete protocol implementation on top of the abstractions already inmaster(theWorkerSessionSessionStatemachine,Termination,WorkerHandle, and theudf_message.proto/udf_service.protosplit).It is deliberately scoped to the session driver and its in-process tests, split out of the larger gRPC-backend change (#55683) so the concurrency-heavy core can be reviewed in isolation. The UDS transport, channel, and dispatcher follow in a second PR.
Added:
GrpcWorkerSession(udf/worker/grpc) — drives theInit -> DataRequest* -> Finish | Cancel/InitResponse -> DataResponse* -> terminatorexchange. Design highlights:SessionStateowned by the baseWorkerSessionalong the protocol-event edges; the only out-of-band flag iscancelRequested(a cancel can be requested before the stream exists).requestLock, withcancelRequestedenforcing the proto invariant "noData/FinishafterCancel".Finishterminal wait is per-poll (reset by each worker event), so the contract is "emit an event everyterminalTimeoutMs", not "finish within it".GrpcWorkerSessionExceptionextendsRuntimeException(notSparkRuntimeException) so theudf-workermodules stay free of aspark-coredependency; the engine integration layer wraps it.GrpcWorkerSessionConcurrencySuite— in-process tests pinning the wire-ordering and fast-fail invariants.grpc/pom.xml— adds thespark-udf-worker-coremodule dependency (the session needscore; the module previously depended only onudf-worker-proto).Class relationship:
GrpcWorkerSessiondepends only on a rawio.grpc.ManagedChannel, not on any transport/channel class — which is what lets it (and its tests) stand alone here, with the channel/dispatcher wiring landing in the follow-up.Why are the changes needed?
The UDF worker protocol defines a language-agnostic execution surface but, prior to this work, had no engine-side driver to run it.
GrpcWorkerSessionis that driver — the piece that turns the protocol abstractions into an actual UDF execution over gRPC bidirectional streaming. Landing it separately from the transport and dispatcher keeps the most correctness-sensitive component (the concurrent stream/state-machine interaction) small and independently testable.Does this PR introduce any user-facing change?
No.
udf/workeris@Experimentaland not wired into any Spark execution path; there are no production callers yet.How was this patch tested?
GrpcWorkerSessionConcurrencySuite(in-process gRPC viaInProcessServerBuilder/InProcessChannelBuilder; no UDS, no subprocess) pins the invariants: a worker-emittedErrorResponse/FinishResponse/CancelResponse/ stream half-close / malformed (RESPONSE_NOT_SET) response beforeInitResponseall failinit()fast rather than hanging on the init timeout; repeatedhasNextafter exhaustion returns immediately;close()concurrent withprocess()settles cleanly. The fast-fail and data-phase paths are re-run under cross-thread (non-directExecutor) delivery so correctness does not silently depend on reentrant delivery.EchoProtocolSuitecontinues to exercise the full wire protocol end-to-end against an in-process echo worker.build/sbt "udf-worker-grpc/testOnly *GrpcWorkerSessionConcurrencySuite *EchoProtocolSuite"— 2 suites, 33 tests, 0 failures.Was this patch authored or co-authored using generative AI tooling?
Yes