chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,959 @@
|
||||
import copy
|
||||
import platform
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional, Type
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
import yaml
|
||||
|
||||
from ray.autoscaler._private.kuberay.autoscaling_config import (
|
||||
GKE_TPU_ACCELERATOR_LABEL,
|
||||
GKE_TPU_TOPOLOGY_LABEL,
|
||||
AutoscalingConfigProducer,
|
||||
_derive_autoscaling_config_from_ray_cr,
|
||||
_get_custom_resources,
|
||||
_get_num_tpus,
|
||||
_get_ray_resources_from_group_spec,
|
||||
_round_up_k8s_quantity,
|
||||
)
|
||||
from ray.autoscaler._private.kuberay.utils import tpu_node_selectors_to_type
|
||||
|
||||
AUTOSCALING_CONFIG_MODULE_PATH = "ray.autoscaler._private.kuberay.autoscaling_config"
|
||||
|
||||
|
||||
def get_basic_ray_cr() -> dict:
|
||||
"""Returns the example Ray CR included in the Ray documentation,
|
||||
modified to include a GPU worker group and a TPU worker group.
|
||||
"""
|
||||
cr_path = str(
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "autoscaler"
|
||||
/ "kuberay"
|
||||
/ "ray-cluster.complete.yaml"
|
||||
)
|
||||
config = yaml.safe_load(open(cr_path).read())
|
||||
gpu_group = copy.deepcopy(config["spec"]["workerGroupSpecs"][0])
|
||||
gpu_group["groupName"] = "gpu-group"
|
||||
gpu_group["template"]["spec"]["containers"][0]["resources"]["limits"].setdefault(
|
||||
"nvidia.com/gpu", 3
|
||||
)
|
||||
gpu_group["maxReplicas"] = 200
|
||||
config["spec"]["workerGroupSpecs"].append(gpu_group)
|
||||
tpu_group = copy.deepcopy(config["spec"]["workerGroupSpecs"][0])
|
||||
tpu_group["groupName"] = "tpu-group"
|
||||
tpu_group["template"]["spec"]["containers"][0]["resources"]["limits"].setdefault(
|
||||
"google.com/tpu", 4
|
||||
)
|
||||
tpu_group["template"]["spec"]["nodeSelector"] = {}
|
||||
tpu_group["template"]["spec"]["nodeSelector"][
|
||||
"cloud.google.com/gke-tpu-topology"
|
||||
] = "2x2x2"
|
||||
tpu_group["template"]["spec"]["nodeSelector"][
|
||||
"cloud.google.com/gke-tpu-accelerator"
|
||||
] = "tpu-v4-podslice"
|
||||
tpu_group["maxReplicas"] = 4
|
||||
tpu_group["numOfHosts"] = 2
|
||||
config["spec"]["workerGroupSpecs"].append(tpu_group)
|
||||
return config
|
||||
|
||||
|
||||
def _get_basic_autoscaling_config() -> dict:
|
||||
"""The expected autoscaling derived from the example Ray CR."""
|
||||
return {
|
||||
"cluster_name": "raycluster-complete",
|
||||
"provider": {
|
||||
"disable_node_updaters": True,
|
||||
"disable_launch_config_check": True,
|
||||
"foreground_node_launch": True,
|
||||
"worker_liveness_check": False,
|
||||
"namespace": "default",
|
||||
"type": "kuberay",
|
||||
},
|
||||
"available_node_types": {
|
||||
"headgroup": {
|
||||
"labels": {},
|
||||
"max_workers": 0,
|
||||
"min_workers": 0,
|
||||
"node_config": {},
|
||||
"resources": {
|
||||
"CPU": 1,
|
||||
"memory": 1000000000,
|
||||
"Custom1": 1,
|
||||
"Custom2": 5,
|
||||
},
|
||||
},
|
||||
"small-group": {
|
||||
"labels": {},
|
||||
"max_workers": 300,
|
||||
"min_workers": 0,
|
||||
"node_config": {},
|
||||
"resources": {
|
||||
"CPU": 1,
|
||||
"memory": 536870912,
|
||||
"Custom2": 5,
|
||||
"Custom3": 1,
|
||||
},
|
||||
},
|
||||
# Same as "small-group" with a GPU resource entry added
|
||||
# and modified max_workers.
|
||||
"gpu-group": {
|
||||
"labels": {},
|
||||
"max_workers": 200,
|
||||
"min_workers": 0,
|
||||
"node_config": {},
|
||||
"resources": {
|
||||
"CPU": 1,
|
||||
"memory": 536870912,
|
||||
"Custom2": 5,
|
||||
"Custom3": 1,
|
||||
"GPU": 3,
|
||||
},
|
||||
},
|
||||
# Same as "small-group" with a TPU resource entry added
|
||||
# and modified max_workers and node_config.
|
||||
"tpu-group": {
|
||||
"labels": {},
|
||||
"max_workers": 8,
|
||||
"min_workers": 0,
|
||||
"node_config": {},
|
||||
"resources": {
|
||||
"CPU": 1,
|
||||
"memory": 536870912,
|
||||
"Custom2": 5,
|
||||
"Custom3": 1,
|
||||
"TPU": 4,
|
||||
"TPU-v4-16-head": 1,
|
||||
},
|
||||
},
|
||||
},
|
||||
"auth": {},
|
||||
"cluster_synced_files": [],
|
||||
"file_mounts": {},
|
||||
"file_mounts_sync_continuously": False,
|
||||
"head_node_type": "headgroup",
|
||||
"head_setup_commands": [],
|
||||
"head_start_ray_commands": [],
|
||||
"idle_timeout_minutes": 1.0,
|
||||
"initialization_commands": [],
|
||||
"max_workers": 508,
|
||||
"setup_commands": [],
|
||||
"upscaling_speed": 1000,
|
||||
"worker_setup_commands": [],
|
||||
"worker_start_ray_commands": [],
|
||||
}
|
||||
|
||||
|
||||
def _get_ray_cr_no_cpu_error() -> dict:
|
||||
"""Incorrectly formatted Ray CR without num-cpus rayStartParam and without resource
|
||||
limits. Autoscaler should raise an error when reading this.
|
||||
"""
|
||||
cr = get_basic_ray_cr()
|
||||
# Verify that the num-cpus rayStartParam is not present for the worker type.
|
||||
assert "num-cpus" not in cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]
|
||||
del cr["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][0][
|
||||
"resources"
|
||||
]["limits"]["cpu"]
|
||||
del cr["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][0][
|
||||
"resources"
|
||||
]["requests"]["cpu"]
|
||||
return cr
|
||||
|
||||
|
||||
def _get_no_cpu_error() -> str:
|
||||
return (
|
||||
"Autoscaler failed to detect `CPU` resources for group small-group."
|
||||
"\nSet the `--num-cpus` rayStartParam and/or "
|
||||
"the CPU resource limit for the Ray container."
|
||||
)
|
||||
|
||||
|
||||
def _get_ray_cr_with_overrides() -> dict:
|
||||
"""CR with memory, cpu, and gpu overrides from rayStartParams."""
|
||||
cr = get_basic_ray_cr()
|
||||
cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]["memory"] = "300000000"
|
||||
# num-gpus rayStartParam with no gpus in container limits
|
||||
cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]["num-gpus"] = "100"
|
||||
# num-gpus rayStartParam overriding gpus in container limits
|
||||
cr["spec"]["workerGroupSpecs"][1]["rayStartParams"]["num-gpus"] = "100"
|
||||
cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]["num-cpus"] = "100"
|
||||
return cr
|
||||
|
||||
|
||||
def _get_autoscaling_config_with_overrides() -> dict:
|
||||
"""Autoscaling config with memory and gpu annotations."""
|
||||
config = _get_basic_autoscaling_config()
|
||||
config["available_node_types"]["small-group"]["resources"]["memory"] = 300000000
|
||||
config["available_node_types"]["small-group"]["resources"]["GPU"] = 100
|
||||
config["available_node_types"]["small-group"]["resources"]["CPU"] = 100
|
||||
config["available_node_types"]["gpu-group"]["resources"]["GPU"] = 100
|
||||
return config
|
||||
|
||||
|
||||
def _get_ray_cr_with_autoscaler_options() -> dict:
|
||||
cr = get_basic_ray_cr()
|
||||
cr["spec"]["autoscalerOptions"] = {
|
||||
"upscalingMode": "Conservative",
|
||||
"idleTimeoutSeconds": 300,
|
||||
}
|
||||
return cr
|
||||
|
||||
|
||||
def _get_ray_cr_with_tpu_custom_resource() -> dict:
|
||||
cr = get_basic_ray_cr()
|
||||
cr["spec"]["workerGroupSpecs"][2]["rayStartParams"][
|
||||
"resources"
|
||||
] = '"{"TPU": 4, "Custom2": 5, "Custom3": 1}"'
|
||||
# remove google.com/tpu k8s resource Pod limit
|
||||
del cr["spec"]["workerGroupSpecs"][2]["template"]["spec"]["containers"][0][
|
||||
"resources"
|
||||
]["limits"]["google.com/tpu"]
|
||||
|
||||
return cr
|
||||
|
||||
|
||||
def _get_ray_cr_with_tpu_k8s_resource_limit_and_custom_resource() -> dict:
|
||||
cr = get_basic_ray_cr()
|
||||
cr["spec"]["workerGroupSpecs"][2]["rayStartParams"][
|
||||
"resources"
|
||||
] = '"{"TPU": 4, "Custom2": 5, "Custom3": 1}"'
|
||||
|
||||
return cr
|
||||
|
||||
|
||||
def _get_ray_cr_with_top_level_labels() -> dict:
|
||||
"""CR with a top-level `labels` field."""
|
||||
cr = get_basic_ray_cr()
|
||||
# This top-level structured labels take priority.
|
||||
cr["spec"]["workerGroupSpecs"][0]["labels"] = {"instance-type": "mx5"}
|
||||
|
||||
# rayStartParams labels field should be ignored.
|
||||
cr["spec"]["workerGroupSpecs"][0]["rayStartParams"]["labels"] = "instance-type=n2"
|
||||
return cr
|
||||
|
||||
|
||||
def _get_autoscaling_config_with_top_level_labels() -> dict:
|
||||
config = _get_basic_autoscaling_config()
|
||||
config["available_node_types"]["small-group"]["labels"] = {"instance-type": "mx5"}
|
||||
return config
|
||||
|
||||
|
||||
def _get_ray_cr_with_invalid_top_level_labels() -> dict:
|
||||
"""CR with a syntactically invalid top-level `labels` field."""
|
||||
cr = get_basic_ray_cr()
|
||||
cr["spec"]["workerGroupSpecs"][0]["labels"] = {"!!invalid-key!!": "some-value"}
|
||||
return cr
|
||||
|
||||
|
||||
def _get_ray_cr_with_top_level_resources() -> dict:
|
||||
"""CR with a top-level `resources` field to test priority."""
|
||||
cr = get_basic_ray_cr()
|
||||
|
||||
# The top-level resources field should take priority.
|
||||
cr["spec"]["workerGroupSpecs"][1]["resources"] = {
|
||||
"CPU": "16",
|
||||
"GPU": "8",
|
||||
"memory": "2Gi",
|
||||
"CustomResource": "99",
|
||||
}
|
||||
# These rayStartParams should be ignored.
|
||||
cr["spec"]["workerGroupSpecs"][1]["rayStartParams"]["num-cpus"] = "1"
|
||||
cr["spec"]["workerGroupSpecs"][1]["rayStartParams"]["memory"] = "100000"
|
||||
cr["spec"]["workerGroupSpecs"][1]["rayStartParams"]["num-gpus"] = "2"
|
||||
cr["spec"]["workerGroupSpecs"][1]["rayStartParams"][
|
||||
"resources"
|
||||
] = '"{"Custom2": 1}"'
|
||||
return cr
|
||||
|
||||
|
||||
def _get_autoscaling_config_with_top_level_resources() -> dict:
|
||||
config = _get_basic_autoscaling_config()
|
||||
|
||||
config["available_node_types"]["gpu-group"]["resources"] = {
|
||||
"CPU": 16,
|
||||
"GPU": 8,
|
||||
"memory": 2147483648,
|
||||
"CustomResource": 99,
|
||||
}
|
||||
return config
|
||||
|
||||
|
||||
def _get_ray_cr_with_top_level_tpu_resource() -> dict:
|
||||
"""CR with a top-level `resources` field for the TPU custom resource."""
|
||||
cr = _get_ray_cr_with_tpu_k8s_resource_limit_and_custom_resource()
|
||||
|
||||
# The top-level field should take priority.
|
||||
cr["spec"]["workerGroupSpecs"][2]["resources"] = {"TPU": "8"}
|
||||
return cr
|
||||
|
||||
|
||||
def _get_ray_cr_with_no_tpus() -> dict:
|
||||
cr = get_basic_ray_cr()
|
||||
# remove TPU worker group
|
||||
cr["spec"]["workerGroupSpecs"].pop(2)
|
||||
|
||||
return cr
|
||||
|
||||
|
||||
def _get_ray_cr_with_only_requests() -> dict:
|
||||
"""CR contains only resource requests"""
|
||||
cr = get_basic_ray_cr()
|
||||
|
||||
for group in [cr["spec"]["headGroupSpec"]] + cr["spec"]["workerGroupSpecs"]:
|
||||
for container in group["template"]["spec"]["containers"]:
|
||||
container["resources"]["requests"] = container["resources"]["limits"]
|
||||
del container["resources"]["limits"]
|
||||
return cr
|
||||
|
||||
|
||||
def _get_ray_cr_with_labels() -> dict:
|
||||
"""CR with labels in rayStartParams of head and worker groups."""
|
||||
cr = get_basic_ray_cr()
|
||||
|
||||
# Pass invalid labels to the head group to test error handling.
|
||||
cr["spec"]["headGroupSpec"]["rayStartParams"]["labels"] = "!!ray.io/node-group=,"
|
||||
# Pass valid labels to each of the worker groups.
|
||||
cr["spec"]["workerGroupSpecs"][0]["rayStartParams"][
|
||||
"labels"
|
||||
] = "ray.io/availability-region=us-central2, ray.io/market-type=spot"
|
||||
cr["spec"]["workerGroupSpecs"][1]["rayStartParams"][
|
||||
"labels"
|
||||
] = "ray.io/accelerator-type=A100"
|
||||
cr["spec"]["workerGroupSpecs"][2]["rayStartParams"][
|
||||
"labels"
|
||||
] = "ray.io/accelerator-type=TPU-V4"
|
||||
return cr
|
||||
|
||||
|
||||
def _get_autoscaling_config_with_labels() -> dict:
|
||||
"""Autoscaling config with parsed labels for each group."""
|
||||
config = _get_basic_autoscaling_config()
|
||||
|
||||
# Since we passed invalid labels to the head group `rayStartParams`,
|
||||
# we expect an empty dictionary in the autoscaling config.
|
||||
config["available_node_types"]["headgroup"]["labels"] = {}
|
||||
config["available_node_types"]["small-group"]["labels"] = {
|
||||
"ray.io/availability-region": "us-central2",
|
||||
"ray.io/market-type": "spot",
|
||||
}
|
||||
config["available_node_types"]["gpu-group"]["labels"] = {
|
||||
"ray.io/accelerator-type": "A100"
|
||||
}
|
||||
config["available_node_types"]["tpu-group"]["labels"] = {
|
||||
"ray.io/accelerator-type": "TPU-V4"
|
||||
}
|
||||
return config
|
||||
|
||||
|
||||
def _get_autoscaling_config_with_options() -> dict:
|
||||
config = _get_basic_autoscaling_config()
|
||||
config["upscaling_speed"] = 1
|
||||
config["idle_timeout_minutes"] = 5.0
|
||||
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
|
||||
|
||||
|
||||
def _get_tpu_group_with_v7x_node_selectors() -> dict[str, Any]:
|
||||
cr = get_basic_ray_cr()
|
||||
tpu_group = cr["spec"]["workerGroupSpecs"][2]
|
||||
tpu_group["template"]["spec"]["nodeSelector"][GKE_TPU_TOPOLOGY_LABEL] = "2x2x2"
|
||||
tpu_group["template"]["spec"]["nodeSelector"][GKE_TPU_ACCELERATOR_LABEL] = "tpu7x"
|
||||
return tpu_group
|
||||
|
||||
|
||||
def _get_ray_cr_with_tpu_v7x() -> dict[str, Any]:
|
||||
cr = get_basic_ray_cr()
|
||||
cr["spec"]["workerGroupSpecs"][2] = _get_tpu_group_with_v7x_node_selectors()
|
||||
return cr
|
||||
|
||||
|
||||
def _get_autoscaling_config_with_v7x() -> dict[str, Any]:
|
||||
config = _get_basic_autoscaling_config()
|
||||
config["available_node_types"]["tpu-group"]["resources"]["TPU-v7x-16-head"] = 1
|
||||
config["available_node_types"]["tpu-group"]["resources"].pop("TPU-v4-16-head", None)
|
||||
return config
|
||||
|
||||
|
||||
def _get_tpu_group_with_v5litepod_node_selectors() -> dict[str, Any]:
|
||||
cr = get_basic_ray_cr()
|
||||
tpu_group = cr["spec"]["workerGroupSpecs"][2]
|
||||
tpu_group["template"]["spec"]["nodeSelector"][GKE_TPU_TOPOLOGY_LABEL] = "2x4"
|
||||
tpu_group["template"]["spec"]["nodeSelector"][
|
||||
GKE_TPU_ACCELERATOR_LABEL
|
||||
] = "tpu-v5-lite-podslice"
|
||||
return tpu_group
|
||||
|
||||
|
||||
def _get_ray_cr_with_tpu_v5litepod() -> dict[str, Any]:
|
||||
cr = get_basic_ray_cr()
|
||||
cr["spec"]["workerGroupSpecs"][2] = _get_tpu_group_with_v5litepod_node_selectors()
|
||||
return cr
|
||||
|
||||
|
||||
def _get_autoscaling_config_with_v5litepod() -> dict[str, Any]:
|
||||
config = _get_basic_autoscaling_config()
|
||||
config["available_node_types"]["tpu-group"]["resources"]["TPU-v5litepod-8-head"] = 1
|
||||
config["available_node_types"]["tpu-group"]["resources"].pop("TPU-v4-16-head", None)
|
||||
return config
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input,output",
|
||||
[
|
||||
# There's no particular discipline to these test cases.
|
||||
("100m", 1),
|
||||
("15001m", 16),
|
||||
("2", 2),
|
||||
("100Mi", 104857600),
|
||||
("1G", 1000000000),
|
||||
],
|
||||
)
|
||||
def test_resource_quantity(input: str, output: int):
|
||||
assert _round_up_k8s_quantity(input) == output, output
|
||||
|
||||
|
||||
PARAM_ARGS = ",".join(
|
||||
[
|
||||
"ray_cr_in",
|
||||
"expected_config_out",
|
||||
"expected_error",
|
||||
"expected_error_message",
|
||||
"expected_log_warning",
|
||||
]
|
||||
)
|
||||
|
||||
TEST_DATA = (
|
||||
[]
|
||||
if platform.system() == "Windows"
|
||||
else [
|
||||
pytest.param(
|
||||
get_basic_ray_cr(),
|
||||
_get_basic_autoscaling_config(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id="basic",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_only_requests(),
|
||||
_get_basic_autoscaling_config(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id="only-requests",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_no_cpu_error(),
|
||||
None,
|
||||
ValueError,
|
||||
_get_no_cpu_error(),
|
||||
None,
|
||||
id="no-cpu-error",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_overrides(),
|
||||
_get_autoscaling_config_with_overrides(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id="overrides",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_autoscaler_options(),
|
||||
_get_autoscaling_config_with_options(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id="autoscaler-options",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_tpu_custom_resource(),
|
||||
_get_basic_autoscaling_config(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id="tpu-custom-resource",
|
||||
),
|
||||
pytest.param(
|
||||
get_basic_ray_cr(),
|
||||
_get_basic_autoscaling_config(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id="tpu-k8s-resource-limit",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_tpu_k8s_resource_limit_and_custom_resource(),
|
||||
_get_basic_autoscaling_config(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id="tpu-k8s-resource-limit-and-custom-resource",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_labels(),
|
||||
_get_basic_autoscaling_config(),
|
||||
None,
|
||||
None,
|
||||
"Ignoring labels: ray.io/accelerator-type=TPU-V4 set in rayStartParams for group 'tpu-group'. Group labels are supported in the top-level Labels field starting in KubeRay v1.5",
|
||||
id="groups-with-raystartparam-labels",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_top_level_labels(),
|
||||
_get_autoscaling_config_with_top_level_labels(),
|
||||
None,
|
||||
None,
|
||||
"Ignoring labels: instance-type=n2 set in rayStartParams for group 'small-group'. Group labels are supported in the top-level Labels field starting in KubeRay v1.5",
|
||||
id="groups-with-top-level-labels",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_invalid_top_level_labels(),
|
||||
_get_basic_autoscaling_config(),
|
||||
ValueError,
|
||||
None,
|
||||
None,
|
||||
id="invalid-top-level-labels",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_tpu_v7x(),
|
||||
_get_autoscaling_config_with_v7x(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id="tpu-v7x",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_tpu_v5litepod(),
|
||||
_get_autoscaling_config_with_v5litepod(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
id="tpu-v5litepod",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
|
||||
@pytest.mark.parametrize(PARAM_ARGS, TEST_DATA)
|
||||
def test_autoscaling_config(
|
||||
ray_cr_in: Dict[str, Any],
|
||||
expected_config_out: Optional[Dict[str, Any]],
|
||||
expected_error: Optional[Type[Exception]],
|
||||
expected_error_message: Optional[str],
|
||||
expected_log_warning: Optional[str],
|
||||
):
|
||||
ray_cr_in["metadata"]["namespace"] = "default"
|
||||
# Reset log_once state to ensure each test case is independent.
|
||||
from ray.util.debug import _logged
|
||||
|
||||
_logged.clear()
|
||||
with mock.patch(f"{AUTOSCALING_CONFIG_MODULE_PATH}.logger") as mock_logger:
|
||||
if expected_error:
|
||||
with pytest.raises(expected_error, match=expected_error_message):
|
||||
_derive_autoscaling_config_from_ray_cr(ray_cr_in)
|
||||
else:
|
||||
assert (
|
||||
_derive_autoscaling_config_from_ray_cr(ray_cr_in) == expected_config_out
|
||||
)
|
||||
if expected_log_warning:
|
||||
mock_logger.warning.assert_called_with(expected_log_warning)
|
||||
else:
|
||||
mock_logger.warning.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
|
||||
def test_cr_image_consistency():
|
||||
"""Verify that the example config uses the same Ray image for all Ray pods."""
|
||||
cr = get_basic_ray_cr()
|
||||
|
||||
group_specs = [cr["spec"]["headGroupSpec"]] + cr["spec"]["workerGroupSpecs"]
|
||||
# Head, CPU group, GPU group, TPU group.
|
||||
assert len(group_specs) == 4
|
||||
|
||||
ray_containers = [
|
||||
group_spec["template"]["spec"]["containers"][0] for group_spec in group_specs
|
||||
]
|
||||
|
||||
# All Ray containers in the example config have "ray-" in their name.
|
||||
assert all("ray-" in ray_container["name"] for ray_container in ray_containers)
|
||||
|
||||
# All Ray images are from the Ray repo.
|
||||
assert all(
|
||||
"rayproject/ray" in ray_container["image"] for ray_container in ray_containers
|
||||
)
|
||||
|
||||
# All Ray images are the same.
|
||||
assert len({ray_container["image"] for ray_container in ray_containers}) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exception", [Exception, requests.HTTPError])
|
||||
@pytest.mark.parametrize("num_exceptions", range(6))
|
||||
def test_autoscaling_config_fetch_retries(exception, num_exceptions):
|
||||
"""Validates retry logic in
|
||||
AutoscalingConfigProducer._fetch_ray_cr_from_k8s_with_retries.
|
||||
"""
|
||||
|
||||
class MockKubernetesHttpApiClient:
|
||||
def __init__(self):
|
||||
self.exception_counter = 0
|
||||
|
||||
def get(self, *args, **kwargs):
|
||||
if self.exception_counter < num_exceptions:
|
||||
self.exception_counter += 1
|
||||
raise exception
|
||||
else:
|
||||
return {"ok-key": "ok-value"}
|
||||
|
||||
class MockAutoscalingConfigProducer(AutoscalingConfigProducer):
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.kubernetes_api_client = MockKubernetesHttpApiClient()
|
||||
self._ray_cr_path = "rayclusters/mock"
|
||||
|
||||
config_producer = MockAutoscalingConfigProducer()
|
||||
# Patch retry backoff period.
|
||||
with mock.patch(
|
||||
"ray.autoscaler._private.kuberay.autoscaling_config.RAYCLUSTER_FETCH_RETRY_S",
|
||||
0,
|
||||
):
|
||||
# If you hit an exception and it's not HTTPError, expect to raise.
|
||||
# If you hit >= 5 exceptions, expect to raise.
|
||||
# Otherwise, don't expect to raise.
|
||||
if (
|
||||
num_exceptions > 0 and exception != requests.HTTPError
|
||||
) or num_exceptions >= 5:
|
||||
with pytest.raises(exception):
|
||||
config_producer._fetch_ray_cr_from_k8s_with_retries()
|
||||
else:
|
||||
out = config_producer._fetch_ray_cr_from_k8s_with_retries()
|
||||
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",
|
||||
"v5litepod-4",
|
||||
id="tpu-v5e-device-test",
|
||||
),
|
||||
pytest.param(
|
||||
"tpu-v5-lite-podslice",
|
||||
"2x4",
|
||||
"v5litepod-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.param(
|
||||
"tpu7x",
|
||||
"2x2x2",
|
||||
"v7x-16",
|
||||
id="tpu-v7x-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",
|
||||
"expected_num_tpus",
|
||||
]
|
||||
)
|
||||
TPU_TEST_DATA = (
|
||||
[]
|
||||
if platform.system() == "Windows"
|
||||
else [
|
||||
pytest.param(
|
||||
get_basic_ray_cr(),
|
||||
4,
|
||||
id="tpu-k8s-resource-limits",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_tpu_custom_resource(),
|
||||
4,
|
||||
id="tpu-custom-resource",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_tpu_k8s_resource_limit_and_custom_resource(),
|
||||
4,
|
||||
id="tpu--k8s-resource-limits-and-custom-resource",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_no_tpus(),
|
||||
0,
|
||||
id="no-tpus-requested",
|
||||
),
|
||||
pytest.param(
|
||||
_get_ray_cr_with_top_level_tpu_resource(),
|
||||
8,
|
||||
id="tpu-top-level-resource",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
|
||||
@pytest.mark.parametrize(TPU_PARAM_ARGS, TPU_TEST_DATA)
|
||||
def test_get_num_tpus(ray_cr_in: Dict[str, Any], expected_num_tpus: int):
|
||||
"""Verify that _get_num_tpus correctly returns the number of requested TPUs."""
|
||||
for worker_group in ray_cr_in["spec"]["workerGroupSpecs"]:
|
||||
group_resources = worker_group.get("resources", {})
|
||||
ray_start_params = worker_group["rayStartParams"]
|
||||
custom_resources = _get_custom_resources(
|
||||
group_resources, ray_start_params, worker_group["groupName"]
|
||||
)
|
||||
k8s_resources = worker_group["template"]["spec"]["containers"][0]["resources"]
|
||||
|
||||
num_tpus = _get_num_tpus(group_resources, custom_resources, k8s_resources)
|
||||
|
||||
if worker_group["groupName"] == "tpu-group":
|
||||
assert num_tpus == expected_num_tpus
|
||||
else:
|
||||
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": 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.param(
|
||||
_get_tpu_group_with_v7x_node_selectors(),
|
||||
False,
|
||||
{
|
||||
"CPU": 1,
|
||||
"memory": 536870912,
|
||||
"Custom2": 5,
|
||||
"Custom3": 1,
|
||||
"TPU": 4,
|
||||
"TPU-v7x-16-head": 1,
|
||||
},
|
||||
id="tpu-group-v7x",
|
||||
),
|
||||
pytest.param(
|
||||
_get_tpu_group_with_v5litepod_node_selectors(),
|
||||
False,
|
||||
{
|
||||
"CPU": 1,
|
||||
"memory": 536870912,
|
||||
"Custom2": 5,
|
||||
"Custom3": 1,
|
||||
"TPU": 4,
|
||||
"TPU-v5litepod-8-head": 1,
|
||||
},
|
||||
id="tpu-group-v5litepod",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@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
|
||||
|
||||
|
||||
@pytest.mark.skipif(platform.system() == "Windows", reason="Not relevant.")
|
||||
def test_top_level_resources_override_warnings():
|
||||
"""
|
||||
Verify all override warnings are logged when a top-level `resources` field is used in
|
||||
addition to specifying those resources in the rayStartParams.
|
||||
"""
|
||||
ray_cr_in = _get_ray_cr_with_top_level_resources()
|
||||
ray_cr_in["metadata"]["namespace"] = "default"
|
||||
|
||||
with mock.patch(f"{AUTOSCALING_CONFIG_MODULE_PATH}.logger") as mock_logger:
|
||||
_derive_autoscaling_config_from_ray_cr(ray_cr_in)
|
||||
|
||||
expected_calls = [
|
||||
mock.call(
|
||||
"'CPU' specified in both the top-level 'resources' field and in 'rayStartParams'. "
|
||||
"Using the value from 'resources': 16."
|
||||
),
|
||||
mock.call(
|
||||
"'GPU' specified in both the top-level 'resources' field and in 'rayStartParams'. "
|
||||
"Using the value from 'resources': 8."
|
||||
),
|
||||
mock.call(
|
||||
"'memory' specified in both the top-level 'resources' field and in 'rayStartParams'. "
|
||||
"Using the value from 'resources': 2Gi."
|
||||
),
|
||||
mock.call(
|
||||
"custom resources specified in both the top-level 'resources' field and in 'rayStartParams'. "
|
||||
"Using the values from 'resources': {'CPU': '16', 'GPU': '8', 'memory': '2Gi', 'CustomResource': '99'}."
|
||||
),
|
||||
]
|
||||
|
||||
# Assert that all expected calls were made, in any order.
|
||||
mock_logger.warning.assert_has_calls(expected_calls, any_order=True)
|
||||
assert mock_logger.warning.call_count == 4
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user