Replace `NodeAffinitySchedulingStrategy` when soft=False with Label Selector API #54940

ryanaoleary · 2025-07-26T10:10:19Z

Why are these changes needed?

The label selector API can now fully replace NodeAffinitySchedulingStrategy when soft=False by specifying a label selector requirement for ray.io/node-id, a default node label which is always added and used by the aforementioned scheduling strategy. This PR replaces uses of NodeAffinitySchedulingStrategy in Ray libraries and tests as part of the implementation plan for Ray Label Selector API proposal. We will continue to replace redundant functionality of NodeAffinitySchedulingStrategy and NodeLabelSchedulingStrategy as fallback label selectors are implemented.

Related issue number

#51564

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

gemini-code-assist

Summary of Changes

Hello @ryanaoleary, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request is a significant step in migrating Ray's internal and external APIs from the older NodeAffinitySchedulingStrategy to the more flexible and standardized label_selector API. The primary goal is to ensure that hard node affinity (where soft=False) is expressed uniformly using node labels, improving consistency and paving the way for future scheduling enhancements.

Highlights

API Migration: I've replaced all instances of NodeAffinitySchedulingStrategy (specifically when soft=False) with the new label_selector API across various Ray libraries and tests. This aligns with the ongoing effort to standardize scheduling mechanisms.
Code Simplification: The changes involve removing direct imports and instantiations of NodeAffinitySchedulingStrategy, leading to cleaner and more concise code where explicit node affinity is required.
Standardized Node Affinity: Node-specific scheduling now consistently leverages the ray.io/node-id label within the label_selector dictionary, as per the new Label Selector API proposal.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request is a great step towards unifying scheduling strategies under the new Label Selector API. The changes are mostly a straightforward replacement of NodeAffinitySchedulingStrategy with label_selector.

I've reviewed the changes and found a few critical issues that need to be addressed:

In python/ray/dashboard/modules/job/job_manager.py, the _get_label_selector function returns a nested dictionary, which will cause incorrect behavior when used.
In python/ray/data/_internal/stats.py, there's a potential NameError as label_selector is not defined in all code paths.
In release/long_running_tests/workloads/many_drivers.py, there are syntax errors in an f-string due to extra parentheses.

These issues could lead to runtime failures. Please see my detailed comments for suggestions on how to fix them. Once these are addressed, the PR should be in good shape.

python/ray/dashboard/modules/job/job_manager.py

python/ray/data/_internal/stats.py

release/long_running_tests/workloads/many_drivers.py

ryanaoleary · 2025-07-26T10:14:14Z

cc: @MengjinYan I replaced the usages in a bunch of files, this PR is quite large (although the changes are almost identical everywhere and pretty straightforward) so I can break it up if necessary. I'll put out a follow-up PR for NodeLabelSchedulingStrategy where appropriate.

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Ryan O'Leary <113500783+ryanaoleary@users.noreply.github.com>

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

…cies

edoakes · 2025-07-30T21:21:33Z

@ryanaoleary looks like there might be some CI failures. I added go label to kick off full run

github-actions · 2025-08-14T00:39:34Z

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

edoakes · 2025-08-14T02:51:58Z

not stale

ryanaoleary · 2025-08-25T17:02:43Z

This PR is blocked because we need to handle the special case where if a node id unavailable, we should directly return an infeasible error rather than block waiting to schedule. I'll work on that fix for label scheduling and then this PR should pass CI.

Replace NodeAffinitySchedulingStrategy soft=False with label selector

14ba17b

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

ryanaoleary requested review from a team as code owners July 26, 2025 10:10

gemini-code-assist bot reviewed Jul 26, 2025

View reviewed changes

ryanaoleary and others added 4 commits July 26, 2025 03:14

Update release/long_running_tests/workloads/many_drivers.py

31afd41

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Ryan O'Leary <113500783+ryanaoleary@users.noreply.github.com>

Update release/long_running_tests/workloads/many_drivers.py

e9f4242

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Ryan O'Leary <113500783+ryanaoleary@users.noreply.github.com>

Update python/ray/dashboard/modules/job/job_manager.py

41bd199

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Ryan O'Leary <113500783+ryanaoleary@users.noreply.github.com>

Add default label selector init

f9c3ec6

Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>

ray-gardener bot added community-contribution Contributed by the community core Issues that should be addressed in Ray Core labels Jul 26, 2025

Merge branch 'master' into replace-label-and-affinity-scheduling-poli…

5d42bb0

…cies

edoakes approved these changes Jul 30, 2025

View reviewed changes

edoakes added the go add ONLY when ready to merge, run all tests label Jul 30, 2025

ryanaoleary mentioned this pull request Aug 7, 2025

[Core] Ray Label Selector API Implementation Tracker #51564

Open

23 tasks

github-actions bot added the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Aug 14, 2025

github-actions bot added unstale A PR that has been marked unstale. It will not get marked stale again if this label is on it. and removed stale The issue is stale. It will be closed within 7 days unless there are further conversation labels Aug 14, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Replace `NodeAffinitySchedulingStrategy` when soft=False with Label Selector API #54940

Replace `NodeAffinitySchedulingStrategy` when soft=False with Label Selector API #54940

Uh oh!

ryanaoleary commented Jul 26, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ryanaoleary commented Jul 26, 2025

Uh oh!

edoakes commented Jul 30, 2025

Uh oh!

github-actions bot commented Aug 14, 2025

Uh oh!

edoakes commented Aug 14, 2025

Uh oh!

ryanaoleary commented Aug 25, 2025

Uh oh!

Uh oh!

Replace NodeAffinitySchedulingStrategy when soft=False with Label Selector API #54940

Are you sure you want to change the base?

Replace NodeAffinitySchedulingStrategy when soft=False with Label Selector API #54940

Uh oh!

Conversation

ryanaoleary commented Jul 26, 2025

Why are these changes needed?

Related issue number

Checks

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ryanaoleary commented Jul 26, 2025

Uh oh!

edoakes commented Jul 30, 2025

Uh oh!

github-actions bot commented Aug 14, 2025

Uh oh!

edoakes commented Aug 14, 2025

Uh oh!

ryanaoleary commented Aug 25, 2025

Uh oh!

Uh oh!

Replace `NodeAffinitySchedulingStrategy` when soft=False with Label Selector API #54940

Replace `NodeAffinitySchedulingStrategy` when soft=False with Label Selector API #54940