Add support for setting lower/upper-bounds for autotuned resources

paasta_tools/cli/schemas/kubernetes_schema.json

@@ -866,6 +866,56 @@
                     "type": "string",
                     "pattern": "^[a-z0-9]([-a-z0-9]*[a-z0-9])?$",
                     "maxLength": 63
+                },
+                "autotune_limits": {
+                    "type": "object",
+                    "properties": {
+                        "cpus": {
+                            "type": "object",
+                            "properties": {
+                                "min": {
+                                    "type": "number",
+                                    "minimum": 0,
+                                    "exclusiveMinimum": true
+                                },
+                                "max": {
+                                    "type": "number",
+                                    "minimum": 0,
+                                    "exclusiveMinimum": true
+                                }
+                            }
+                        },
+                        "mem": {
+                            "type": "object",
+                            "properties": {
+                                "min": {
+                                    "type": "integer",
+                                    "minimum": 32,
+                                    "exclusiveMinimum": false
+                                },
+                                "max": {
+                                    "type": "integer",
+                                    "minimum": 32,
+                                    "exclusiveMinimum": false
+                                }
+                            }
+                        },
+                        "disk": {
+                            "type": "object",
+                            "properties": {
+                                "min": {
+                                    "type": "integer",
+                                    "minimum": 128,
+                                    "exclusiveMinimum": false
+                                },
+                                "max": {
+                                    "type": "integer",
+                                    "minimum": 128,
+                                    "exclusiveMinimum": false
+                                }
+                            }
+                        }
+                    }
                 }
             }
         }

paasta_tools/config_utils.py

@@ -1,3 +1,4 @@
+import copy
 import logging
 import os
 import subprocess
@@ -7,10 +8,12 @@
 from typing import List
 from typing import Optional
 from typing import Set
+from typing import Tuple
 
 import ruamel.yaml as yaml
 
 from paasta_tools.cli.cmds.validate import validate_schema
+from paasta_tools.utils import AUTO_SOACONFIG_SUBDIR
 from paasta_tools.utils import DEFAULT_SOA_DIR
 
 
@@ -259,3 +262,88 @@ def commit_to_remote(self, extra_message: str = ""):
             _push_to_remote(self.branch)
         else:
             log.info("No files changed, no push required.")
+
+    def _clamp_recommendations(
+        self, merged_recommendation: Dict[str, Any], config: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        clamped_recomendation = copy.deepcopy(merged_recommendation)
+        for limit_type, limits in config.get("autotune_limits", {}).items():
+            log.debug(f"Processing {limit_type} autotune limits...")
+            min_value = limits.get("min")
+            max_value = limits.get("max")
+            unclamped_resource_value = clamped_recomendation.get(limit_type)
+
+            # no autotune data present, but min value present
+            if not unclamped_resource_value and min_value:
+                # use the min value since this is likely an autogenerated service where we know we have a given minimum CPU
+                # that we'd like to allocate
+                log.debug(
+                    f"No {limit_type} autotune data found, using autotune limit lower bound ({min_value})."
+                )
+                clamped_recomendation[limit_type] = min_value
+
+            # otherwise, we can do some pretty rote clamping of resource values
+            elif unclamped_resource_value is not None:
+                if min_value and unclamped_resource_value < min_value:
+                    log.debug(
+                        f"{limit_type} autotune config under configured limit ({min_value}), using autotune limit lower bound."
+                    )
+                    clamped_recomendation[limit_type] = min_value
+                if max_value and unclamped_resource_value > max_value:
+                    log.debug(
+                        f"{limit_type} autotune config over configured limit ({min_value}), using autotune limit upper bound."
+                    )
+                    clamped_recomendation[limit_type] = max_value
+            else:
+                log.debug(
+                    f"No {limit_type} autotune data or limits found, will continue using PaaSTA defaults."
+                )
+
+        return clamped_recomendation
+
+    def merge_recommendations(
+        self, recs: Dict[Tuple[str, str], Dict[str, Any]]
+    ) -> Dict[Tuple[str, str], Dict[str, Any]]:
+        """
+        :param recs: Dictionary of (service, instance_type_cluster) -> recommendations.
+        NOTE: instance_type_cluster is something like "kubernetes-pnw-prod".
+        :returns: a dictionary of the same format, with the previous recommendations merged in and autotune_limits applied.
+        """
+        merged_recs = {}
+        for (
+            service,
+            instance_type_cluster,
+        ), recommendations_by_instance in recs.items():
+            log.info(f"Starting to process {service}/{instance_type_cluster}.yaml...")
+            log.debug(
+                f"Getting current autotune configs for {service}/{instance_type_cluster}.yaml"
+            )
+            existing_recommendations = self.get_existing_configs(
+                service=service,
+                file_name=instance_type_cluster,
+                sub_dir=AUTO_SOACONFIG_SUBDIR,
+            )
+
+            log.debug(
+                f"Getting current configs for {service}/{instance_type_cluster}.yaml..."
+            )
+            existing_configs = self.get_existing_configs(
+                service=service,
+                file_name=instance_type_cluster,
+            )
+
+            for instance_name, recommendation in recommendations_by_instance.items():
+                log.debug(
+                    f"Merging recommendations for {instance_name} in {service}/{AUTO_SOACONFIG_SUBDIR}/{instance_type_cluster}.yaml..."
+                )
+                existing_recommendations.setdefault(instance_name, {})
+                existing_recommendations[instance_name].update(recommendation)
+
+                existing_recommendations[instance_name] = self._clamp_recommendations(
+                    merged_recommendation=existing_recommendations[instance_name],
+                    config=existing_configs[instance_name],
+                )
+            merged_recs[(service, instance_type_cluster)] = existing_recommendations
+            log.info(f"Done processing {service}/{instance_type_cluster}.yaml.")
+
+        return merged_recs

paasta_tools/contrib/rightsizer_soaconfigs_update.py

@@ -19,6 +19,9 @@
 from paasta_tools.utils import format_git_url
 from paasta_tools.utils import load_system_paasta_config
 
+
+log = logging.getLogger(__name__)
+
 NULL = "null"
 SUPPORTED_CSV_KEYS = (
     "cpus",
@@ -276,7 +279,7 @@ def get_recommendations_by_service_file(
             result["cluster"],
         )  # e.g. (foo, marathon-norcal-stagef)
         instance_type = result["cluster"].split("-", 1)[0]
-        rec: Union[KubernetesRecommendation, CassandraRecommendation]
+        rec: Union[KubernetesRecommendation, CassandraRecommendation] = {}
         if instance_type == "cassandracluster":
             rec = get_cassandra_recommendation_from_result(result, keys_to_apply)
         elif instance_type == "kubernetes":
@@ -317,17 +320,17 @@ def main(args):
         validation_schema_path=AUTO_SOACONFIG_SUBDIR,
     )
     with updater:
-        for (service, extra_info), instance_recommendations in results.items():
-            existing_recommendations = updater.get_existing_configs(
-                service, extra_info, AUTO_SOACONFIG_SUBDIR
+        for (
+            service,
+            instance_type_cluster,
+        ), instance_recommendations in updater.merge_recommendations(results).items():
+            log.info(
+                f"Writing configs for {service} to {AUTO_SOACONFIG_SUBDIR}/{instance_type_cluster}.yaml..."
             )
-            for instance_name, recommendation in instance_recommendations.items():
-                existing_recommendations.setdefault(instance_name, {})
-                existing_recommendations[instance_name].update(recommendation)
             updater.write_configs(
                 service,
-                extra_info,
-                existing_recommendations,
+                instance_type_cluster,
+                instance_recommendations,
                 AUTO_SOACONFIG_SUBDIR,
                 HEADER_COMMENT,
             )

tests/test_config_utils.py

@@ -317,3 +317,62 @@ def test_auto_config_updater_commit(mock_push, mock_commit, did_commit, updater)
         mock_push.assert_called_once_with(updater.branch)
     else:
         assert mock_push.call_count == 0
+
+
+def test_auto_config_updater_merge_recommendations_limits(updater):
+    service = "foo"
+    conf_file = "notk8s-euwest-prod"
+    limited_instance = "some_instance"
+    noop_instance = "other_instance"
+    autotune_data = {
+        limited_instance: {"cpus": 0.1, "mem": 167, "disk": 256, "cpu_burst_add": 0.5}
+    }
+    user_data = {
+        limited_instance: {
+            "cmd": "ls",
+            "autotune_limits": {
+                "cpus": {"min": 1, "max": 2},
+                "mem": {"min": 1024, "max": 2048},
+                "disk": {"min": 512, "max": 1024},
+            },
+        },
+        noop_instance: {"cmd": "ls"},
+    }
+
+    recs = {
+        (service, conf_file): {
+            limited_instance: {
+                "mem": 1,
+                "disk": 700,
+                "cpus": 3,
+            },
+            noop_instance: {
+                "cpus": 100,
+                "mem": 10000,
+                "disk": 2048,
+            },
+        }
+    }
+
+    with mock.patch.object(
+        updater,
+        "get_existing_configs",
+        autospec=True,
+        side_effect=[autotune_data, user_data],
+    ):
+        assert updater.merge_recommendations(recs) == {
+            (service, conf_file): {
+                limited_instance: {
+                    "mem": 1024,  # use lower bound
+                    "disk": 700,  # unchanged
+                    "cpus": 2,  # use upper bound
+                    "cpu_burst_add": 0.5,  # no updated rec or limit, leave alone
+                },
+                # this instances recommendations should be left untouched
+                noop_instance: {
+                    "cpus": 100,
+                    "mem": 10000,
+                    "disk": 2048,
+                },
+            }
+        }