[Train / Autoscaler] Fix 60s worker group startup stall when the elastic reservation changes after a resize decision

[pseudo-rnd-thoughts]

Description

test_elastic_e2e started timing out on CI about 45% of the time after #66002, which pins Train V2 workers to autoscaling coordinator reservations. The root cause is that elastic worker-group starts stall for the full 60s (RAY_TRAIN_WORKER_GROUP_START_TIMEOUT_S) whenever the reservation changes between the resize decision and pinning. This was due to the scaling policy required that exacting the resize labels be available (if len(label_selectors) != num_workers:) and if too many or too few then the whole thing would stall until the startup would timeout causing a new resize attempt that would pass.

More detail on the two cases:

1. The reservation grows past the decision. The elastic policy requests max_workers bundles, so its reservation keeps growing as nodes join. _try_get_selectors required len(label_selectors) == num_workers exactly. A ResizeDecision(14) followed by more nodes above 14 joining therefore never counted as ready. It polled for 60s, raised WorkerGroupStartupTimeoutError and burned a retry. In the failing CI run, this cost 60 of the 180 seconds:

   Deciding to resize the worker group from 8 -> 14 workers.
   The worker group startup timed out after 60.0 seconds waiting for 14 workers.
   Detected ready resources for 32 workers ... Started training worker group of size 32
   

   Cursor Bugbot flagged this on [train] Pin Train V2 workers to autoscaling coordinator reservations … #66002 ("Extra reserved slots block pinning"), but the comment was never addressed.

2. The reservation shrinks below the decision. make_decision_for_non_running_worker_group sized the group from the cached reservation (up to 1s old), but pinning uses recompute=True. After a node dies, the decision still counts the dead node's slots, so pinning waits 60s for workers that don't exist. The rewritten e2e test below exposed this issue: after removing one node, the controller decided on 12 workers when only 8 remained.

[gemini-code-assist]

Code Review

This pull request refactors the autoscaling coordinator's reservation logic by introducing a helper function _sum_bundles and updating _compute_reservations to be more flexible. It also improves the elastic scaling policy by ensuring worker counts are verified against a fresh view of reserved resources to avoid stale data issues. I have identified a potential AttributeError in the new check that should be addressed.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

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

[iamjustinhsu]

Thanks!

[justinvyu]

Thanks for tracking this down. Both stalls come from the same thing: the ResizeDecision.num_workers and the label selector pins are computed from different reservation snapshots. The size comes from the cached view in the policy, and the pins come from a recompute=True poll in _try_get_selectors. Any change between the two becomes a 60s timeout.

The approach in this PR is to let the pinning step tolerate a mismatch, but that adds some complexity around using a larger reservation for a smaller number of desired workers. Could we instead just take one snapshot and derive the label selector pins along with the ResizeDecision.num_workers value, so that by construction there can't be a mismatch?

Proposal for this PR: compute elastic's pins in the scaling policy and put them on the decision

@dataclass
class ResizeDecision(ScalingDecision):
    num_workers: int
    resources_per_worker: Dict[str, float]
    # Per-worker node pins taken from the same reservation snapshot that chose
    # `num_workers`. None = not pinned by the policy (TPU, zero-resource
    # workers, or a policy that doesn't compute pins yet -- see below).
    label_selectors: Optional[List[LabelSelector]] = None

• Elastic: ElasticScalingPolicy reads the reservation once (recompute=True), counts workers from it, and builds selectors from that same snapshot with _reserved_resources_to_bundle_label_selectors. The "grew" and "shrank" cases can't happen anymore to lead to mismatched len(label_selectors).
• Controller: _start_worker_group uses decision.label_selectors when they're there. It keeps the existing poll as a fallback for the fixed policy. This fallback can be removed after the followup to unify the FixedScalingPolicy below:

if can_pin_to_reservation:
    if decision.label_selectors is not None:
        label_selector = decision.label_selectors
    else:
        # Fixed policy (for now): poll until the reservation covers num_workers.
        label_selector = self._scaling_policy.get_reserved_bundle_label_selectors(
            num_workers
        )
        if label_selector is None:
            raise WorkerGroupStartupTimeoutError(num_workers=num_workers)

Benefit of this is that we don't need logic reconciling M reserved worker resources with N target workers.

Follow-up: Unify the FixedScalingPolicy to also reserve resources from the coordinator and package label selector pins in the ResizeDecision

• Today, elastic and fixed scaling policies differ in the way they reserve resources:
  • Fixed: immediately assume that num_workers is ready (not necessarily true in the cluster upscaling case) and return a ResizeDecision. Rely on a 60 second timeout when starting the worker group before retrying.
  • Elastic: Return NoopDecision until the coordinator returns enough reserved resources for min_workers before returning a ResizeDecision.

This can also be an opportunity to unify these two codepaths paths to both request resources from the coordinator up front (and determine the label selectors) and wait for all resources to be ready before attempting to launch the worker group with ResizeDecision.

I attached a full sketch PR as a reference, but I think it can be broken down into 2 PRs with the split described above. #66646

[pseudo-rnd-thoughts]

Closing in favor of #66646