[core] Placement Group bundle label selectors are not respected

[andrewsykim]

What happened + What you expected to happen

I am testing the new JaxTrainer which leverages bundle label selectors for placement group bundles. I am encountering a bug where the bundle label selector is not being respected. This breaks multi-host training on my RayCluster that has a 4x4 slice (4 nodes) and a couple of single-host nodes.

Here are the label selectors in the placement group table from Ray Dashboard which look correct:

The node labels look correct as well. Here's the node labels from a node in the 4x4 topology:

And here are the labels of the single-host node:

Both the placement gruop label selector and Raylet labels look correct. However, when I run my job, I can see from the logs that 2 train workers run on the single host node and the other 2 run on the 4x4 slice. Based on the label selectors I expect all 4 train workers to run on the 4x4 slice.

I checked the job task table and also confirmed a bunch of tasks running on the single-host workers:

Versions / Dependencies

nightly version of Ray on v6e TPUs on GKE

Reproduction script

Deploy a RayCluster with 2 worker groups (4x4 slice and single-host pool). Deploy a training job to the 4x4 slice.

Issue Severity

None

[andrewsykim]

cc @Yevet @MengjinYan

[andrewsykim]

ref #55162

[matthewdeng]

Is this reproducible with just PlacementGroups?

I.e. would the following request which shouldn't fit on the 4x4 slice hang?

num_bundles = 5 # greater than 4

pg = placement_group(
    bundles=[{"TPU":4}] * num_bundles,
    bundle_label_selector=[{"ray.io/tpu-slice-name":"tpu-group-0"}] * num_bundles,
)

ray.get(pg)

[MengjinYan]

Took a look at the issue and found that we didn't check the label selectors during bundle scheduling. Sent a PR to add the logic and updated the corresponding test.

[edoakes]

Took a look at the issue and found that we didn't check the label selectors during bundle scheduling. Sent a PR to add the logic and updated the corresponding test.

Make sure to link this issue in the PR please

[MengjinYan]

Took a look at the issue and found that we didn't check the label selectors during bundle scheduling. Sent a PR to add the logic and updated the corresponding test.

Make sure to link this issue in the PR please

Thanks for reminder! I already added this issue the PR when I sent it out last night: [Core] Add Logic to Check Label Selector in PG Scheduling #55599](#55599).