From b475f0e25660a26fcff70700ac33e3e24f1c3726 Mon Sep 17 00:00:00 2001
From: Akshay Chitneni <akshayadatta@gmail.com>
Date: Tue, 24 Oct 2023 13:16:39 -0700
Subject: [PATCH] Ability to configure cull_idle_timeout with kernelSpec
 (#1342)

Co-authored-by: Akshay Chitneni <achitneni@apple.com>
---
 .../services/kernels/kernelmanager.py         |  5 ++-
 tests/services/kernels/test_cull.py           | 39 +++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/jupyter_server/services/kernels/kernelmanager.py b/jupyter_server/services/kernels/kernelmanager.py
index 312abeb73a..28b0c6e45f 100644
--- a/jupyter_server/services/kernels/kernelmanager.py
+++ b/jupyter_server/services/kernels/kernelmanager.py
@@ -645,6 +645,9 @@ async def cull_kernel_if_idle(self, kernel_id):
             await ensure_async(self.shutdown_kernel(kernel_id))
             return
 
+        kernel_spec_metadata = kernel.kernel_spec.metadata
+        cull_idle_timeout = kernel_spec_metadata.get("cull_idle_timeout", self.cull_idle_timeout)
+
         if hasattr(
             kernel, "last_activity"
         ):  # last_activity is monkey-patched, so ensure that has occurred
@@ -657,7 +660,7 @@ async def cull_kernel_if_idle(self, kernel_id):
             dt_now = utcnow()
             dt_idle = dt_now - kernel.last_activity
             # Compute idle properties
-            is_idle_time = dt_idle > timedelta(seconds=self.cull_idle_timeout)
+            is_idle_time = dt_idle > timedelta(seconds=cull_idle_timeout)
             is_idle_execute = self.cull_busy or (kernel.execution_state != "busy")
             connections = self._kernel_connections.get(kernel_id, 0)
             is_idle_connected = self.cull_connected or not connections
diff --git a/tests/services/kernels/test_cull.py b/tests/services/kernels/test_cull.py
index 88402a1e04..50ecbf2b96 100644
--- a/tests/services/kernels/test_cull.py
+++ b/tests/services/kernels/test_cull.py
@@ -12,6 +12,12 @@
 CULL_TIMEOUT = 30 if platform.python_implementation() == "PyPy" else 5
 CULL_INTERVAL = 1
 
+sample_kernel_json_with_metadata = {
+    "argv": ["cat", "{connection_file}"],
+    "display_name": "Test kernel",
+    "metadata": {"cull_idle_timeout": 0},
+}
+
 
 @pytest.fixture(autouse=True)
 def suppress_deprecation_warnings():
@@ -24,6 +30,21 @@ def suppress_deprecation_warnings():
         yield
 
 
+@pytest.fixture
+def jp_kernelspec_with_metadata(jp_data_dir):
+    """Configures some sample kernelspecs in the Jupyter data directory."""
+    kenrel_spec_name = "sample_with_metadata"
+    sample_kernel_dir = jp_data_dir.joinpath("kernels", kenrel_spec_name)
+    sample_kernel_dir.mkdir(parents=True)
+    # Create kernel json file
+    sample_kernel_file = sample_kernel_dir.joinpath("kernel.json")
+    kernel_json = sample_kernel_json_with_metadata.copy()
+    sample_kernel_file.write_text(json.dumps(kernel_json))
+    # Create resources text
+    sample_kernel_resources = sample_kernel_dir.joinpath("resource.txt")
+    sample_kernel_resources.write_text("resource")
+
+
 @pytest.mark.parametrize(
     "jp_server_config",
     [
@@ -73,6 +94,24 @@ async def test_cull_idle(jp_fetch, jp_ws_fetch):
     assert culled
 
 
+async def test_cull_idle_disable(jp_fetch, jp_ws_fetch, jp_kernelspec_with_metadata):
+    r = await jp_fetch("api", "kernels", method="POST", allow_nonstandard_methods=True)
+    kernel = json.loads(r.body.decode())
+    kid = kernel["id"]
+
+    # Open a websocket connection.
+    ws = await jp_ws_fetch("api", "kernels", kid, "channels")
+
+    r = await jp_fetch("api", "kernels", kid, method="GET")
+    model = json.loads(r.body.decode())
+    assert model["connections"] == 1
+    culled = await get_cull_status(kid, jp_fetch)  # connected, should not be culled
+    assert not culled
+    ws.close()
+    culled = await get_cull_status(kid, jp_fetch)  # not connected, should not be culled
+    assert not culled
+
+
 # Pending kernels was released in Jupyter Client 7.1
 # It is currently broken on Windows (Jan 2022). When fixed, we can remove the Windows check.
 # See https://github.com/jupyter-server/jupyter_server/issues/672