[joblib] Add opt-in autoscaling to the default 'ray' actor pool

[OneSizeFitsQuorum]

Add opt-in autoscaling to the joblib actor pool

Closes #64959.

What

An opt-in autoscale=True mode on the existing "ray" joblib backend:

from ray.util.joblib import register_ray
register_ray(autoscale=True, max_size=64)

import joblib
with joblib.parallel_backend("ray", n_jobs=8):
    search.fit(X, y)

Set it once at register_ray — all subsequent parallel_backend("ray") calls
use autoscaling. Per-call params (max_size, min_size, etc.) override the
register-time defaults.

Why

Ray maintainer edoakes commented on this PR: "Instead, could we make the pool
of actors autoscale? That would be preferable over supporting two backends."

This PR does exactly that — folds autoscaling into the default "ray" actor
pool (one backend), and withdraws the earlier ray_tasks task-backend prototype
(pre-release, never merged to master).

The #22048 deadlock was caused by the eager create-all + ray.get(ping)
readiness barrier, not by num_cpus per se. Removing the barrier and creating
actors lazily (capped at max_size) lets pending num_cpus=1 actors drive the
autoscaler instead of deadlocking construction. The default path
(autoscale=False) is byte-identical — no behavioral change unless opted in.

How

All changes are in python/ray/util/multiprocessing/pool.py (~130 lines) +
python/ray/util/joblib/ray_backend.py (~20 lines) +
python/ray/util/joblib/__init__.py (~10 lines):

• Slot model: _actor_pool becomes a fixed-length Optional slot list
  (len = max_size, stable) so all existing chunking/fan-out/round-robin code
  is unchanged. _run_batch lazy-creates None slots on dispatch.
• Barrier removal: _start_actor_pool autoscale branch skips the eager
  create-all + ping barrier; actors are created on demand.
• Idle reaper: a daemon thread reaps idle actors past idle_timeout_s down
  to min_size, reusing the existing _stop_actor graceful-terminate path.
• num_cpus: autoscale=True → num_cpus=1 (drives the autoscaler);
  autoscale=False → num_cpus=0 (back-compat).
• register_ray factory: register_ray(autoscale=True, max_size=64) sets
  defaults once via a factory closure; per-call params override. RayBackend
  stores the params in __init__ and forwards them to Pool in configure.
• joblib 1.2.0: uses the existing apply_async → AsyncResult →
  ResultThread path. No 1.5 dependency.

Tests

python/ray/tests/test_joblib_autoscale.py (5 tests):

1. Default-path byte-identity (autoscale=False → len == processes, all slots non-None).
2. Lazy creation (initial_size=0 → no actors at startup, lazy-create on dispatch).
3. Idle reaping (tiny idle_timeout_s → actors reaped to 0).
4. Performance benchmark (@pytest.mark.skip): warm-actor pool.map vs
   per-task cold dispatch — asserts the actor pool is faster. Skipped in CI
   (timing is environment-dependent); see local validation below.
5. AutoscalingCluster e2e: pending num_cpus=1 actors drive the real Ray
   autoscaler to add worker nodes (cluster CPU grows 2 → 6 → 10). Uses the
   same autoscaler KubeRay uses.

Local validation (not in CI)

Run on py311, joblib==1.2.0, main-repo clean _raylet.so build overlaid
with this PR's three changed files:

• Perf (criterion 1): warm-actor pool.map dispatch is ~5.1× faster
  than per-task cold dispatch (median 0.076s vs 0.388s, N=20, 3 rounds).
  Pool and task are measured separately: a fixed pool of num_cpus=1 actors
  saturates the cluster's CPUs, so concurrent num_cpus=1 tasks would
  deadlock — hence the separate measurement and the @skip in CI.
• AutoscalingCluster e2e: test_autoscaling_cluster_e2e PASSED (65s).
  Autoscaler log: head 2 CPU → Adding 2 node(s) of type cpu_worker →
  Resized to 6 CPUs → Resized to 10 CPUs.
• K8s/KubeRay e2e (scale up + scale down): KubeRay operator + a
  RayCluster (head 2 CPU, worker group minReplicas:0 / maxReplicas:3 / 2 CPU each, enableInTreeAutoscaling). The PR's three changed files were
  patched into the head pod; Pool(autoscale=True, max_size=6, idle_timeout_s=10) submitted num_cpus=1 actors.
  • Scale up: pending placements drove the KubeRay autoscaler to add
    worker pods — kubectl get pods showed 0 → 2 Running within ~10s;
    RESULTS_OK: True (50 batches correct), CLUSTER_CPUS: 6 (head 2 +
    2 worker×2).
  • Scale down: after work completes and actors idle, the PR's reaper
    (idle_timeout_s=10) reaps idle actors → worker pods go idle → the
    KubeRay autoscaler scales worker pods 2 → 0 and CLUSTER_CPUS drops
    6 → 2 (head only) within ~100s. Full 0 → 2 → 0 cycle captured.
    The joblib backend layer (RayBackend.configure) is the same path
    validated in SC3; SC4 focuses on the KubeRay provider (RayCluster CR →
    worker pod scale up + down).

Scope

~370 additions across 7 files. The core change is ~130 lines in pool.py
(one file) — no new module, no controller, no adapter, no second
ResultThread. The earlier 1336-line elastic-actor prototype and KubeRay
harness were removed (preserved locally, not in the PR).

Not in scope

• ray.util.ActorPool migration (pull-based contract; separate PR).
• KubeRay e2e harness (not a bazel target; external).
• ray_tasks task backend (withdrawn; pre-release, never merged to master).
  The effective_n_jobs hardening from that work is retained.

Duplicate-work check

No open issue/PR proposes autoscaling the joblib actor pool. Sibling: #31128
(ray.util.multiprocessing.Pool fixed-size, same root-cause family). #22048
set PoolActor to num_cpus=0 (context for the barrier this PR removes).

AI assistance

This change was prepared with AI assistance. The implementation, tests, and
this description were reviewed end-to-end before submission.

[cursor]

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

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

[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]

Instead, could we make the pool of actors autoscale? That would preferable over supporting two backends, and would likely have better and more predictable performance

[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.