Merge pull request #3830 from Yelp/add_autotune_limits_to_service_sha…

…rd_update

Add support for autotune limits in service_shard_update

paasta_tools/contrib/service_shard_update.py

@@ -39,6 +39,13 @@ def parse_args():
         action="store_true",
         dest="verbose",
     )
+    parser.add_argument(
+        "-d",
+        "--dry-run",
+        help="Do not commit changes to git",
+        action="store_true",
+        dest="dry_run",
+    )
     parser.add_argument(
         "--source-id",
         help="String to attribute the changes in the commit message.",
@@ -115,6 +122,48 @@ def parse_args():
         type=int,
         dest="timeout_server_ms",
     )
+    parser.add_argument(
+        "--autotune-min-cpus",
+        help="Minimum number of CPUs Autotune should give the shard",
+        required=False,
+        type=float,
+        dest="autotune_min_cpus",
+    )
+    parser.add_argument(
+        "--autotune-max-cpus",
+        help="Maximum number of CPUs Autotune should give the shard",
+        required=False,
+        type=float,
+        dest="autotune_max_cpus",
+    )
+    parser.add_argument(
+        "--autotune-min-mem",
+        help="Minimum amount of memory Autotune should give the shard",
+        required=False,
+        type=int,
+        dest="autotune_min_mem",
+    )
+    parser.add_argument(
+        "--autotune-max-mem",
+        help="Maximum amount of memory Autotune should give the shard",
+        required=False,
+        type=int,
+        dest="autotune_max_mem",
+    )
+    parser.add_argument(
+        "--autotune-min-disk",
+        help="Minimum amount of disk Autotune should give the shard",
+        required=False,
+        type=int,
+        dest="autotune_min_disk",
+    )
+    parser.add_argument(
+        "--autotune-max-disk",
+        help="Maximum amount of disk Autotune should give the shard",
+        required=False,
+        type=int,
+        dest="autotune_max_disk",
+    )
     return parser.parse_args()
 
 
@@ -195,9 +244,11 @@ def main(args):
                     instance_config = {
                         "deploy_group": f"{deploy_prefix}.{args.shard_name}",
                         "min_instances": args.min_instance_count,
-                        "max_instances": args.prod_max_instance_count
-                        if deploy_prefix == "prod"
-                        else args.non_prod_max_instance_count,
+                        "max_instances": (
+                            args.prod_max_instance_count
+                            if deploy_prefix == "prod"
+                            else args.non_prod_max_instance_count
+                        ),
                         "env": {
                             "PAASTA_SECRET_BUGSNAG_API_KEY": "SECRET(bugsnag_api_key)",
                         },
@@ -217,6 +268,40 @@ def main(args):
                         instance_config["cpus"] = args.cpus
                     if args.mem is not None:
                         instance_config["mem"] = args.mem
+                    if any(
+                        (
+                            args.autotune_min_cpus,
+                            args.autotune_max_cpus,
+                            args.autotune_min_mem,
+                            args.autotune_max_mem,
+                            args.autotune_min_disk,
+                            args.autotune_max_disk,
+                        )
+                    ):
+                        limit_config = {}
+                        limit_config["cpus"] = {
+                            "min": args.autotune_min_cpus,
+                            "max": args.autotune_max_cpus,
+                        }
+                        limit_config["mem"] = {
+                            "min": args.autotune_min_mem,
+                            "max": args.autotune_max_mem,
+                        }
+                        limit_config["disk"] = {
+                            "min": args.autotune_min_disk,
+                            "max": args.autotune_max_disk,
+                        }
+
+                        # remove any None values to keep the config clean
+                        for resource in list(limit_config):
+                            for key in list(limit_config[resource]):
+                                if limit_config[resource][key] is None:
+                                    del limit_config[resource][key]
+                            if len(limit_config[resource]) == 0:
+                                del limit_config[resource]
+
+                        if len(limit_config) > 0:
+                            instance_config["autotune_limits"] = limit_config
                     # If the service config does not contain definitions for the shard in each ecosystem
                     # Add the missing definition and write to the corresponding config
                     if args.shard_name not in config_file.keys():
@@ -247,7 +332,7 @@ def main(args):
             log.info(f"{args.shard_name} is in smartstack config already, skipping.")
 
         # Only commit to remote if changes were made
-        if changes_made:
+        if changes_made and not args.dry_run:
             updater.commit_to_remote()
             trigger_deploys(args.service)
         else: