feat(ray): Implement dynamic scale-in for RaySwordfishActor

[huleilei]

Changes Made

This commit implements the dynamic scaling down (scale-in) functionality for RaySwordfishActor to release idle resources.

• Idle Worker Retirement: Implemented retire_idle_ray_workers in RayWorkerManager. It identifies workers that have been idle for longer than DAFT_AUTOSCALING_DOWNSCALE_IDLE_SECONDS (default: 60s) and releases them.
• Scheduler Integration: The SchedulerActor loop now periodically checks for idle workers and triggers retirement while maintaining a minimum survivor count (DAFT_AUTOSCALING_MIN_SURVIVOR_WORKERS, default: 1).
• Blacklist Mechanism: Added a pending_release_blacklist to prevent the autoscaler from immediately respawning workers that were just released.

Configuration

New environment variables added:

• DAFT_AUTOSCALING_DOWNSCALE_ENABLED: Enable/disable downscaling (default: true).
• DAFT_AUTOSCALING_DOWNSCALE_IDLE_SECONDS: Seconds a worker must be idle before retirement (default: 60).
• DAFT_AUTOSCALING_MIN_SURVIVOR_WORKERS: Minimum number of workers to keep alive (default: 1).
• DAFT_AUTOSCALING_PENDING_RELEASE_EXCLUDE_SECONDS: TTL for blacklisted worker IDs (default: 120).

Related Issues

#5683

[greptile-apps]

Greptile Summary

This PR implements dynamic scale-in for RaySwordfishActor to release idle Ray workers and reduce cluster costs. The implementation adds idle worker retirement logic with configurable thresholds, maintains a minimum survivor count, and uses a blacklist to prevent immediate respawning of released workers.

Key changes:

• Added retire_idle_ray_workers() method to WorkerManager that identifies workers idle for longer than DAFT_AUTOSCALING_DOWNSCALE_IDLE_SECONDS (default: 60s) and releases them
• Integrated downscaling into the scheduler loop with periodic checks every 1 second when no scale-up is needed
• Implemented pending_release_blacklist mechanism with TTL to prevent autoscaler from immediately respawning just-released workers
• Added head node protection to prevent retirement of workers on the Ray head node (which cannot be scaled down by Ray autoscaler)
• Final cleanup on job completion releases all idle workers and clears autoscaling requests
• Configuration via environment variables: DAFT_AUTOSCALING_DOWNSCALE_ENABLED (default: true), DAFT_AUTOSCALING_MIN_SURVIVOR_WORKERS (default: 1), DAFT_AUTOSCALING_DOWNSCALE_IDLE_SECONDS (default: 60), DAFT_AUTOSCALING_PENDING_RELEASE_EXCLUDE_SECONDS (default: 120)

The implementation is well-tested with comprehensive unit tests for the retirement logic, blacklist mechanism, and integration with the scheduler.

Confidence Score: 4/5

• Safe to merge with minor edge case considerations around idle detection timing
• Implementation is solid with proper head node protection, blacklist mechanism, and comprehensive testing. Score reflects a minor semantic gap where scheduler counts idle workers differently than worker manager filters them (by duration), which could cause fewer retirements than expected but isn't a correctness issue. The defensive programming (checking active tasks before release, clearing blacklist on scale-up demand) and extensive test coverage demonstrate quality work.
• No files require special attention. The most complex logic in worker_manager.rs is well-structured with appropriate safety checks.

Important Files Changed

Filename	Overview
src/daft-distributed/src/python/ray/worker_manager.rs	Implements retire_idle_ray_workers method with blacklist mechanism, head node protection, and idle duration filtering. Good defensive logic to prevent premature worker respawn.
src/daft-distributed/src/scheduling/scheduler/scheduler_actor.rs	Integrates downscaling into scheduler loop with environment-based configuration. Counts idle candidates based on empty active tasks, but retire_idle_ray_workers applies additional idle duration filter which may cause mismatch.
src/daft-distributed/src/python/ray/worker.rs	Adds ActorState enum, is_idle(), idle_duration(), and release() methods. Clean implementation with proper state transitions and safety checks.

Sequence Diagram

sequenceDiagram
    participant S as SchedulerActor
    participant WM as WorkerManager
    participant W as RaySwordfishWorker
    participant Ray as Ray Autoscaler
    
    Note over S: Main scheduler loop (every 1s tick)
    
    S->>WM: worker_snapshots()
    WM-->>S: List of worker states
    
    alt Has pending tasks requiring scale-up
        S->>WM: try_autoscale(bundles)
        WM->>WM: Clear pending_release_blacklist
        WM->>Ray: request_resources(bundles)
    else No scale-up needed AND downscale_enabled
        S->>S: Count idle workers (empty active_task_details)
        S->>S: Calculate num_to_retire = min(idle_count, allowed_to_retire)
        alt num_to_retire > 0
            S->>WM: retire_idle_ray_workers(num_to_retire, false)
            WM->>Ray: get_head_node_id()
            Ray-->>WM: head_node_id
            WM->>WM: Filter candidates: skip head node, check idle duration
            WM->>WM: Sort by longest idle duration
            WM->>WM: Select top N candidates
            loop For each selected worker
                WM->>W: release()
                W->>W: Check no active tasks
                W->>W: shutdown()
                W->>W: Set state to Released
                WM->>WM: Add worker_id to pending_release_blacklist
            end
            WM->>Ray: clear_autoscaling_requests()
        end
    end
    
    Note over S: On job completion (loop exit)
    alt downscale_enabled
        S->>WM: retire_idle_ray_workers(all_workers, true)
        Note over WM: force_all_when_cluster_idle=true
        WM->>Ray: clear_autoscaling_requests()
        WM->>WM: Release all idle workers
    end

Loading

[greptile-apps]

Additional Comments (3)

1. src/daft-distributed/src/python/ray/worker_manager.rs, line 287-340 (link)

   logic: Holding the mutex lock while calling worker.release(py) and Python operations can cause significant lock contention. The state mutex is held from line 288 through line 340, during which Python GIL operations occur (lines 334-340). This blocks other operations like submit_tasks_to_workers unnecessarily.

   Consider releasing the lock before Python operations:

2. src/daft-distributed/src/python/ray/worker.rs, line 138-146 (link)

   style: The release method silently returns early if there are inflight tasks without setting state or logging. This could lead to confusion during debugging.

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

3. src/daft-distributed/src/scheduling/scheduler/scheduler_actor.rs, line 138-144 (link)

   style: Environment variable parsing with defaults lacks documentation. Consider adding comments explaining these configuration options and their defaults.

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

_{8 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 48.11594% with 179 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.83%. Comparing base (7bec778) to head (fab6e3d).

Files with missing lines	Patch %	Lines
.../daft-distributed/src/python/ray/worker_manager.rs	0.00%	124 Missing ⚠️
src/daft-distributed/src/python/ray/worker.rs	0.00%	36 Missing ⚠️
...ibuted/src/scheduling/scheduler/scheduler_actor.rs	77.14%	8 Missing ⚠️
...aft-distributed/src/scheduling/scheduler/linear.rs	0.00%	6 Missing ⚠️
daft/runners/flotilla.py	20.00%	4 Missing ⚠️
src/daft-distributed/src/scheduling/worker.rs	99.24%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5903      +/-   ##
==========================================
- Coverage   72.91%   72.83%   -0.09%     
==========================================
  Files         973      973              
  Lines      126166   126494     +328     
==========================================
+ Hits        91995    92132     +137     
- Misses      34171    34362     +191     

Files with missing lines	Coverage Δ
...ft-distributed/src/scheduling/scheduler/default.rs	88.99% <100.00%> (+0.02%)	⬆️
...c/daft-distributed/src/scheduling/scheduler/mod.rs	88.04% <ø> (ø)
src/daft-distributed/src/scheduling/worker.rs	86.74% <99.24%> (+12.50%)	⬆️
daft/runners/flotilla.py	46.85% <20.00%> (-0.79%)	⬇️
...aft-distributed/src/scheduling/scheduler/linear.rs	87.50% <0.00%> (-2.25%)	⬇️
...ibuted/src/scheduling/scheduler/scheduler_actor.rs	89.22% <77.14%> (-0.96%)	⬇️
src/daft-distributed/src/python/ray/worker.rs	0.00% <0.00%> (ø)
.../daft-distributed/src/python/ray/worker_manager.rs	0.00% <0.00%> (ø)

... and 6 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[huleilei]

@colin-ho @universalmind303 help me review when you are convenient. Thanks

[srilman]

Overall this makes sense to me and I think would be useful. However, I feel like it would make more sense for the worker manager to have the logic to determine what workers should be retired vs within the scheduler and then passing into the worker manager.

[srilman]

Actually I think these configurations would be useful to have in the daft.set_runner_ray configuration on top of these environment variables

[madvart]

@huleilei - Checking to see if you are ok working on changes recommended by @srilman

[huleilei]

I'm

@huleilei - Checking to see if you are ok working on changes recommended by @srilman

Sorry, I have been reviewing this PR recently. Thank you

[codspeed-hq]

Merging this PR will not alter performance

✅ 40 untouched benchmarks
⏩ 10 skipped benchmarks¹

---

_{Comparing huleilei:hll/auto (55953ef) with main (766f0f8)}

Footnotes

1. 10 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[huleilei]

@srilman @madvart Please help me review when you have time, thanks.

[DogerW666]

This PR is also useful for us.

[srilman]

Overall, LGTM. Had 1 small thing, but once thats addressed plus the merge conflict, we can merge

[srilman]

Rather than having this code here, can we move it to the WorkerManager directly? Plus, rather than a specific retire_idle_ray_workers, can we have a generic retire_idle_workers

[madvart]

@huleilei - Thanks for your patience and work on this. Once you are able to resolve the last comment and the merge conflicts, we can merge this. Thanks again.

[huleilei]

@madvart @srilman Hello, I have made the revisions. Thank you for reviewing them again. Best regards

[huleilei]

@madvart @srilman Hello, I have made the revisions. Thank you for reviewing them again. Best regards