Kishu: Time-Traveling for Computational Notebooks

Data cleaning and visualization operations are oftentimes irreversible (e.g., df = df.drop_col(’a’)) and/or partially performed with unobservable side-effects (e.g., a mid-cell execution error), and the user may want to return to the previous state if the execution was erroneous or undesirable (Groups, 2024; Brachmann and Spoth, 2020). To enable interruption-free time-traveling, we can checkpoint the state delta to storage after each operation such that the session state prior to performing the operation can be returned to via loading the appropriate deltas. We empirically study this use case in § 7.5.1.

Reactive execution allows users to quickly investigate alternative paths in data science workflows: when a previously executed cell is re-executed (e.g., train model with different hyperparameters), all its dependent cells (e.g., plotting cells) are too (Koop and Patel, 2017; Shankar et al., 2022). If we can efficiently persist all different variations of objects in different execution paths (i.e., as incremental deltas w.r.t. the shared state), users can efficiently evaluate each variation of values against each other: to switch paths, only the (small portion of) data differing between paths need to be updated via loading the appropriate deltas. We empirically study this use case in § 7.5.2.

Tools such as CRIU (CRIU, 2023) and DMTCP (Ansel et al., 2009) can be used to create memory snapshots of notebook processes, which contain all data in the session state. (Incremental Checkpointing) Subsequent snapshots can be made incrementally w.r.t. previous snapshots such that only dirty memory pages are stored. However, memory page granularity is coarse and incurs high checkpoint storage overhead (§ 7.3), and OS-level checkpointing is limited to single processes, failing on notebooks utilizing multiple/remote processors (e.g., Spark (Zaharia et al., 2010) and Ray (Moritz et al., 2018) pipelines) or moving data out of CPU (e.g., onto the GPU). (Complete Checkout) Memory snapshots must be entirely loaded to restore the notebook state and additionally require the existing notebook process to be killed (otherwise, a PID conflict occurs) before restoration. This process is both not seamless and incurs high data loading costs (§ 7.5).

Application-level tools such as Dill’s dumpsession (Foundation, 2023g) and ElasticNotebook (Li et al., [n.d.]) can be used to create checkpoint files by serializing data (e.g., into bytestrings) in the session state. (Complete Checkpointing) Notably, no existing application-level tool supports incremental checkpointing: each checkpoint must independently contain sufficient data for restoring the session state. (Complete Checkout) Checkpoint files need to be entirely loaded (via data deserialization) for session restoration; despite capable of restoring into an existing notebook process, these tools do not utilize variables that are already present within the kernel which may not need to be updated for faster checkout (§ 7.5).

If we can detect the delta between session states at a finer granularity (compared to OS-level tools) and utilize existing data within the kernel to speedup checkout (compared to Application-level dumps), we can achieve efficient and generalized application-level incremental checkpointing and checkout, as follows: (Incremental and Generalized Checkpointing) We utilize application-level information to track state deltas at a finer-than-memory-page-granularity for storage-efficient incremental checkpointing. For generalizability, we can utilize an object’s reduction (i.e., __reduce__ (Foundation, 2023f)) as storage instructions to handle checkpointing multiprocessing and off-CPU workloads. (Incremental Checkout) If we know the contents of the target session state to restore to (e.g.., in terms of a state snapshot (Bernstein and Goodman, 1983)) and accurately compute its difference from the current state, we can perform incremental checkout by only loading and updating data that differ between the current and target state.

Specifically, we compute state differences at the Co-variable granularity, which we describe and motivate in the next session.

System-level checkpointing tools (CRIU, 2023; Ansel et al., 2009) track state deltas via dirty memory pages. While fast, the delta granularity is insufficient for tracking Python notebook session states, as (1) modifying any object, regardless of size (e.g., x += 1) will cause the entire page (e.g., 4KB) holding the object to become dirty, and (2) Python data structures (e.g., lists) are often constructed and stored in a fragmented manner, leading to simple operations (e.g., mapping elements in a list in-place) creating multiple dirty pages, leading to high incremental checkpoint costs.

Notebook systems such as IPyFlow (Macke et al., 2020), Dataflow (Koop and Patel, 2017), and ElasticNotebook (Li et al., [n.d.]) track session state evolution at the variable-level via provenance-based code analysis (Foundation, 2023a). While capable of producing fine-grained deltas, the tracker has to be either conservative on identifying changed variables (e.g., w.r.t. control flows and external function calls (Li et al., [n.d.])) causing many false positives and large deltas, or perform extensive live instrumentation for resolution (Macke et al., 2020), which can result in high overhead (§ 7.6).

In order to avoid low granularity of memory-page tracking and potential efficiency issues of provenance tracking, we propose to use live object comparison to track updates to Co-variables—connected components of objects (w.r.t. pointer references). Our intuition is that (live) tracking and checkpointing/checking out individual objects are expensive and risky (i.e., may break shared references (Li et al., [n.d.])), respectively. However, at Co-variable granularity, we can achieve low-overhead tracking by reasoning in a principled way based on access patterns which Co-variables were updated by each cell execution, and correctly store/load Co-variables during incremental checkpointing/checkout as if they are independent data tables (§ 4). We depict this idea in LABEL:fig:background_variable_blob: {ls,obj} is a Co-variable (red) as the objects reachable from these variables overlap, i.e., &ls[1] = &obj.foo. {df} is another Co-variable (blue), and there is no way to reach objects under df from objects under ls or obj via references. Notably, Co-variables are the minimum granularity which data in the session state can be stored and loaded without risking breaking shared references (e.g., unlike variable-level KV-stores (shove, 2024; 0xnurl, 2024)). We formally describe Co-variables and how we correctly and efficiently capture state delta at this granularity in § 4.

Suppose a data analyst is performing text mining (LABEL:fig:background_example). They load the corpus (Cell 1), define category lists (Cell 2), and sort texts by sentiment into the lists (Cell 3). The analyst incrementally checkpoints the state after each cell execution.

(1) Incremental Checkpointing: The analyst wishes to map the lists to clean the contained text and tests with sad_ls (Cell 4, blue). Due to its interleaved construction (with other lists), the list sad_ls is fragmented; incrementally checkpointing the state at memory page granularity for Cell 4 (w.r.t. Cell 3) results in all pages overlapping with sad_ls being copied. However, a Co-variable granularity incremental checkpoint stores only (the bytestring of) sad_ls.

(2) Incremental Checkout: The analyst decided to undo the mapping function in Cell 4 as the results were unsatisfactory. Returning to the state of Cell 3 by (completely) loading a memory snapshot is slow as it requires reloading the corpus. However, by identifying that the states between Cells 3 and 4 differ only by sad_ls, we can only load the value of sad_ls from Cell 3 (red) to replace value of sad_ls from Cell 4 (blue) to perform incremental checkout while keeping the rest of the session state untouched.

At the start of the notebook session, Kishu patches the session namespace in order to monitor accesses to its contents between successive cell executions (§ 4.3). It identifies the candidate Co-variables to check for updates through tracking user-referenced variable names, and passes the candidates to the Delta Detector to compute the Co-variable granularity state delta.

The Delta Detector computes the state delta based on the candidates identified from the Patched Namespace (i.e., which of the candidate Co-variables were actually updated by the cell execution). We discuss the Kishu’s delta detection in (§ 4).

The Checkpoint Graph is a tree-like structure analogous to Git’s commit graph (Git, 2024a), in which Kishu writes, stores, and versions incremental checkpoints consisting of the updated Co-variables (i.e., the state delta) of each cell execution (§ 5.1). The incremental checkpoints stored in the Checkpoint Graph are used by the State Loader to perform incremental checkout.

The State Loader restores to a session state upon requested checkout. It first identifies the difference between the current session state (i.e., existing items in the namespace) and the target state according to the Checkpoint Graph, then loads only the necessary data from the Checkpoint Graph for replacing the variables that need updating (§ 5.2). If required data for checkout is missing (e.g., Kishu failed to serialize it into the Checkpoint Graph) or fails to load, the Data Restorer is invoked to restore the data.

The Data Restorer is a mechanism that utilizes fallback recomputation to restore missing data for checkout (e.g., Kishu failed to serialize the data during prior checkpointing). It reconstructs missing data by combining loading dependent data and cell re-runs according to the Checkpoint Graph. (§ 5.3)

When initializing a notebook session, Kishu will be attached to the kernel; it will patch the namespace and initialize the Checkpoint Graph on storage.

After each cell execution, the Delta Detector utilizes the Patched Namespace to identify the updated Co-variables, then stores them as a new incremental checkpoint (i.e., node) in the Checkpoint Graph.

Kishu will restore a previous session state on request. The State Restorer first identifies the difference between the current and target state at Co-variable granularity according to the Checkpoint Graph (i.e., which variables need to be restored), then loads only the necessary data according to the difference to complete the restoration. If necessary, the Data Restorer reconstructs data that is missing (i.e., failed to serialize during checkpointing earlier) or failed to load via fallback recomputation.

In Python and the Jupyter Notebook ecosystem, variables and objects are 2 distinct concepts: A variable is a named entity from which one or more objects are reachable - for example, for a list ls=[1,2,3], the list name (ls) is a variable and each of the elements (1, 2, 3) is an object. We define reachability according to references, i.e., object y is reachable from variable x if y can be accessed from x through a chain of references. Some common reachability patterns include subscripting (e.g., y = x[0]), class member (e.g., y = x.attr), and attribution (e.g., y = x.__dict__). Given our distinction between variables and objects and our definition of reachability, we now define the Co-variable as follows:

A Co-variable can consist of one name (e.g., a primitive, x = 1) or multiple names from which the same object can be reached (i.e., shared references). LABEL:fig:checkpoint_variable_update shows an example—the string object ’b’ is reachable from both list ls and object obj via subscript and class member respectively, hence {ls,obj} is a Co-variable. Co-variables are self-contained by definition, i.e., there are no inter-Co-variable references. Co-variables can be modified by cell executions:

For example, the Co-variable {ls,obj} in LABEL:fig:checkpoint_variable_update is modified node-wise with “ls[0] = ’e’” (bottom), and is modified edge-wise with “obj.foo = ls[2]” (bottom-right). Co-variables can also be created (or deleted) via split and merge (right): the Co-variable {obj,ls} is deleted via a split (as obj and ls longer share references), and the Co-variable {obj, st} is created through a merge. For brevity, we collectively refer to Co-variable modifications, creations, and deletions as updates—the set of Co-variables updated by a cell execution form the execution’s state delta.

Kishu uses VarGraphs to detect Co-variable membership and updates. The VarGraph is a graph structure constructed from each variable in the namespace that captures its reachable objects. An example is shown in LABEL:fig:checkpoint_id_graph: each node in a variable’s VarGraph corresponds to a reachable object, containing the (1) object type, (2) memory address, and one of (3) pointers to other reachable objects (i.e., children) for non-primitive objects, or (4) value for primitive objects. For example, the node for the list reachable from ls contains 3 child pointers to the 3 nodes for strings ’a’, ’b’, and ’c’, and the node for string ’b’ holds its value ’b’. ²²2The VarGraph is inspired by ElasticNotebook’s ID graph (Li et al., [n.d.]) which captures reachable objects’ memory addresses; VarGraphs uniquely contain datatypes and primitive values for additional robustness (e.g., detecting a different primitive in the same address).

Co-variable membership is determined by intersecting VarGraphs. For example, in LABEL:fig:checkpoint_id_graph, ls and obj form a Co-variable as the node ’b’ is in both graphs (red).

Co-variable updates is determined by comparing VarGraphs before and after cell executions. A graph structure modification and/or a node attribute change (e.g., object memory address or type) indicates an update to the Co-variable.

As Kishu constructs VarGraphs following object rechability, it detects Co-variable updates with no false negatives (verified empirically in § 7.2.1). However, Kishu’s update detection is conservative: there may be false positives if objects are dynamically generated (e.g., datatype objects) with a different memory address each time during VarGraph construction/object traversal, or cannot be traversed into (i.e., lacking referencing instructions, e.g., generators (Foundation, 2023b), which Kishu assumes to be updated on access).

Cell executions in Jupyter Notebook interact with the global namespace (i.e., globals()). Therefore, if Kishu can capture variable references in the cell execution, it can reason about which Co-variables were possibly updated (and which ones were definitely not), as follows:

A Co-variable being accessed indicates the possibility of it being updated (e.g., a subscription access ls[0] = ’e’). Kishu patches the accessor, setter, and deletion methods of the global namespace (LABEL:fig:checkpoint_detect_fast) to capture variable (hence Co-variable) accesses. Kishu identifies candidate Co-variables for update checking via the captured accessed variables: if a Co-variable’s members $(x_{1},...,x_{i})$ overlaps with the accessed variables of the cell execution, then the Co-variable may have been updated (e.g., {ls,obj} in LABEL:fig:checkpoint_detect_fast) and Kishu will have to verify the update by (1) re-generating VarGraphs for its member variables, (2) comparing the VarGraphs with those before the cell execution to identify modifications, and (3) intersecting the VarGraphs amongst variables of accessed Co-variables to identify merges and splits. Otherwise, the Co-variable surely wasn’t updated and Kishu skips its check for this cell execution (e.g., {df} in LABEL:fig:checkpoint_detect_fast, greyed out). We provide a short proof of contradiction:

As only a small portion of variables are accessed per cell in a typical data science notebook, Kishu significantly reduces delta detection overhead with this approach (empirically verified in § 7.6.1).

As Kishu patches the global namespace of the notebook session, it is impossible for the user to use variables from within the notebook (i.e., to modify objects) undetected. Hence, Kishu will not misidentify Co-variables possibly updated via references.

The user may still use non-referencing methods such as direct C-pointer-based modifications, but these cases are rare in notebooks (found in 0 out of 60 surveyed notebooks (Li et al., [n.d.])) Some libraries, such as NumPy (Team, 2024a), do perform memory-based updates (e.g., via slicing). However, the objects are supported by Kishu as their memory-based updates are still invoked via referencing (e.g., arr[0,1] += 1, empirically verified in § 7.2.1).

The Checkpoint Graph is a directed tree of (incremental) checkpoints representing the branch-based evolution of session states. It is grown (i.e., via adding nodes) with each checkpoint Kishu performs, and nodes are timestamped according to the completion time of the corresponding cell execution $t$ (which we refer as CE $t$ and node $t$ for simplicity). The Checkpoint Graph maintains a head node, which tracks the user’s current state. Each node $t$ contains the state delta consisting of (only) Co-variables updated by CE $t$. When Co-variables are stored in the Checkpoint Graph, they are versioned according to their corresponding CE $t$:

Versioned Co-variables are analogous to versioned datasets: the same Co-variable (w.r.t. variable membership, i.e., $\mathcal{X}=\{x_{1},...,x_{i}\}$) can take on multiple values throughout a notebook session being updated by different cell executions. LABEL:fig:checkpoint_graph show an example: CE $t_{3}$ creates the Co-variable {plot}, hence it is stored in node $t_{3}$ (red) as the Versioned Co-variable $((plot),t_{3})$.

Kishu creates a new node $t$ in the Checkpoint Graph after each CE $t$. The new node contains (1) the state delta of CE $t$ consisting of versioned Co-variables, (2) the code of CE $t$, and (3) the versioned Co-variables (possibly stored in previous checkpoints) accessed by CE $t$ (§ 4.3). For example, the node $t_{3}$ in LABEL:fig:restore_diff contains the code of CE $t_{3}$ (“plot=gmm.result()“) and its dependency on the versioned Co-variable from node $t_{2}$ $((gmm),t_{2})$ (dashed line), analogous to the transformation, transaction, and dependencies in database versioning. The new node $t$ is written into the Checkpoint Graph under the head node $s$, and an edge (i.e., a parent-child relationship) is added from the head node $s$ to the new node $t$ (which is now the new head node).

If Kishu cannot write an updated Co-variable in the state delta into the Checkpoint Graph (e.g., it contains an unserializable object such as a generator (Foundation, 2023b) or hash (Foundation, 2023c)), Kishu simply skips its storage. Instead, upon checkout, the missing (unserializable) Co-variable will be restored through fallback recomputation enabled by the cell code and dependencies stored in the Checkpoint Graph node, which we discuss in § 5.3.

When checking out to the state at node $t$, The State Restorer (§ 3.1) performs the following steps:

1. 1.

   Load the appropriate Versioned Co-variables from nodes (i.e., node $t$ and ancestors of node $t$) to update diverged Co-variables between the state of the current head node $s$ and node $t$.

2. 2.

   Update/re-generate VarGraphs (§ 4.2) for updated Co-variables.

3. 3.

   Move the head from node $s$ to the checked out node $t$.

Notably, the next cell execution will create a node in a new branch rooted at $t$ in the Checkpoint Graph, e.g., the graph in LABEL:fig:restore_diff is generated through the sequence $t_{1}\rightarrow t_{2}\rightarrow t_{3}\rightarrow(\text{checkout to }\,t_{1})\rightarrow t_{4}\rightarrow t_{5}$. If during checkout, a required Versioned Co-variable is missing (i.e., due to serialization failure, § 5.1) or fails to load (i.e., deserialization failure), Kishu restores it via fallback recomputation.

As each Checkpoint Graph node $t$ contains the code of CE $t$ and (2) which Versioned Co-variables $(\mathcal{X}_{j},t_{j})$ (stored in previous deltas, $t_{j}<t$) CE $t$ accessed (§ 5.1), any Versioned Co-variable $(\mathcal{X},t)$ in the state delta of node $t$ can be recomputed by (1) loading the accessed Versioned Co-variables from previous checkpoints into a temporary namespace, then re-running CE $t$. For example, in LABEL:fig:restore_fallback, suppose Versioned Co-variable ({plot}, $t_{3}$) (green) fails to load when checking out to session state $t_{3}$. The versioned Co-variable (({gmm}, $t_{2}$) is required to rerun CE $t_{3}$ (red); therefore, ({gmm}, $t_{2}$) is loaded from the parent node $t_{2}$, and after rerunning CE $t_{3}$ on the input ({gmm}, $t_{2}$), ({plot}, $t_{3}$) is restored.

Kishu’s fallback recomputation is dynamic and recursive—if another Co-variable is missing or fails to load when retrieving the inputs (i.e., accessed Co-variables) for performing fallback recomputation, fallback recomputation can be recursively performed for that Co-variable. For example, if ({gmm}, $t_{2}$) from node $t_{2}$ fails to load (as part of fallback recomputation for ({plot}, $t_{3}$) from node $t_{3}$), it itself can be recomputed by loading ({gmm}, $t_{1}$) from node $t_{1}$ and rerunning CE $t_{2}$ (blue).

Kishu guarantees exact restoration for all serializable Co-variables during checkout, i.e., they will have the same bytestring representation before and after checkout if there are no hidden serialization errors (§ 6.2). While Kishu is capable of restoring problematic Co-variables via fallback recomputation, it currently does not support exactly restoring Co-variables that both (1) fail to store or load and (2) are created through non-deterministic means (e.g., random generators). This is a limitation similarly present in Spark (Zaharia et al., 2010) and Ray’s (Moritz et al., 2018) lineage-based fault tolerance; however, in our case, unserializable objects are sufficiently rare in data science libraries (§ 7.2), hence we consider this limitation to be acceptable.

For seamless integration, Kishu is implemented as a separate application from the notebook process, and can be used without altering the base Jupyter application. When a session is initialized, Kishu places hooks into the kernel (pre_run_cell and post_run_cell (IPython, 2024b)) and patches the namespace (user_ns (IPython, 2024a)) (§ 3.1), allowing the standalone Kishu process to perform delta detection, write data in the namespace to storage, and overwrite data in the namespace upon checkout transparently.

The Pickle protocol (i.e., __reduce__ (Foundation, 2023f)) is employed for (1) object serialization and (2) constructing VarGraphs (hence identifying Co-variables), i.e., an object y is reachable from another object x if pickle(x) includes y. As Pickle is the de-facto standard (in Python) observed by almost all data science libraries (e.g., NumPy, PyTorch (et al., 2018)), Kishu can be used for almost all use cases. Furthermore, Kishu’s per-Co-variable storage enables mixing and matching serialization libraries for coverage. Currently, Kishu will try CloudPickle (cloudpipe, 2024) first, then use Dill (Foundation, 2023g) as a fallback for Co-variables containing objects that CloudPickle fails on.

Currently, Kishu uses SQLite (SQLite, 2024) to store Versioned Co-variables in the Checkpoint Graph. However, any storage mechanism can be used in its place—even in-memory ones if the user wants to maximize checkpointing/checkout efficiency.

In Kishu, the session state is formally defined according to Co-variables in the user namespace (user_ns), which contains key-value pairs of variable names to their reachable objects. The session state does not include (1) local/module/hidden variables and (2) external items outside of the notebook session (e.g., on-disk files), which we do not aim to checkpoint nor restore.

Certain object classes may contain incorrect serialization instructions, which, despite being able to be stored/loaded to/from storage, result in silent errors. Kishu currently assumes that the serialization instructions are correctly implemented for all objects in terms of equality before and after pickling (i.e., x == unpickle(pickle(x)) and does not prevent these silent errors. However, these cases are rare (§ 7.2.1), and Kishu provides a blocklist file for users to force fallback recomputation for Co-variables containing objects belonging to these classes.

We select a total of 8 data science notebooks from Kaggle Grandmaster-level users or tutorials of established tools (e.g., Ray). LABEL:tbl:workload reports our selected notebooks’ dataset sizes and runtimes.

Additionally, we select 146 data science library classes commonly used in data science notebooks, on which we evaluate the robustness of Kishu’s modification detection and time-traveling correctness. LABEL:tbl:library reports a classification of the libraries.

We evaluate Kishu against existing tools capable of enabling time-travelling on notebooks to various degrees:

• •

  CRIU (CRIU, 2023): Performs a system-level memory dump of the process hosting the notebook session. The session state is restored by loading the memory dump and reviving the process.

• •

  CRIU-Incremental (CRIU, 2023): CRIU with snapshot deduplication, storing only dirty memory pages in subsequent snapshots.

• •

  DumpSession (Foundation, 2023h): An application-level checkpointing tool that serializes the entire session state into one single file.

• •

  ElasticNotebook (Li et al., [n.d.]): An application-level notebook migration tool that balances data serialization and cell recomputation to achieve optimized session replication times.

We additionally compare the overhead of Kishu’s update detection mechanism with the following methods:

• •

  IPyFlow (Shankar et al., 2022): A hybrid dynamic-static (i.e., AST analysis with live symbol resolution) for obtaining sub-variable (i.e., symbols, e.g., ls[x]) level granularity to perform reactive cell executions.

• •

  Kishu (Check all): Always perform update detection for all Co-variables in the session state after each cell execution, regardless of whether they were accessed by the user in the previous cell.

We consider these methods regarding session state delta tracking time (§ 7.6) to gauge the effectiveness of our overhead reduction via identifying candidates for updated Co-variables.

Experiments are performed on an Ubuntu server with 2 AMD EPYC 7552 48-Core Processors and 1TB RAM. The checkpoints for all methods are written to a mounted NFS, with disk read and write speeds of 519.8 MB/s and 358.9 MB/s, respectively.

We measure (1) (incremental) checkpoint time as time spent checkpointing (including both tracking and data writing) after each cell execution, (2) checkout time as time to restore the state from checkpoint files, and (3) tracking overhead as time spent after each cell execution tracking Co-variable granularity updates. We clear the page cache between runs.

Our implementation of Kishu, experiment notebooks, full list of data science library classes we test with, and scripts can be found in our Github repository.⁴⁴4https://anonymous.4open.science/r/kishu-vldb-E724