[core][joblib] Support adjustable actor capacity in ray.util.multiprocessing.Pool

[OneSizeFitsQuorum]

Summary

Closes #31128.

ray.util.multiprocessing.Pool previously fixed its actor count at construction time (defaulting to the CPUs then visible to Ray). This prevented a Pool from releasing actors after a burst and, in configurations such as a zero-CPU head node, from expressing pending actor demand to the Ray cluster autoscaler.

This PR adds opt-in adjustable capacity between min_size and max_size. Supplying any capacity option selects the adjustable scheduler. A Pool without capacity options keeps the previous fixed scheduler and creates processes actors, preserving the existing default path. A lower min_size allows actors above that floor to be released after idle_timeout_s:

with Pool(
    processes=64,
    min_size=1,
    idle_timeout_s=60,
    ray_remote_args={"num_cpus": 1},
) as pool:
    results = pool.map(f, items)

The Ray Joblib backend forwards the same capacity options and maxtasksperchild. Joblib's n_jobs remains the concurrency ceiling, so max_size may lower but never raise it.

Pool capacity and cluster autoscaling remain separate. Actor resource requirements do not change implicitly: callers use ray_remote_args to request CPUs, GPUs, or custom resources when pending actors should create cluster demand.

Design

There are two deliberately separate scheduling paths:

• Pools without capacity options use the original fixed, round-robin actor scheduler.
• Pools with an explicit capacity option use the adjustable slot scheduler described below.

This keeps the common fixed Pool behavior unchanged and limits the new lifecycle machinery to callers that opt into elasticity. Non-default max_concurrency, max_restarts, max_task_retries, and get_if_exists remain supported by the fixed scheduler. They are rejected when combined with adjustable capacity because those Ray actor policies make bounded physical actor ownership ambiguous.

For adjustable pools, Ray actor mailboxes remain the task queue and ObjectRefs remain the result protocol. Batches are submitted directly to actors, including actors that are waiting for resources or running their initializer. There is no driver-side task dispatcher, task migration, recovery poller, operation registry, or batch replay. The only new background thread retires idle actors.

Each adjustable slot follows this lifecycle:

stateDiagram-v2
    [*] --> EMPTY
    EMPTY --> STARTING: create actor
    STARTING --> ACTIVE: readiness succeeds
    STARTING --> DRAINING: close or task limit
    ACTIVE --> DRAINING: close, task limit, idle timeout, or resource handoff
    DRAINING --> EMPTY: actor exit confirmed
    STARTING --> EMPTY: actor death confirmed
    ACTIVE --> EMPTY: actor death confirmed

Loading

A slot owns its actor until Ray confirms the actor has exited. STARTING actors may already have accepted mailbox work, while ACTIVE only means readiness has been observed. Generation checks prevent delayed callbacks from modifying a later actor that reuses the same slot.

The reaper and readiness, batch, and termination ObjectRef callbacks observe the actor set through weak references. A callback retains only a slot index and generation—not the slot or actor handle—then resolves the current slot after confirming the actor set still exists. An abandoned Pool therefore does not remain alive solely because an actor is permanently pending.

An idle ACTIVE actor normally waits for idle_timeout_s before retirement. The one normal exception is resource handoff: when a STARTING actor already owns accepted work but cannot acquire resources, an ACTIVE actor with no outstanding work drains immediately. Because accepted mailbox work cannot migrate between actors, this releases Pool resources that might otherwise strand the pending work until the idle timeout. The capacity floor is rechecked before each such retirement.

The scheduler prefers an already-ready actor while it is at most one batch more loaded than a STARTING actor. This intentionally avoids imposing actor startup latency on small tasks; a short warm burst can consequently use less than the eventual maximum parallelism.

maxtasksperchild counts accepted batches. Once an actor reaches the limit, its slot stops accepting work, drains accepted batches, and becomes reusable only after actor-exit confirmation.

The Joblib backend constructs the Ray Pool directly instead of mutating Joblib's process-global PicklingPool bases. It owns backend arguments instead of relying on version-specific Joblib base-class storage, then merges them with arguments supplied to configure(). Pool options such as initializer, initargs, ray_address, capacity settings, and maxtasksperchild therefore reach the Ray Pool on both Joblib 1.2 and 1.5.

Failure semantics and lifecycle boundaries

• close() rejects new submissions and preserves accepted work; terminate() may abort accepted work; join() waits for actor cleanup.
• A user-function exception fails its result without poisoning a live actor. A confirmed ActorDiedError releases the slot so capacity can be restored when policy requires it.
• ActorUnavailableError is local to the affected ObjectRef and therefore fails that result or Joblib batch. It does not prove that the actor exited, so the slot retains the original handle and is not reused; it also does not permanently poison the Pool. A later call may succeed, report unavailability again, or confirm actor death.
• The Pool never replays an accepted batch. For Joblib, a failed batch follows Joblib's native behavior: the current Parallel call aborts and reports the error. Any Ray-level replay requires an explicit non-default max_task_retries, which remains on the fixed scheduler rather than adjustable capacity.
• The adjustable scheduler does not add quarantine health checks or recovery polling. A persistently unavailable actor can continue to occupy its bounded slot until Ray reports its death or the caller shuts down the Pool.
• An accepted batch is not migrated from one actor mailbox to another. An infeasible resource request or initializer that never returns can therefore remain pending until the caller terminates the Pool. A result timeout only bounds the caller's wait; it does not cancel the work.
• A sequence of submissions, such as chunks from one map_async() call, is not a distributed transaction. Earlier chunks may already be accepted if a later submission fails.
• The Pool does not reconnect actor handles after Ray session replacement, add workload-independent backpressure, or impose an independent deadline on join().

Validation

The focused autoscaling suite contains 29 test functions (39 parameterized cases) covering:

• capacity validation and rejection of ambiguous adjustable/advanced actor-option combinations;
• selection and behavior of the unchanged default fixed scheduler;
• scale-up, minimum capacity, idle retirement, scale-from-zero, and reuse;
• resource handoff to pending work and permanently infeasible requests;
• graceful close with pending actors, actor-death recovery, and maxtasksperchild recycling;
• deterministic slot invariants across startup, activity, draining, and confirmed exit;
• callback-registration failure, ambiguous termination, exit-before-reuse, and deterministically interleaved close/terminate convergence;
• garbage collection of an abandoned elastic Pool, including the actor set and idle reaper while an actor is permanently pending;
• batch-local actor unavailability without unsafe slot reuse or Pool poisoning;
• exception values as successful data;
• Joblib capacity forwarding, stored backend Pool arguments, recycling, failure handling, backend reuse, and argument filtering;
• pending demand from a zero-CPU head followed by worker arrival.

The latest local source-overlay validation used an installed Ray wheel for compiled extensions and the PR's Python modules and tests. The focused suite was invoked with --noconftest and an equivalent local shutdown_only fixture because the checkout's repository-wide conftest imports build-only modules that do not match the installed wheel. The complete suite passed with the repository-pinned Joblib 1.2.0; the Joblib integration regressions also passed with Joblib 1.5.3:

pytest.main(
    ["--noconftest", "-q", "python/ray/tests/test_joblib_autoscale.py"],
    plugins=[LocalIsolation()],
)
Joblib 1.2.0: 39 passed in 167.11s

python -m pytest --noconftest -q \
  python/ray/tests/test_joblib_autoscale.py::test_joblib_forwards_backend_pool_arguments \
  python/ray/tests/test_joblib_autoscale.py::test_joblib_respects_capacity_and_maxtasksperchild \
  python/ray/tests/test_joblib_autoscale.py::test_joblib_backend_can_be_reused_after_task_failure
Joblib 1.5.3: 3 passed in 24.77s

PYTHONPATH=python/ray/tests python -m pytest --noconftest -q python/ray/tests/test_multiprocessing.py
16 passed in 35.67s

pre-commit run --files doc/source/ray-more-libs/joblib.rst doc/source/ray-more-libs/multiprocessing.rst python/ray/tests/test_joblib_autoscale.py python/ray/util/joblib/ray_backend.py python/ray/util/multiprocessing/pool.py
Passed

python -m py_compile python/ray/util/multiprocessing/pool.py python/ray/util/joblib/ray_backend.py python/ray/tests/test_joblib_autoscale.py
Passed

git diff --check
Passed

The wheel supplied _raylet and other compiled extensions while the workspace supplied the changed pure-Python modules. Built-source CI remains authoritative for repository-wide and documentation integration.

Duplicate-work check

Issue #31128 and open PRs in the Joblib/Pool autoscaling area were checked. No competing open implementation was found.

AI assistance and human accountability

AI assistance was used to develop, review, test, and document this change.

• The human submitter has reviewed every changed line and can defend the final design and implementation.
• The human submitter has independently confirmed the final test results before requesting maintainer approval.

These items intentionally remain unchecked for the human submitter to complete.

[gemini-code-assist]

Code Review

This pull request introduces an elastic Ray task backend ("ray_tasks") for joblib, allowing each batch to run as a short-lived Ray task instead of using a fixed-size actor pool. This enables better integration with the Ray autoscaler and allows idle workers to scale down. The changes include documentation, unit tests, and the backend implementation itself. Feedback focuses on ensuring thread safety when modifying the in-flight task list concurrently from background threads, wrapping the RayBatchedCalls import in a try-except block to handle potential import failures, and making effective_n_jobs more robust when Ray is not yet initialized or when CPU resources are missing from the cluster state.

[edoakes]

@OneSizeFitsQuorum do you use the joblib backend?

[OneSizeFitsQuorum]

@edoakes Yes, I’m actively using the joblib backend. My goal is to take tasks that are currently parallelized locally with joblib and scale them out to a Kubernetes-backed Ray cluster by allowing the worker pool to grow with demand. I found that the current pool cannot scale elastically on Kubernetes, which led me to investigate how this could be improved.

I agree that making the actor pool autoscale would be preferable if we can retain a single backend. However, the current pool uses fixed num_cpus=0 actors and follows fixed-size multiprocessing.Pool semantics, so this may affect resource accounting, startup behavior, actor lifetime, and potentially existing default behavior.

I’m happy to explore an autoscaling actor pool, potentially as an opt-in mode first. Would you prefer implementing it in ray.util.multiprocessing.Pool, or keeping it internal to the joblib backend?

[edoakes]

I’m happy to explore an autoscaling actor pool, potentially as an opt-in mode first. Would you prefer implementing it in ray.util.multiprocessing.Pool, or keeping it internal to the joblib backend?

I would suggest prototyping/validating it in whatever way is easiest. Could be in the multiprocessing pool, actor pool, or just hand-written. Once you have validated the behavior/performance, we can decide the best abstraction & layering. Off the top of my head, I would think it probably makes sense to have the actor pool support autoscaling optionally (and use it for the joblib backend), and then the multiprocessing pool could just be a fixed-size actor pool.

[OneSizeFitsQuorum]

@edoakes Thanks for the idea! I'll test it out and evaluate it.

[gemini-code-assist]

Code Review

This pull request introduces an experimental opt-in autoscaling mode (autoscale=True) for the Ray joblib backend and ray.util.multiprocessing.Pool, allowing the actor pool to grow on demand and shrink when idle. The feedback highlights several critical issues in the implementation: a strong reference cycle in the background reaper thread that prevents garbage collection of the Pool instance, a bug in task count tracking when replacing actors that reach maxtasksperchild, and a race condition where active actors executing long-running tasks could be prematurely reaped. To address these, it is recommended to use a static method with a weak reference for the reaper thread and to track active object references to ensure only truly idle actors are reaped.

[OneSizeFitsQuorum]

Additional adversarial validation

I ran two manual adversarial test matrices against the latest pull-based
implementation (3b1712452d). These supplement the focused tests listed in the
PR description.

Local Ray state-machine matrix

Environment: Python 3.11, local Ray cluster with 2 CPUs, using the clean
main-checkout _raylet.so and the PR Python sources.

Scenario	Stress condition	Result
Graceful close with backlog	map_async, 80 one-item batches, max_size=8, then immediate close() and join()	All results returned in order; dispatcher exited
Terminate with multiple results	Eight independent asynchronous maps, 80 total batches, termination with actors running and work queued	Every result became ready with an error; no blocked result thread
Lazy iterator stop race	Repeated imap races against both close() and terminate()	No stranded queue entries or blocked next() calls
Capacity below max_size plus recycling	2 CPUs, max_size=12, maxtasksperchild=2, 100 one-item batches	Completed in order using ready actors; pending actors did not receive work
Immediate idle reaping	idle_timeout_s=0, six consecutive bursts with long-running batches	Running actors were not reaped; all bursts completed
Concurrent driver threads	Eight driver threads sharing one Pool, 400 tasks	No lost or duplicated batches
Exact-multiple and empty imap	Empty iterator and 10 one-item chunks with only 2 CPUs but max_size=8	All results returned and the dynamic ResultThread exited
Unexpected actor death	Kill a ready actor before follow-up work	Follow-up result completed or failed within the timeout; no hang
Dropped Pool handle	Drop the last Pool reference with four asynchronous results outstanding	Results completed; Pool and dispatcher were garbage-collected

The temporary adversarial harness ran 10 test instances successfully after
the fixes below. The permanent correctness suite then passed:

17 passed, 1 deselected

This matrix found two related iterator lifecycle bugs:

1. An autoscaling imap whose length was an exact multiple of chunksize
   could leave its dynamic ResultThread alive indefinitely.
2. Sending a plain end sentinel when exhaustion was detected could race ahead
   of ObjectRefs that the pull dispatcher had not registered yet, causing
   premature completion or a blocked iterator.

The final implementation sends the known final chunk count together with a
wakeup, so the ResultThread waits for late-dispatched ObjectRefs without
leaking on empty input. The regression is
test_autoscale_imap_exact_chunks_stops_result_thread.

KubeRay control-plane matrix

Environment:

• kind single-node cluster with 8 host CPUs
• KubeRay operator 1.6.2
• Ray 2.52.0
• head: 2 CPUs
• worker group: 0–3 replicas, 2 CPUs per worker
• actual schedulable Ray capacity: 6 CPUs because Kubernetes system/operator
  pods consume part of the node

Scenario	Expected property	Result
Head-only workload	max_size=2 must not create workers	Stayed at 2 Ray CPUs and 0 workers
Partial scale-up	Demand for four actors should add one worker	Ray CPU reached 4; workload completed
Cold scale-up	Start from zero workers and submit sustained demand	Ray CPU changed 2 → 6
max_size above cluster capacity	max_size=12 with only 6 schedulable Ray CPUs must not strand work on pending actors	120 ordered results completed on 6 ready actors across 3 Ray nodes
Actor scale-to-zero and reuse	Reap every Pool actor, then submit another burst to the same Pool	Both bursts completed with newly created actors
Shared-Pool client concurrency	Eight driver threads submit 200 tasks concurrently	All task groups completed without loss or duplication
Graceful close during demand	Close immediately after submitting 100 batches	All batches completed; join() and dispatcher exited
Terminate during demand	Terminate with running actors and a large backlog	Returned in 0.12 s; result failed bounded with ActorDiedError
joblib concurrency semantics	backend max_size=12, call-level n_jobs=4	Exactly 4 actors processed 80 tasks across head and one worker
Actor churn	maxtasksperchild=1 with 40 one-item batches	40 tasks completed on 40 actor IDs
Independent clients	Two separate driver processes each create a Pool and submit concurrently	Both completed; one used head-only capacity while the other also used a worker

Observed Kubernetes lifecycle:

worker pods:
0
  -> 3 requested (2 Running + 1 Pending due to the kind CPU limit)
  -> 2
  -> 1
  -> 0

Ray CPUs:          2 -> 6 -> 2
desired workers:   0 -> 3 -> 0
pending demands:   none at completion

Two deployment observations are worth calling out:

• joblib must be installed consistently on the head and worker images. A
  head-only installation correctly failed worker deserialization with
  ModuleNotFoundError; installing joblib 1.2.0 on each active worker made the
  real joblib scenario pass.
• On this static kind node, the third requested worker could not be scheduled.
  Pool demand disappeared promptly after completion, but KubeRay cleaned up
  the already-allocated pending pod in stages. This delayed pod convergence
  beyond the 60-second Ray-node idle timeout, but it did converge to zero and
  did not block task completion. A Kubernetes Cluster Autoscaler would
  normally add node capacity; hard-quota clusters should set max_size
  accordingly.

No additional Pool correctness issue was found in the KubeRay matrix.

[OneSizeFitsQuorum]

@edoakes @robertnishihara This PR is now ready for review.

After several iterations, adjustable actor capacity is now integrated as a native part of ray.util.multiprocessing.Pool, rather than being modeled as separate fixed and elastic implementations. The existing default behavior remains unchanged.

I’ve also put significant effort into simplifying the implementation, organizing the main path ahead of compatibility handling, and keeping the code and documentation readable.

Finally, the behavior has been validated across Pool and Joblib integration scenarios, including scaling, idle retirement, pending-resource handoff, actor recycling, failures, and compatibility.

I’d appreciate another review when you have time. Thanks!

[gemini-code-assist]

Code Review

This pull request introduces dynamic actor capacity management (autoscaling) for the Ray backend in Joblib and multiprocessing.Pool, allowing the pool to scale actors between min_size and max_size and release idle actors after idle_timeout_s. The feedback highlights critical deadlock risks where self._pool_lock is held during blocking submit operations, which would prevent closing or terminating the pool. Additionally, the reviewer suggests using isinstance() instead of direct type comparisons for PEP 8 compliance, and warns that treating transient RayActorErrors as terminal could permanently fail the pool unnecessarily.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

^{Reviewed by Cursor Bugbot for commit 2147276. Configure here.}