From 5a1ffc510c83dcb1c0c6e7278e86f34f9bd2fe34 Mon Sep 17 00:00:00 2001
From: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
Date: Thu, 14 Dec 2023 08:37:58 +0100
Subject: [PATCH] fix: handle cell timeouts

Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
---
 nbrunner/nb_runner.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/nbrunner/nb_runner.py b/nbrunner/nb_runner.py
index 137357d..0d025d8 100644
--- a/nbrunner/nb_runner.py
+++ b/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