diff --git a/python/ray/autoscaler/_private/kuberay/utils.py b/python/ray/autoscaler/_private/kuberay/utils.py index beb18b0b8952..1b045b422b66 100644 --- a/python/ray/autoscaler/_private/kuberay/utils.py +++ b/python/ray/autoscaler/_private/kuberay/utils.py @@ -11,6 +11,7 @@ "tpu-v5-lite-device": "v5e", "tpu-v5-lite-podslice": "v5e", "tpu-v5p-slice": "v5p", + "tpu-v6e-slice": "v6e", } @@ -102,9 +103,9 @@ def tpu_node_selectors_to_type(topology: str, accelerator: str) -> Optional[str] # Reduce e.g. "2x2x2" to 8 chip_dimensions = [int(chip_count) for chip_count in topology.split("x")] num_chips = reduce(lambda x, y: x * y, chip_dimensions) - default_num_cores_per_chip = 2 - if generation == "v5e": - default_num_cores_per_chip = 1 + default_num_cores_per_chip = 1 + if generation == "v4" or generation == "v5p": + default_num_cores_per_chip = 2 num_cores = num_chips * default_num_cores_per_chip return f"{generation}-{num_cores}" return None diff --git a/python/ray/tests/kuberay/test_autoscaling_config.py b/python/ray/tests/kuberay/test_autoscaling_config.py index 12b5c239f4a6..699df522eb1a 100644 --- a/python/ray/tests/kuberay/test_autoscaling_config.py +++ b/python/ray/tests/kuberay/test_autoscaling_config.py @@ -16,6 +16,8 @@ _get_custom_resources, ) +from ray.autoscaler._private.kuberay.utils import tpu_node_selectors_to_type + AUTOSCALING_CONFIG_MODULE_PATH = "ray.autoscaler._private.kuberay.autoscaling_config" @@ -402,6 +404,75 @@ def _fetch_ray_cr_from_k8s(self) -> Dict[str, Any]: assert out == {"ok-key": "ok-value"} +TPU_TYPES_ARGS = ",".join( + [ + "accelerator", + "topology", + "expected_tpu_type", + ] +) +TPU_TYPES_DATA = ( + [] + if platform.system() == "Windows" + else [ + pytest.param( + "tpu-v4-podslice", + None, + None, + id="tpu-none-topology", + ), + pytest.param( + None, + "2x2x2", + None, + id="tpu-none-accelerator", + ), + pytest.param( + "tpu-v4-podslice", + "2x2x2", + "v4-16", + id="tpu-v4-test", + ), + pytest.param( + "tpu-v5-lite-device", + "2x2", + "v5e-4", + id="tpu-v5e-device-test", + ), + pytest.param( + "tpu-v5-lite-podslice", + "2x4", + "v5e-8", + id="tpu-v5e-podslice-test", + ), + pytest.param( + "tpu-v5p-slice", + "2x2x4", + "v5p-32", + id="tpu-v5p-test", + ), + pytest.param( + "tpu-v6e-slice", + "16x16", + "v6e-256", + id="tpu-v6e-test", + ), + ] +) + + +@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.") +@pytest.mark.parametrize(TPU_TYPES_ARGS, TPU_TYPES_DATA) +def test_tpu_node_selectors_to_type( + accelerator: str, topology: str, expected_tpu_type: str +): + """Verify that tpu_node_selectors_to_type correctly returns TPU type from + TPU nodeSelectors. + """ + tpu_type = tpu_node_selectors_to_type(topology, accelerator) + assert expected_tpu_type == tpu_type + + TPU_PARAM_ARGS = ",".join( [ "ray_cr_in",