[SPARK-3540] Add reboot-slaves functionality to the ec2 script

Tested on a real cluster.

Author: Reynold Xin <rxin@apache.org>

Closes #2404 from rxin/ec2-reboot-slaves and squashes the following commits:

00a2dbd [Reynold Xin] Allow rebooting slaves.

1 files changed, 15 insertions, 1 deletions

diff --git a/ec2/spark_ec2.py b/ec2/spark_ec2.py
index bfd07593b9..5682e96aa8 100755
--- a/ec2/spark_ec2.py
+++ b/ec2/spark_ec2.py
@@ -52,7 +52,7 @@ class UsageError(Exception):
 def parse_args():
     parser = OptionParser(
         usage="spark-ec2 [options] <action> <cluster_name>"
-        + "\n\n<action> can be: launch, destroy, login, stop, start, get-master",
+        + "\n\n<action> can be: launch, destroy, login, stop, start, get-master, reboot-slaves",
         add_help_option=False)
     parser.add_option(
         "-h", "--help", action="help",
@@ -950,6 +950,20 @@ def real_main():
         subprocess.check_call(
             ssh_command(opts) + proxy_opt + ['-t', '-t', "%s@%s" % (opts.user, master)])
 
+    elif action == "reboot-slaves":
+        response = raw_input(
+            "Are you sure you want to reboot the cluster " +
+            cluster_name + " slaves?\n" +
+            "Reboot cluster slaves " + cluster_name + " (y/N): ")
+        if response == "y":
+            (master_nodes, slave_nodes) = get_existing_cluster(
+                conn, opts, cluster_name, die_on_error=False)
+            print "Rebooting slaves..."
+            for inst in slave_nodes:
+                if inst.state not in ["shutting-down", "terminated"]:
+                    print "Rebooting " + inst.id
+                    inst.reboot()
+
     elif action == "get-master":
         (master_nodes, slave_nodes) = get_existing_cluster(conn, opts, cluster_name)
         print master_nodes[0].public_dns_name