[KYUUBI #7684][SPARK] Keep the python operation context private to the operation - #7685
[KYUUBI #7684][SPARK] Keep the python operation context private to the operation#7685moelhoussein wants to merge 3 commits into
Conversation
…QL does Spark local properties are thread local. The python engine runs every operation of a session on one pinned Py4J thread and clears the operation context with the empty string, while SparkOperation runs each operation on its own thread and clears with null. Both of those differences matter. AuthZUtils.getAuthzUgi skips the session user only when the property is null, so an empty user reaches UserGroupInformation.createRemoteUser, which rejects it with IllegalArgumentException: Null user. With session user signing enabled the same empty value fails verification as a blank identifier. Clearing with null makes a cleared context read as absent, and falls back to the engine user the way the SQL path already does. Holding the worker lock across withLocalProperties gives the context the isolation that the SQL path gets from running on its own thread, so a concurrent operation of the same session cannot observe the cleared window. AuthzSessionUserSuite covers the empty and null cases. The interleaving between two operations did not reproduce locally, so it has no test. Co-authored-by: Cursor <cursoragent@cursor.com>
…tion context Both defects are observable from the engine module, where a test can talk to the session's python worker directly: the cleared value lives on the worker's pinned thread, and the worker lock is what keeps one operation's context from being visible to another. Each test fails on the current behaviour, the first with `''` instead of `None`, the second because a second thread reaches the worker while an operation holds its context. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
Pull request overview
Fixes cross-operation context corruption for Spark Python operations sharing one worker.
Changes:
- Serializes the complete Python operation context lifecycle.
- Clears Spark local properties with
null. - Adds regression and authorization contract tests.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
ExecutePython.scala |
Makes context setup, execution, and cleanup atomic. |
ExecutePythonSuite.scala |
Tests cleanup and worker-lock isolation. |
AuthzSessionUserSuite.scala |
Verifies null versus empty-user authorization behavior. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| // Clear with null, not the empty string. AuthZUtils.getAuthzUgi only skips the property | ||
| // when it is null, so an empty user reaches UserGroupInformation.createRemoteUser and | ||
| // fails with `Null user`, and fails session user signing as a blank identifier. |
There was a problem hiding this comment.
This is only true when session-user signing is disabled. When signing is enabled, verifyKyuubiSessionUser runs first and rejects both null and empty values. Could you clarify that the outer lock, rather than the null value, protects the signing-enabled path?
There was a problem hiding this comment.
This test verifies existing AuthZUtils behavior, but it does not exercise the ExecutePython change and still passes when the patch is reverted. Since ExecutePythonSuite already verifies that the property is cleared with null, could we remove this suite or fold the relevant assertion into that regression test?
|
Both failures in this job come from the newly added This CI job compiles the embedded test JVM against Spark 3.5.8 while The later Could you update the tests to work with the existing cross-binary CI matrix, either by using matching Spark/PySpark versions for the embedded worker or by testing through the actual Spark binary, and rerun the checks? |
…ySpark version The suite runs the python worker inside the test JVM, so it cannot pass in the cross-binary CI jobs, where the worker imports PySpark from a Spark 4.x SPARK_HOME while the JVM holds Spark 3.5 classes and Py4J fails at import. Cancel the tests when the two do not match, resolving SPARK_HOME the way createSessionPythonWorker does. PySparkTests keeps the cross-binary python coverage, since it drives an engine from SPARK_HOME so both sides match. Also correct the comment on the context clear, since verifyKyuubiSessionUser rejects a blank user, so with signing enabled the worker lock is what protects that path, and drop AuthzSessionUserSuite, which only covered AuthZUtils. Co-authored-by: Cursor <cursoragent@cursor.com>
|
CI on The one red job is unrelated to this change. I do not have rights to re-run it. Could you kick that job when you get a chance? |
Fixes #7684.
Why are the changes needed?
Python operations of one session all run on a single pinned Py4J thread, and Spark local properties are thread local, so the per-operation context in
ExecutePython.withLocalPropertiesis not private to the operation the way it is for SQL. Two problems follow from that.The teardown clears
kyuubi.session.userwith the empty string, whileSparkOperation.withLocalPropertiesclears withnull.AuthZUtils.getAuthzUgiskips the property only when it is null, so an empty user reachesUserGroupInformation.createRemoteUserand fails withIllegalArgumentException: Null user. Withkyuubi.session.user.sign.enabledit fails earlier, asAccessControlException: Invalid user identifier []. The empty string was already the outlier in its own method, since the samefinallyblock clears the signing keys withnullthroughclearSessionUserSign. Null is not what saves the signed path, whereverifyKyuubiSessionUsertreats null as blank and reportsInvalid user identifier [null]just the same. There the lock is what keeps a concurrent operation from reading a cleared context.SessionPythonWorker.runCodetakes the worker lock per call, so the 7 to 12 internal round-trips that apply and clear the context are not atomic. Another operation of the same session can set or clear these properties in between, which loses the session user and also lets one operation'ssetJobGrouptag another operation's jobs, so itscancelJobGroupcancels them.This patch clears with
nulland holds the worker lock across the whole body ofwithLocalProperties, which substitutes for the isolation the SQL path gets from running each operation on its own thread. Throughput should be unaffected, becauserunCode(statement)already holds the lock for the whole execution and same-session operations are already serialized behind it. Cancellation is unaffected, becauseExecutePython.cleanupinterrupts the worker with a signal and never takes this lock.How was this patch tested?
Two new tests in
ExecutePythonSuite, both of which fail on master. The first runs an operation and then probes the worker thread with a barerunCode, which observes exactly what the teardown left behind. Master leaves'', so it fails with"['']" did not equal "[None]". The second holds an operation's context and checks that a second thread cannot reach the worker. On master it can, so it fails withotherOperationRanCode.get() was true.Both tests start the python worker inside the test JVM, where its PySpark reaches this JVM's Spark classes over Py4J, so they cancel when the PySpark under
SPARK_HOMEdoes not match the engine's Spark. That is the case in the jobs that verify a Spark 3.5 build on a Spark 4.x binary, which carry no 3.5 PySpark to point the worker at.PySparkTestsstill covers python in those jobs, since it drives an engine fromSPARK_HOMEand both sides match there.Local results are
ExecutePythonSuite2 of 2 andPySparkTests7 of 8. The single failure isexecutePython support timeout, which also fails on unmodified master in the same environment (one pass in six runs on master, zero in six on this branch). It fails in the second half of that test, wherebad_codeafter a timeout returns no error, which points at response handling on the worker stream after the interrupt rather than at the context.Was this patch assisted by generative AI tooling?
Assisted-by: Claude Opus 5