cxmanage: Add --retry COUNT option

This option will cause cxmanage to retry the command a fixed number
of times on failed hosts, instead of prompting. This will probably
be useful for very large clusters where failures are more likely to
occur.

For example, to try the config reset command 3 times (2 retries):

cxmanage --retry 2 config reset

2 files changed, 31 insertions, 13 deletions

diff --git a/cxmanage/controller.py b/cxmanage/controller.py
index 1f7b0f9..b47ab0b 100644
--- a/cxmanage/controller.py
+++ b/cxmanage/controller.py
@@ -54,13 +54,14 @@ class Controller:
     cxmanage. Scripts or UIs can build on top of this to provide an user
     interface. """
 
-    def __init__(self, verbosity=0, max_threads=1,
-            image_class=Image, target_class=Target):
+    def __init__(self, verbosity=0, max_threads=1, image_class=Image,
+            target_class=Target, retries=None):
         self.tftp = None
         self.targets = []
         self.images = []
         self.verbosity = verbosity
         self.max_threads = max_threads
+        self.retries = retries
         self.target_class = target_class
         self.image_class = image_class
         self.work_dir = tempfile.mkdtemp(prefix="cxmanage-")
@@ -532,6 +533,7 @@ class Controller:
     def _retry_command(self, name, *args):
         """ Run a generic retrying command on all targets """
         targets = self.targets
+        retries = self.retries
 
         while True:
             # Get results and errors
@@ -543,18 +545,31 @@ class Controller:
                     print "Command completed successfully.\n"
                 return False
             else:
-                # Print errors and retry prompt, update target list
+                # Print errors
                 self._print_errors(errors)
-                sys.stdout.write("Retry on failed hosts? (y/n): ")
-                sys.stdout.flush()
-                while True:
-                    command = raw_input().strip().lower()
-                    if command in ['y', 'yes']:
-                        print
-                        break
-                    elif command in ['n', 'no']:
-                        print
+
+                # Decide whether or not to retry
+                if retries == None:
+                    sys.stdout.write("Retry on failed hosts? (y/n): ")
+                    sys.stdout.flush()
+                    while True:
+                        command = raw_input().strip().lower()
+                        if command in ['y', 'yes']:
+                            print
+                            break
+                        elif command in ['n', 'no']:
+                            print
+                            return True
+                else:
+                    if retries == 0:
                         return True
+                    if retries == 1:
+                        print "Retrying command 1 more time...\n"
+                    else:
+                        print "Retrying command %i more times...\n" % retries
+                    retries -= 1
+
+                # Update target list
                 targets = [x for x in targets if x.address in errors]
 
     def _run_command(self, targets, name, *args):
diff --git a/scripts/cxmanage b/scripts/cxmanage
index d83ba01..8a9d7c9 100755
--- a/scripts/cxmanage
+++ b/scripts/cxmanage
@@ -62,6 +62,8 @@ def build_parser():
             metavar='THREAD_COUNT', help='Number of threads to use')
     p['parser'].add_argument('--force-yes', action='store_true',
             help='Don\'t stop at confirmation dialogue')
+    p['parser'].add_argument('--retry', help='Retry command on multiple times',
+            type=int, default=None, metavar='COUNT')
     p['verbosity'] = p['parser'].add_mutually_exclusive_group()
     p['verbosity'].add_argument('-v', '--verbose', action='store_true',
             help='Verbose output')
@@ -245,7 +247,8 @@ def main():
     elif args.quiet:
         verbosity = 0
 
-    controller = Controller(verbosity=verbosity, max_threads=args.threads)
+    controller = Controller(verbosity=verbosity, max_threads=args.threads,
+            retries=args.retry)
 
     set_tftp(controller, args)
     add_targets(controller, args)