REP: Topology Aware Schedulling

[aaronscalene]

This REP proposes topology aware scheduling for placement groups. This allows functionality such as rack aware placement group scheduling, enabling rack aware fault tolerance. Motivated by new GB300 racks that are connected through the NVLink domain, necessitating topology aware scheduling to take advantage of faster memory bandwidth.

Requested reviewers: @edoakes @MengjinYan @Sparks0219

[andrewsykim]

For AZ topology scheduling, I think the two main scenarios are something like:

• For training, STRICT_PACK in single AZ to avoid cross AZ network transfer (incurs additional cost)
• For inference: STRICT_SPREAD across AZs for better fault tolerance

The inference case has a bit more nuance in that if you're doing multi-host serving, you want each group of nodes to be STRICT_PACK but each "replica" of the model to be STRICT_SPREAD. This is maybe irrelevant if we always run a placement group per model replica?

[andrewsykim]

Either way, we should add the first scenario (strict pack in single AZ) as a use-case as it comes up a lot

[aaronscalene]

Sounds good, I added an example of Strick Pack AZ + Strict Pack Rack hierarchical topology. For your inference case, I am not entirely sure about the model replica implementation, but I believe what you can do would be to define a group of bundles for each replica, and have a hierarchical scheduling to STRICT_SPREAD each group of bundles across AZs. You can then probably target each group of bundles using their bundle ids (probably should discuss more on this).

[andrewsykim]

Worth nothing that in KubeRay we started to add better metadata for multi-host scenarios. Specifically if you set replicas > 0 and numOfHosts > 1, we set the following labels on every Pod:

ray.io/worker-group-replica-index=<replica index>
ray.io/replica-host-index=<host-index>

So in the NVL72 case, you set numOfHosts=18 and each "replica" is a rack. Then rack_id can just inherit whatever value is used for ray.io/worker-group-replica-index. This is a QoL improvement to the cluster operator because they don't need to manually set unique labels for every rack.

[andrewsykim]

@ryanaoleary do you know if it's possible now to set Ray labels based on Pod labels?

For example something like:

  workerGroupSpecs:
  - replicas: 1
    minReplicas: 1
    maxReplicas: 10
    numOfHosts: 18
    groupName: gb200
    labels:
      rack-id: ${WORKER_GROUP_REPLICA_INDEX}
    rayStartParams: {}
    template:
      spec:
        containers:
        - name: ray-worker
          image: rayproject/ray:nightly
          resources:
            limits:
              cpu: "2"
              memory: "4Gi"
            requests:
              cpu: "2"
              memory: "4Gi"
        env:
        - name: WORKER_GROUP_REPLICA_INDEX
          valueFrom:
            fieldRef:
              fieldPath: metadata.labels['worker-group-replica-index']

[ryanaoleary]

The above example fails with:

Warning  InvalidRayClusterSpec  13s   raycluster-controller  The RayCluster spec is invalid default/raycluster-manual-label-test: invalid label value for key 'rack-id' in gb200 group: '${WORKER_GROUP_REPLICA_INDEX}', error: a valid label must be an empty string or consist of alphanumeric characters, '-', '_' or '.', and must start and end with an alphanumeric character (e.g. 'MyValue',  or 'my_value',  or '12345', regex used for validation is '(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])?')

I think it's not currently supported, if we wanted the label to be set by default when the env var is present we could add it here: https://github.com/ray-project/ray/blob/7ddc3faa761ab533eaff081be4db7dcea683ea56/python/ray/_private/resource_and_label_spec.py#L271. We do currently set a label based on the TPU_NAME var only for TPU, that serves a similar function (unique per multi-host replica): https://github.com/ray-project/ray/blob/7ddc3faa761ab533eaff081be4db7dcea683ea56/python/ray/_private/accelerators/tpu.py#L750.

I think we should set Ray node labels by default for k8s labels on the pod with the prefix ray.io/. I'll put out this change since we previously talked about adding this support.

[andrewsykim]

Is that only failing due to validation though? If we skip validation does the env var render correctly at runtime?

[ryanaoleary]

Sorry missed this comment - tested with skipping the KubeRay validation (specifically: validateRayGroupLabels)and it fails validation at the kube-apiserver level when trying to create the Pod:

3s          Normal    CreatedService            raycluster/raycluster-manual-label-test       Created service default/raycluster-manual-label-test-head-svc
3s          Normal    CreatedService            raycluster/raycluster-manual-label-test       Created service default/raycluster-manual-label-test-headless
1s          Warning   FailedToCreateWorkerPod   raycluster/raycluster-manual-label-test       Failed to create worker Pod for the cluster default/raycluster-manual-label-test, Pod "raycluster-manual-label-test-gb200-worker-jlxz8" is invalid
: metadata.labels: Invalid value: "${WORKER_GROUP_REPLICA_INDEX}": a valid label must be an empty string or consist of alphanumeric characters, '-', '_' or '.', and must start and end with an alphanumeric character (e.g. 'MyValue',  or 'my_value',  or '12345', regex used for validation is '(([A-Za-z0-9][-A-Za-z0-9_.]*)?[A-Za-z0-9])?')

I'm not sure if we can skip/disable the above validation.

it works if we overwrite the rayStartCommand though and skip the KubeRay validation:

Amended worker group spec to:

  workerGroupSpecs:
  - replicas: 1
    minReplicas: 1
    maxReplicas: 10
    numOfHosts: 18
    groupName: gb200
    labels: {}
    rayStartParams: {}
    template:
      metadata:
        annotations:
          ray.io/overwrite-container-cmd: "true"
      spec:
        containers:
        - name: ray-worker
          image: rayproject/ray:nightly
          
          command: ["/bin/bash", "-c"]
          # Inject the downward API variable natively into the ray start arguments
          args: ["ulimit -n 65536; $KUBERAY_GEN_RAY_START_CMD --labels=\"rack-id=${WORKER_GROUP_REPLICA_INDEX}\""]
          env:
          - name: WORKER_GROUP_REPLICA_INDEX
            valueFrom:
              fieldRef:
                fieldPath: metadata.labels['ray.io/worker-group-replica-index']

Ran this test on the worker group with the above spec:

import ray
ray.init()
print([n.get('Labels', {}).get('rack-id') for n in ray.nodes()])

labels set on the nodes:

['0', '0', None, '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0', '0']

[ryanaoleary]

Can we add a section discussing how autoscaler support will be implemented. I'm interested in how we plan to:

1. Set default labels for multi-host / topology aware groups on the Ray node and in the autoscaling config.
2. Integrate the label domain key into the autoscaler - will this be supported the same as bundle_label_selector? (i.e. after a bundle_label_selector has been chosen with a value that's the same across all bundles, does it just follow the same path as though a bundle_label_selector was provided explicitly?)

Unrelated but I'm also wondering if a future step will include supporting topology_strategy in the fallback_strategy argument.

[connorp-ui]

This would likely be easier to use if we put the topology strategy itself in the bundle then have a way to have topology labels for a group of bundles.

[ryanaoleary]

will topology strategy be a Dict or List of dicts to support hierarchical scheduling?

[ryanaoleary]

LGTM with a few clarifying questions worth resolving, happy to help with implementing/reviewing this soon - especially pieces like SPREAD and hierarchical scheduling support