fix: handle cell timeouts

Signed-off-by: Panos Vagenas <[email protected]>

nbrunner/nb_runner.py

@@ -7,7 +7,11 @@
 import deepsearch as ds
 import nbformat
 from deepsearch.cps.client.components.elastic import ElasticProjectDataCollectionSource
-from nbconvert.preprocessors import CellExecutionError, ExecutePreprocessor
+from nbconvert.preprocessors import (
+    CellExecutionError,
+    CellTimeoutError,
+    ExecutePreprocessor,
+)
 from rich.console import Console
 from rich.style import Style
 from rich.table import Table
@@ -55,7 +59,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