elastic-kernel 0.0.28__tar.gz → 0.0.30__tar.gz

{elastic_kernel-0.0.28 → elastic_kernel-0.0.30}/MANIFEST.in

@@ -1,3 +1,3 @@
 include README.md
-include elastic_notebook_slim/LICENSE
+include elastic_notebook/LICENSE
 recursive-include elastic_kernel *.py *.json

{elastic_kernel-0.0.28/elastic_kernel.egg-info → elastic_kernel-0.0.30}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: elastic-kernel
-Version: 0.0.28
+Version: 0.0.30
 Summary: An IPython Kernel that automatically saves and restores Jupyter Notebook execution states.
 Author-email: ryutarom <m2002r1028@icloud.com>
 License:                                  Apache License
@@ -249,12 +249,26 @@ ElasticKernel is a custom IPython kernel that **automatically checkpoints your n
 
 > 🇯🇵 日本語版は [README.ja.md](README.ja.md) を参照してください。
 
-<!--
-TODO: Add a ~15s demo GIF at the top showing the core value:
-  heavy computation → kernel restart/crash → variables are still there.
-A demo GIF here dramatically improves shareability.
-![ElasticKernel demo](docs/assets/demo.gif)
--->
+## Demo
+
+The same workflow in both clips: define a variable `a = 1`, then **restart the kernel**. The difference is what happens next.
+
+<table>
+<tr>
+<th align="center">Standard kernel (<code>ipykernel</code>)</th>
+<th align="center">ElasticKernel</th>
+</tr>
+<tr>
+<td width="50%"><video src="https://github.com/user-attachments/assets/9ba4d267-20e8-4b06-a3ea-869c19687f81" controls muted></video></td>
+<td width="50%"><video src="https://github.com/user-attachments/assets/9f3c7acc-2ae9-4f98-873c-60f37f1a434e" controls muted></video></td>
+</tr>
+<tr>
+<td align="center">❌ After the restart, <code>a</code> is <b>gone</b> — <code>%whos</code> reports <i>"Interactive namespace is empty."</i></td>
+<td align="center">✅ After the restart, <code>a</code> is <b>automatically restored</b> — <code>%whos</code> still shows <code>a&nbsp;&nbsp;int&nbsp;&nbsp;1</code>.</td>
+</tr>
+</table>
+
+> If the videos don't play inline, click them to view: [ipykernel](docs/assets/ipykernel.mp4) · [ElasticKernel](docs/assets/elastickernel.mp4).
 
 ## Why ElasticKernel?
 

{elastic_kernel-0.0.28 → elastic_kernel-0.0.30}/README.md

@@ -11,12 +11,26 @@ ElasticKernel is a custom IPython kernel that **automatically checkpoints your n
 
 > 🇯🇵 日本語版は [README.ja.md](README.ja.md) を参照してください。
 
-<!--
-TODO: Add a ~15s demo GIF at the top showing the core value:
-  heavy computation → kernel restart/crash → variables are still there.
-A demo GIF here dramatically improves shareability.
-![ElasticKernel demo](docs/assets/demo.gif)
--->
+## Demo
+
+The same workflow in both clips: define a variable `a = 1`, then **restart the kernel**. The difference is what happens next.
+
+<table>
+<tr>
+<th align="center">Standard kernel (<code>ipykernel</code>)</th>
+<th align="center">ElasticKernel</th>
+</tr>
+<tr>
+<td width="50%"><video src="https://github.com/user-attachments/assets/9ba4d267-20e8-4b06-a3ea-869c19687f81" controls muted></video></td>
+<td width="50%"><video src="https://github.com/user-attachments/assets/9f3c7acc-2ae9-4f98-873c-60f37f1a434e" controls muted></video></td>
+</tr>
+<tr>
+<td align="center">❌ After the restart, <code>a</code> is <b>gone</b> — <code>%whos</code> reports <i>"Interactive namespace is empty."</i></td>
+<td align="center">✅ After the restart, <code>a</code> is <b>automatically restored</b> — <code>%whos</code> still shows <code>a&nbsp;&nbsp;int&nbsp;&nbsp;1</code>.</td>
+</tr>
+</table>
+
+> If the videos don't play inline, click them to view: [ipykernel](docs/assets/ipykernel.mp4) · [ElasticKernel](docs/assets/elastickernel.mp4).
 
 ## Why ElasticKernel?
 

{elastic_kernel-0.0.28 → elastic_kernel-0.0.30}/elastic_kernel/command.py

@@ -1,8 +1,6 @@
 import sys
 from pathlib import Path
 
-from setuptools.command.install import install
-
 
 def install_kernel():
     try:
@@ -24,13 +22,6 @@ def install_kernel():
     return True
 
 
-class PostInstallCommand(install):
-    def run(self):
-        install.run(self)
-        print("=== Elastic Kernel: Installing Jupyter kernel ===")
-        install_kernel()
-
-
 def main():
     if len(sys.argv) > 1 and sys.argv[1] == "install":
         install_kernel()

{elastic_kernel-0.0.28 → elastic_kernel-0.0.30}/elastic_kernel/kernel.py

@@ -5,25 +5,14 @@ import os
 import time
 import traceback
 import urllib
+import urllib.parse
+import urllib.request
 from datetime import datetime, timedelta, timezone
 
 from ipykernel.ipkernel import IPythonKernel
 
 from elastic_notebook import ElasticNotebook
-
-
-class JSTFormatter(logging.Formatter):
-    """日本時間（JST）用のログフォーマッター"""
-
-    def converter(self, timestamp):
-        dt = datetime.fromtimestamp(timestamp)
-        return dt.astimezone(timezone(timedelta(hours=9)))  # UTC+9
-
-    def formatTime(self, record, datefmt=None):
-        dt = self.converter(record.created)
-        if datefmt:
-            return dt.strftime(datefmt)
-        return dt.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]  # マイクロ秒を3桁まで表示
+from elastic_notebook.core.common.logging_setup import setup_logger
 
 
 class ElasticKernel(IPythonKernel):
@@ -69,6 +58,9 @@ class ElasticKernel(IPythonKernel):
         # ===========================================
 
         # ElasticNotebookをロードする
+        # 生成に失敗した場合は None のままにし、以降は「追跡なしの素のカーネル」として
+        # 動作を継続する（毎セルで AttributeError を吐き続けないようにする）。(D-13)
+        self.elastic_notebook = None
         try:
             self.elastic_notebook = ElasticNotebook(
                 shell=self.shell,
@@ -77,9 +69,14 @@ class ElasticKernel(IPythonKernel):
             self.logger.info("ElasticNotebook successfully loaded.")
         except Exception as e:
             self.logger.error(f"Error loading ElasticNotebook: {e}")
+            self.logger.error(
+                "Continuing without checkpoint tracking (plain kernel mode)."
+            )
 
         # チェックポイントファイルをロードする
-        if os.path.exists(self.checkpoint_file_path):
+        if self.elastic_notebook is not None and os.path.exists(
+            self.checkpoint_file_path
+        ):
             self.logger.info("Checkpoint file exists. Loading checkpoint.")
             try:
                 start_time = datetime.now(timezone(timedelta(hours=9))).strftime(
@@ -227,33 +224,14 @@ class ElasticKernel(IPythonKernel):
         """
         ロガーの設定
         """
-        # ロガーの設定
-        self.logger = logging.getLogger("ElasticKernelLogger")
-
-        # 環境変数からログレベルを取得
-        log_level_str = os.environ.get("ELASTIC_KERNEL_LOG_LEVEL", "INFO").upper()
-        log_level = getattr(logging, log_level_str, logging.INFO)
-        self.logger.setLevel(log_level)
-
-        formatter = JSTFormatter(
-            "[%(asctime)s %(name)s %(filename)s:%(lineno)d %(levelname)s] %(message)s",
-            "%Y-%m-%d %H:%M:%S.%f",
-        )
-
-        # ローテーティングファイルハンドラー
-        rotating_file_handler = logging.handlers.RotatingFileHandler(
-            self.log_file_path,
-            maxBytes=5 * 1024 * 1024,
-            backupCount=5,  # 5MBのログサイズでローテーション、5世代保存
-        )
-        rotating_file_handler.setLevel(log_level)
-        rotating_file_handler.setFormatter(formatter)
-        self.logger.addHandler(rotating_file_handler)
+        self.logger = setup_logger("ElasticKernelLogger", self.log_file_path)
 
     def __del_from_user_ns_hidden(self):
         """
         %whoで表示されるようにするために復元した変数をself.shell.user_ns_hiddenから削除する
         """
+        if self.elastic_notebook is None:
+            return
         variable_snapshots = set(
             self.elastic_notebook.dependency_graph.variable_snapshots
         )
@@ -294,24 +272,34 @@ class ElasticKernel(IPythonKernel):
 
         self.logger.debug(f"Executing Code:\n{code}")
 
+        skip_record = self.__skip_record(code)
+
         pre_execution_user_ns = (
-            set(self.shell.user_ns.keys()) if not self.__skip_record(code) else None
+            set(self.shell.user_ns.keys()) if not skip_record else None
         )
-        start_time = time.time() if not self.__skip_record(code) else None
+        start_time = time.time() if not skip_record else None
 
         result = await super().do_execute(
             code, silent, store_history, user_expressions, allow_stdin
         )
 
-        if not self.__skip_record(code):
+        # Q3: 実行が成功したセルのみ記録する。失敗したセルを記録すると、依存グラフに
+        # 「実行されたセル」として残り、復元のたびに再計算（再実行）されてしまうため。
+        execution_succeeded = isinstance(result, dict) and result.get("status") == "ok"
+
+        if skip_record:
+            self.logger.debug("Skipping record event")
+        elif not execution_succeeded:
+            self.logger.debug("Skipping record event for failed cell execution")
+        elif self.elastic_notebook is None:
+            self.logger.debug("Skipping record event (ElasticNotebook unavailable)")
+        else:
             cell_runtime = time.time() - start_time
             self.logger.debug(f"Cell runtime: {cell_runtime}")
             self.elastic_notebook.record_event(
                 code, pre_execution_user_ns, start_time, cell_runtime
             )
             self.logger.debug("Recording event")
-        else:
-            self.logger.debug("Skipping record event")
 
         return result
 
@@ -319,6 +307,12 @@ class ElasticKernel(IPythonKernel):
         """
         カーネル終了時に呼び出されるメソッド
         """
+        # ElasticNotebook の生成に失敗していた場合はチェックポイントを保存できないので
+        # スキップして通常のシャットダウンを継続する。(D-13)
+        if self.elastic_notebook is None:
+            self.logger.info("ElasticNotebook unavailable; skipping checkpoint save.")
+            return super().do_shutdown(restart)
+
         try:
             start_time = datetime.now(timezone(timedelta(hours=9))).strftime(
                 "%Y-%m-%dT%H:%M:%S.%f%z"

{elastic_kernel-0.0.28 → elastic_kernel-0.0.30/elastic_kernel.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: elastic-kernel
-Version: 0.0.28
+Version: 0.0.30
 Summary: An IPython Kernel that automatically saves and restores Jupyter Notebook execution states.
 Author-email: ryutarom <m2002r1028@icloud.com>
 License:                                  Apache License
@@ -249,12 +249,26 @@ ElasticKernel is a custom IPython kernel that **automatically checkpoints your n
 
 > 🇯🇵 日本語版は [README.ja.md](README.ja.md) を参照してください。
 
-<!--
-TODO: Add a ~15s demo GIF at the top showing the core value:
-  heavy computation → kernel restart/crash → variables are still there.
-A demo GIF here dramatically improves shareability.
-![ElasticKernel demo](docs/assets/demo.gif)
--->
+## Demo
+
+The same workflow in both clips: define a variable `a = 1`, then **restart the kernel**. The difference is what happens next.
+
+<table>
+<tr>
+<th align="center">Standard kernel (<code>ipykernel</code>)</th>
+<th align="center">ElasticKernel</th>
+</tr>
+<tr>
+<td width="50%"><video src="https://github.com/user-attachments/assets/9ba4d267-20e8-4b06-a3ea-869c19687f81" controls muted></video></td>
+<td width="50%"><video src="https://github.com/user-attachments/assets/9f3c7acc-2ae9-4f98-873c-60f37f1a434e" controls muted></video></td>
+</tr>
+<tr>
+<td align="center">❌ After the restart, <code>a</code> is <b>gone</b> — <code>%whos</code> reports <i>"Interactive namespace is empty."</i></td>
+<td align="center">✅ After the restart, <code>a</code> is <b>automatically restored</b> — <code>%whos</code> still shows <code>a&nbsp;&nbsp;int&nbsp;&nbsp;1</code>.</td>
+</tr>
+</table>
+
+> If the videos don't play inline, click them to view: [ipykernel](docs/assets/ipykernel.mp4) · [ElasticKernel](docs/assets/elastickernel.mp4).
 
 ## Why ElasticKernel?
 

{elastic_kernel-0.0.28 → elastic_kernel-0.0.30}/elastic_kernel.egg-info/SOURCES.txt

@@ -12,16 +12,16 @@ elastic_kernel.egg-info/dependency_links.txt
 elastic_kernel.egg-info/entry_points.txt
 elastic_kernel.egg-info/requires.txt
 elastic_kernel.egg-info/top_level.txt
+elastic_notebook/LICENSE
 elastic_notebook/__init__.py
 elastic_notebook/elastic_notebook.py
 elastic_notebook/algorithm/__init__.py
-elastic_notebook/algorithm/baseline.py
 elastic_notebook/algorithm/optimizer_exact.py
 elastic_notebook/algorithm/selector.py
 elastic_notebook/core/__init__.py
 elastic_notebook/core/common/__init__.py
 elastic_notebook/core/common/checkpoint_file.py
-elastic_notebook/core/common/profile_graph_size.py
+elastic_notebook/core/common/logging_setup.py
 elastic_notebook/core/common/profile_migration_speed.py
 elastic_notebook/core/common/profile_variable_size.py
 elastic_notebook/core/graph/__init__.py
@@ -32,7 +32,6 @@ elastic_notebook/core/io/__init__.py
 elastic_notebook/core/io/adapter.py
 elastic_notebook/core/io/filesystem_adapter.py
 elastic_notebook/core/io/migrate.py
-elastic_notebook/core/io/pickle.py
 elastic_notebook/core/io/recover.py
 elastic_notebook/core/mutation/__init__.py
 elastic_notebook/core/mutation/fingerprint.py
@@ -43,4 +42,18 @@ elastic_notebook/core/notebook/checkpoint.py
 elastic_notebook/core/notebook/find_input_vars.py
 elastic_notebook/core/notebook/find_output_vars.py
 elastic_notebook/core/notebook/restore_notebook.py
-elastic_notebook/core/notebook/update_graph.py
+elastic_notebook/core/notebook/update_graph.py
+tests/test_checkpoint_file.py
+tests/test_checkpoint_func.py
+tests/test_find_input_vars.py
+tests/test_find_output_vars.py
+tests/test_graph.py
+tests/test_id_graph.py
+tests/test_integration.py
+tests/test_logging_setup.py
+tests/test_migrate_resume.py
+tests/test_object_hash.py
+tests/test_optimizer.py
+tests/test_profile_migration_speed.py
+tests/test_profile_variable_size.py
+tests/test_skip_record.py

{elastic_kernel-0.0.28 → elastic_kernel-0.0.30}/elastic_kernel.egg-info/top_level.txt

@@ -1,4 +1,2 @@
-dist
-docs
 elastic_kernel
 elastic_notebook

elastic_kernel-0.0.30/elastic_notebook/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

{elastic_kernel-0.0.28 → elastic_kernel-0.0.30}/elastic_notebook/algorithm/optimizer_exact.py

@@ -1,6 +1,8 @@
 # This file has been modified from the original ElasticNotebook.
 # Original: https://github.com/illinoisdata/ElasticNotebook
 
+from typing import Union
+
 import networkx as nx
 import numpy as np
 from networkx.algorithms.flow import shortest_augmenting_path
@@ -23,9 +25,6 @@ class OptimizerExact(Selector):
         self.active_oes = None
         self.compute_graph = None
 
-        # CEs required to recompute a variables last modified by a given CE.
-        self.recomputation_ces = {}
-
         self.idx = 0
 
     def get_new_idx(self) -> int:
@@ -36,7 +35,12 @@ class OptimizerExact(Selector):
         self.idx += 1
         return idx
 
-    def dfs(self, current: str, visited: set, recompute_ces: str):
+    def dfs(
+        self,
+        current: Union[CellExecution, VariableSnapshot, None],
+        visited: set,
+        recompute_ces: set,
+    ):
         """
         Perform DFS on the Application History Graph for finding the CEs required to recompute a variable.
         Args:
@@ -71,7 +75,7 @@ class OptimizerExact(Selector):
                 self.dfs(ce, set(), recompute_ces)
                 self.recomputation_ces[ce] = recompute_ces
 
-    def select_vss(self, notebook_name=None, optimizer_name=None) -> set:
+    def select_vss(self, notebook_name=None, optimizer_name=None) -> tuple:
         self.find_prerequisites()
 
         # Construct flow graph for computing mincut

{elastic_kernel-0.0.28 → elastic_kernel-0.0.30}/elastic_notebook/algorithm/selector.py

@@ -1,20 +1,6 @@
 # This file has been modified from the original ElasticNotebook.
 # Original: https://github.com/illinoisdata/ElasticNotebook
 
-from enum import Enum
-
-
-class OptimizerType(Enum):
-    """
-    String representations of all implemented optimizers.
-    """
-
-    EXACT = "exact"
-    EXACT_C = "exact_c"
-    EXACT_R = "exact_r"
-    MIGRATE_ALL = "migrate_all"
-    RECOMPUTE_ALL = "recompute_all"
-
 
 class Selector:
     """
@@ -32,13 +18,17 @@ class Selector:
         self.overlapping_vss = None
         self.migration_speed_bps = migration_speed_bps
 
-    def select_vss(self) -> set:
+        # CEs required to recompute a variable last modified by a given CE.
+        # Declared on the base so callers (e.g. checkpoint()) can rely on it.
+        self.recomputation_ces: dict = {}
+
+    def select_vss(self, notebook_name=None, optimizer_name=None) -> tuple:
         """
-        Classes that inherit from the `Selector` class (such as `Optimizer` and various baselines)
-            should override `select_vss`.
+        Classes that inherit from the `Selector` class (such as `OptimizerExact`) should
+        override `select_vss`.
 
         Returns:
-            set(VariableSnapshot): a subset of active VSs selected to migrate based on various heuristics and
-            algorithms.
+            Tuple[set(VariableSnapshot), set(CellExecution)]: the VSs selected to migrate and
+            the CEs selected to recompute.
         """
         raise NotImplementedError()