diff --git a/paasta_tools/contrib/service_shard_update.py b/paasta_tools/contrib/service_shard_update.py index 9ec91c9858..a48031c68d 100644 --- a/paasta_tools/contrib/service_shard_update.py +++ b/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: