spark-test commented on code in PR #59713:
URL: https://github.com/apache/spark/pull/59713#discussion_r4352743294
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDTFSuite.scala:
##########
@@ -118,6 +130,85 @@ class PythonUDTFSuite extends SharedSparkSession {
Some(StructType.fromDDL("v ARRAY<VARIANT>"))
)
+ test("SPARK-54055: local session eviction does not stop an active Python
UDTF planner") {
+ assume(shouldTestPythonUDFs && !Utils.isWindows)
+ withTempDir { tempDir =>
+ val uuid = UUID.randomUUID().toString
+ val started = new File(tempDir, "started")
+ val resume = new File(tempDir, "resume")
+ val script = s"""
+ |import os
+ |import time
+ |from pyspark.sql.types import StructType
+ |from pyspark.sql.udtf import AnalyzeResult
+ |
+ |class BlockingAnalyze:
+ | @staticmethod
+ | def analyze():
+ | with open("${started.getCanonicalPath}", "w") as started:
+ | started.write(str(os.getpid()))
+ | deadline = time.monotonic() + 30
+ | while not os.path.exists("${resume.getCanonicalPath}"):
+ | if time.monotonic() > deadline:
+ | raise TimeoutError("Planner was not resumed")
+ | time.sleep(0.01)
+ | return AnalyzeResult(StructType().add("value", "integer"))
+ |
+ | def eval(self):
+ | yield 1,
+ |""".stripMargin
+ val udtf = createUserDefinedPythonTableFunction("BlockingAnalyze",
script, None)
+ val loader = new MutableURLClassLoader(Array.empty,
getClass.getClassLoader)
+ val session = new IsolatedSessionState(
+ uuid, loader, loader, new HashMap[String, Long], new HashMap[String,
Long],
+ new HashMap[String, Long], None)
+ IsolatedSessionState.sessions.put(uuid, session)
+ val pool =
ThreadUtils.newDaemonSingleThreadExecutor("python-planner-eviction")
+ implicit val executionContext: ExecutionContext =
ExecutionContext.fromExecutor(pool)
+ val planning = Future {
+ spark.withActive {
+ JobArtifactSet.withActiveJobArtifactState(JobArtifactState(uuid,
None)) {
+ val runner = new UserDefinedPythonTableFunctionAnalyzeRunner(
+ udtf.name, udtf.func, Seq.empty, Seq.empty,
spark.sessionState.sqlParser)
+ runner.runInPython(useDaemon = true)
+ }
+ }
+ }
+ try {
+ eventually(timeout(10.seconds), interval(10.milliseconds)) {
+ assert(started.isFile)
+ assert(Files.readString(started.toPath).nonEmpty)
+ }
+ val child =
ProcessHandle.of(Files.readString(started.toPath).toLong).get()
+ val daemon = child.parent().get()
+ val directory = new File(SparkFiles.getRootDirectory(), uuid)
+ assert(directory.isDirectory)
+
+ // No executor tasks hold this session, but the driver is still
planning its query.
+ session.markEvicted()
+ assert(!directory.exists())
+ assert(!IsolatedSessionState.sessions.containsKey(uuid))
+ assert(child.isAlive)
+ assert(daemon.isAlive)
+ Files.writeString(resume.toPath, "resume")
+ assert(ThreadUtils.awaitResult(planning, 10.seconds).schema ==
+ StructType.fromDDL("value int"))
+ } finally {
+ Files.writeString(resume.toPath, "resume")
+ pool.shutdown()
+ try {
+ assert(pool.awaitTermination(10, TimeUnit.SECONDS))
+ } finally {
+ pool.shutdownNow()
+ if (IsolatedSessionState.sessions.containsKey(uuid)) {
+ session.markEvicted()
+ }
+ SparkEnv.get.removePythonWorkerFactories(uuid)
Review Comment:
**Nit (P3):** `removePythonWorkerFactories` filters out `driverOwned` keys,
and the factory this test creates comes from `PythonPlannerRunner` with
`driverOwned = true`, so this cleanup call does nothing for it. The planner
daemon (whose working directory the test just deleted) and its pooled worker
stay alive in the shared SparkEnv until the suite's SparkContext stops, even
though this line reads as if it retires them. The test needs a way to stop and
remove the driver-owned factory for `uuid`, and could then assert that the
daemon observed above has exited.
**Verification:**
- **Regression:** After the test's cleanup, the planner daemon process
observed during the test has exited and SparkEnv holds no factory for the
test's session UUID.
- **Compatibility:** Executor session eviction still leaves an active
driver-owned planner running, as the test already asserts.
##########
core/src/main/scala/org/apache/spark/SparkEnv.scala:
##########
@@ -370,8 +389,19 @@ class SparkEnv (
daemonModule: String,
envVars: Map[String, String],
worker: PythonWorker): Unit = {
+ destroyPythonWorker(
+ pythonExec, workerModule, daemonModule, envVars, worker, driverOwned =
false)
+ }
+
+ private[spark] def destroyPythonWorker(
+ pythonExec: String,
+ workerModule: String,
+ daemonModule: String,
+ envVars: Map[String, String],
+ worker: PythonWorker,
+ driverOwned: Boolean): Unit = {
synchronized {
- val key = PythonWorkersKey(pythonExec, workerModule, daemonModule,
envVars)
+ val key = PythonWorkersKey(pythonExec, workerModule, daemonModule,
envVars, driverOwned)
Review Comment:
**Non-blocking (P2):** When the factory for `key` has already been removed
by `removePythonWorkerFactories`, this `pythonWorkers.get(key).foreach(...)` is
a silent no-op, and the same holds for the `releasePythonWorker` overload
below. A reachable sequence is: `PythonRunner.MonitorThread` wins
`releasedOrClosed.compareAndSet` for an interrupted task (so the task
completion listener skips `worker.stop()`), the task then completes, and
`Executor`'s finally block calls `isolatedSession.release()` on an evicted
session. `cleanup()` removes and stops the factory before the monitor's
`destroyPythonWorker` gets the SparkEnv lock. `factory.stop()` kills the daemon
(or the simple workers), but it never closes this `PythonWorker`'s JVM channel
and selector, so they leak and nothing is logged. At the merge base the factory
was never removed at runtime, so this call always reached `stopWorker`, which
ends with `worker.stop()`. The PR description says workers that finish after
their factory wa
s removed are closed, but nothing does that. Falling back to `worker.stop()`
when the lookup misses, for both destroy and release, would make that hold.
**Verification:**
- **Regression:** After the session's factories are removed, destroying and
releasing one of its workers each close that worker's channel. Neither is
silently ignored.
- **Behavior:** With the factory still registered, release still returns the
worker to that factory's pool and destroy still goes through the factory's
stopWorker.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
