elastic-kernel 0.0.2__py3-none-any.whl

elastic_notebook/core/common/profile_variable_size.py

@@ -0,0 +1,66 @@
+import sys
+import types
+
+
+def get_total_size(data):
+    """
+    Compute the estimated total size of a variable.
+    """
+
+    def get_memory_size(obj, is_initialize, visited):
+        # same memory space should be calculated only once
+        obj_id = id(obj)
+        if obj_id in visited:
+            return 0
+        visited.add(obj_id)
+
+        try:
+            total_size = sys.getsizeof(obj)
+        except Exception:
+            total_size = float("inf")
+
+        obj_type = type(obj)
+        if obj_type in [int, float, str, bool, type(None)]:
+            # if the original obj is not primitive, then the size is already included
+            if not is_initialize:
+                return 0
+        else:
+            if obj_type in [list, tuple, set]:
+                for e in obj:
+                    total_size = total_size + get_memory_size(e, False, visited)
+            elif obj_type is dict:
+                for k, v in obj.items():
+                    total_size = total_size + get_memory_size(k, False, visited)
+                    total_size = total_size + get_memory_size(v, False, visited)
+            # function, method, class
+            elif obj_type in [
+                types.FunctionType,
+                types.MethodType,
+                types.BuiltinFunctionType,
+                types.ModuleType,
+            ] or isinstance(
+                obj, type
+            ):  # True if obj is a class
+                pass
+            # custom class instance
+            elif isinstance(type(obj), type):
+                # if obj has no builtin size and has additional pointers
+                # if obj has builtin size, all the additional memory space is already added
+                if not hasattr(obj, "__sizeof__") and hasattr(obj, "__dict__"):
+                    for k, v in getattr(obj, "__dict__").items():
+                        total_size = total_size + get_memory_size(k, False, visited)
+                        total_size = total_size + get_memory_size(v, False, visited)
+            else:
+                raise NotImplementedError("Not handled", obj)
+        return total_size
+
+    return get_memory_size(data, True, set())
+
+
+def profile_variable_size(x) -> int:
+    """
+    Profiles the size of variable x. Notably, this should recursively find the size of lists, sets and dictionaries.
+    Args:
+        x: The variable to profile.
+    """
+    return get_total_size(x)

elastic_notebook/core/graph/cell_execution.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2021-2022 University of Illinois
+from typing import List
+
+
+class CellExecution:
+    """
+    A cell execution (object) corresponds to a cell execution (action, i.e. press play) in the notebook session.
+    """
+
+    def __init__(
+        self,
+        cell_num: int,
+        cell: str,
+        cell_runtime: float,
+        start_time: float,
+        src_vss: List,
+        dst_vss: List,
+    ):
+        """
+        Create an operation event from cell execution metrics.
+        Args:
+            cell_num (int): The nth cell execution of the current session.
+            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.
+            src_vss (List[VariableSnapshot]): Nodeset containing input VSs of the cell execution.
+            dst_vss (List[VariableSnapshot]): Nodeset containing output VSs of the cell execution.
+        """
+        self.cell_num = cell_num
+        self.cell = cell
+        self.cell_runtime = cell_runtime
+        self.start_time = start_time
+
+        self.src_vss = src_vss
+        self.dst_vss = dst_vss

elastic_notebook/core/graph/graph.py

@@ -0,0 +1,75 @@
+from collections import defaultdict
+from typing import List
+
+from elastic_notebook.core.graph.cell_execution import CellExecution
+from elastic_notebook.core.graph.variable_snapshot import VariableSnapshot
+
+
+class DependencyGraph:
+    """
+    A dependency graph is a snapshot of the history of a notebook instance.
+    Nodesets and operation events are the nodes and edges of the dependency graph.
+    """
+
+    def __init__(self):
+        """
+        Create a new dependency graph. Called when the magic extension of elastic notebook is loaded with %load_ext.
+        """
+        # Cell executions.
+        self.cell_executions = []
+
+        # Dict of variable snapshots.
+        # Keys are variable names, while values are lists of the actual VSs.
+        # i.e. {"x": [(x, 1), (x, 2)], "y": [(y, 1), (y, 2), (y, 3)]}
+        self.variable_snapshots = defaultdict(list)
+
+    def create_variable_snapshot(
+        self, variable_name: str, deleted: bool
+    ) -> VariableSnapshot:
+        """
+        Creates a new variable snapshot for a given variable.
+        Args:
+            variable_name (str): variable_name
+            deleted (bool): Whether this VS is created for the deletion of a variable, i.e. 'del x'.
+        """
+
+        # Assign a version number to the VS.
+        if variable_name in self.variable_snapshots:
+            version = len(self.variable_snapshots[variable_name])
+        else:
+            version = 0
+
+        # Create a new VS instance and store it in the graph.
+        vs = VariableSnapshot(variable_name, version, deleted)
+        self.variable_snapshots[variable_name].append(vs)
+        return vs
+
+    def add_cell_execution(
+        self, cell, cell_runtime: float, start_time: float, src_vss: List, dst_vss: List
+    ):
+        """
+        Create a cell execution from captured metrics.
+        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.
+            src_vss (List): List containing input VSs of the cell execution.
+            dst_vss (List): List containing output VSs of the cell execution.
+        """
+
+        # Create a cell execution.
+        ce = CellExecution(
+            len(self.cell_executions), cell, cell_runtime, start_time, src_vss, dst_vss
+        )
+
+        # Add the newly created cell execution to the graph.
+        self.cell_executions.append(ce)
+
+        # Set the newly created cell execution as dependent on its input variable snapshots.
+        for src_vs in src_vss:
+            src_vs.input_ces.append(ce)
+
+        # Set the newly created cell execution as the parent of its output variable snapshots.
+        for dst_vs in dst_vss:
+            dst_vs.output_ce = ce

elastic_notebook/core/graph/variable_snapshot.py

@@ -0,0 +1,31 @@
+class VariableSnapshot:
+    """
+    A variable snapshot in the dependency graph corresponds to a version of a variable.
+    I.e. if variable 'x' has been assigned 3 times (x = 1, x = 2, x = 3), then 'x' will have 3 corresponding
+    variable snapshots.
+    """
+
+    def __init__(self, name: str, version: int, deleted: bool):
+        """
+        Create a variable snapshot from variable properties.
+        Args:
+            name (str): Variable name.
+            version (int): The nth update to the corresponding variable.
+            deleted (bool): Whether this VS is created for the deletion of a variable, i.e. 'del x'.
+        """
+        self.name = name
+        self.version = version
+
+        # Whether this VS corresponds to a deleted variable.
+        # i.e. if this VS was created for 'del x' we set this to true so this variable is explicitly not considered
+        # for migration.
+        self.deleted = deleted
+
+        # Cell executions accessing this variable snapshot (i.e. require this variable snapshot to run).
+        self.input_ces = []
+
+        # The unique cell execution creating this variable snapshot.
+        self.output_ce = None
+
+        # Size of variable pointed to by VS; estimated at migration time.
+        self.size = 0

elastic_notebook/core/io/adapter.py

@@ -0,0 +1,18 @@
+from pathlib import Path
+
+
+class Adapter:
+    def __init__(self):
+        pass
+
+    def read_all(self, path: Path):
+        raise NotImplementedError()
+
+    def create(self, path: Path):
+        raise NotImplementedError()
+
+    def write_all(self, path: Path, buf):
+        raise NotImplementedError()
+
+    def remove(self, path: Path):
+        raise NotImplementedError()

elastic_notebook/core/io/filesystem_adapter.py

@@ -0,0 +1,30 @@
+import gc
+from pathlib import Path
+
+import dill
+
+from elastic_notebook.core.io.adapter import Adapter
+
+
+class FilesystemAdapter(Adapter):
+    def __init__(self):
+        super().__init__()
+
+    def read_all(self, path: Path):
+        """
+        The following (read then decode) is faster vs. directly returning dill.load when network speed is low.
+        """
+        gc.disable()
+        contents_bytestring = open(path, "rb").read()
+        contents = dill.loads(contents_bytestring)
+        gc.enable()
+        return contents
+
+    def create(self, path: Path):
+        path.touch()
+
+    def write_all(self, path: Path, buf):
+        dill.dump(buf, open(path, "wb"))
+
+    def remove(self, path: Path):
+        path.unlink()

elastic_notebook/core/io/migrate.py

@@ -0,0 +1,98 @@
+from collections import defaultdict
+from pathlib import Path
+
+import dill
+
+# from elastic_notebook.core.io.filesystem_adapter import FilesystemAdapter
+from ipykernel.zmqshell import ZMQInteractiveShell
+
+from elastic_notebook.core.common.checkpoint_file import CheckpointFile
+from elastic_notebook.core.graph.graph import DependencyGraph
+
+# Default checkpoint location if a file path isn't specified.
+FILENAME = "./notebook.pickle"
+
+
+def migrate(
+    graph: DependencyGraph,
+    shell: ZMQInteractiveShell,
+    vss_to_migrate: set,
+    vss_to_recompute: set,
+    ces_to_recompute: set,
+    udfs,
+    recomputation_ces,
+    overlapping_vss,
+    filename: str,
+):
+    """
+    Writes the graph representation of the notebook, migrated variables, and instructions for recomputation as the
+    specified file.
+
+    Args:
+        graph (DependencyGraph): dependency graph representation of the notebook.
+        shell (ZMQInteractiveShell): interactive Jupyter shell storing the state of the current session.
+        vss_to_migrate (set): set of VSs to migrate.
+        vss_to_recompute (set): set of VSs to recompute.
+        ces_to_recompute (set): set of CEs to recompute post-migration.
+        filename (str): the location to write the checkpoint to.
+        udfs (set): set of user-declared functions.
+    """
+    # Retrieve variables
+    variables = defaultdict(list)
+    for vs in vss_to_migrate:
+        variables[vs.output_ce].append(vs)
+
+    # construct serialization order list.
+    temp_dict = {}
+    serialization_order = []
+    for vs1, vs2 in overlapping_vss:
+        if vs1 in temp_dict:
+            temp_dict[vs1].add(vs2)
+        elif vs2 in temp_dict:
+            temp_dict[vs2].add(vs1)
+        else:
+            # create new entry
+            new_set = {vs1, vs2}
+            temp_dict[vs1] = new_set
+            temp_dict[vs2] = new_set
+
+    for vs in vss_to_migrate:
+        if vs not in temp_dict:
+            temp_dict[vs] = {vs}
+
+    for v in temp_dict.values():
+        serialization_order.append(list(v))
+
+    # Construct checkpoint JSON.
+    # adapter = FilesystemAdapter()
+    metadata = (
+        CheckpointFile()
+        .with_dependency_graph(graph)
+        .with_variables(variables)
+        .with_vss_to_migrate(vss_to_migrate)
+        .with_vss_to_recompute(vss_to_recompute)
+        .with_ces_to_recompute(ces_to_recompute)
+        .with_recomputation_ces(recomputation_ces)
+        .with_serialization_order(serialization_order)
+        .with_udfs(udfs)
+    )
+
+    if filename:
+        print("Checkpoint saved to:", filename)
+        write_path = filename
+    else:
+        write_path = FILENAME
+
+    with open(Path(write_path), "wb") as output_file:
+        dill.dump(metadata, output_file)
+        for vs_list in serialization_order:
+            obj_list = []
+            for vs in vs_list:
+                obj_list.append(shell.user_ns[vs.name])
+            dill.dump(obj_list, output_file)
+    # Write the JSON file to the specified location. Uses the default location if a file path isn't specified.
+    # if filename:
+    #    print("Checkpoint saved to:", filename)
+    #    adapter.write_all(Path(filename), metadata)
+    # else:
+    # adapter.write_all(Path(FILENAME), metadata)

elastic_notebook/core/io/pickle.py

@@ -0,0 +1,71 @@
+# /usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2021-2022 University of Illinois
+
+import hashlib
+
+# import polars as pl
+import inspect
+import mmap
+import pickle
+import types
+
+import dill
+import matplotlib.pyplot as plt
+import networkx
+import pandas as pd
+import seaborn
+from scipy import sparse
+
+
+def is_picklable_fast(obj):
+    if type(obj) in {types.GeneratorType, mmap.mmap, hashlib.sha256}:
+        return False
+    return True
+
+
+def is_picklable(obj):
+    """
+    Checks whether an object is pickleable.
+    """
+    if is_exception(obj) or inspect.ismodule(obj):
+        return True
+    try:
+        # This function can crash.
+        return _is_picklable_dill(obj)
+    except Exception:
+        try:
+            # Double check with function from pickle module
+            return _is_picklable_raw(obj)
+        except Exception:
+            return False
+
+
+def is_exception(obj):
+    """
+    List of objects which _is_picklable_dill returns false (or crashes) but are picklable.
+    """
+    if hasattr(obj, "__module__") and getattr(obj, "__module__", None).split(".")[
+        0
+    ] in {plt.__name__, seaborn.__name__, networkx.__name__, pd.__name__}:
+        return True
+    exceptions = [pd.core.frame.DataFrame, sparse.csr.csr_matrix]
+    return type(obj) in exceptions
+
+
+def _is_picklable_raw(obj):
+    try:
+        # dumps can be slow for large objects that can be pickled
+        pickle.dumps(obj)
+    except Exception:
+        return False
+    return True
+
+
+def _is_picklable_dill(obj):
+    # compared to _is_picklable_raw, this may be faster
+    # however, dill's correctness is worrying because
+    #   it currently considers Pandas DataFrame as not
+    #   picklable, which is not true
+    return dill.pickles(obj)

elastic_notebook/core/io/recover.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2021-2022 University of Illinois
+
+import _pickle
+import dill
+
+from elastic_notebook.core.graph.graph import DependencyGraph
+
+
+def resume(filename):
+    """
+    Resumes the notebook from the checkpoint file.
+    Args:
+        filename (str): location of the checkpoint file.
+    Returns:
+        tuple: (dependency_graph, fingerprint_dict, udfs, recomputation_ces, overlapping_vss)
+    """
+    try:
+        with open(filename, "rb") as output_file:
+            try:
+                # 新しい形式で保存されたデータを読み込む
+                graph, user_ns, udfs, metadata = dill.load(output_file)
+
+                # fingerprint_dictを再構築
+                fingerprint_dict = {}
+                for var_name in user_ns:
+                    if var_name in graph.variable_snapshots:
+                        fingerprint_dict[var_name] = (
+                            None,
+                            set(),
+                            None,
+                        )  # 簡易的な指紋情報
+
+                # metadataから情報を取得
+                recomputation_ces = metadata.get("recomputation_ces", {})
+                overlapping_vss = metadata.get("overlapping_vss", [])
+
+                return graph, fingerprint_dict, udfs, recomputation_ces, overlapping_vss
+
+            except _pickle.UnpicklingError as e:
+                print(f"Warning: Checkpoint file is corrupted: {e}")
+                return DependencyGraph(), {}, set(), {}, []
+            except EOFError as e:
+                print(f"Warning: Checkpoint file is incomplete: {e}")
+                return DependencyGraph(), {}, set(), {}, []
+
+    except Exception as e:
+        print(f"Error loading checkpoint: {e}")
+        return DependencyGraph(), {}, set(), {}, []

elastic_notebook/core/mutation/fingerprint.py

@@ -0,0 +1,184 @@
+import time
+from collections.abc import Iterable
+from types import FunctionType
+
+import pandas as pd
+
+from elastic_notebook.core.mutation.id_graph import (
+    construct_id_graph,
+    is_root_equals,
+    is_structure_equals,
+)
+from elastic_notebook.core.mutation.object_hash import (
+    DataframeObj,
+    ImmutableObj,
+    ModuleObj,
+    NoneObj,
+    NpArrayObj,
+    NxGraphObj,
+    ScipyArrayObj,
+    TorchTensorObj,
+    UncomparableObj,
+    UnserializableObj,
+    construct_object_hash,
+)
+
+BASE_TYPES = [
+    str,
+    int,
+    float,
+    bool,
+    type(None),
+    FunctionType,
+    ImmutableObj,
+    UncomparableObj,
+    NoneObj,
+    NxGraphObj,
+    TorchTensorObj,
+    ModuleObj,
+    UnserializableObj,
+    NpArrayObj,
+    ScipyArrayObj,
+    DataframeObj,
+]
+
+
+def base_typed(obj, visited):
+    """
+    Recursive reflection method to convert any object property into a comparable form.
+    From: https://stackoverflow.com/questions/1227121/compare-object-instances-for-equality-by-their-attributes
+    """
+    T = type(obj)
+    from_numpy = T.__module__ == "numpy"
+
+    if T in BASE_TYPES or callable(obj) or (from_numpy and not isinstance(T, Iterable)):
+        return obj
+
+    visited.add(id(obj))
+
+    if isinstance(obj, Iterable):
+        return obj
+    d = obj if T is dict else obj.__dict__
+
+    comp_dict = {}
+    for k, v in d.items():
+        if id(v) not in visited:
+            comp_dict[k] = base_typed(v, visited)
+
+    return comp_dict
+
+
+def deep_equals(*args):
+    """
+    Extended equality comparison which compares objects recursively by their attributes, i.e., it also works for
+    certain user-defined objects with no equality (__eq__) defined.
+    """
+    return all(
+        base_typed(args[0], set()) == base_typed(other, set()) for other in args[1:]
+    )
+
+
+def construct_fingerprint(obj, profile_dict):
+    """
+    Construct a fingerprint of the object (ID graph + hash).
+    """
+    start = time.time()
+    id_graph, id_set = construct_id_graph(obj)
+    end = time.time()
+    profile_dict["idgraph"] += end - start
+
+    start = time.time()
+    object_representation = construct_object_hash(obj, deepcopy=True)
+    end = time.time()
+    profile_dict["representation"] += end - start
+
+    return [id_graph, id_set, object_representation]
+
+
+def compare_fingerprint(
+    fingerprint_list, new_obj, profile_dict, input_variables_id_graph_union
+):
+    """
+    Check whether an object has been changed by comparing it to its previous fingerprint.
+    """
+    changed = False
+    overwritten = False
+    uncomparable = False
+
+    # Hack: check for pandas dataframes and series: if the flag has been flipped back on, the object has been changed.
+    if isinstance(new_obj, pd.DataFrame):
+        for _, col in new_obj.items():
+            if col.__array__().flags.writeable:
+                changed = True
+                break
+
+    elif isinstance(new_obj, pd.Series):
+        if new_obj.__array__().flags.writeable:
+            changed = True
+
+    # ID graph check: check whether the structure of the object has changed (i.e. its relation with other objects)
+    start = time.time()
+
+    id_graph, id_set = construct_id_graph(new_obj)
+
+    if id_set != fingerprint_list[1] or not is_structure_equals(
+        id_graph, fingerprint_list[0]
+    ):
+        # Distinguish between overwritten variables and modified variables (i.e., x = 1 vs. x[0] = 1)
+        if not is_root_equals(id_graph, fingerprint_list[0]):
+            overwritten = True
+        changed = True
+        fingerprint_list[0] = id_graph
+        fingerprint_list[1] = id_set
+
+    end = time.time()
+    profile_dict["idgraph"] += end - start
+
+    # Value check via object hash: check whether the object's value has changed
+    if not changed:
+        start = time.time()
+        try:
+            new_repr = construct_object_hash(new_obj, deepcopy=False)
+
+            # Variable is uncomparable
+            if isinstance(new_repr, UncomparableObj):
+                if id_set.intersection(input_variables_id_graph_union):
+                    changed = True
+                uncomparable = True
+                fingerprint_list[2] = UncomparableObj()
+            else:
+                if not deep_equals(new_repr, fingerprint_list[2]):
+                    # Variable has equality defined; the variable has been modified.
+                    if (
+                        "__eq__" in type(new_repr).__dict__.keys()
+                        or "eq" in type(new_repr).__dict__.keys()
+                    ):
+                        changed = True
+                    else:
+                        # Object is uncomparable
+                        if id_set.intersection(input_variables_id_graph_union):
+                            changed = True
+                        uncomparable = True
+                        fingerprint_list[2] = UncomparableObj()
+        except Exception as e:
+            print(e)
+            # Variable is uncomparable
+            if id_set.intersection(input_variables_id_graph_union):
+                changed = True
+            uncomparable = True
+            fingerprint_list[2] = UncomparableObj()
+
+    # Update the object hash if either:
+    # 1. the object has been completely overwritten
+    # 2. the object has been modified, and is of a comparable type (i.e., hashable or unhashable but has equality
+    # defined)
+    if overwritten or (
+        changed
+        and not uncomparable
+        and not isinstance(fingerprint_list[2], UncomparableObj)
+    ):
+        fingerprint_list[2] = construct_object_hash(new_obj, deepcopy=True)
+    end = time.time()
+    profile_dict["representation"] += end - start
+
+    return changed, overwritten