Add v6e TPU Head Resource Autoscaling Support

[ryanaoleary]

Why are these changes needed?

This PR adds tpu-v6e-slice to the list of known TPU accelerators, enabling the KubeRay autoscaler to add a TPU-v6e...-Head resource to the autoscaling resource config. This PR also adds a unit test to cover test_tpu_node_selectors_to_type.

Related issue number

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 :(

[ryanaoleary]

cc: @kevin85421 @andrewsykim

[ryanaoleary]

This PR was manually tested as follows:

1. Build Ray from source with these changes.
2. Create an autoscaling RayCluster CR with the following TPU worker group:

- replicas: 0
    minReplicas: 0
    maxReplicas: 10
    numOfHosts: 4
    groupName: tpu-group
    rayStartParams: {}
    template:
      spec:
        securityContext:
          runAsUser: 0
        containers:
        - name: ray-worker
          image: $DEV-IMAGE
          imagePullPolicy: Always
          resources:
            limits:
              cpu: "24"
              ephemeral-storage: 10Gi
              google.com/tpu: "4"
              memory: 40G
            requests:
              cpu: "24"
              ephemeral-storage: 10Gi
              google.com/tpu: "4"
              memory: 40G
          env:
          - name: NODE_IP
            valueFrom:
              fieldRef:
                fieldPath: status.hostIP
          - name: VBAR_CONTROL_SERVICE_URL
            value: $(NODE_IP):8353
          - name: JAX_PLATFORMS
            value: tpu,cpu
          - name: ENABLE_PJRT_COMPATIBILITY
            value: "true"
          ports:
          - containerPort: 8081
            name: mxla
        nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
            cloud.google.com/gke-tpu-topology: 4x4

3. Verify the TPU head resource is added to the autoscaler resources after a replica is scaled up:

Usage:
 0.0/96.0 CPU
 0.0/16.0 TPU
 0.0/1.0 TPU-v6e-16-head
 0B/150.87GiB memory
 0B/45.02GiB object_store_memory
 0.0/4.0 tpu-group-0

Demands:
 (no resource demands)
2024-10-22 21:23:50,547 - INFO - The autoscaler took 0.046 seconds to complete the update iteration.
2024-10-22 21:23:50,547 INFO autoscaler.py:470 -- The autoscaler took 0.046 seconds to complete the update iteration.

[andrewsykim]

LGTM, just a minor nit

[ryanaoleary]

@kevin85421 @hongchaodeng pinging this again to see if it can get reviewed/approved by a code-owner, with v6e TPUs in private preview in GKE it would be good to ensure the Ray autoscaler support is there

[kevin85421]

How can I determine the exact value of default_num_cores_per_chip to verify that this logic is correct? I briefly used Ctrl + F to search for some keywords in https://cloud.google.com/kubernetes-engine/docs/how-to/tpus#run, but I couldn't find the information.

[ryanaoleary]

The relation between cores to number of chips for each TPU generation in one spot is viewable in the Cloud TPU documentation for each version. Under System Architecture it states the number of TensorCores per TPU chip.

[kevin85421]

I have already pinged my colleague to merge this PR.