elements: Add terminate prototype and reboot nodes

This adds functions to terminate the prototype gracefully or
forcefully, and to reboot all nodes in the experiment. This is useful
when running the experiment interactively, so swapping out and back in
is not always required.

1 files changed, 14 insertions, 3 deletions

diff --git a/rumba/elements/experimentation.py b/rumba/elements/experimentation.py
index 1c27dff..d0910a6 100644
--- a/rumba/elements/experimentation.py
+++ b/rumba/elements/experimentation.py
@@ -503,7 +503,7 @@ class Experiment(object):
         raise Exception('prototype_name() method not implemented')
 
     @abc.abstractmethod
-    def _terminate_prototype(self):
+    def _terminate_prototype(self, force=False):
         raise Exception('terminate_prototype() method not implemented')
 
     def swap_in(self):
@@ -521,8 +521,6 @@ class Experiment(object):
         Swap the experiment out of the testbed.
         """
         start = time.time()
-        # Terminate prototype gracefully
-        self._terminate_prototype()
         for node in self.nodes:
             if node.ssh_config.client is not None:
                 node.ssh_config.client.close()
@@ -533,6 +531,19 @@ class Experiment(object):
         end = time.time()
         logger.info("Swap-out took %.2f seconds", end - start)
 
+    def terminate_prototype(self, force=False):
+        """
+        Terminate the prototype in the experiment.
+        """
+        self._terminate_prototype()
+
+    def reboot_nodes(self):
+        """
+        Reboot all nodes in the experiment.
+        """
+        for node in self.nodes:
+            node.execute_command('reboot', as_root=True)
+
     @abc.abstractmethod
     def export_dif_bandwidth(self, filename, dif):
         raise Exception('Export DIF bandwidth method not implemented')