fix: handle cell timeouts (#43)

Signed-off-by: Panos Vagenas <[email protected]>
DS4SD · Dec 14, 2023 · c85f982 · c85f982
1 parent d7b1a54
commit c85f982
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 6 deletions.
diff --git a/nbrunner/nb_runner.py b/nbrunner/nb_runner.py
@@ -7,6 +7,7 @@
 import deepsearch as ds
 import nbformat
 from deepsearch.cps.client.components.elastic import ElasticProjectDataCollectionSource
+from nbclient.exceptions import CellTimeoutError
 from nbconvert.preprocessors import CellExecutionError, ExecutePreprocessor
 from rich.console import Console
 from rich.style import Style
@@ -55,7 +56,7 @@ def execute_notebook(self, run_id, notebook_path):
         )
         try:
             out = ep.preprocess(nb, {"metadata": {"path": notebook_path.parent}})
-        except CellExecutionError as e:
+        except (CellExecutionError, CellTimeoutError) as e:
             print(f"=> Error during {run_id}; check {output_filename} for details")
             print(e.traceback)
             out = None

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -31,6 +31,7 @@ spacy = "^3.5.2"
 pydantic = "^2.0.3"
 python-dotenv = "^1.0.0"
 pydantic-settings = "^2.1.0"
+nbclient = "^0.9.0"
 
 [tool.poetry.group.dev.dependencies]
 black = {extras = ["jupyter"], version = "^22.1.0"}

diff --git a/requirements.txt b/requirements.txt
@@ -78,7 +78,7 @@ mistune==3.0.2 ; python_version >= "3.8" and python_version < "3.11"
 mols2grid==2.0.0 ; python_version >= "3.8" and python_version < "3.11"
 monotonic==1.6 ; python_version >= "3.8" and python_version < "3.11"
 murmurhash==1.0.10 ; python_version >= "3.8" and python_version < "3.11"
-nbclient==0.8.0 ; python_version >= "3.8" and python_version < "3.11"
+nbclient==0.9.0 ; python_version >= "3.8" and python_version < "3.11"
 nbconvert==7.9.2 ; python_version >= "3.8" and python_version < "3.11"
 nbformat==5.9.2 ; python_version >= "3.8" and python_version < "3.11"
 nest-asyncio==1.5.8 ; python_version >= "3.8" and python_version < "3.11"