diff --git a/paasta_tools/nrtsearchservice_tools.py b/paasta_tools/nrtsearchservice_tools.py index 6b0d7abb3f..9612a0d7ef 100644 --- a/paasta_tools/nrtsearchservice_tools.py +++ b/paasta_tools/nrtsearchservice_tools.py @@ -10,6 +10,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any +from typing import Dict from typing import List from typing import Mapping from typing import Optional @@ -28,6 +30,7 @@ class NrtsearchServiceDeploymentConfigDict(LongRunningServiceConfigDict, total=False): replicas: int + serverSets: List[Dict[str, Any]] class NrtsearchServiceDeploymentConfig(LongRunningServiceConfig): diff --git a/paasta_tools/nrtsearchserviceeks_tools.py b/paasta_tools/nrtsearchserviceeks_tools.py index 29b3e4bbad..ac6835c00f 100644 --- a/paasta_tools/nrtsearchserviceeks_tools.py +++ b/paasta_tools/nrtsearchserviceeks_tools.py @@ -2,6 +2,8 @@ import service_configuration_lib +from paasta_tools.autoscaling.utils import MetricsProviderDict +from paasta_tools.long_running_service_tools import METRICS_PROVIDER_PROMQL from paasta_tools.nrtsearchservice_tools import NrtsearchServiceDeploymentConfig from paasta_tools.utils import DEFAULT_SOA_DIR from paasta_tools.utils import BranchDictV2 @@ -13,6 +15,64 @@ class NrtsearchServiceEksDeploymentConfig(NrtsearchServiceDeploymentConfig): config_filename_prefix = "nrtsearchserviceeks" + def _get_autoscalable_server_set(self) -> Optional[dict]: + """Return the non-primary serverSet that has targetGpuUtilization > 0.""" + server_sets = self.config_dict.get("serverSets", []) + for server_set in server_sets: + if server_set.get("primary", False): + continue + autoscaling = server_set.get("autoscaling") + if not autoscaling: + continue + if autoscaling.get("targetGpuUtilization", 0) > 0: + return server_set + return None + + def get_autoscaling_metrics_provider( + self, provider_type: str + ) -> Optional[MetricsProviderDict]: + if provider_type != METRICS_PROVIDER_PROMQL: + return None + server_set = self._get_autoscalable_server_set() + if server_set is None: + return None + + instance_name = "replica" + deployment_name = self.get_sanitised_deployment_name() + namespace = self.get_namespace() + paasta_cluster = self.get_cluster() + service = self.get_service() + + metrics_query = ( + "avg(" + "DCGM_FI_DEV_GPU_UTIL" + " * on(kube_pod, kube_namespace) group_left()" + " (kube_pod_labels{" + f"label_paasta_yelp_com_service='{service}'," + f"label_yelp_com_paasta_instance='{instance_name}'," + f"paasta_cluster='{paasta_cluster}'" + "})" + ")" + ) + return MetricsProviderDict( + type=METRICS_PROVIDER_PROMQL, + metrics_query=metrics_query, + series_query=( + f"kube_deployment_labels{{" + f"deployment='{deployment_name}'," + f"paasta_cluster='{paasta_cluster}'," + f"namespace='{namespace}'" + f"}}" + ), + setpoint=1.0, + ) + + def get_sanitised_deployment_name(self) -> str: + return f"{self.instance}-replica-dep" + + def namespace_custom_prometheus_metric_name(self, metric_name: str) -> str: + return f"{self.get_sanitised_deployment_name()}-gpu-prom" + def load_nrtsearchserviceeks_instance_config( service: str, diff --git a/paasta_tools/setup_prometheus_adapter_config.py b/paasta_tools/setup_prometheus_adapter_config.py index 0f3f982206..92cecf8a35 100755 --- a/paasta_tools/setup_prometheus_adapter_config.py +++ b/paasta_tools/setup_prometheus_adapter_config.py @@ -21,6 +21,7 @@ from typing import Dict from typing import List from typing import Optional +from typing import Union from typing import cast import ruamel.yaml as yaml @@ -65,6 +66,7 @@ from paasta_tools.long_running_service_tools import METRICS_PROVIDER_UWSGI from paasta_tools.long_running_service_tools import METRICS_PROVIDER_UWSGI_V2 from paasta_tools.long_running_service_tools import METRICS_PROVIDER_WORKER_LOAD +from paasta_tools.nrtsearchserviceeks_tools import NrtsearchServiceEksDeploymentConfig from paasta_tools.paasta_service_config_loader import PaastaServiceConfigLoader from paasta_tools.utils import DEFAULT_SOA_DIR from paasta_tools.utils import get_services_for_cluster @@ -88,6 +90,7 @@ K8S_INSTANCE_TYPE_CLASSES = ( KubernetesDeploymentConfig, EksDeploymentConfig, + NrtsearchServiceEksDeploymentConfig, ) @@ -204,7 +207,9 @@ def _minify_promql(query: str) -> str: def create_instance_scaling_rule( service: str, - instance_config: KubernetesDeploymentConfig, + instance_config: Union[ + KubernetesDeploymentConfig, NrtsearchServiceEksDeploymentConfig + ], metrics_provider_config: MetricsProviderDict, paasta_cluster: str, ) -> Optional[PrometheusAdapterRule]: @@ -284,7 +289,9 @@ def create_instance_scaling_rule( def create_instance_active_requests_scaling_rule( service: str, - instance_config: KubernetesDeploymentConfig, + instance_config: Union[ + KubernetesDeploymentConfig, NrtsearchServiceEksDeploymentConfig + ], metrics_provider_config: MetricsProviderDict, paasta_cluster: str, metric_name: str, @@ -394,7 +401,9 @@ def create_instance_active_requests_scaling_rule( def create_instance_uwsgi_scaling_rule( service: str, - instance_config: KubernetesDeploymentConfig, + instance_config: Union[ + KubernetesDeploymentConfig, NrtsearchServiceEksDeploymentConfig + ], metrics_provider_config: MetricsProviderDict, paasta_cluster: str, metric_name: str, @@ -496,7 +505,9 @@ def create_instance_uwsgi_scaling_rule( def create_instance_uwsgi_v2_scaling_rule( service: str, - instance_config: KubernetesDeploymentConfig, + instance_config: Union[ + KubernetesDeploymentConfig, NrtsearchServiceEksDeploymentConfig + ], metrics_provider_config: MetricsProviderDict, paasta_cluster: str, metric_name: str, @@ -567,7 +578,9 @@ def create_instance_uwsgi_v2_scaling_rule( def create_instance_worker_load_scaling_rule( service: str, - instance_config: KubernetesDeploymentConfig, + instance_config: Union[ + KubernetesDeploymentConfig, NrtsearchServiceEksDeploymentConfig + ], metrics_provider_config: MetricsProviderDict, paasta_cluster: str, metric_name: str, @@ -640,7 +653,9 @@ def create_instance_worker_load_scaling_rule( def create_instance_piscina_scaling_rule( service: str, - instance_config: KubernetesDeploymentConfig, + instance_config: Union[ + KubernetesDeploymentConfig, NrtsearchServiceEksDeploymentConfig + ], metrics_provider_config: MetricsProviderDict, paasta_cluster: str, metric_name: str, @@ -737,7 +752,9 @@ def create_instance_piscina_scaling_rule( def create_instance_gunicorn_scaling_rule( service: str, - instance_config: KubernetesDeploymentConfig, + instance_config: Union[ + KubernetesDeploymentConfig, NrtsearchServiceEksDeploymentConfig + ], metrics_provider_config: MetricsProviderDict, paasta_cluster: str, metric_name: str, @@ -845,7 +862,9 @@ def create_instance_gunicorn_scaling_rule( def create_instance_arbitrary_promql_scaling_rule( service: str, - instance_config: KubernetesDeploymentConfig, + instance_config: Union[ + KubernetesDeploymentConfig, NrtsearchServiceEksDeploymentConfig + ], metrics_provider_config: MetricsProviderDict, paasta_cluster: str, metric_name: str, @@ -903,7 +922,9 @@ def create_instance_arbitrary_promql_scaling_rule( def get_rules_for_service_instance( service_name: str, - instance_config: KubernetesDeploymentConfig, + instance_config: Union[ + KubernetesDeploymentConfig, NrtsearchServiceEksDeploymentConfig + ], paasta_cluster: str, ) -> List[PrometheusAdapterRule]: """ @@ -964,6 +985,16 @@ def create_prometheus_adapter_config( ) } ) + services.update( + { + service_name + for service_name, _ in get_services_for_cluster( + cluster=paasta_cluster, + instance_type="nrtsearchserviceeks", + soa_dir=str(soa_dir), + ) + } + ) for service_name in services: config_loader = PaastaServiceConfigLoader( service=service_name, soa_dir=str(soa_dir) @@ -976,7 +1007,13 @@ def create_prometheus_adapter_config( rules.extend( get_rules_for_service_instance( service_name=service_name, - instance_config=instance_config, + instance_config=cast( + Union[ + KubernetesDeploymentConfig, + NrtsearchServiceEksDeploymentConfig, + ], + instance_config, + ), paasta_cluster=paasta_cluster, ) ) diff --git a/tests/test_setup_prometheus_adapter_config.py b/tests/test_setup_prometheus_adapter_config.py index 5c85d8636e..caffecbbc9 100644 --- a/tests/test_setup_prometheus_adapter_config.py +++ b/tests/test_setup_prometheus_adapter_config.py @@ -11,6 +11,7 @@ from paasta_tools.long_running_service_tools import METRICS_PROVIDER_UWSGI from paasta_tools.long_running_service_tools import METRICS_PROVIDER_UWSGI_V2 from paasta_tools.long_running_service_tools import METRICS_PROVIDER_WORKER_LOAD +from paasta_tools.nrtsearchserviceeks_tools import NrtsearchServiceEksDeploymentConfig from paasta_tools.setup_prometheus_adapter_config import _minify_promql from paasta_tools.setup_prometheus_adapter_config import ( create_instance_active_requests_scaling_rule, @@ -485,3 +486,112 @@ def test_create_instance_arbitrary_promql_scaling_rule_with_custom_resources(): "metricsQuery": "foo", "seriesQuery": "bar", } + + +def test_nrtsearch_gpu_autoscaling_via_service_instance(): + """Test that nrtsearch configs with targetGpuUtilization generate adapter rules + through the standard get_rules_for_service_instance path.""" + instance_config = NrtsearchServiceEksDeploymentConfig( + service="nrtsearch", + cluster="pnw-prod", + instance="my-cluster", + config_dict={ + "serverSets": [ + {"name": "indexing", "primary": True, "count": 1}, + { + "name": "search", + "primary": False, + "autoscaling": { + "minInstances": 2, + "maxInstances": 10, + "targetCpuUtilization": 80, + "targetGpuUtilization": 70, + "stabilizationWindow": 600, + }, + }, + ], + }, + branch_dict=None, + soa_dir="/mock/soa", + ) + + rules = get_rules_for_service_instance( + service_name="nrtsearch", + instance_config=instance_config, + paasta_cluster="pnw-prod", + ) + + assert len(rules) == 1 + rule = rules[0] + assert rule["name"]["as"] == "my-cluster-replica-dep-gpu-prom" + assert "DCGM_FI_DEV_GPU_UTIL" in rule["metricsQuery"] + assert "my-cluster-replica-dep" in rule["seriesQuery"] + assert "pnw-prod" in rule["seriesQuery"] + assert rule["resources"] == { + "overrides": { + "namespace": {"resource": "namespace"}, + "deployment": {"group": "apps", "resource": "deployments"}, + }, + } + + +def test_nrtsearch_no_gpu_config_returns_no_rules(): + """Test that nrtsearch configs without targetGpuUtilization produce no rules.""" + instance_config = NrtsearchServiceEksDeploymentConfig( + service="nrtsearch", + cluster="pnw-prod", + instance="my-cluster", + config_dict={ + "serverSets": [ + { + "name": "search", + "primary": False, + "autoscaling": { + "minInstances": 2, + "maxInstances": 10, + "targetCpuUtilization": 80, + }, + }, + ], + }, + branch_dict=None, + soa_dir="/mock/soa", + ) + + rules = get_rules_for_service_instance( + service_name="nrtsearch", + instance_config=instance_config, + paasta_cluster="pnw-prod", + ) + + assert rules == [] + + +def test_nrtsearch_primary_server_set_skipped(): + """Test that primary serverSets are never used for GPU autoscaling.""" + instance_config = NrtsearchServiceEksDeploymentConfig( + service="nrtsearch", + cluster="pnw-prod", + instance="my-cluster", + config_dict={ + "serverSets": [ + { + "name": "indexing", + "primary": True, + "autoscaling": { + "targetGpuUtilization": 70, + }, + }, + ], + }, + branch_dict=None, + soa_dir="/mock/soa", + ) + + rules = get_rules_for_service_instance( + service_name="nrtsearch", + instance_config=instance_config, + paasta_cluster="pnw-prod", + ) + + assert rules == []