diff options
-rwxr-xr-x | examples/script-example.py | 6 | ||||
-rw-r--r-- | rumba/storyboard.py | 205 | ||||
-rw-r--r-- | tools/scriptgenerator.py | 2 |
3 files changed, 151 insertions, 62 deletions
diff --git a/examples/script-example.py b/examples/script-example.py index 316a1d1..5e14823 100755 --- a/examples/script-example.py +++ b/examples/script-example.py @@ -93,7 +93,11 @@ if __name__ == '__main__': client_a.add_node(node_b) client_b.add_node(node_b) client_c.add_node(node_b) - story.parse_script_file('example-script.rsb') + script_file = os.path.join( + os.path.dirname(__file__), + 'example-script.rsb' + ) + story.parse_script_file(script_file) log.flush_log() with ExperimentManager(exp): exp.swap_in() diff --git a/rumba/storyboard.py b/rumba/storyboard.py index b0ae320..b98cc6f 100644 --- a/rumba/storyboard.py +++ b/rumba/storyboard.py @@ -338,9 +338,11 @@ class StoryBoard(SBEntity): self.active_clients = self.process_dict.values() self.start_time = None self.commands_list = {} - self.script = script self.node_map = {} self._build_nodes_lists() + # The following must be last, because it needs the info from + # _build_nodes_list + self._script = script if script is not None else _Script(self) def _build_nodes_lists(self): """Populates server_nodes and client_nodes lists""" @@ -373,6 +375,11 @@ class StoryBoard(SBEntity): self._build_nodes_lists() def set_experiment(self, experiment): + """ + Set the storyboard's underlying experiment instance + + @param experiment: the experiment instance + """ if not isinstance(experiment, model.Experiment): raise TypeError('Experiment instance required.') self.experiment = experiment @@ -391,6 +398,9 @@ class StoryBoard(SBEntity): """ Utility method to simultaneously add a server to a sb and a node to the server. + + @param server: the server to be added to the storyboard + @param node: the node upon which the server should run """ if self.experiment is None: raise ValueError("Cannot add a server before " @@ -416,8 +426,8 @@ class StoryBoard(SBEntity): """ if self.experiment is None: raise ValueError("An experiment is needed to schedule commands.") - if self.script is None: - self.script = Script(self) + if self._script is None: + self._script = _Script(self) if isinstance(node, str): node = self.node_map[node] if node not in self.experiment.nodes: @@ -426,7 +436,7 @@ class StoryBoard(SBEntity): if isinstance(command, str): command = [command] action = functools.partial(self.run_command, node, command) - self.script.add_event(Event(action, ev_time=t)) + self._script.add_event(Event(action, ev_time=t)) def run_command(self, node, command): """ @@ -446,6 +456,25 @@ class StoryBoard(SBEntity): command = [command] node.execute_commands(command) + def add_event(self, event): + """ + Add an event to this script, provided either as an Event + instance or as a string as read from a .rsb script. + + :param event: the event to add + :type event: (Event or str) + """ + self._script.add_event(event) + + def del_event(self, event): + """ + Remove an event from this storyboard + + :param event: the event (or id thereof) to remove + :type event: (Event or str) + """ + self._script.del_event(event) + def run_client_of(self, server, duration=None, node=None, proc_id=None): """ Runs a random client of the specified server @@ -492,7 +521,7 @@ class StoryBoard(SBEntity): self.process_dict[process.id] = process process.run() action = functools.partial(self.kill_process, process.id) - self.script.add_event(Event(action, ev_time=(self.cur_time + duration))) + self._script.add_event(Event(action, ev_time=(self.cur_time + duration))) def start_client_of(self, server, duration=None, node=None, proc_id=None): """ @@ -503,10 +532,10 @@ class StoryBoard(SBEntity): is not specified, it will be randomly generated according to the server parameters (client apps and their nodes). - If the client app must be shutdown manually or the duration - parameter is None, the client process will _not_ be stopped - automatically, and will continue running unless otherwise - killed. + Note that this method, as opposed to + :meth:`rumba.storyboard.run_client_of`, will not generate an event + to stop the client after the duration is expired. In most cases, + :meth:`rumba.storyboard.run_client_of` is the way to go. @param server: the server of which one client should be run @param duration: the duration of the client process @@ -525,10 +554,10 @@ class StoryBoard(SBEntity): If the node parameter is not given, it will be chosen at random among the client default nodes. - If the app must be shutdown manually or the duration - parameter is None, it will _not_ be stopped - automatically, and will continue running unless otherwise - killed. + Note that this method, as opposed to + :meth:`rumba.storyboard.run_client`, will not generate an event + to stop the client after the duration is expired. In most cases, + :meth:`rumba.storyboard.run_client` is the way to go. @param client: the client which should be run @param duration: the duration of the client process @@ -551,21 +580,21 @@ class StoryBoard(SBEntity): del self.process_dict[proc_id] def periodic_check(self, t): - self.script.check_for_ready_ev(t) - self.script.run_ready() - - def parse_script(self, buffer): - self.script = Script(self) - self.script.parse(buffer) + self._script.check_for_ready_ev(t) + self._script.run_ready() - def parse_script_file(self, filename): - self.script = Script(self) - self.script.parse_file(filename) + def generate_script(self, clean=True): + """ + Randomly generate a script for this experiment based on the + parameters provided to the instances of servers, nodes and clients. - def generate_script(self): + @param clean: if True, discard the current script before + generating a new one. + """ if self.experiment is None: raise ValueError('Cannot generate script without an experiment') - self.script = Script(self) + if clean: + self._script = _Script(self) t = self.SCRIPT_RESOLUTION marker = 5 last_marker = 0 @@ -585,7 +614,7 @@ class StoryBoard(SBEntity): p, c ) - self.script.add_event(start) + self._script.add_event(start) t += self.SCRIPT_RESOLUTION def _make_process_events(self, t, d, n, p, c): @@ -602,7 +631,7 @@ class StoryBoard(SBEntity): def start(self): if self.experiment is None: raise ValueError("Cannot run sb with no experiment.") - if self.script is None: + if self._script is None: self.generate_script() logger.info('Starting storyboard execution') self._build_nodes_lists() @@ -668,12 +697,12 @@ class StoryBoard(SBEntity): if not isinstance(dif, model.ShimEthDIF): raise Exception("Not a Shim Ethernet DIF.") - if self.script is None: - self.script = Script(self) + if self._script is None: + self._script = _Script(self) for ipcp in dif.ipcps: action = functools.partial(ipcp.node.set_link_state, ipcp, state) - self.script.add_event(Event(action, ev_time=t)) + self._script.add_event(Event(action, ev_time=t)) def set_link_up(self, t, dif): self.set_link_state(t, dif, 'up') @@ -685,14 +714,14 @@ class StoryBoard(SBEntity): if self.experiment is None: raise ValueError("An experiment is needed to schedule commands.") - if self.script is None: - self.script = Script(self) + if self._script is None: + self._script = _Script(self) for ipcp in node.ipcps: if not isinstance(ipcp, model.ShimEthIPCP): continue action = functools.partial(ipcp.node.set_link_state, ipcp, state) - self.script.add_event(Event(action, ev_time=t)) + self._script.add_event(Event(action, ev_time=t)) def set_node_up(self, t, node): self.set_node_state(t, node, 'up') @@ -700,6 +729,72 @@ class StoryBoard(SBEntity): def set_node_down(self, t, node): self.set_node_state(t, node, 'down') + def write_script(self, buffer): + """ + Writes the script on a string buffer, at the current position + + @param buffer: a string buffer. + """ + self._script.write(buffer) + + def write_script_to_file(self, filename, clean=True): + """ + Writes the script to a file, overwriting content. + + @param filename: the name of the destination file + @param clean: if True, current file contents will be overwritten. + """ + mode = 'w' + if not clean: + mode += '+' + with open(filename, mode) as f: + self.write_script(f) + + def write_script_string(self): + """ + Writes the script into a string and returns it. + + @return: the script as a string. + """ + s = StringIO() + self.write_script(s) + return s + + def parse_script(self, buffer, clean=True): + """ + Reads a script from a buffer, at the current position. + + @param buffer: the buffer to read from. + @param clean: if True, discard the current script before reading. + """ + if clean: + self._script = _Script(self) + self._script.parse(buffer) + + def parse_script_file(self, filename, clean=True): + """ + Reads a script from a file. + + @param filename: the file to read from. + @param clean: if True, discard the current script before reading. + """ + if clean: + self._script = _Script(self) + with open(filename, 'r') as f: + self.parse_script(f, clean) + + def parse_script_string(self, string, clean=True): + """ + Reads a script from a string. + + @param string: the string to read from. + @param clean: if True, discard the current script before reading. + """ + if clean: + self._script = _Script(self) + buffer = StringIO(string) + self.parse_script(buffer, clean) + class Event(object): @@ -815,12 +910,11 @@ class Event(object): return self._repr -class Script(object): +class _Script(object): def __init__(self, storyboard): - # Brute force 'dump all in memory' approach to avoid - # iterating through the events a lot of times - # at runtime + if storyboard is None: + raise ValueError("storyboard must not be None") self.events_by_id = {} self._waiting_events = {} self._events_ready = [] @@ -828,17 +922,24 @@ class Script(object): self._nodes = {} self._servers = {} self._clients = {} - self._testbed = None - self._experiment = None - self._storyboard = None + self._storyboard = storyboard self._entities = {} - self._parse_entities(storyboard) + self._parse_entities() - def _parse_entities(self, storyboard): - self._storyboard = storyboard - self._experiment = self._storyboard.experiment + @property + def _experiment(self): + return self._storyboard.experiment + + @property + def _testbed(self): + exp = self._experiment + if exp is None: + return None + else: + return exp.testbed + + def _parse_entities(self): self._nodes = self._storyboard.node_map - self._testbed = self._experiment.testbed self._servers = self._storyboard.server_apps self._clients = self._storyboard.client_apps self._processes = {} @@ -1009,28 +1110,12 @@ class Script(object): except ValueError as e: raise ValueError(str(e) + ' -> @ line %s' % (index,)) - def parse_file(self, filename): - with open(filename, 'r') as f: - self.parse(f) - - def parse_string(self, s): - buffer = StringIO(s) - self.parse(buffer) - def write(self, buffer): ev_list = list(self.events_by_id.values()) ev_list.sort(key=lambda x: x.time if x.time is not None else float('+inf')) for event in ev_list: buffer.write(repr(event) + '\n') - def write_to_file(self, filename): - with open(filename, 'w') as f: - self.write(f) - - def write_string(self): - s = StringIO() - self.write(s) - return s def _parse_conditions(self, *conditions): """ diff --git a/tools/scriptgenerator.py b/tools/scriptgenerator.py index 059163e..cc3e1ea 100644 --- a/tools/scriptgenerator.py +++ b/tools/scriptgenerator.py @@ -45,7 +45,7 @@ def main(duration, exp, run=False, script='generated_script.txt'): f.write('################################################\n') f.write('# SCRIPT GENERATED WITH RUMBA SCRIPT GENERATOR #\n') f.write('################################################\n') - story.script.write(f) + story.write_script(f) if run: with ExperimentManager(exp, swap_out_strategy=PAUSE_SWAPOUT): |