You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kyuubi.apache.org by ch...@apache.org on 2026/09/02 11:18:46 UTC
(kyuubi) branch branch-1.10 updated: [KYUUBI #7684][SPARK] Keep the python operation context private to the operation
This is an automated email from the ASF dual-hosted git repository.
pan3793 pushed a commit to branch branch-1.10
in repository https://gitbox.apache.org/repos/asf/kyuubi.git
The following commit(s) were added to refs/heads/branch-1.10 by this push:
new 6f43a5aad5 [KYUUBI #7684][SPARK] Keep the python operation context private to the operation
6f43a5aad5 is described below
commit 6f43a5aad5cc5b5e974ef0e29fb457dd3ff86274
Author: Mohamed <81...@users.noreply.github.com>
AuthorDate: Wed Sep 2 19:17:53 2026 +0800
[KYUUBI #7684][SPARK] Keep the python operation context private to the operation
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.withLocalProperties` is not private to the operation the way it is for SQL. Two problems follow from that.
The teardown clears `kyuubi.session.user` with the empty string, while `SparkOperation.withLocalProperties` clears with `null`. `AuthZUtils.getAuthzUgi` skips the property only when it is null, so an empty user reaches `UserGroupInformation.createRemoteUser` and fails with `IllegalArgumentException: Null user`. With `kyuubi.session.user.sign.enabled` it fails earlier, as `AccessControlException: Invalid user identifier []`. The empty string was already the outlier in its own method, s [...]
`SessionPythonWorker.runCode` takes 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's `setJobGroup` tag another operation's jobs, so its `cancelJobGroup` cancels them.
This patch clears with `null` and holds the worker lock across the whole body of `withLocalProperties`, which substitutes for the isolation the SQL path gets from running each operation on its own thread. Throughput should be unaffected, because `runCode(statement)` already holds the lock for the whole execution and same-session operations are already serialized behind it. Cancellation is unaffected, because `ExecutePython.cleanup` interrupts the worker with a signal and never takes t [...]
### 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 bare `runCode`, 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 with `otherOperationRanCode.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_HOME` does 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. `PySparkTests` still covers python in those jobs, since it drives an engine from `SPARK_HOME` and both sides match there.
Local results are `ExecutePythonSuite` 2 of 2 and `PySparkTests` 7 of 8. The single failure is `executePython 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, where `bad_code` after 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
Closes #7685 from moelhoussein/python-session-user-race.
Closes #7684
da9e7d9e2 [MElHoussein] [KYUUBI #7684][SPARK] Gate the python tests on the SPARK_HOME PySpark version
1fef42e06 [MElHoussein] [KYUUBI #7684][SPARK] Add regression tests for the python operation context
ed177ee47 [MElHoussein] [KYUUBI #7684][SPARK] Clear the python operation context like SQL does
Lead-authored-by: Mohamed <81...@users.noreply.github.com>
Co-authored-by: MElHoussein <ME...@geico.com>
Signed-off-by: Cheng Pan <ch...@apache.org>
(cherry picked from commit 084a14f9d4c0f0a5e017753143ca96514058cba2)
Signed-off-by: Cheng Pan <ch...@apache.org>
---
.../engine/spark/operation/ExecutePython.scala | 73 ++++++-----
.../spark/operation/ExecutePythonSuite.scala | 144 +++++++++++++++++++++
2 files changed, 188 insertions(+), 29 deletions(-)
diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala
index 6cf525924b..1f6f85062b 100644
--- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala
+++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecutePython.scala
@@ -143,35 +143,46 @@ class ExecutePython(
}
override protected def withLocalProperties[T](f: => T): T = {
- try {
- // to prevent the transferred set job group python code broken
- val jobDesc = s"Python statement: $statementId"
- // for python, the boolean value is capitalized
- val pythonForceCancel = if (forceCancel) "True" else "False"
- worker.runCode(
- "spark.sparkContext.setJobGroup" +
- s"('$statementId', '$jobDesc', $pythonForceCancel)",
- internal = true)
- setSparkLocalProperty(KYUUBI_SESSION_USER_KEY, session.user)
- setSparkLocalProperty(KYUUBI_STATEMENT_ID_KEY, statementId)
- schedulerPool match {
- case Some(pool) =>
- setSparkLocalProperty(SPARK_SCHEDULER_POOL_KEY, pool)
- case None =>
- }
- if (isSessionUserSignEnabled) {
- setSessionUserSign()
- }
+ // Spark local properties are thread local, and every python operation of a session shares
+ // one worker thread, so the context below is not private to this operation the way it is for
+ // SQL. Holding the worker lock across apply, run and clear keeps a concurrent operation of
+ // the same session from setting or clearing these properties in between our round-trips.
+ worker.withLockRequired {
+ try {
+ // to prevent the transferred set job group python code broken
+ val jobDesc = s"Python statement: $statementId"
+ // for python, the boolean value is capitalized
+ val pythonForceCancel = if (forceCancel) "True" else "False"
+ worker.runCode(
+ "spark.sparkContext.setJobGroup" +
+ s"('$statementId', '$jobDesc', $pythonForceCancel)",
+ internal = true)
+ setSparkLocalProperty(KYUUBI_SESSION_USER_KEY, session.user)
+ setSparkLocalProperty(KYUUBI_STATEMENT_ID_KEY, statementId)
+ schedulerPool match {
+ case Some(pool) =>
+ setSparkLocalProperty(SPARK_SCHEDULER_POOL_KEY, pool)
+ case None =>
+ }
+ if (isSessionUserSignEnabled) {
+ setSessionUserSign()
+ }
- f
- } finally {
- setSparkLocalProperty(KYUUBI_SESSION_USER_KEY, "")
- setSparkLocalProperty(KYUUBI_STATEMENT_ID_KEY, "")
- setSparkLocalProperty(SPARK_SCHEDULER_POOL_KEY, "")
- // using cancelJobGroup for pyspark, see details in pyspark/context.py
- worker.runCode(s"spark.sparkContext.cancelJobGroup('$statementId')", internal = true)
- if (isSessionUserSignEnabled) {
- clearSessionUserSign()
+ f
+ } finally {
+ // Clear with null, like SparkOperation does. AuthZUtils.getAuthzUgi skips the property
+ // only when it is null, so an empty user reaches UserGroupInformation.createRemoteUser
+ // and fails with `Null user`. With session user signing enabled, verifyKyuubiSessionUser
+ // rejects a blank user either way, so there the worker lock is what keeps a concurrent
+ // operation from seeing the cleared context.
+ setSparkLocalProperty(KYUUBI_SESSION_USER_KEY, null)
+ setSparkLocalProperty(KYUUBI_STATEMENT_ID_KEY, null)
+ setSparkLocalProperty(SPARK_SCHEDULER_POOL_KEY, null)
+ // using cancelJobGroup for pyspark, see details in pyspark/context.py
+ worker.runCode(s"spark.sparkContext.cancelJobGroup('$statementId')", internal = true)
+ if (isSessionUserSignEnabled) {
+ clearSessionUserSign()
+ }
}
}
}
@@ -194,7 +205,7 @@ case class SessionPythonWorker(
new BufferedReader(new InputStreamReader(workerProcess.getInputStream), 1)
private val lock = new ReentrantLock()
- private def withLockRequired[T](block: => T): T = Utils.withLockRequired(lock)(block)
+ def withLockRequired[T](block: => T): T = Utils.withLockRequired(lock)(block)
/**
* Run the python code and return the response. This method maybe invoked internally,
@@ -203,6 +214,10 @@ case class SessionPythonWorker(
* please make sure the internal python code simple and set internal flag, to be aware of the
* internal python code failure.
*
+ * The lock is held for one call only. A sequence of calls that must not be seen half applied,
+ * such as the operation context in `ExecutePython.withLocalProperties`, has to be wrapped in
+ * [[withLockRequired]] by the caller.
+ *
* @param code the python code
* @param internal whether is internal python code
* @return the python response
diff --git a/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/ExecutePythonSuite.scala b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/ExecutePythonSuite.scala
new file mode 100644
index 0000000000..451405fa2b
--- /dev/null
+++ b/externals/kyuubi-spark-sql-engine/src/test/scala/org/apache/kyuubi/engine/spark/operation/ExecutePythonSuite.scala
@@ -0,0 +1,144 @@
+/*
+ * 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.kyuubi.engine.spark.operation
+
+import java.nio.charset.StandardCharsets.UTF_8
+import java.nio.file.{Files, Paths}
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.util.Try
+
+import org.apache.kyuubi.Utils
+import org.apache.kyuubi.config.KyuubiReservedKeys.{KYUUBI_SESSION_USER_KEY, KYUUBI_STATEMENT_ID_KEY}
+import org.apache.kyuubi.engine.spark.KyuubiSparkUtil.SPARK_ENGINE_RUNTIME_VERSION
+import org.apache.kyuubi.engine.spark.WithSparkSQLEngine
+import org.apache.kyuubi.operation.OperationHandle
+import org.apache.kyuubi.session.Session
+import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TProtocolVersion
+import org.apache.kyuubi.tags.PySparkTest
+
+@PySparkTest
+class ExecutePythonSuite extends WithSparkSQLEngine {
+
+ override def withKyuubiConf: Map[String, String] = Map.empty
+
+ /**
+ * Every python operation of a session runs on the worker's pinned thread, so a bare `runCode`
+ * outside any operation observes exactly the Spark local properties that the last operation's
+ * teardown left on that thread.
+ */
+ private def probe(worker: SessionPythonWorker, expr: String): String = {
+ val response = worker.runCode(s"print(repr($expr))")
+ assert(response.isDefined, s"no response for $expr")
+ val content = response.get.content
+ assert(
+ content.status === PythonResponse.OK_STATUS,
+ s"probe failed: ${content.getEname()}: ${content.getEvalue()}")
+ content.data.getOrElse("text/plain", "").toString.trim
+ }
+
+ /**
+ * The version of the PySpark that the worker imports, resolved the way
+ * `createSessionPythonWorker` resolves `SPARK_HOME` and `execute_python.py` extends `sys.path`.
+ * `None` when it cannot be read, so that only a known mismatch cancels a test.
+ */
+ private lazy val workerPySparkVersion: Option[String] =
+ Try(sys.env.getOrElse("SPARK_HOME", ExecutePython.defaultSparkHome)).toOption
+ .map(Paths.get(_, "python", "pyspark", "version.py"))
+ .filter(Files.exists(_))
+ .flatMap { path =>
+ """__version__[^'"]*["']([^'"]+)["']""".r
+ .findFirstMatchIn(new String(Files.readAllBytes(path), UTF_8)).map(_.group(1))
+ }
+
+ private def withSessionAndWorker(f: (Session, SessionPythonWorker) => Unit): Unit = {
+ // The worker's PySpark talks to this JVM's Spark classes over Py4J, so it cannot attach to a
+ // different Spark, as in the CI jobs that run a Spark 3.5 engine on a Spark 4.x binary.
+ assume(
+ workerPySparkVersion.forall(SPARK_ENGINE_RUNTIME_VERSION === _),
+ s"pyspark ${workerPySparkVersion.getOrElse("")} under SPARK_HOME cannot attach to" +
+ s" Spark $SPARK_ENGINE_RUNTIME_VERSION")
+
+ val sessionManager = engine.backendService.sessionManager
+ val handle = sessionManager.openSession(
+ TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V11,
+ Utils.currentUser,
+ "",
+ "127.0.0.1",
+ Map.empty)
+ val session = sessionManager.getSession(handle)
+ ExecutePython.init()
+ val worker = ExecutePython.createSessionPythonWorker(spark, session)
+ try {
+ f(session, worker)
+ } finally {
+ worker.close()
+ sessionManager.closeSession(handle)
+ }
+ }
+
+ test("KYUUBI #7684 - the operation context is cleared with null, not the empty string") {
+ withSessionAndWorker { (session, worker) =>
+ val operation =
+ new ExecutePython(session, "print(1)", false, 0L, worker, OperationHandle())
+ operation.run()
+
+ // An empty string here passes the null check in AuthZUtils.getAuthzUgi and reaches
+ // UserGroupInformation.createRemoteUser, which rejects it with `Null user`.
+ assert(probe(worker, s"spark.sparkContext.getLocalProperty('$KYUUBI_SESSION_USER_KEY')")
+ === "None")
+ assert(probe(worker, s"spark.sparkContext.getLocalProperty('$KYUUBI_STATEMENT_ID_KEY')")
+ === "None")
+ }
+ }
+
+ test("KYUUBI #7684 - the operation context is atomic with respect to the shared worker") {
+ withSessionAndWorker { (session, worker) =>
+ val operation = new ProbeExecutePython(session, worker)
+ val otherOperationRanCode = new AtomicBoolean(false)
+ val other = new Thread(new Runnable {
+ override def run(): Unit = {
+ worker.runCode("1")
+ otherOperationRanCode.set(true)
+ }
+ })
+
+ operation.withContext {
+ // The context of this operation is applied. Another operation of the same session shares
+ // this worker, so if it can run code here it can also set or clear these properties here.
+ other.start()
+ other.join(TimeUnit.SECONDS.toMillis(5))
+ assert(
+ !otherOperationRanCode.get(),
+ "another operation reached the worker while this operation held the context")
+ }
+
+ other.join(TimeUnit.SECONDS.toMillis(10))
+ }
+ }
+}
+
+/**
+ * Gives a test access to the operation context without running a python statement, so that the
+ * body of `withLocalProperties` does not take the worker lock on its own.
+ */
+private class ProbeExecutePython(session: Session, worker: SessionPythonWorker)
+ extends ExecutePython(session, "print(1)", false, 0L, worker, OperationHandle()) {
+ def withContext[T](f: => T): T = withLocalProperties(f)
+}