elastic-kernel 0.0.2__py3-none-any.whl

elastic_notebook/core/notebook/restore_notebook.py

@@ -0,0 +1,91 @@
+import sys
+import time
+from io import StringIO
+
+from ipykernel.zmqshell import ZMQInteractiveShell
+from IPython import get_ipython
+
+from elastic_notebook.core.graph.graph import DependencyGraph
+
+
+def restore_notebook(
+    graph: DependencyGraph,
+    shell: ZMQInteractiveShell,
+    variables: dict,
+    ces_to_recompute: set,
+    write_log_location=None,
+    notebook_name=None,
+    optimizer_name=None,
+):
+    """
+    Restores the notebook. Declares variables back into the kernel and recomputes the CEs to restore non-migrated
+    variables.
+    Args:
+        graph (DependencyGraph): dependency graph representation of the notebook.
+        shell (ZMQInteractiveShell): interactive Jupyter shell storing the state of the current session.
+        variables (Dict): Mapping from OEs to lists of variables defined in those OEs.
+        oes_to_recompute (set): OEs to recompute to restore non-migrated variables.
+        write_log_location (str): location to write component runtimes to. For experimentation only.
+        notebook_name (str): notebook name. For experimentation only.
+        optimizer_name (str): optimizer name. For experimentation only.
+    """
+
+    # Recompute OEs following the order they were executed in.
+    recompute_start = time.time()
+    for ce in graph.cell_executions:
+        if ce in ces_to_recompute:
+            # Rerun cell code; suppress stdout when rerunning.
+            print("Rerunning cell", ce.cell_num + 1)
+
+            # 出力を完全に抑制するための設定
+            ipython = get_ipython()
+
+            # 元のdisplay_pubを保存
+            original_display_pub = ipython.display_pub
+
+            # display_pubを一時的に無効化
+            ipython.display_pub = None
+
+            # stdoutとstderrをリダイレクト
+            original_stdout = sys.stdout
+            original_stderr = sys.stderr
+            sys.stdout = StringIO()
+            sys.stderr = StringIO()
+
+            try:
+                # セルを実行（出力は抑制される）
+                cell_output = ipython.run_cell(ce.cell)
+                cell_output.raise_error()
+            except Exception as e:
+                # エラーが発生した場合は元のstderrに出力
+                sys.stderr = original_stderr
+                raise e
+            finally:
+                # 元の設定を復元
+                ipython.display_pub = original_display_pub
+                sys.stdout = original_stdout
+                sys.stderr = original_stderr
+
+        # Define output variables in the CE.
+        for pair in variables[ce.cell_num]:
+            print("Declaring variable", pair[0].name)
+            shell.user_ns[pair[0].name] = pair[1]
+
+    recompute_end = time.time()
+
+    if write_log_location:
+        with open(
+            write_log_location
+            + "/output_"
+            + notebook_name
+            + "_"
+            + optimizer_name
+            + ".txt",
+            "a",
+        ) as f:
+            f.write(
+                "Recompute stage took - "
+                + repr(recompute_end - recompute_start)
+                + " seconds"
+                + "\n"
+            )

elastic_notebook/core/notebook/update_graph.py

@@ -0,0 +1,46 @@
+from elastic_notebook.core.graph.graph import DependencyGraph
+
+
+def update_graph(
+    cell: str,
+    cell_runtime: float,
+    start_time: float,
+    input_variables: set,
+    created_and_modified_variables: set,
+    deleted_variables: set,
+    graph: DependencyGraph,
+):
+    """
+    Updates the graph according to the newly executed cell and its input and output variables.
+    Args:
+         cell (str): Raw cell cell.
+         cell_runtime (float): Cell runtime.
+         start_time (time): Time of start of cell execution. Note that this is different from when the cell was
+             queued.
+         input_variables (set): Set of input variables of the cell.
+         created_and_modified_variables (set): set of created and modified variables.
+         deleted_variables (set): set of deleted variables
+         graph (DependencyGraph): Dependency graph representation of the notebook.
+    """
+
+    # Retrieve input variable snapshots
+    input_vss = set(
+        graph.variable_snapshots[variable][-1] for variable in input_variables
+    )
+
+    # Create output variable snapshots
+    output_vss_create = set(
+        graph.create_variable_snapshot(k, False) for k in created_and_modified_variables
+    )
+    output_vss_delete = set(
+        graph.create_variable_snapshot(k, True) for k in deleted_variables
+    )
+
+    # Add the newly created OE to the graph.
+    graph.add_cell_execution(
+        cell,
+        cell_runtime,
+        start_time,
+        input_vss,
+        output_vss_create.union(output_vss_delete),
+    )

elastic_notebook/elastic_notebook.py

@@ -0,0 +1,336 @@
+from __future__ import print_function
+
+import time
+import types
+from os.path import dirname
+
+from IPython import get_ipython
+from IPython.core.interactiveshell import InteractiveShell
+from pympler import asizeof
+
+from elastic_notebook.algorithm.baseline import MigrateAllBaseline, RecomputeAllBaseline
+from elastic_notebook.algorithm.optimizer_exact import OptimizerExact
+from elastic_notebook.algorithm.selector import OptimizerType
+from elastic_notebook.core.common.profile_graph_size import profile_graph_size
+from elastic_notebook.core.common.profile_migration_speed import profile_migration_speed
+from elastic_notebook.core.graph.graph import DependencyGraph
+from elastic_notebook.core.io.recover import resume
+from elastic_notebook.core.mutation.fingerprint import (
+    compare_fingerprint,
+    construct_fingerprint,
+)
+from elastic_notebook.core.mutation.object_hash import UnserializableObj
+from elastic_notebook.core.notebook.checkpoint import checkpoint
+from elastic_notebook.core.notebook.find_input_vars import find_input_vars
+from elastic_notebook.core.notebook.find_output_vars import find_created_deleted_vars
+from elastic_notebook.core.notebook.restore_notebook import restore_notebook
+from elastic_notebook.core.notebook.update_graph import update_graph
+
+
+class ElasticNotebook:
+    """
+    Magics class for Elastic Notebook. Enable this in the notebook by running '%load_ext ElasticNotebook'.
+    Enables efficient checkpointing of intermediate notebook state via balancing migration and recomputation
+    costs.
+    """
+
+    def __init__(self, shell: InteractiveShell):
+        self.shell = shell
+
+        # Initialize the dependency graph for capturing notebook state.
+        self.dependency_graph = DependencyGraph()
+
+        # Migration properties.
+        self.migration_speed_bps = 100000
+        self.alpha = 1
+        self.selector = OptimizerExact(migration_speed_bps=self.migration_speed_bps)
+
+        # Dictionary of object fingerprints. For detecting modified references.
+        self.fingerprint_dict = {}
+
+        # Set of user-declared functions.
+        self.udfs = set()
+
+        # Flag if migration speed has been manually set. In this case, skip profiling of migration speed at checkpoint
+        # time.
+        self.manual_migration_speed = False
+
+        # Location to log runtimes to. For experiments only.
+        self.write_log_location = None
+
+        # Strings for determining log filename. For experiments only.
+        self.optimizer_name = ""
+        self.notebook_name = ""
+
+        # Total elapsed time spent inferring cell inputs and outputs.
+        # For measuring overhead.
+        self.total_recordevent_time = 0
+
+        # Dict for recording overhead of profiling operations.
+        self.profile_dict = {"idgraph": 0, "representation": 0}
+
+        # デバッグフラグ
+        self.debug = False
+
+        # マイグレーションと再計算の変数リスト
+        self._vss_to_migrate = []
+        self._vss_to_recompute = []
+
+    @property
+    def vss_to_migrate(self):
+        """マイグレーション対象の変数リストを取得"""
+        return self._vss_to_migrate
+
+    @property
+    def vss_to_recompute(self):
+        """再計算対象の変数リストを取得"""
+        return self._vss_to_recompute
+
+    def update_migration_lists(self, vss_to_migrate, vss_to_recompute):
+        """マイグレーションと再計算の変数リストを更新"""
+        self._vss_to_migrate = [vs.name for vs in vss_to_migrate]
+        self._vss_to_recompute = [vs.name for vs in vss_to_recompute]
+
+    def __str__(self):
+        """文字列表現を定義"""
+        return f"マイグレーション対象: {self.vss_to_migrate}\n再計算対象: {self.vss_to_recompute}"
+
+    def set_debug(self, debug):
+        self.debug = debug
+
+    def record_event(self, cell):
+        if self.debug:
+            print("Recording event...")
+
+        pre_execution = set(self.shell.user_ns.keys())
+
+        # Create id trees for output variables
+        for var in self.dependency_graph.variable_snapshots.keys():
+            if var not in self.fingerprint_dict and var in self.shell.user_ns:
+                self.fingerprint_dict[var] = construct_fingerprint(
+                    self.shell.user_ns[var], self.profile_dict
+                )
+
+        # Find input variables (variables potentially accessed) of the cell.
+        input_variables, function_defs = find_input_vars(
+            cell,
+            set(self.dependency_graph.variable_snapshots.keys()),
+            self.shell,
+            self.udfs,
+        )
+        # Union of ID graphs of input variables. For detecting modifications to unserializable variables.
+        input_variables_id_graph_union = set()
+        for var in input_variables:
+            if var in self.fingerprint_dict:
+                input_variables_id_graph_union = input_variables_id_graph_union.union(
+                    self.fingerprint_dict[var][1]
+                )
+
+        # Run the cell.
+        start_time = time.time()
+        try:
+            cell_output = get_ipython().run_cell(cell)
+            cell_output.raise_error()
+            # traceback_list = []
+        except Exception:
+            pass
+            # _, _, tb = sys.exc_info()
+            # traceback_list = traceback.extract_tb(tb).format()
+        cell_runtime = time.time() - start_time
+        post_execution = set(self.shell.user_ns.keys())
+        infer_start = time.time()
+
+        # Find created and deleted variables by computing difference between namespace pre and post execution.
+        created_variables, deleted_variables = find_created_deleted_vars(
+            pre_execution, post_execution
+        )
+
+        # Remove stored ID graphs for deleted variables.
+        for var in deleted_variables:
+            del self.fingerprint_dict[var]
+            if var in self.udfs:
+                self.udfs.remove(var)
+
+        # Find modified variables by comparing ID graphs and object hashes.
+        modified_variables = set()
+        for k, v in self.fingerprint_dict.items():
+            changed, overwritten = compare_fingerprint(
+                self.fingerprint_dict[k],
+                self.shell.user_ns[k],
+                self.profile_dict,
+                input_variables_id_graph_union,
+            )
+            if changed:
+                modified_variables.add(k)
+
+            # In the case of non-overwrite modification, the variable is additionally considered as accessed.
+            if changed and not overwritten:
+                input_variables.add(k)
+
+            # A user defined function has been overwritten.
+            elif overwritten and k in self.udfs:
+                self.udfs.remove(k)
+
+            # Select unserializable variables are assumed to be modified if accessed.
+            if (
+                not changed
+                and not overwritten
+                and isinstance(self.fingerprint_dict[k][2], UnserializableObj)
+            ):
+                if self.fingerprint_dict[k][1].intersection(
+                    input_variables_id_graph_union
+                ):
+                    modified_variables.add(k)
+
+        # Create ID graphs for output variables
+        for var in created_variables:
+            self.fingerprint_dict[var] = construct_fingerprint(
+                self.shell.user_ns[var], self.profile_dict
+            )
+
+        # Record newly defined UDFs
+        for udf in function_defs:
+            if udf in self.shell.user_ns and isinstance(
+                self.shell.user_ns[udf], types.FunctionType
+            ):
+                self.udfs.add(udf)
+
+        # Update the dependency graph.
+        update_graph(
+            cell,
+            cell_runtime,
+            start_time,
+            input_variables,
+            created_variables.union(modified_variables),
+            deleted_variables,
+            self.dependency_graph,
+        )
+
+        # Update total recordevent time tally.
+        infer_end = time.time()
+        self.total_recordevent_time += infer_end - infer_start
+
+    def set_migration_speed(self, migration_speed):
+        try:
+            if float(migration_speed) > 0:
+                self.migration_speed_bps = float(migration_speed)
+                self.manual_migration_speed = True
+            else:
+                pass
+                # print("Migration speed is not positive.")
+        except ValueError:
+            print("Migration speed is not a number.")
+        self.selector.migration_speed_bps = self.migration_speed_bps
+
+    def set_optimizer(self, optimizer):
+        self.optimizer_name = optimizer
+
+        if optimizer == OptimizerType.EXACT.value:
+            self.selector = OptimizerExact(self.migration_speed_bps)
+            self.alpha = 1
+        elif optimizer == OptimizerType.EXACT_C.value:
+            self.selector = OptimizerExact(self.migration_speed_bps)
+            self.alpha = 20
+        elif optimizer == OptimizerType.EXACT_R.value:
+            self.selector = OptimizerExact(self.migration_speed_bps)
+            self.alpha = 0.05
+        elif optimizer == OptimizerType.MIGRATE_ALL.value:
+            self.selector = MigrateAllBaseline(self.migration_speed_bps)
+        elif optimizer == OptimizerType.RECOMPUTE_ALL.value:
+            self.selector = RecomputeAllBaseline(self.migration_speed_bps)
+
+    def set_write_log_location(self, filename):
+        self.write_log_location = filename
+
+    def set_notebook_name(self, name):
+        self.notebook_name = name
+
+    def checkpoint(self, filename):
+        """チェックポイントを作成"""
+        if self.debug:
+            print("Checkpointing...")
+
+        # Write overhead metrics to file (for experiments).
+        if self.write_log_location:
+            with open(
+                self.write_log_location
+                + "/output_"
+                + self.notebook_name
+                + "_"
+                + self.optimizer_name
+                + ".txt",
+                "a",
+            ) as f:
+                f.write(
+                    "comparison overhead - "
+                    + repr(
+                        asizeof.asizeof(self.dependency_graph)
+                        + asizeof.asizeof(self.fingerprint_dict)
+                    )
+                    + " bytes"
+                    + "\n"
+                )
+                f.write(
+                    "notebook overhead - "
+                    + repr(asizeof.asizeof(self.shell.user_ns))
+                    + " bytes"
+                    + "\n"
+                )
+                f.write(
+                    "Dependency graph storage overhead - "
+                    + repr(profile_graph_size(self.dependency_graph))
+                    + " bytes"
+                    + "\n"
+                )
+                f.write(
+                    "Cell monitoring overhead - "
+                    + repr(self.total_recordevent_time)
+                    + " seconds"
+                    + "\n"
+                )
+
+        # Profile the migration speed to filename.
+        if not self.manual_migration_speed:
+            self.migration_speed_bps = profile_migration_speed(
+                dirname(filename), alpha=self.alpha
+            )
+            self.selector.migration_speed_bps = self.migration_speed_bps
+
+        # Checkpoint the notebook.
+        return checkpoint(
+            self.dependency_graph,
+            self.shell,
+            self.fingerprint_dict,
+            self.selector,
+            self.udfs,
+            filename,
+            self.profile_dict,
+            self.write_log_location,
+            self.notebook_name,
+            self.optimizer_name,
+            self,  # 自身のインスタンスを渡す
+        )
+
+    def load_checkpoint(self, filename):
+        if self.debug:
+            print("Loading checkpoint...")
+
+        # start_time = time.time()
+        (
+            self.dependency_graph,
+            variables,
+            vss_to_migrate,
+            vss_to_recompute,
+            self.udfs,
+        ) = resume(filename)
+
+        # Recompute missing VSs and redeclare variables into the kernel.
+        restore_notebook(
+            self.dependency_graph,
+            self.shell,
+            variables,
+            self.write_log_location,
+            self.notebook_name,
+            self.optimizer_name,
+        )
+        # print("Checkpoint load time:", time.time() - start_time)