[SPARK-56413] Add gRPC UDF execution protocol#55657
[SPARK-56413] Add gRPC UDF execution protocol#55657haiyangsun-db wants to merge 2 commits intoapache:masterfrom
Conversation
haiyangsun-db
changed the title
| // with no branch set. | ||
| oneof control { | ||
| Init init = 1; | ||
| PayloadChunk payload = 2; |
There was a problem hiding this comment.
what is the reason that payload does not require a confirmation from the worker?
There was a problem hiding this comment.
Payloadchunks are an addition to the Init message, we can assume it's part of the Init, only used when payload is large. It doesn't require a response. Init + (payload_chunk)* maps to one InitResponse
| // the engine: the engine forwards [[payload]] and [[format]] | ||
| // unchanged, and the worker decodes them per the format the client | ||
| // and worker have agreed on. | ||
| message UdfPayload { |
There was a problem hiding this comment.
how about payload language? I assume because the worker is already tied to a specific language, so it does not need to know what language this UDF payload is in?
There was a problem hiding this comment.
exactly, that does not matter. worker is already provided by the worker spec, engine doesn't care what language it is
| // a typed field number from the reserved range right after this | ||
| // block and is removed from [[session_conf]]. [[timezone]] below | ||
| // is an example of a key that has already been promoted. | ||
| map<string, string> session_conf = 6; |
There was a problem hiding this comment.
why would this and task_context above have no optional prefix?
There was a problem hiding this comment.
map field cannot be optional in protobuf. Leave it to be an empty map when it is not set.
|
|
||
| // (Optional) Session timezone, promoted out of [[session_conf]] | ||
| // because every eval needs it for timestamp encoding/decoding. | ||
| optional string timezone = 7; |
There was a problem hiding this comment.
is string the canonical type to represent the timezone? I am afraid all kinds of conversion errors may happen with no schema/enum enforcement.
There was a problem hiding this comment.
this is convention from Spark, timezone is a string in spark.
| message Heartbeat {} | ||
|
|
||
| // Acknowledgment for [[Heartbeat]]. | ||
| message HeartbeatAck {} |
There was a problem hiding this comment.
| message HeartbeatAck {} | |
| message HeartbeatResponse {} |
just for some consistency?
|
|
||
| // (Optional) Session timezone, promoted out of [[session_conf]] | ||
| // because every eval needs it for timestamp encoding/decoding. | ||
| optional string timezone = 7; |
There was a problem hiding this comment.
We should specify the exact format in which the timezone will be reported since its a string
| // Packed by the client side of the protocol; opaque to the | ||
| // wire protocol. Left unset whenever the worker's built-in | ||
| // decoders are sufficient. | ||
| optional bytes input_encoder = 6; |
There was a problem hiding this comment.
AFAIK there are no use cases for the custom in/output encoders at the moment. Should we maybe only add them when they are needed?
| def cancel(): Unit | ||
|
|
||
| /** Closes this session and releases resources. */ | ||
| override def close(): Unit |
There was a problem hiding this comment.
I think we need to clarify the exact semantics of close/finish and cancel within the background of how we could implement calling them in Spark.
From my current understanding, finish would indicate to the UDF worker that no more input batches are to be send. Therefore, the worker would finish processing the batches it has already received/buffered and then response with a FinishResponse. From the Spark side (e.g. in a operator), we could send multiple batches and then call finish, meaning the batches and the finish message would sit in the client or server-side gRPC buffer. If the Spark task now gets canceled, we have to wait for the UDF worker to finish as the finish message was already sent. Is this a scenario, would we like to support more eager cancellation? This also somewhat depends on the processing time per batch and the buffer size. Alternatively, we can force worker termination via the system-level primitives (SIGTERM/SIGKILL).
Another concern is a potential race condition between finish and cancel calls. Cancellations are most likely going to be implemented using a taskInterruptListener on the Spark task running a UDF. The callback invoked from this listener does not necessarily share the whole context of the operator execution. Therefore, it might not know if finish has already been called/queued. Would it make sense for the cancel/finish calls from the WorkerSession to implement a no-op in this case? The session has all the state and can know whether the session has already previously been canceled/finished.
There was a problem hiding this comment.
The same concern exists between init and cancel. If we implement cancellation via a taskInterruptListener, cancel might be called before init was called.
3 participants
What changes were proposed in this pull request?
Adds
udf_protocol.proto, the gRPC wire contract between the Spark engine and aUDF worker process, as described in SPIP. Sits next to the existing
worker_spec.proto.Defines a
Workerservice with two RPCs:Execute(stream UdfRequest) returns (stream UdfResponse)— one bidirectionalstream per UDF execution. Lifecycle on the stream:
Init→ 0..NDataRequest/DataResponse→ exactly oneFinishorCancel.PayloadChunkstreams oversized UDF bodies.Manage(WorkerRequest) returns (WorkerResponse)— unary, worker-scoped(heartbeat, graceful shutdown).
UdfPayloadcarries the engine-opaque callable bytes plus aformattag,an
eval_typeworker-dispatch hint, and optional input/output encoders.Initcarriesdata_format, schemas,session_conf,task_context, andtimezone(the first graduate fromsession_conf); a reserved field rangeabsorbs future graduates.
Also fixes two typos in
common.proto(exachanged/bidrectional).Out of scope
No planning info on the wire (no execution-shape / cardinality enum, no
chained-UDF metadata). Both can be added additively later.
Why are the changes needed?
Spark Connect's UDF support today is Python-only and tied to a Python-specific
socket protocol. Onboarding other client languages requires a structured,
language-neutral wire contract. This PR lands the proto layer; engine and
worker implementations will follow.
Does this PR introduce any user-facing change?
No. Wire contract only; not yet wired into any end-to-end path.
How was this patch tested?
Verified the proto compiles with
protocagainstcommon.protoandworker_spec.proto, and inspected the generated descriptor for field-numberand oneof correctness. End-to-end conformance tests will land with the
engine-side client and first worker implementation.
Was this patch authored or co-authored using generative AI tooling?
Yes