Add prometheus-adapter rule generation for nrtsearch GPU autoscaling

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

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,

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,
                     )
                 )

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 == []