ray-project · edoakes · Jun 23, 2025 · Jun 18, 2025 · davidxia · Jun 18, 2025
@@ -28,6 +28,9 @@
 MAX_RAYCLUSTER_FETCH_TRIES = 5
 RAYCLUSTER_FETCH_RETRY_S = 5
 
+GKE_TPU_TOPOLOGY_LABEL = "cloud.google.com/gke-tpu-topology"
+GKE_TPU_ACCELERATOR_LABEL = "cloud.google.com/gke-tpu-accelerator"
+
 # Logical group name for the KubeRay head group.
 # Used as the name of the "head node type" by the autoscaler.
 _HEAD_GROUP_NAME = "headgroup"
@@ -264,15 +267,27 @@ def _get_ray_resources_from_group_spec(
         resource labels on worker 0 of each replica:
             worker 0: resources = {"TPU": 4, "TPU-v4-16-head": 1}
         """
-        topology = group_spec["template"]["spec"]["nodeSelector"][
-            "cloud.google.com/gke-tpu-topology"
-        ]
-        accelerator = group_spec["template"]["spec"]["nodeSelector"][
-            "cloud.google.com/gke-tpu-accelerator"
-        ]
-        accelerator_type = utils.tpu_node_selectors_to_type(topology, accelerator)
-        if accelerator_type:
-            resources[f"TPU-{accelerator_type}-head"] = 1
+        if (
+            "nodeSelector" in group_spec["template"]["spec"]
+            and GKE_TPU_TOPOLOGY_LABEL in group_spec["template"]["spec"]["nodeSelector"]
+            and GKE_TPU_ACCELERATOR_LABEL
+            in group_spec["template"]["spec"]["nodeSelector"]
+        ):
+            topology = group_spec["template"]["spec"]["nodeSelector"][
+                GKE_TPU_TOPOLOGY_LABEL
+            ]
+            accelerator = group_spec["template"]["spec"]["nodeSelector"][
+                GKE_TPU_ACCELERATOR_LABEL
+            ]
+            accelerator_type = utils.tpu_node_selectors_to_type(topology, accelerator)
+            if accelerator_type:
+                resources[f"TPU-{accelerator_type}-head"] = 1
+        else:
+            logger.error(
+                f"Pods using TPUs require both `{GKE_TPU_TOPOLOGY_LABEL}` and `{GKE_TPU_ACCELERATOR_LABEL}` node selectors. "
+                "See https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/tpu.html#configuring-ray-pods-for-tpu-usage "
+                "and https://cloud.google.com/kubernetes-engine/docs/how-to/tpus."
+            )
 
     if memory is not None:
         resources["memory"] = memory

diff --git a/python/ray/tests/kuberay/test_autoscaling_config.py b/python/ray/tests/kuberay/test_autoscaling_config.py
@@ -3,18 +3,21 @@
 import platform
 import requests
 import sys
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Type
 from unittest import mock
 import yaml
 
 import pytest
 
 from ray.autoscaler._private.kuberay.autoscaling_config import (
+    GKE_TPU_ACCELERATOR_LABEL,
+    GKE_TPU_TOPOLOGY_LABEL,
     _derive_autoscaling_config_from_ray_cr,
     AutoscalingConfigProducer,
     _round_up_k8s_quantity,
     _get_num_tpus,
     _get_custom_resources,
+    _get_ray_resources_from_group_spec,
 )
 from ray.autoscaler._private.kuberay.utils import tpu_node_selectors_to_type
 
@@ -242,6 +245,27 @@ def _get_autoscaling_config_with_options() -> dict:
     return config
 
 
+def _get_tpu_group_with_no_node_selectors() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    tpu_group = cr["spec"]["workerGroupSpecs"][2]
+    tpu_group["template"]["spec"].pop("nodeSelector", None)
+    return tpu_group
+
+
+def _get_tpu_group_without_accelerator_node_selector() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    tpu_group = cr["spec"]["workerGroupSpecs"][2]
+    tpu_group["template"]["spec"]["nodeSelector"].pop(GKE_TPU_ACCELERATOR_LABEL, None)
+    return tpu_group
+
+
+def _get_tpu_group_without_topology_node_selector() -> dict[str, Any]:
+    cr = get_basic_ray_cr()
+    tpu_group = cr["spec"]["workerGroupSpecs"][2]
+    tpu_group["template"]["spec"]["nodeSelector"].pop(GKE_TPU_TOPOLOGY_LABEL, None)
+    return tpu_group
+
+
 @pytest.mark.parametrize(
     "input,output",
     [
@@ -344,7 +368,7 @@ def test_resource_quantity(input: str, output: int):
 def test_autoscaling_config(
     ray_cr_in: Dict[str, Any],
     expected_config_out: Optional[Dict[str, Any]],
-    expected_error: Optional[Exception],
+    expected_error: Optional[Type[Exception]],
     expected_error_message: Optional[str],
     expected_log_warning: Optional[str],
 ):
@@ -552,5 +576,113 @@ def test_get_num_tpus(ray_cr_in: Dict[str, Any], expected_num_tpus: int):
             assert num_tpus is None
 
 
+RAY_RESOURCES_PARAM_ARGS = ",".join(
+    [
+        "group_spec",
+        "is_head",
+        "expected_resources",
+    ]
+)
+RAY_RESOURCES_TEST_DATA = (
+    []
+    if platform.system() == "Windows"
+    else [
+        pytest.param(
+            get_basic_ray_cr()["spec"]["headGroupSpec"],
+            True,
+            {
+                "CPU": 1,
+                "memory": 1000000000,
+                "Custom1": 1,
+                "Custom2&qu
8000
ot;: 5,
+            },
+            id="head-group",
+        ),
+        pytest.param(
+            get_basic_ray_cr()["spec"]["workerGroupSpecs"][0],
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+            },
+            id="cpu-group",
+        ),
+        pytest.param(
+            get_basic_ray_cr()["spec"]["workerGroupSpecs"][1],
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "GPU": 3,
+            },
+            id="gpu-group",
+        ),
+        pytest.param(
+            get_basic_ray_cr()["spec"]["workerGroupSpecs"][2],
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+                "TPU-v4-16-head": 1,
+            },
+            id="tpu-group",
+        ),
+        pytest.param(
+            _get_tpu_group_with_no_node_selectors(),
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+            },
+            id="tpu-group-no-node-selectors",
+        ),
+        pytest.param(
+            _get_tpu_group_without_accelerator_node_selector(),
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+            },
+            id="tpu-group-no-accelerator-node-selector",
+        ),
+        pytest.param(
+            _get_tpu_group_without_topology_node_selector(),
+            False,
+            {
+                "CPU": 1,
+                "memory": 536870912,
+                "Custom2": 5,
+                "Custom3": 1,
+                "TPU": 4,
+            },
+            id="tpu-group-no-topology-node-selector",
+        ),
+    ]
+)
+
+
+@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
+@pytest.mark.parametrize(RAY_RESOURCES_PARAM_ARGS, RAY_RESOURCES_TEST_DATA)
+def test_get_ray_resources_from_group_spec(
+    group_spec: Dict[str, Any],
+    is_head: bool,
+    expected_resources: Dict[str, Any],
+):
+    assert _get_ray_resources_from_group_spec(group_spec, is_head) == expected_resources
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))