From 9a8147a32a0c9197eaa8a21d77399dd71ebf1309 Mon Sep 17 00:00:00 2001
From: ClaytonAstrom <castromxiii@gmail.com>
Date: Wed, 25 Oct 2023 04:03:40 -0500
Subject: [PATCH] Allow for custom kernel_spec_manager class (#1404)

* Allow for custom kernel spec manager class

* Added startup_timeout config

---------

Co-authored-by: castrom <Clayton.Astrom@morningstar.com>
---
 docs/customize.md      | 18 ++++++++++++++++++
 voila/app.py           |  6 ++++--
 voila/configuration.py |  8 ++++++++
 voila/execute.py       | 14 +++++++++++++-
 4 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/docs/customize.md b/docs/customize.md
index d9cc7cc49..954ec5ade 100644
--- a/docs/customize.md
+++ b/docs/customize.md
@@ -639,3 +639,21 @@ By using the [layout customization system](https://jupyterlab.readthedocs.io/en/
    }
 }
 ```
+
+## Custom kernel_spec_manager class
+
+By default, Voilà uses `jupyter_client`'s KernelSpecManager. To change this class, add the following to your config like so:
+
+```py
+import CustomKernelSpecManager
+
+c.VoilaConfiguration.kernel_spec_manager_class = CustomKernelSpecManager
+```
+
+## Kernel startup_timeout
+
+By default, Voilà's grace period for kernel startup time is 60 seconds. It can be adjusted as such, in seconds
+
+```sh
+voila --VoilaExecutor.startup_timeout=60
+```
diff --git a/voila/app.py b/voila/app.py
index eb60b3059..78a49c5e3 100644
--- a/voila/app.py
+++ b/voila/app.py
@@ -35,7 +35,6 @@
 import jinja2
 import tornado.ioloop
 import tornado.web
-from jupyter_client.kernelspec import KernelSpecManager
 from jupyter_core.paths import jupyter_config_path, jupyter_path
 from jupyter_server.base.handlers import FileFindHandler, path_regex
 from jupyter_server.config_manager import recursive_update
@@ -169,6 +168,7 @@ class Voila(Application):
         "template": "VoilaConfiguration.template",
         "theme": "VoilaConfiguration.theme",
         "classic_tree": "VoilaConfiguration.classic_tree",
+        "kernel_spec_manager_class": "VoilaConfiguration.kernel_spec_manager_class",
     }
     classes = [VoilaConfiguration, VoilaExecutor, VoilaExporter]
     connection_dir_root = Unicode(
@@ -544,7 +544,9 @@ def init_settings(self) -> Dict:
         # default server_url to base_url
         self.server_url = self.server_url or self.base_url
 
-        self.kernel_spec_manager = KernelSpecManager(parent=self)
+        self.kernel_spec_manager = self.voila_configuration.kernel_spec_manager_class(
+            parent=self
+        )
 
         # we create a config manager that load both the serverconfig and nbconfig (classical notebook)
         read_config_path = [
diff --git a/voila/configuration.py b/voila/configuration.py
index 880359bb1..03ba6a44e 100644
--- a/voila/configuration.py
+++ b/voila/configuration.py
@@ -158,6 +158,14 @@ def _valid_file_blacklist(self, proposal):
         """,
     )
 
+    kernel_spec_manager_class = Type(
+        config=True,
+        default_value="jupyter_client.kernelspec.KernelSpecManager",
+        klass="jupyter_client.kernelspec.KernelSpecManager",
+        help="""The kernel spec manager class.  Allows for setting a custom kernel spec manager for finding and running kernels
+        """,
+    )
+
     http_header_envs = List(
         Unicode(),
         [],
diff --git a/voila/execute.py b/voila/execute.py
index 1e584116e..6071725ee 100644
--- a/voila/execute.py
+++ b/voila/execute.py
@@ -10,7 +10,7 @@
 from nbclient.client import NotebookClient
 from nbclient.exceptions import CellExecutionError
 from nbconvert.preprocessors.clearoutput import ClearOutputPreprocessor
-from traitlets import Bool, Unicode
+from traitlets import Bool, Unicode, Integer
 
 
 def strip_code_cell_warnings(cell):
@@ -50,6 +50,18 @@ class VoilaExecutor(NotebookClient):
         help=("Whether to send tracebacks to clients on exceptions."),
     )
 
+    startup_timeout: int = Integer(
+        60,
+        config=True,
+        help=(
+            """
+            The time to wait (in seconds) for the kernel to start.
+            If kernel startup takes longer, a RuntimeError is
+            raised.
+            """
+        ),
+    )
+
     def execute(self, nb, resources, km=None):
         try:
             result = super().execute()