Add tests for preheated mode

voila-dashboards · Oct 4, 2021 · 75f7fd3 · 75f7fd3
1 parent 73fcea8
commit 75f7fd3
Show file tree

Hide file tree

Showing 7 changed files with 231 additions and 5 deletions.
diff --git a/tests/app/preheat_activation_test.py b/tests/app/preheat_activation_test.py
@@ -0,0 +1,103 @@
+import pytest
+import time
+import asyncio
+
+
+@pytest.fixture
+def voila_args(notebook_directory, voila_args_extra):
+    return [
+        f"{notebook_directory}/pre_heat.ipynb",
+        "--VoilaTest.log_level=INFO",
+        "--VoilaConfiguration.preheat_kernel=true",
+    ] + voila_args_extra
+
+
+NOTEBOOK_EXECUTION_TIME = 2
+
+
+async def send_request(sc, url, wait=0):
+    await asyncio.sleep(wait)
+    real_time = time.time()
+    response = await sc.fetch(url)
+    real_time = time.time() - real_time
+    html_text = response.body.decode("utf-8")
+    return real_time, html_text
+
+
+async def test_request_before_kernel_heated(http_server_client, base_url):
+    """
+    We send a request to server immediately, a normal kernel
+    will be used so it should take more than `NOTEBOOK_EXECUTION_TIME`
+    to display result.
+    """
+    time, text = await send_request(sc=http_server_client, url=base_url)
+    assert "hello world" in text
+    assert time > NOTEBOOK_EXECUTION_TIME
+
+
+async def test_render_time_with_preheated_kernel(http_server_client, base_url):
+    """
+    We wait for kernel properly heated before sending request,
+    it should take under 0.5s to return result
+    """
+    time, text = await send_request(
+        sc=http_server_client, url=base_url, wait=NOTEBOOK_EXECUTION_TIME + 2
+    )
+    assert "hello world" in text
+    assert time < 0.5
+
+
+async def test_render_time_with_multiple_requests(http_server_client, base_url):
+    """
+    We send a second request just after the first one, so the
+    pool is not filled yet and a normal kernel is used instead
+    """
+    time_list = []
+    for wait in [NOTEBOOK_EXECUTION_TIME + 1, 0]:
+
+        time, _ = await send_request(sc=http_server_client, url=base_url, wait=wait)
+        time_list.append(time)
+
+    assert max(time_list) > 0.5  # Render time for a normal kernel
+    assert min(time_list) < 0.5  # Render time for a preheated kernel
+
+
+async def test_request_with_query(http_server_client, base_url):
+    """
+    We sent request with query parameter, preheat kernel should
+    be disable is this case.
+    """
+    url = f"{base_url}?foo=bar"
+    time, _ = await send_request(
+        sc=http_server_client, url=url, wait=NOTEBOOK_EXECUTION_TIME + 1
+    )
+    assert time > 0.5
+
+
+async def test_request_with_theme_parameter(http_server_client, base_url):
+    """
+    We sent request with theme parameter, preheat kernel should
+    be used if requested theme is same as theme used for prerendered.
+    """
+    wait = NOTEBOOK_EXECUTION_TIME + 1
+
+    url = f"{base_url}?voila-theme=light"
+    time, _ = await send_request(sc=http_server_client, url=url, wait=wait)
+    assert time < 0.5
+
+    url = f"{base_url}?voila-theme=dark"
+    time, _ = await send_request(sc=http_server_client, url=url, wait=wait)
+    assert time > 0.5
+
+
+async def test_request_with_template_parameter(http_server_client, base_url):
+    """
+    We sent request with theme parameter, preheat kernel should
+    be used if requested theme is same as theme used for prerendered.
+    """
+    wait = NOTEBOOK_EXECUTION_TIME + 1
+
+    url = f"{base_url}?voila-template=lab"
+    time, _ = await send_request(sc=http_server_client, url=url, wait=wait)
+    assert time < 0.5
+
diff --git a/tests/app/preheat_configuration_test.py b/tests/app/preheat_configuration_test.py
@@ -0,0 +1,53 @@
+import pytest
+import time
+import asyncio
+import os
+
+BASE_DIR = os.path.dirname(__file__)
+
+@pytest.fixture
+def voila_config_file_paths_arg():
+    path = os.path.join(BASE_DIR, '..', 'configs', 'preheat')
+    return '--VoilaTest.config_file_paths=[%r]' % path
+
+
+@pytest.fixture
+def voila_args(notebook_directory, voila_args_extra, voila_config_file_paths_arg):
+    return [
+        f"{notebook_directory}/pre_heat.ipynb",
+        "--VoilaTest.log_level=INFO", voila_config_file_paths_arg
+    ] + voila_args_extra 
+
+
+NOTEBOOK_EXECUTION_TIME = 2
+NUMBER_PREHEATED_KERNEL = 2
+
+async def send_request(sc, url, wait=0):
+    await asyncio.sleep(wait)
+    real_time = time.time()
+    response = await sc.fetch(url)
+    real_time = time.time() - real_time
+    html_text = response.body.decode("utf-8")
+    return real_time, html_text
+
+
+async def test_refill_kernel_asynchronously(http_server_client, base_url):
+    await asyncio.sleep(NUMBER_PREHEATED_KERNEL*NOTEBOOK_EXECUTION_TIME + 1)
+    fast = []
+    slow = []
+    for i in range(5*NUMBER_PREHEATED_KERNEL):
+        time, _ = await send_request(sc=http_server_client, url=base_url)
+        if time < 0.5:
+            fast.append(time)
+        else:
+            slow.append(time)
+
+    assert len(fast) == 7
+    assert len(slow) == 3
+
+
+async def test_env_variable_defined_in_kernel(http_server_client, base_url):
+    await asyncio.sleep(NUMBER_PREHEATED_KERNEL*NOTEBOOK_EXECUTION_TIME + 1)
+    _, text = await send_request(sc=http_server_client, url=base_url)
+    assert "bar" in text
+
diff --git a/tests/configs/preheat/voila.json b/tests/configs/preheat/voila.json
@@ -0,0 +1,14 @@
+{
+  "VoilaConfiguration": {
+    "preheat_kernel": true
+  },
+  "VoilaKernelManager": {
+    "kernel_pools_size": {
+      "python3": 2
+    },
+    "fill_delay": 0,
+    "kernel_env_variables": {
+      "foo": "bar"
+    }
+  }
+}
diff --git a/tests/notebooks/pre_heat.ipynb b/tests/notebooks/pre_heat.ipynb
@@ -0,0 +1,56 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7efe3e8f-4b8d-4fd3-99d1-b06d79b88ae2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "time.sleep(2)\n",
+    "print('hello world')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "782377cd-2bc0-4d0d-8b63-74dcb7e9d645",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.getenv('foo')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0e689ec5-708c-4cac-98ba-02b00411e41d",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/voila/handler.py b/voila/handler.py
@@ -165,7 +165,6 @@ def should_use_rendered_notebook(
             return False
         if theme is not None and rendered_theme != theme:
             return False
-
         args_list = [
             key for key in request_args if key not in ['voila-template', 'voila-theme']
         ]

diff --git a/voila/notebook_renderer.py b/voila/notebook_renderer.py
@@ -30,7 +30,7 @@
 
 class NotebookRenderer(LoggingConfigurable):
     def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+        super().__init__()
         self.root_dir = kwargs.get('root_dir', [])
         self.notebook_path = kwargs.get('notebook_path', [])  # should it be []
         self.template_paths = kwargs.get('template_paths', [])

diff --git a/voila/voila_kernel_manager.py b/voila/voila_kernel_manager.py
@@ -52,7 +52,7 @@ class VoilaKernelManager(base_class):
             """
 
             kernel_pools_size = Dict(
-                {'python3': 3},
+                {'python3': 1},
                 config=True,
                 help='Mapping from kernel name to the number of started kernels to keep on standby.'
             )
@@ -167,8 +167,9 @@ async def shutdown_all(self, *args, **kwargs):
                             kid = await fut
                         except Exception:
                             pass
-                        if kid in self:
-                            await ensure_async(self.shutdown_kernel(kid, *args, **kwargs)) 
+                        else:
+                            if kid in self:
+                                await ensure_async(self.shutdown_kernel(kid, *args, **kwargs)) 
 
             async def _initialize(self, kernel_id_future: str) -> str:
                 """Run any configured initialization code in the kernel"""