cxmanage: Allow --retry option to work with any command

2 files changed, 118 insertions, 89 deletions

diff --git a/cxmanage/controller.py b/cxmanage/controller.py
index f73e2c9..6d2e452 100644
--- a/cxmanage/controller.py
+++ b/cxmanage/controller.py
@@ -56,6 +56,8 @@ class Controller:
 
     def __init__(self, verbosity=0, max_threads=1, image_class=Image,
             target_class=Target, retries=None):
+        if retries == None:
+            retries = "prompt"
         self.tftp = None
         self.targets = []
         self.images = []
@@ -193,7 +195,11 @@ class Controller:
         # Get IP info
         if self.verbosity >= 1:
             print "Getting IP addresses..."
-        results, errors = self._run_command(targets, "get_ipinfo", self.tftp)
+        retries = self.retries
+        if retries == "prompt":
+            retries = 0
+        results, errors = self._run_command_on_targets(targets, retries,
+                "get_ipinfo", self.tftp)
 
         # Add all resulting targets
         for target in targets:
@@ -207,7 +213,6 @@ class Controller:
             for target in self.targets:
                 print target.address
             print
-        self._print_errors(targets, errors)
 
         return len(errors) > 0
 
@@ -217,14 +222,20 @@ class Controller:
         """ Send the given power command to all targets """
         if self.verbosity >= 1:
             print "Sending power %s command..." % mode
-        return self._retry_command("set_power", mode)
+
+        results, errors = self._run_command(True, "set_power", mode)
+
+        if self.verbosity >= 1 and len(errors) == 0:
+            print "Command completed successfully.\n"
+        return len(errors) > 0
+
 
     def power_status(self):
         """ Retrieve power status from all targets in group """
 
         if self.verbosity >= 1:
             print "Getting power status..."
-        results, errors = self._run_command(self.targets, "get_power")
+        results, errors = self._run_command(False, "get_power")
 
         # Print results
         if len(results) > 0:
@@ -238,8 +249,8 @@ class Controller:
                     print "%s: %s" % (target.address.ljust(16), result)
             print
 
-        # Print errors
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -247,13 +258,18 @@ class Controller:
         """ Set the power policy for all targets """
         if self.verbosity >= 1:
             print "Setting power policy to %s..." % mode
-        return self._retry_command("set_power_policy", mode)
+
+        results, errors = self._run_command(True, "set_power_policy", mode)
+
+        if self.verbosity >= 1 and len(errors) == 0:
+            print "Command completed successfully.\n"
+        return len(errors) > 0
 
     def power_policy_status(self):
         """ Get power policy status for all targets """
         if self.verbosity >= 1:
             print "Getting power policy status..."
-        results, errors = self._run_command(self.targets, "get_power_policy")
+        results, errors = self._run_command(False, "get_power_policy")
 
         # Print results
         if len(results) > 0:
@@ -264,8 +280,8 @@ class Controller:
                             results[target.address])
             print
 
-        # Print errors
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -273,20 +289,30 @@ class Controller:
         """ Send an MC reset command to all targets """
         if self.verbosity >= 1:
             print "Sending MC reset command..."
-        return self._retry_command("mc_reset")
+
+        results, errors = self._run_command(True, "mc_reset")
+
+        if self.verbosity >= 1 and len(errors) == 0:
+            print "Command completed successfully.\n"
+        return len(errors) > 0
 
     def update_firmware(self, partition_arg="INACTIVE"):
         """ Send firmware update commands to all targets """
         if self.verbosity >= 1:
             print "Updating firmware..."
-        return self._retry_command("update_firmware", self.tftp, self.images,
-                partition_arg)
+
+        results, errors = self._run_command(True, "update_firmware", self.tftp,
+                self.images, partition_arg)
+
+        if self.verbosity >= 1 and len(errors) == 0:
+            print "Command completed successfully.\n"
+        return len(errors) > 0
 
     def firmware_info(self):
         """ Print firmware info for all targets """
         if self.verbosity >= 1:
             print "Getting firmware info..."
-        results, errors = self._run_command(self.targets, "get_firmware_info")
+        results, errors = self._run_command(False, "get_firmware_info")
 
         for target in self.targets:
             if target.address in results:
@@ -303,7 +329,8 @@ class Controller:
                     print "In Use             : %s" % partition.in_use
                     print
 
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -311,7 +338,7 @@ class Controller:
         """ Get sensor readings from all targets """
         if self.verbosity >= 1:
             print "Getting sensor readings..."
-        results, errors = self._run_command(self.targets, "get_sensors", name)
+        results, errors = self._run_command(False, "get_sensors", name)
 
         if len(results) > 0:
             # Get sensor names
@@ -361,7 +388,8 @@ class Controller:
 
                 print
 
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -369,8 +397,7 @@ class Controller:
         """ Get IP info from all targets """
         if self.verbosity >= 1:
             print "Getting IP info..."
-        results, errors = self._run_command(self.targets, "get_ipinfo",
-                self.tftp)
+        results, errors = self._run_command(False, "get_ipinfo", self.tftp)
 
         # Print results
         if len(results) > 0:
@@ -382,7 +409,8 @@ class Controller:
                         print "Node %i: %s" % entry
                     print
 
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -390,21 +418,20 @@ class Controller:
         """ Get mac addresses from all targets """
         if self.verbosity >= 1:
             print "Getting MAC addresses..."
-        results, errors = self._run_command(self.targets, "get_macaddrs",
-                self.tftp)
+        results, errors = self._run_command(False, "get_macaddrs", self.tftp)
 
         # Print results
         if len(results) > 0:
             for target in self.targets:
                 if target.address in results:
-                    print "Mac addresses from %s" % target.address
+                    print "MAC addresses from %s" % target.address
                     macaddrs = results[target.address]
                     for entry in macaddrs:
                         print "Node %i, Port %i: %s" % entry
-                    if target != self.targets[-1] or len(errors) > 0:
-                        print
+                    print
 
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -412,7 +439,12 @@ class Controller:
         """ Send config reset command to all targets """
         if self.verbosity >= 1:
             print "Sending config reset command..."
-        return self._retry_command("config_reset", self.tftp)
+
+        results, errors = self._run_command(True, "config_reset", self.tftp)
+
+        if self.verbosity >= 1 and len(errors) == 0:
+            print "Command completed successfully.\n"
+        return len(errors) > 0
 
     def config_boot(self, boot_args):
         """ Send config boot command to all targets """
@@ -426,14 +458,19 @@ class Controller:
 
         if self.verbosity >= 1:
             print "Setting boot order..."
-        return self._retry_command("set_boot_order", self.tftp, boot_args)
+
+        results, errors = self._run_command(True, "set_boot_order", self.tftp,
+                boot_args)
+
+        if self.verbosity >= 1 and len(errors) == 0:
+            print "Command completed successfully.\n"
+        return len(errors) > 0
 
     def config_boot_status(self):
         """ Get boot order from all targets """
         if self.verbosity >= 1:
             print "Getting boot orders..."
-        results, errors = self._run_command(self.targets, "get_boot_order",
-                self.tftp)
+        results, errors = self._run_command(False, "get_boot_order", self.tftp)
 
         # Print results
         if len(results) > 0:
@@ -444,8 +481,8 @@ class Controller:
                             ",".join(results[target.address]))
             print
 
-        # Print errors
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -453,7 +490,7 @@ class Controller:
         """ Get basic SoC info from all targets """
         if self.verbosity >= 1:
             print "Getting SoC info..."
-        results, errors = self._run_command(self.targets, "info_basic")
+        results, errors = self._run_command(False, "info_basic")
 
         # Print results
         if len(results) > 0:
@@ -468,8 +505,8 @@ class Controller:
                             time.ctime(int(result.timestamp)))
                     print
 
-        # Print errors
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -477,8 +514,7 @@ class Controller:
         """ Print u-boot environment for all targets """
         if self.verbosity >= 1:
             print "Getting u-boot environments..."
-        results, errors = self._run_command(self.targets, "get_ubootenv",
-                self.tftp)
+        results, errors = self._run_command(False, "get_ubootenv", self.tftp)
 
         # Print results
         if len(results) > 0:
@@ -490,8 +526,8 @@ class Controller:
                         print "%s=%s" % (variable, ubootenv.variables[variable])
                     print
 
-        # Print errors
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -499,8 +535,7 @@ class Controller:
         """ Dump info from all targets """
         if self.verbosity >= 1:
             print "Getting all info..."
-        results, errors = self._run_command(self.targets, "info_dump",
-                self.tftp)
+        results, errors = self._run_command(False, "info_dump", self.tftp)
 
         # Print results
         if len(results) > 0:
@@ -510,8 +545,8 @@ class Controller:
                     print results[target.address]
                     print
 
-        # Print errors
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
@@ -519,7 +554,7 @@ class Controller:
         """ Run an arbitrary ipmitool command on all targets """
         if self.verbosity >= 1:
             print "Running IPMItool command..."
-        results, errors = self._run_command(self.targets, "ipmitool_command",
+        results, errors = self._run_command(True, "ipmitool_command",
                 ipmitool_args)
 
         # Print results
@@ -530,55 +565,20 @@ class Controller:
                     print results[target.address]
                     print
 
-        # Print errors
-        self._print_errors(self.targets, errors)
+        if self.verbosity >= 1 and len(errors) > 0:
+            print "Some errors occured during the command.\n"
 
         return len(errors) > 0
 
-    def _retry_command(self, name, *args):
-        """ Run a generic retrying command on all targets """
-        targets = self.targets
+    def _run_command(self, retry_prompt, name, *args):
+        """ Run a command on all targets """
         retries = self.retries
+        if retries == "prompt" and not retry_prompt:
+            retries = 0
+        return self._run_command_on_targets(self.targets, retries, name, *args)
 
-        while True:
-            # Get results and errors
-            results, errors = self._run_command(targets, name, *args)
-
-            if len(errors) == 0:
-                # No errors encountered, print message and return
-                if self.verbosity >= 1:
-                    print "Command completed successfully.\n"
-                return False
-            else:
-                # Print errors
-                self._print_errors(targets, errors)
-
-                # 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):
-        """ Run a target command with multiple threads
+    def _run_command_on_targets(self, targets, retries, name, *args):
+        """ Run a multi-threaded command on the specified targets.
 
         Returns (results, errors) which map addresses to their results """
 
@@ -627,6 +627,34 @@ class Controller:
             indicator.stop()
             print "\n"
 
+        # Handle errors
+        should_retry = False
+        if len(errors) > 0:
+            self._print_errors(targets, errors)
+            if retries == "prompt":
+                sys.stdout.write("Retry command on failed hosts? (y/n): ")
+                sys.stdout.flush()
+                while True:
+                    command = raw_input().strip().lower()
+                    if command in ['y', 'yes']:
+                        should_retry = True
+                        break
+                    elif command in ['n', 'no']:
+                        print
+                        break
+            elif retries >= 1:
+                should_retry = True
+                if retries == 1:
+                    print "Retrying command 1 more time..."
+                elif retries > 1:
+                    print "Retrying command %i more times..." % retries
+                retries -= 1
+        if should_retry:
+            targets = [x for x in targets if x.address in errors]
+            new_results, errors = self._run_command_on_targets(targets,
+                    retries, name, *args)
+            results.update(new_results)
+
         return results, errors
 
     def _print_errors(self, targets, errors):
diff --git a/scripts/cxmanage b/scripts/cxmanage
index b01ea7b..7bb4c0a 100755
--- a/scripts/cxmanage
+++ b/scripts/cxmanage
@@ -333,7 +333,8 @@ def add_targets(controller, args):
     # Add hosts to controller
     if args.all_nodes:
         if controller.add_fabrics(hosts, args.user, args.password):
-            print "ERROR: Failed to get IP addresses. Aborting.\n"
+            print "ERROR: Failed to get IP info from all hosts. Aborting."
+            print
             sys.exit(1)
     else:
         for host in hosts: