[serve][3/n] Gang scheduling -- core scheduling engine

[jeffreywang88]

Description

Implements the core gang scheduling logic: Gang-scheduled deployments atomically reserve placement groups for groups of replicas and start them together, ensuring all members of a gang are co-scheduled or none are.

Approach

• Scheduler (deployment_scheduler.py)

  • Added schedule_gang_placement_groups to DeploymentScheduler.
  • The default scheduler now creates named gang placement groups and assigns replica ranks within each gang.
  • Replica scheduling checks for a gang placement group first, and falls back to per-replica placement if none exists.
  • Gang reservation results are passed to the deployment state machine.

• State Machine (deployment_state.py)

  • Introduced a new step in the update loop to reserve gang placement groups.
  • Added _add_replicas_with_gang_scheduling() to start replicas with gang context (gang_id, rank, world_size, member_replica_ids).
  • If any replica in a gang fails during startup, all replicas in that gang are stopped.
  • Gracefully handles placement group removal failures for shared gang placement groups.

• Replica (replica.py)

  • Extended ReplicaMetadata to include GangContext.
  • ActorReplicaWrapper now stores and exposes gang context and passes it through check_ready().

Related issues

RFC: #60873
Precedent: #61205
Original PR: #60802

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request introduces the core logic for gang scheduling in Ray Serve, a significant feature that enables atomic scheduling of replica groups. The changes span across the deployment scheduler, state machine, and replica logic, and are well-supported by a comprehensive suite of new tests.

The implementation correctly introduces a new step for reserving placement groups, updates the scheduling logic to utilize them, and handles gang-level failures to ensure atomicity. The approach is robust and the code is generally of high quality.

I have one suggestion to improve error logging to prevent potential resource leaks from being silently ignored.

[abrarsheikh]

let's add a few more integration tests, ignore if any of these already exist

1. serve.delete() on a gang app, verify PGs are cleaned up?
2. Multiple gang deployments in one app
3. One replica in a gang fails during startup. Both replicas in that gang are stopped; no partial gang left running.
4. Running gang replica fails health check or crashes. Whole gang is torn down and restarted

[jeffreywang88]

@abrarsheikh I addressed your comments regarding to tests -- ready for another pass.

[cursor]

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

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

[cursor]

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

^{Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

[cursor]

Startup gang cleanup misses PENDING_MIGRATION siblings

Medium Severity

The gang sibling cleanup in _check_startup_replicas iterates over {original_state, ReplicaState.RUNNING} but omits ReplicaState.PENDING_MIGRATION. This is inconsistent with the health-check gang cleanup in check_and_update_replicas, which correctly iterates over [ReplicaState.RUNNING, ReplicaState.PENDING_MIGRATION]. If a gang member succeeds startup and transitions to RUNNING, then gets migrated to PENDING_MIGRATION, and its sibling subsequently fails startup, the migrating sibling won't be stopped — leaving a partial gang running.

Additional Locations (1)

• python/ray/serve/_private/deployment_state.py#L3617-L3633

 

[jeffreywang88]

This is resolved in #61216

[abrarsheikh]

nits

[jeffreywang88]

@abrarsheikh I addressed your comments except for #61206 (comment) and also adjusted startup failure counting logic for gangs.

Previously, in gang scheduling, each replica startup failure is counted towards the threshold, but I think counting failure per gang makes more sense. Startup failure occurs when there's an allocation (e.g. insufficient resources) or initialization (e.g. actor initialization) issue, and replicas in a gang could run into issues with the same root cause, and therefore the previous approach will inflate the failure count by gang_size.

[jeffreywang88]

Addressed all comments :)