Wait for resources to become available between actor pool stages

[ayushdg]

Description

Ray actor pool executor goes stage by stage kills all actors and the end of the stage and restarts new set of actors for the next stage.

The executor gauges how many actors it can start by looking at cluster.available_resources.
But given that ray.kill is async and it may take time for the available resources to update, the default sleep of 0.2s may not be enough for all machines to free up existing processes.
This may lead to fewer resources being spun up for future stages or no resources available at all leading to errors.

The new proposed solution adds the following:

1. interval -> Poll resources every interval seconds
2. timeout -> If we do not get all intended resources by timeout s the stage starts with whatever is available.

Here's the new process:

1. Calculate the available resources since pipeline execution - Update with a higher value if more become available later.
2. Calculate the number of actors a stage could have with these max seen resources.
   a. If a stage now has enough resources to get these actors, start the stage.
   b. If a stage doesn't have enough resources to get these actors, sleep for interval seconds and recheck.
   c. If a stage doesn't have enough resources to get these actors even after timeout, start as many actors as possible with the given shape or error out if even 1 cannot be started.

Usage

# Add snippet demonstrating usage

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[ayushdg]

/ok to test 0f04448

[ayushdg]

@greptile review

[greptile-apps]

Greptile Summary

This PR adds resource high-water tracking and bounded polling between Ray actor-pool stages so asynchronous actor teardown can release capacity before the next pool is sized. Major changes:

• Adds configurable resource wait timeout and polling interval behavior.
• Reuses the observed resource baseline when sizing ordinary, shuffle, RAFT, and LSH actor pools.
• Adds tests for polling, timeout degradation, unavailable resources, reservation accounting, and LSH iterations.

Confidence Score: 4/5

The PR should not merge until polling is bounded by the configured timeout.

The new loop checks its deadline before sleeping for an unrestricted polling interval, allowing valid configuration values to delay degradation or failure well beyond the promised timeout.

Files Needing Attention: nemo_curator/backends/ray_actor_pool/utils.py

Important Files Changed

Filename	Overview
nemo_curator/backends/ray_actor_pool/executor.py	Integrates resource-baseline tracking and configurable waiting into every actor-pool stage path.
nemo_curator/backends/ray_actor_pool/utils.py	Implements resource snapshots, high-water tracking, and timeout-based pool sizing, but polling can substantially exceed the configured timeout.
tests/backends/ray_actor_pool/test_executor.py	Covers the principal resource polling and fallback paths but does not exercise an interval longer than the timeout.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Finish stage] --> B[Kill stage actors]
  B --> C[Update resource high-water baseline]
  C --> D[Calculate intended actor count]
  D --> E[Poll currently available resources]
  E --> F{Enough for intended pool?}
  F -->|Yes| G[Start intended actor pool]
  F -->|No| H{Timeout reached?}
  H -->|No| I[Sleep for polling interval]
  I --> E
  H -->|Yes, at least one actor fits| J[Start degraded actor pool]
  H -->|Yes, no actor fits| K[Raise TimeoutError]

Loading

_{Reviews (1): Last reviewed commit: "Reduce interval and timeout defaults" | Re-trigger Greptile}

[greptile-apps]

Polling overshoots resource timeout

When resource_wait_interval_s exceeds the time remaining before resource_wait_timeout_s, this unbounded sleep delays degradation or failure beyond the configured timeout; for example, a 5-second timeout with a 60-second interval blocks for about 60 seconds.

Suggested change

	time.sleep(interval)
	time.sleep(min(interval, max(0.0, deadline - time.monotonic())))