better cmd args and interactive user confirmation

scale/scale.py

@@ -5,6 +5,7 @@
 from update_nodes import update_unschedulable
 from cluster_update import gce_cluster_control
 from settings import settings
+from utils import confirm
 
 import logging
 import argparse
@@ -16,7 +17,7 @@
 scale_logger = logging.getLogger("scale")
 
 
-def shutdown_empty_nodes(nodes, k8s, cluster):
+def shutdown_empty_nodes(nodes, k8s, cluster, test=False):
     """
     Search through all nodes and shut down those that are unschedulable
     and devoid of non-critical pods
@@ -25,31 +26,44 @@ def shutdown_empty_nodes(nodes, k8s, cluster):
     """
     for node in nodes:
         if k8s.get_pods_number_on_node(node) == 0 and node.spec.unschedulable:
-            scale_logger.info(
-                "Shutting down empty node: %s", node.metadata.name)
-            cluster.shutdown_specified_node(node.metadata.name)
+            if confirm(("Shutting down empty node: %s" % node.metadata.name)):
+                scale_logger.info(
+                    "Shutting down empty node: %s", node.metadata.name)
+                if not test:
+                    cluster.shutdown_specified_node(node.metadata.name)
 
 
-def resize_for_new_nodes(new_total_nodes, k8s, cluster):
+def shutdown_empty_nodes_test(nodes, k8s, cluster):
+    shutdown_empty_nodes(nodes, k8s, cluster, True)
+
+
+def resize_for_new_nodes(new_total_nodes, k8s, cluster, test=False):
     """create new nodes to match new_total_nodes required
     only for scaling up
 
     TODO: Add gcloud python client for sizing up
     instead of CLI call"""
-    scale_logger.info("Resizing up to: %d nodes", new_total_nodes)
-    cluster.add_new_node(
-        new_total_nodes)
+    if confirm(("Resizing up to: %d nodes" % new_total_nodes)):
+        scale_logger.info("Resizing up to: %d nodes", new_total_nodes)
+        if not test:
+            cluster.add_new_node(new_total_nodes)
+
+
+def resize_for_new_nodes_test(new_total_nodes, k8s, cluster):
+    resize_for_new_nodes(new_total_nodes, k8s, cluster, True)
 
 
-def scale(options, context, test=0):
+def scale(options, context):
     """Update the nodes property based on scaling policy
     and create new nodes if necessary"""
+
+    # ONLY GCE is supported for scaling at this time
     cluster = gce_cluster_control(options)
-    if test == 1:
+    if options.test_k8s:
         k8s = k8s_control_test(options, context)
     else:
         k8s = k8s_control(options, context)
-        # ONLY GCE is supported for scaling at this time
+
     scale_logger.info("Scaling on cluster %s", k8s.get_cluster_name())
 
     nodes = []  # a list of nodes that are NOT critical
@@ -61,16 +75,21 @@ def scale(options, context, test=0):
     scale_logger.info("Total nodes in the cluster: %i", len(k8s.nodes))
     scale_logger.info("Found %i critical nodes; recommending additional %i nodes for service",
                       len(k8s.nodes) - len(nodes), goal)
-
-    update_unschedulable(max(len(nodes) - goal, 0), nodes, k8s)
+    if confirm(("Updating unschedulable flags to ensure %i nodes are unschedulable", max(len(nodes) - goal, 0))):
+        update_unschedulable(max(len(nodes) - goal, 0), nodes, k8s)
 
     if len(k8s.critical_node_names) + goal > len(k8s.nodes):
         scale_logger.info("Resize the cluster to %i nodes to satisfy the demand", (
             len(k8s.critical_node_names) + goal))
-        if test == 0:
+        if options.test_cloud:
+            resize_for_new_nodes_test(
+                len(k8s.critical_node_names) + goal, k8s, cluster)
+        else:
             resize_for_new_nodes(
                 len(k8s.critical_node_names) + goal, k8s, cluster)
-    if test == 0:
+    if options.test_cloud:
+        shutdown_empty_nodes_test(nodes, k8s, cluster)
+    else:
         # CRITICAL NODES SHOULD NOT BE SHUTDOWN
         shutdown_empty_nodes(nodes, k8s, cluster)
 
@@ -79,9 +98,13 @@ def scale(options, context, test=0):
     parser.add_argument(
         "-v", "--verbose", help="Show verbose output (debug)", action="store_true")
     parser.add_argument(
-        "--test", help="Run the script in test mode, no real action", action="store_true")
+        "-T", "--test", help="Run the script in TEST mode, log expected behavior, no real action will be taken", action="store_true")
+    parser.add_argument(
+        "--test-k8s", help="Run the script to test kubernetes actions: log expected commands to kubernetes, no real action on node specs", action="store_true")
     parser.add_argument(
-        "--testcloud", help="Run the script to test cloud apis, no real action on actual nodes", action="store_true")
+        "--test-cloud", help="Run the script to test cloud actions: log expected commands to the cloud provider, no real action on actual VM pool", action="store_true")
+    parser.add_argument(
+        "-y", help="Run the script without user interactive confirmation", action="store_true")
     parser.add_argument(
         "-c", "--context", required=True, help="A unique segment in the context name to specify which to use to instantiate Kubernetes")
     args = parser.parse_args()
@@ -90,15 +113,22 @@ def scale(options, context, test=0):
     else:
         scale_logger.setLevel(logging.INFO)
 
-    t = 0
+    # Retrieve settings from the environment
+    options = settings()
+
     if args.test:
-        t = 1
         scale_logger.warning(
             "Running in test mode, no action will actually be taken")
-    if args.testcloud:
-        t = 2
+    elif args.test_cloud:
+        options.test_k8s = False
+        scale_logger.warning(
+            "Running in test cloud mode, no action on VM pool")
+    elif args.test_k8s:
+        options.test_cloud = False
         scale_logger.warning(
-            "Running in test cloud mode, not all action will actually be taken")
+            "Running in test kubernetes mode, no action on node specs")
 
-    options = settings()
-    scale(options, args.context, t)
+    if args.y:
+        options.yes = True
+
+    scale(options, args.context)

scale/settings.py

@@ -42,5 +42,9 @@ def __init__(self):
         self.omit_namespaces = os.environ.get(
             "OMIT_NAMESPACES", "kube-system").split(self.env_delimiter)
 
+        self.test_cloud = True
+        self.test_k8s = True
+        self.yes = False
+
         # only used for debugging
         self.default_context = os.environ.get("DEFAULT_CONTEXT", "dev")

scale/utils.py

@@ -49,3 +49,43 @@ def check_list_intersection(list1, list2):
     if list1 is None or list2 is None:
         return False
     return len(set(list1).intersection(set(list2))) != 0
+
+
+def confirm(prompt=None, resp=False):
+    """prompts for yes or no response from the user. Returns True for yes and
+    False for no.
+
+    'resp' should be set to the default value assumed by the caller when
+    user simply types ENTER.
+
+    >>> confirm(prompt='Create Directory?', resp=True)
+    Create Directory? [y]|n: 
+    True
+    >>> confirm(prompt='Create Directory?', resp=False)
+    Create Directory? [n]|y: 
+    False
+    >>> confirm(prompt='Create Directory?', resp=False)
+    Create Directory? [n]|y: y
+    True
+
+    """
+
+    if prompt is None:
+        prompt = 'Confirm'
+
+    if resp:
+        prompt = '%s [%s]|%s: ' % (prompt, 'y', 'n')
+    else:
+        prompt = '%s [%s]|%s: ' % (prompt, 'n', 'y')
+
+    while True:
+        ans = input(prompt)
+        if not ans:
+            return resp
+        if ans not in ['y', 'Y', 'n', 'N']:
+            print('please enter y or n.')
+            continue
+        if ans == 'y' or ans == 'Y':
+            return True
+        if ans == 'n' or ans == 'N':
+            return False