-
Notifications
You must be signed in to change notification settings - Fork 7.1k
[Serve] Add label_selector and bundle_label_selector to Serve API
#57694
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
[Serve] Add label_selector and bundle_label_selector to Serve API
#57694
Conversation
Signed-off-by: Ryan O'Leary <[email protected]>
|
cc: @MengjinYan @eicherseiji I think this change can help enable TPU use-cases with Ray LLM, since it'll allow users to target the desired slice/topology based on labels like these: |
MengjinYan
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The current change looks good.
At the same time, I have a general question, probably also to @eicherseiji as well: I'm not familiar with the serve codebase. But from the look at the code, the change in the PR seems not cover the whole code path from the replica config to actually creating the placement group where we need to apply the bundle label selector (e.g. CreatePlacementGroupRequest, DeploymentTargetState, DeploymentVersion, etc.).
Wondering if we should include the change to the rest of the code path in this PR as well?
|
Hi @ryanaoleary! Seems like it may be useful, but could you go into more detail about the problem this solved for you? I.e. is it typical to have multiple TPU node pools/heterogeneous compute in a TPU cluster? |
Signed-off-by: Ryan O'Leary <[email protected]>
The use-case would be for TPU multi-slice, where it's important that the Ray nodes reserved with some group of resources (i.e. a bundle of TPUs) are a part of the same actual TPU slice so that they can benefit from the high speed ICI interconnects. There are GKE webhooks that set some TPU information as env vars when the Pods are created, including topology and a unique identifier for the slice which we set as Ray node labels. They also inject For a RayCluster with multiple TPU slices (of the same topology or different), we currently only schedule using |
|
@ryanaoleary I see, thanks! Is there a reason we can't extend/re-use the |
In addition to the TPU support, in general, we want to have label support in all library APIs so that user can do scheduling based on node labels as well. |
I think just configuring |
|
This pull request has been automatically marked as stale because it has not had 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. |
|
@ryanaoleary I will review this PR next week |
|
thanks for working on this! Really looking forward to using label selectors in Serve. |
Signed-off-by: ryanaoleary <[email protected]>
Signed-off-by: ryanaoleary <[email protected]>
@abrarsheikh Done, I updated the PR description with an outline that I think covers all these situations. |
|
I think the logic of your implementation is correct afaict, I am writing down my understanding of how this stuff is supposed to work, so that we both are on the same page Decision Flow Chart For Scheduler
|
Co-authored-by: Abrar Sheikh <[email protected]> Signed-off-by: Ryan O'Leary <[email protected]>
Co-authored-by: Abrar Sheikh <[email protected]> Signed-off-by: Ryan O'Leary <[email protected]>
Co-authored-by: Abrar Sheikh <[email protected]> Signed-off-by: Ryan O'Leary <[email protected]>
Co-authored-by: Abrar Sheikh <[email protected]> Signed-off-by: Ryan O'Leary <[email protected]>
Co-authored-by: Abrar Sheikh <[email protected]> Signed-off-by: Ryan O'Leary <[email protected]>
Co-authored-by: Abrar Sheikh <[email protected]> Signed-off-by: Ryan O'Leary <[email protected]>
Co-authored-by: Abrar Sheikh <[email protected]> Signed-off-by: Ryan O'Leary <[email protected]>
Co-authored-by: Abrar Sheikh <[email protected]> Signed-off-by: Ryan O'Leary <[email protected]>
Co-authored-by: Abrar Sheikh <[email protected]> Signed-off-by: Ryan O'Leary <[email protected]>
Signed-off-by: ryanaoleary <[email protected]>
Signed-off-by: ryanaoleary <[email protected]>
Signed-off-by: ryanaoleary <[email protected]>
Signed-off-by: ryanaoleary <[email protected]>
Signed-off-by: ryanaoleary <[email protected]>
| return False | ||
|
|
||
| # !in operator | ||
| if selector_value.startswith("!in(") and selector_value.endswith(")"): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
should we trim the selector_value before passing through the string parsers?
| # Actor: Use Actor label selector | ||
| if "label_selector" in scheduling_request.actor_options: | ||
| primary_labels = [ | ||
| scheduling_request.actor_options["label_selector"] or {} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
is None allowed in actor_options["label_selector"]?
| fallback_labels = [fallback.get("label_selector", {}) or {}] | ||
| placement_candidates.append( | ||
| (scheduling_request.required_resources, fallback_labels) | ||
| ) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can there be a fallback with label_selector == None?
Also should we skip adding to placement_candidates if fallback_labels is empty?
| @pytest.mark.parametrize( | ||
| "node_value, selector_value, expected", | ||
| [ | ||
| # Equals operator | ||
| ("us-west", "us-west", True), | ||
| ("us-east", "us-west", False), | ||
| (None, "us-west", False), | ||
| # Not equals (!) operator | ||
| ("us-west", "!us-east", True), | ||
| ("us-east", "!us-east", False), | ||
| (None, "!us-east", True), | ||
| # In operator | ||
| ("A100", "in(A100, H100)", True), | ||
| ("T4", "in(A100, H100)", False), | ||
| (None, "in(A100, H100)", False), | ||
| ("value", "in(value)", True), | ||
| # Not in operator | ||
| ("T4", "!in(A100, H100)", True), | ||
| ("A100", "!in(A100, H100)", False), | ||
| (None, "!in(A100, H100)", True), | ||
| # Invalid types | ||
| ("A100", None, False), | ||
| ("A100", 123, False), | ||
| ], | ||
| ) | ||
| def test_match_label_selector_value(node_value, selector_value, expected): | ||
| assert match_label_selector_value(node_value, selector_value) == expected | ||
|
|
||
|
|
||
| @pytest.mark.parametrize( | ||
| "node_labels, selector, expected", | ||
| [ | ||
| # Match all | ||
| ( | ||
| {"region": "us-west", "gpu": "A100"}, | ||
| {"region": "us-west", "gpu": "A100"}, | ||
| True, | ||
| ), | ||
| # Partial match | ||
| ( | ||
| {"region": "us-west", "gpu": "T4"}, | ||
| {"region": "us-west", "gpu": "A100"}, | ||
| False, | ||
| ), | ||
| # Multiple operators | ||
| ( | ||
| {"region": "us-west", "env": "prod"}, | ||
| {"region": "in(us-west, us-east)", "env": "!dev"}, | ||
| True, | ||
| ), | ||
| # Missing keys in node labels | ||
| ({"region": "us-west"}, {"gpu": "!in(H100)"}, True), | ||
| ({"region": "us-west"}, {"gpu": "in(H100)"}, False), | ||
| ], | ||
| ) | ||
| def test_match_label_selector(node_labels, selector, expected): | ||
| assert match_label_selector(node_labels, selector) == expected |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
| @pytest.mark.parametrize( | |
| "node_value, selector_value, expected", | |
| [ | |
| # Equals operator | |
| ("us-west", "us-west", True), | |
| ("us-east", "us-west", False), | |
| (None, "us-west", False), | |
| # Not equals (!) operator | |
| ("us-west", "!us-east", True), | |
| ("us-east", "!us-east", False), | |
| (None, "!us-east", True), | |
| # In operator | |
| ("A100", "in(A100, H100)", True), | |
| ("T4", "in(A100, H100)", False), | |
| (None, "in(A100, H100)", False), | |
| ("value", "in(value)", True), | |
| # Not in operator | |
| ("T4", "!in(A100, H100)", True), | |
| ("A100", "!in(A100, H100)", False), | |
| (None, "!in(A100, H100)", True), | |
| # Invalid types | |
| ("A100", None, False), | |
| ("A100", 123, False), | |
| ], | |
| ) | |
| def test_match_label_selector_value(node_value, selector_value, expected): | |
| assert match_label_selector_value(node_value, selector_value) == expected | |
| @pytest.mark.parametrize( | |
| "node_labels, selector, expected", | |
| [ | |
| # Match all | |
| ( | |
| {"region": "us-west", "gpu": "A100"}, | |
| {"region": "us-west", "gpu": "A100"}, | |
| True, | |
| ), | |
| # Partial match | |
| ( | |
| {"region": "us-west", "gpu": "T4"}, | |
| {"region": "us-west", "gpu": "A100"}, | |
| False, | |
| ), | |
| # Multiple operators | |
| ( | |
| {"region": "us-west", "env": "prod"}, | |
| {"region": "in(us-west, us-east)", "env": "!dev"}, | |
| True, | |
| ), | |
| # Missing keys in node labels | |
| ({"region": "us-west"}, {"gpu": "!in(H100)"}, True), | |
| ({"region": "us-west"}, {"gpu": "in(H100)"}, False), | |
| ], | |
| ) | |
| def test_match_label_selector(node_labels, selector, expected): | |
| assert match_label_selector(node_labels, selector) == expected | |
| @pytest.mark.parametrize( | |
| "node_value, selector_value, expected", | |
| [ | |
| # Equals operator | |
| ("us-west", "us-west", True), | |
| ("us-east", "us-west", False), | |
| (None, "us-west", False), | |
| # Equals with empty strings | |
| ("", "", True), | |
| ("value", "", False), | |
| ("", "value", False), | |
| # Not equals (!) operator | |
| ("us-west", "!us-east", True), | |
| ("us-east", "!us-east", False), | |
| (None, "!us-east", True), | |
| # Not equals with empty string | |
| ("value", "!", True), # !<empty> means not equal to empty | |
| ("", "!", False), # empty string equals empty, so !empty is False | |
| ("", "!value", True), # empty string != "value" | |
| # In operator | |
| ("A100", "in(A100, H100)", True), | |
| ("T4", "in(A100, H100)", False), | |
| (None, "in(A100, H100)", False), | |
| ("value", "in(value)", True), | |
| # In operator without spaces | |
| ("A100", "in(A100,H100)", True), | |
| ("H100", "in(A100,H100)", True), | |
| ("T4", "in(A100,H100)", False), | |
| # In operator with many values | |
| ("c", "in(a, b, c, d, e)", True), | |
| ("z", "in(a, b, c, d, e)", False), | |
| # In operator with empty string value not in list | |
| ("", "in(A100, H100)", False), | |
| # Not in operator | |
| ("T4", "!in(A100, H100)", True), | |
| ("A100", "!in(A100, H100)", False), | |
| (None, "!in(A100, H100)", True), | |
| # Not in operator without spaces | |
| ("T4", "!in(A100,H100)", True), | |
| ("A100", "!in(A100,H100)", False), | |
| # Not in with single value | |
| ("T4", "!in(A100)", True), | |
| ("A100", "!in(A100)", False), | |
| # Not in with empty string | |
| ("", "!in(A100, H100)", True), # empty string not in list | |
| # Invalid selector types | |
| ("A100", None, False), | |
| ("A100", 123, False), | |
| (None, None, False), | |
| ("value", [], False), | |
| ("value", {}, False), | |
| # Node value that looks like operator syntax (edge case) | |
| ("in(x)", "in(x)", True), | |
| ("!value", "!value", True), | |
| # Case sensitivity | |
| ("A100", "a100", False), | |
| ("A100", "in(a100, h100)", False), | |
| ("A100", "!a100", True), | |
| ], | |
| ) | |
| def test_match_label_selector_value(node_value, selector_value, expected): | |
| assert match_label_selector_value(node_value, selector_value) == expected | |
| @pytest.mark.parametrize( | |
| "node_labels, selector, expected", | |
| [ | |
| # Empty cases | |
| ({}, {}, True), # Empty selector matches everything | |
| ({"region": "us-west"}, {}, True), # Empty selector matches any node | |
| ({}, {"region": "us-west"}, False), # Missing key with equals fails | |
| ({}, {"region": "!us-west"}, True), # Missing key with not-equals succeeds | |
| # Match all | |
| ( | |
| {"region": "us-west", "gpu": "A100"}, | |
| {"region": "us-west", "gpu": "A100"}, | |
| True, | |
| ), | |
| # Single key match | |
| ({"region": "us-west"}, {"region": "us-west"}, True), | |
| ({"region": "us-east"}, {"region": "us-west"}, False), | |
| # Partial match (one key matches, one doesn't) | |
| ( | |
| {"region": "us-west", "gpu": "T4"}, | |
| {"region": "us-west", "gpu": "A100"}, | |
| False, | |
| ), | |
| # Multiple operators | |
| ( | |
| {"region": "us-west", "env": "prod"}, | |
| {"region": "in(us-west, us-east)", "env": "!dev"}, | |
| True, | |
| ), | |
| # Multiple operators - failure case | |
| ( | |
| {"region": "us-west", "env": "dev"}, | |
| {"region": "in(us-west, us-east)", "env": "!dev"}, | |
| False, | |
| ), | |
| # Missing keys in node labels | |
| ({"region": "us-west"}, {"gpu": "!in(H100)"}, True), | |
| ({"region": "us-west"}, {"gpu": "in(H100)"}, False), | |
| ({"region": "us-west"}, {"gpu": "!H100"}, True), # Missing key with ! succeeds | |
| # Extra keys in node_labels that selector doesn't care about | |
| ( | |
| {"region": "us-west", "gpu": "A100", "env": "prod", "zone": "a"}, | |
| {"region": "us-west"}, | |
| True, | |
| ), | |
| # All keys matching with mixed operators | |
| ( | |
| {"region": "us-west", "gpu": "A100", "env": "prod"}, | |
| {"region": "us-west", "gpu": "in(A100, H100)", "env": "!dev"}, | |
| True, | |
| ), | |
| # First key matches, second doesn't | |
| ( | |
| {"region": "us-west", "gpu": "T4"}, | |
| {"region": "in(us-west, us-east)", "gpu": "in(A100, H100)"}, | |
| False, | |
| ), | |
| # First key doesn't match | |
| ( | |
| {"region": "eu-west", "gpu": "A100"}, | |
| {"region": "in(us-west, us-east)", "gpu": "in(A100, H100)"}, | |
| False, | |
| ), | |
| # Empty string values | |
| ({"region": ""}, {"region": ""}, True), | |
| ({"region": "us-west"}, {"region": ""}, False), | |
| # Complex multi-key scenario | |
| ( | |
| {"cloud": "aws", "region": "us-west", "gpu": "A100", "spot": "true"}, | |
| {"cloud": "!gcp", "region": "in(us-west, us-east)", "gpu": "A100"}, | |
| True, | |
| ), | |
| ], | |
| ) | |
| def test_match_label_selector(node_labels, selector, expected): | |
| assert match_label_selector(node_labels, selector) == expected |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
just add more params
| "labels": {"region": "us-east", "gpu-type": "H100"}, | ||
| } | ||
| cluster.add_node(**node2_config) | ||
| node_2_id = ray.get(_get_node_id.options(resources={"worker2": 1}).remote()) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
this is clever, nice :)
|
|
||
| serve.shutdown() | ||
|
|
||
|
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
let's add a few more tests here
- STRICT_PACK PG + bundle_label_selector (Pack Scheduling Mode)
- STRICT_SPREAD PG Strategy
- placement_group_fallback_strategy → NotImplementedError
- Non-STRICT_PACK PG Forces Spread Mode (when RAY_SERVE_USE_PACK_SCHEDULING_STRATEGY=1)
- Multiple Bundles with bundle_label_selector
- Actor label_selector + fallback with Multiple Fallbacks
Why are these changes needed?
This PR adds the
label_selectoroption to the supported list of Actor options for a Serve deployment. Additionally, we addbundle_label_selectorto specify label selectors for bundles whenplacement_group_bundlesare specified for the deployment. These two options are already supported for Tasks/Actors and placement groups respectively.Example use case:
The expected behaviors of these new fields is as follows:
Pack scheduling enabled
PACK/STRICT_PACK PG strategy:
Standard PG without bundle_label_selector or fallback:
PG node label selector provided:
PG node label selector and fallback:
Same as above but when scheduling tries the following:
SPREAD/STRICT_SPREAD PG strategy:
Spread scheduling enabled
Related issue number
#51564
Checks
git commit -s) in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.