[SPARK-52971] [PYTHON] Limit idle Python worker queue size #51684
Conversation
| mockWorkers.foreach(_.stop()) | ||
| worker3.stop() |
There was a problem hiding this comment.
Shall we surround the above with try block and put these in the final clause to make sure all the workers are cleaned up?
There was a problem hiding this comment.
Thanks, made this change.
| @@ -33,6 +33,12 @@ import org.apache.spark.util.ThreadUtils | |||
| // Tests for PythonWorkerFactory. | |||
| class PythonWorkerFactorySuite extends SparkFunSuite with SharedSparkContext { | |||
|
|
|||
| private def getIdleWorkerCount(factory: PythonWorkerFactory): Int = { | |||
| val field = factory.getClass.getDeclaredField("idleWorkers") | |||
There was a problem hiding this comment.
I guess it's ok to make idleWorkers as private[spark]? cc @HyukjinKwon
There was a problem hiding this comment.
but let's comment that this exposed for testing purpose.
There was a problem hiding this comment.
Thank you, changed to private[spark] and added comment.
core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala
Outdated
Show resolved
Hide resolved
| // Visible for testing | ||
| private[spark] val idleWorkers = new mutable.Queue[PythonWorker]() | ||
| @GuardedBy("self") | ||
| private val idleWorkerPoolSize = authHelper.conf.get(PYTHON_FACTORY_IDLE_WORKER_MAX_POOL_SIZE) |
There was a problem hiding this comment.
why do we access the spark conf via authHelper?
There was a problem hiding this comment.
Thanks, changed everywhere to use conf instead of authHelper.conf when needed.
| // Visible for testing | ||
| private[spark] val idleWorkers = new mutable.Queue[PythonWorker]() | ||
| @GuardedBy("self") | ||
| private val idleWorkerPoolSize = authHelper.conf.get(PYTHON_FACTORY_IDLE_WORKER_MAX_POOL_SIZE) |
There was a problem hiding this comment.
could you adjust the variable name to reflect that it's the maximum size?
| @@ -484,6 +487,15 @@ private[spark] class PythonWorkerFactory( | |||
| self.synchronized { | |||
| lastActivityNs = System.nanoTime() | |||
| idleWorkers.enqueue(worker) | |||
| if (idleWorkerPoolSize.exists(idleWorkers.size > _)) { | |||
There was a problem hiding this comment.
should this be >=? Otherwise you'll have idleWorkerPoolSize + 1 worker in the pool
There was a problem hiding this comment.
Yes, but should not be a problem at the moment.
If we are at maxIdleWorkerPoolSize workers in the queue and we add 1 more, we will go over the limit and then remove the worker.
Not a problem now but could be an issue in the future if the queue implementation is changed to one that actually enforces this and throws an exception.
Changed.
There was a problem hiding this comment.
Makes sense, I see that the worker was enqueued anyways before you were doing this check, so in either case it will be at most workers
| @@ -484,6 +487,15 @@ private[spark] class PythonWorkerFactory( | |||
| self.synchronized { | |||
| lastActivityNs = System.nanoTime() | |||
| idleWorkers.enqueue(worker) | |||
| if (idleWorkerPoolSize.exists(idleWorkers.size > _)) { | |||
There was a problem hiding this comment.
Makes sense, I see that the worker was enqueued anyways before you were doing this check, so in either case it will be at most workers
|
Thanks! merging to master. |
What changes were proposed in this pull request?
Makes the number of idle workers in the
PythonWorkerFactorypool 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.0How 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