[SPARK-52971][PYTHON] Limit idle Python worker queue size

### What changes were proposed in this pull request?

Makes the number of idle workers in the `PythonWorkerFactory` pool configurable.

### Why are the changes needed?

Without limiting the maximum queue size, the idle worker pool can grow unbounded. Allows better control over number of workers allowed.

### Does this PR introduce _any_ user-facing change?

Yes, adds a new optional configuration entry: `spark.python.factory.idleWorkerMaxPoolSize`, from Spark 4.1.0

### How was this patch tested?

This patch adds two new test to verify behavior with and without the worker limit configuration.

### Was this patch authored or co-authored using generative AI tooling?

No

Closes #51684 from craiuconstantintiberiu/SPARK-52971-idle-pool.

Authored-by: Tibi Craiu <[email protected]>
Signed-off-by: Takuya Ueshin <[email protected]>

Author: tibicrai-db

core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala

@@ -34,7 +34,7 @@ import org.apache.spark._
 import org.apache.spark.errors.SparkCoreErrors
 import org.apache.spark.internal.Logging
 import org.apache.spark.internal.LogKeys._
-import org.apache.spark.internal.config.Python.{PYTHON_UNIX_DOMAIN_SOCKET_DIR, PYTHON_UNIX_DOMAIN_SOCKET_ENABLED}
+import org.apache.spark.internal.config.Python.PYTHON_FACTORY_IDLE_WORKER_MAX_POOL_SIZE
 import org.apache.spark.security.SocketAuthHelper
 import org.apache.spark.util.{RedirectThread, Utils}
 
@@ -98,8 +98,9 @@ private[spark] class PythonWorkerFactory(
     !Utils.isWindows && useDaemonEnabled
   }
 
-  private val authHelper = new SocketAuthHelper(SparkEnv.get.conf)
-  private val isUnixDomainSock = authHelper.conf.get(PYTHON_UNIX_DOMAIN_SOCKET_ENABLED)
+  private val conf = SparkEnv.get.conf
+  private val authHelper = new SocketAuthHelper(conf)
+  private val isUnixDomainSock = authHelper.isUnixDomainSock
 
   @GuardedBy("self")
   private var daemon: Process = null
@@ -111,7 +112,11 @@ private[spark] class PythonWorkerFactory(
   @GuardedBy("self")
   private var daemonSockPath: String = _
   @GuardedBy("self")
-  private val idleWorkers = new mutable.Queue[PythonWorker]()
+  // Visible for testing
+  private[spark] val idleWorkers = new mutable.Queue[PythonWorker]()
+  @GuardedBy("self")
+  private val maxIdleWorkerPoolSize =
+    conf.get(PYTHON_FACTORY_IDLE_WORKER_MAX_POOL_SIZE)
   @GuardedBy("self")
   private var lastActivityNs = 0L
   new MonitorThread().start()
@@ -127,7 +132,7 @@ private[spark] class PythonWorkerFactory(
   def create(): (PythonWorker, Option[ProcessHandle]) = {
     if (useDaemon) {
       self.synchronized {
-        // Pull from idle workers until we one that is alive, otherwise create a new one.
+        // Pull from idle workers until we get one that is alive, otherwise create a new one.
         while (idleWorkers.nonEmpty) {
           val worker = idleWorkers.dequeue()
           daemonWorkers.get(worker).foreach { workerHandle =>
@@ -203,8 +208,7 @@ private[spark] class PythonWorkerFactory(
       blockingMode: Boolean): (PythonWorker, Option[ProcessHandle]) = {
     var serverSocketChannel: ServerSocketChannel = null
     lazy val sockPath = new File(
-      authHelper.conf.get(PYTHON_UNIX_DOMAIN_SOCKET_DIR)
-        .getOrElse(System.getProperty("java.io.tmpdir")),
+      authHelper.sockDir,
       s".${UUID.randomUUID()}.sock")
     try {
       if (isUnixDomainSock) {
@@ -307,8 +311,7 @@ private[spark] class PythonWorkerFactory(
         if (isUnixDomainSock) {
           workerEnv.put(
             "PYTHON_WORKER_FACTORY_SOCK_DIR",
-            authHelper.conf.get(PYTHON_UNIX_DOMAIN_SOCKET_DIR)
-              .getOrElse(System.getProperty("java.io.tmpdir")))
+            authHelper.sockDir)
           workerEnv.put("PYTHON_UNIX_DOMAIN_ENABLED", "True")
         } else {
           workerEnv.put("PYTHON_WORKER_FACTORY_SECRET", authHelper.secret)
@@ -483,6 +486,15 @@ private[spark] class PythonWorkerFactory(
     if (useDaemon) {
       self.synchronized {
         lastActivityNs = System.nanoTime()
+        if (maxIdleWorkerPoolSize.exists(idleWorkers.size >= _)) {
+          val oldestWorker = idleWorkers.dequeue()
+          try {
+            stopWorker(oldestWorker)
+          } catch {
+            case e: Exception =>
+              logWarning("Failed to stop evicted worker", e)
+          }
+        }
         idleWorkers.enqueue(worker)
       }
     } else {

core/src/main/scala/org/apache/spark/internal/config/Python.scala

@@ -127,4 +127,15 @@ private[spark] object Python {
       .timeConf(TimeUnit.SECONDS)
       .checkValue(_ >= 0, "The interval should be 0 or positive.")
       .createWithDefault(0)
+
+  val PYTHON_FACTORY_IDLE_WORKER_MAX_POOL_SIZE =
+    ConfigBuilder("spark.python.factory.idleWorkerMaxPoolSize")
+      .doc("Maximum number of idle Python workers to keep. " +
+        "If unset, the number is unbounded. " +
+        "If set to a positive integer N, at most N idle workers are retained; " +
+        "least-recently used workers are evicted first.")
+      .version("4.1.0")
+      .intConf
+      .checkValue(_ > 0, "If set, the idle worker max size must be > 0.")
+      .createOptional
 }

core/src/test/scala/org/apache/spark/api/python/PythonWorkerFactorySuite.scala

@@ -19,6 +19,7 @@ package org.apache.spark.api.python
 
 import java.net.SocketTimeoutException
 
+import scala.collection.mutable
 // scalastyle:off executioncontextglobal
 import scala.concurrent.ExecutionContext.Implicits.global
 // scalastyle:on executioncontextglobal
@@ -56,4 +57,56 @@ class PythonWorkerFactorySuite extends SparkFunSuite with SharedSparkContext {
     // Timeout ensures that the test fails in 5 minutes if createSimplerWorker() doesn't return.
     ThreadUtils.awaitReady(createFuture, 5.minutes)
   }
+
+  test("idle worker pool is unbounded when idleWorkerMaxPoolSize is not set") {
+    sc.conf.remove("spark.python.factory.idleWorkerMaxPoolSize")
+
+    val factory = new PythonWorkerFactory("python3", "pyspark.worker", Map.empty, true)
+
+    assert(factory.idleWorkers.size === 0)
+
+    val mockWorkers: mutable.Queue[PythonWorker] = mutable.Queue.empty
+    try {
+      (1 to 3).foreach { _ =>
+        val mockChannel = java.nio.channels.SocketChannel.open()
+        mockChannel.configureBlocking(false)
+        mockWorkers.enqueue(PythonWorker(mockChannel))
+      }
+      mockWorkers.foreach(factory.releaseWorker)
+      assert(factory.idleWorkers.size === 3)
+
+    } finally {
+      mockWorkers.foreach(factory.stopWorker)
+    }
+  }
+
+  test("idle worker pool is bounded when idleWorkerMaxPoolSize is set") {
+    sc.conf.set("spark.python.factory.idleWorkerMaxPoolSize", "2")
+
+    val factory = new PythonWorkerFactory("python3", "pyspark.worker", Map.empty, true)
+
+    assert(factory.idleWorkers.size === 0)
+    val mockWorkers: mutable.Queue[PythonWorker] = mutable.Queue.empty
+    try {
+      (1 to 2).foreach { _ =>
+        val mockChannel = java.nio.channels.SocketChannel.open()
+        mockChannel.configureBlocking(false)
+        mockWorkers.enqueue(PythonWorker(mockChannel))
+      }
+      mockWorkers.foreach(factory.releaseWorker)
+      assert(factory.idleWorkers.size === 2)
+
+
+      val worker3 = {
+        val mockChannel = java.nio.channels.SocketChannel.open()
+        mockChannel.configureBlocking(false)
+        PythonWorker(mockChannel)
+      }
+      mockWorkers.enqueue(worker3)
+      factory.releaseWorker(worker3)
+      assert(factory.idleWorkers.size === 2)
+    } finally {
+      mockWorkers.foreach(factory.stopWorker)
+    }
+  }
 }