[data] Support multiple datasets in a cluster (2/2): partition cluster resources by subcluster label

[TimothySeah]

Summary

The end goal is to support 2 ray data datasets in 1 cluster with subcluster label scheduling. In such a setup, we have 2 datasets and 2 trainers sharing the same AutoscalingCoordinator. The previous PR in this stack (#63331) made sure that each dataset's tasks ended up in the correct subcluster. This PR ensures that all requesters, whether they are trainers or datasets, only request and receive resources in their subcluster.

To this end, the main change was to AutoscalingCoordinator._tick, which is called at regular intervals. AutoscalingCoordinator._tick calls 3 helper methods, which this PR changes as follows:

1. merge_and_send_requests: each autoscaling request now includes the subcluster label of the requester
2. update_cluster_node_resources: we now group cluster nodes by subcluster
3. _reallocate_resources: we now update OngoingRequests with their subcluster-scoped resources.

I also changed the try_trigger_scaling method, which creates datasets' autoscaling requests. Before this change, this method tried to scale up every node in the cluster. Now, it only scales up the relevant subcluster. Note that this only applies to dataset requesters; trainer requesters attempt scaleup by requesting resource bundles with their corresponding label selectors (which includes subcluster labels), so I didn't need to touch that path.

API summary

To use subcluster scheduling, the user must set the __subcluster__ label in their compute config

- name: training_node
  instance_type: p4d.24xlarge
  max_workers: 2
  min_workers: 2
  use_spot: false
  labels:
    __subcluster__: training

- name: validation_node
  instance_type: p4d.24xlarge
  max_workers: 1
  min_workers: 1
  use_spot: false
  labels:
    __subcluster__: validation

and the label_selector on their dataset

# Option 1: set using ray.train.DataConfig
trainer = TorchTrainer(
	train_fn,
	datasets={"train": train_ds, "val": val_ds},
	dataset_config=ray.train.DataConfig(
		datasets_to_split=["train"],
		data_execution_options=DataExecutionOptions(
			per_dataset_execution_options={
				"train": ExecutionOptions(
					label_selector={"__subcluster__": "train"}
				),
				"val": ExecutionOptoins(
					label_selector={"__subcluster__": "validation"}
				)
			}
		)
	)

# Option 2: set directly using DataContext.ExecutionOptions
train_ds.context.label_selector = {"__subcluster__": "training"}

Testing

Ran multitenancy stress test based on this PR (PR: #63737, test: https://buildkite.com/ray-project/release/builds/95982).

[gemini-code-assist]

Code Review

This pull request implements label_selector and subcluster_label_key support across Ray Data and Ray Train, allowing users to constrain task and actor placement to specific labeled subsets of a cluster. The changes include updates to ExecutionOptions, the AutoscalingCoordinator for resource bucketing, and broad propagation of these selectors through physical operators, planners, and data source implementations. Feedback was provided regarding the merge_label_selector utility, suggesting that it should always return a new dictionary to resolve a contradiction in its docstring and prevent potential mutation bugs.

[justinvyu]

Great work!

[justinvyu]

Thanks!

[cursor]

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

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

[justinvyu]

when is request_resources(label_selectors) used? what happens if you do request_resources(label_selectors, subcluster_selector)? Does one overwrite the other?

Is that meant to be "non-subcluster related labels"?

Can we also raise an error to explicitly disallow one requester trying to request bundles from multiple subclusters?

around here:

  if subcluster_selector and label_selectors:
      req_subcluster = subcluster_selector.get(SUBCLUSTER_LABEL_KEY)
      for i, sel in enumerate(label_selectors):
          bundle_subcluster = sel.get(SUBCLUSTER_LABEL_KEY)
          if bundle_subcluster is not None and bundle_subcluster != req_subcluster:
              raise ValueError(
                  f"Bundle {i} label_selector targets subcluster "
                  f"{bundle_subcluster!r}, but requester is registered to "
                  f"{req_subcluster!r}. Per-bundle cross-subcluster "
                  f"allocation is not supported."
              )

[TimothySeah]

when is request_resources(label_selectors) used?
Is that meant to be "non-subcluster related labels"?

Check out #58845 and #63287. In the former PR, my goal was to support placing Ray Train workers on nodes with particular attributes. These would usually be subcluster labels, but could also be nodes within a subcluster. For example, we may want to place Ray Train workers on gpu nodes within the training subcluster, as opposed to Ray Data workers for the training dataset on the cpu nodes within the training subcluster. However, I forgot to update the AutoscalingCoordinator to scale up these nodes if they don't currently exist, which @liulehui added in the latter PR.

Right now, there are two types of requesters - datasets and ray train. Datasets will always request the subcluster using subcluster_selector, while Ray Train will always request all desired node attributes - including the subcluster - together using label_selectors. I agree it's a bit clunky/confusing though so I am open to suggestions on how to clean up this separation.

what happens if you do request_resources(label_selectors, subcluster_selector)? Does one overwrite the other?
Can we also raise an error to explicitly disallow one requester trying to request bundles from multiple subclusters?

subcluster_selector takes precedence: https://github.com/ray-project/ray/pull/63375/changes#diff-23e42254510d06fc2e4595cb52c69872e0b16f6c52932f06b502d63548e72067R361. I also implemented your ValueError suggestion, so now we should raise an error before we even get to this point.

[justinvyu]

Thanks! Can you update the PR description?