Skip to content

[train] Pin elastic worker groups to the reservation snapshot that sized them - #66762

Draft
justinvyu wants to merge 1 commit into
masterfrom
justinvyu/train-reservation-pins
Draft

justinvyu wants to merge 1 commit into
masterfrom
justinvyu/train-reservation-pins

Conversation

@justinvyu

@justinvyu justinvyu commented Oct 5, 2026 •

Copy link
Copy Markdown
Contributor

Description

test_elastic_e2e has been timing out on CI since #66002, which pins Train V2 workers to AutoscalingCoordinator reservations. The elastic worker group start waits the full 60s (RAY_TRAIN_WORKER_GROUP_START_TIMEOUT_S) whenever the reservation changes between the scaling decision and pinning. The decision is sized from the cached reservation, while the pins come from a separate recompute=True poll that requires an exact match, so:

  1. The reservation grows past the decision. Elastic requests max_workers bundles, so after ResizeDecision(14) the reservation can reach 32 slots, and len(label_selectors) != num_workers never passes.
  2. The reservation shrinks below the decision. After a node dies, the cached view still counts it, and pinning waits for workers that don't exist.

This PR has the elastic policy take the size and the pins from the same snapshot, so they can't disagree:

  • Pins on ResizeDecision. A new ResizeDecision.label_selectors field holds one node pin per worker, built from the reservation that chose num_workers using the same _floor_slots math. The STRICT_PACK/STRICT_SPREAD layout check from [train] Pin Train V2 workers to autoscaling coordinator reservations … #66002 is kept and now lives in ScalingPolicy._get_reserved_label_selectors.
  • Elastic decides and pins from a fresh (recompute=True) view. In the non-running path it reads the fresh view only once the cached one looks ready, so waiting doesn't trigger a full coordinator tick every loop. The running path reads fresh every elastic_resize_monitor_interval_s.
  • Controller applies the decision's pins. Decisions without pins (FixedScalingPolicy) still use the existing poll, get_reserved_bundle_label_selectors, now built on the same helper. The stacked PR above this one moves the fixed policy to the same model and removes the poll.

Compared with polling for an exact match, this also avoids an extra restart: the group starts at the size the reservation supports now, not at a stale size that elastic would resize one monitor interval later.

Related issues

Supersedes #66601 (agreed with its author), which fixes the same flake by picking a subset of a larger reservation with _compute_reservations and adding a relative _FIT_EPSILON to the coordinator's shared _bundle_can_fit_on_node. This PR doesn't change the coordinator or add a Train → Ray Data _compute_reservations dependency.

Part of a stack:

  1. This PR: pins on ResizeDecision + elastic.
  2. Fixed scaling policy waits for its reservation in the decision (stacked on this one).

Related, independent: early stop when a pinned node dies (linked in a comment).

Testing

All passing locally (macOS, ray_dev_py311):

pytest python/ray/train/v2/tests/test_autoscaling_coordinator_client.py \
       python/ray/train/v2/tests/test_elastic_scaling_policy.py \
       python/ray/train/v2/tests/test_controller.py \
       python/ray/train/v2/tests/test_controller_callback_behaviour.py \
       python/ray/train/v2/tests/test_failure_policy.py      # 93 passed
pre-commit run --files <changed files>                      # clean

test_scheduling.py, test_state.py and test_elastic_e2e.py were run on the top of the stack (see the stacked PR).

AI assistance (Claude Code) was used to write this change.

🤖 Generated with Claude Code

…zed them

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Signed-off-by: Justin Yu <justinvyu@anyscale.com>
@justinvyu
justinvyu added this pull request to stack #66764 October 5, 2026 23:41
@justinvyu justinvyu changed the title justinvyu/train reservation pins [train] Pin elastic worker groups to the reservation snapshot that sized them Oct 5, 2026
@justinvyu

Copy link
Copy Markdown
Contributor Author

Related, independent fix for the remaining test_elastic_e2e stall (a pin made while the GCS still reports a just-removed node as alive): #66765. Can merge in either order.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant