Add parameters to shift_nodepool_capacity script

experiment/maintenance/shift_nodepool_capacity.py

@@ -20,22 +20,17 @@
 # shift_nodepool_capacity.py pool-to-drain pool-to-grow shrink_increment:grow_increment num_to_add
 #
 # EG:
-# shift_nodepool_capacity.py default-pool pool-n1-highmem-8-300gb 2:1 5
+# shift_nodepool_capacity.py 5 2:1 default-pool pool-n1-highmem-8-300gb
 #
 # for nodefs on the prow builds cluster
 # USE AT YOUR OWN RISK.
 
+import argparse
 import sys
 import subprocess
 import json
 import math
 
-# xref prow/Makefile get-build-cluster-credentials
-# TODO: perhaps make these configurable
-CLUSTER = 'prow'
-ZONE = 'us-central1-f'
-PROJECT = 'k8s-prow-builds'
-
 
 def get_pool_sizes(project, zone, cluster):
     """returns a map of node pool name to size using the gcloud cli."""
@@ -92,18 +87,39 @@ def prompt_confirmation():
     print('Confirmed.')
 
 
-def main():
+# xref prow/Makefile get-build-cluster-credentials
+def parse_args(args):
+    parser = argparse.ArgumentParser()
+    parser.add_argument('nodes', type=int,
+                        help='Number of Nodes to add.')
+    parser.add_argument('ratio', type=str,
+                        help='ShrinkIncrement:GrowIncrement, Ex 2:1.')
+    parser.add_argument('shrink', type=str,
+                        help='Pool name to drain nodes from.')
+    parser.add_argument('grow', type=str,
+                        help='Pool name to grow nodes into.')
+    parser.add_argument('--cluster', type=str, default="prow",
+                        help='Name of GCP cluster.')
+    parser.add_argument('--zone', type=str, default='us-central1-f',
+                        help='GCP zonal location of the cluster.')
+    parser.add_argument('--project', type=str, default='k8s-prow-builds',
+                        help='GCP Project that the cluster exists within.')
+    return parser.parse_args(args)
+
+
+def main(options):
     # parse cli
-    nodes_to_add = int(sys.argv[-1])
+    nodes_to_add = options.nodes
 
-    ratio = sys.argv[-2].split(':')
+    ratio = options.ratio.split(':')
     shrink_increment, grow_increment = int(ratio[0]), int(ratio[1])
 
-    pool_to_grow = sys.argv[-3]
-    pool_to_shrink = sys.argv[-4]
+    pool_to_grow = options.grow
+    pool_to_shrink = options.shrink
 
     # obtain current pool sizes
-    pool_sizes = get_pool_sizes(PROJECT, ZONE, CLUSTER)
+    project, zone, cluster = options.project, options.zone, options.cluster
+    pool_sizes = get_pool_sizes(project, zone, cluster)
     pool_to_grow_initial = pool_sizes[pool_to_grow]
     pool_to_shrink_initial = pool_sizes[pool_to_shrink]
 
@@ -120,7 +136,7 @@ def main():
         'Shifting NodePool capacity for project = "{project}",'
         'zone = "{zone}", cluster = "{cluster}"'
         ).format(
-            project=PROJECT, zone=ZONE, cluster=CLUSTER,
+            project=project, zone=zone, cluster=cluster,
         ))
     print('')
     print((
@@ -163,7 +179,7 @@ def main():
             shrink_increment=shrink_increment, pool_to_shrink=pool_to_shrink,
         ))
         new_size = max(pool_to_shrink_initial - (i*shrink_increment + shrink_increment), 0)
-        resize_nodepool(pool_to_shrink, new_size, PROJECT, ZONE, CLUSTER)
+        resize_nodepool(pool_to_shrink, new_size, project, zone, cluster)
         print('')
 
         # ditto for growing, modulo the cap
@@ -172,11 +188,11 @@ def main():
             num_to_add=num_to_add, pool_to_grow=pool_to_grow,
         ))
         new_size = pool_to_grow_initial + (i*grow_increment + num_to_add)
-        resize_nodepool(pool_to_grow, new_size, PROJECT, ZONE, CLUSTER)
+        resize_nodepool(pool_to_grow, new_size, project, zone, cluster)
         print('')
 
     print('')
     print('Done')
 
 if __name__ == '__main__':
-    main()
+    main(parse_args(sys.argv[1:]))