Add tests for multiple notebooks case

tests/app/preheat_activation_test.py

@@ -11,7 +11,7 @@ def preheat_mode():
 
 @pytest.fixture
 def voila_notebook(notebook_directory):
-    return os.path.join(notebook_directory, 'pre_heat.ipynb')
+    return os.path.join(notebook_directory, 'preheat', 'pre_heat.ipynb')
 
 
 NOTEBOOK_EXECUTION_TIME = 2

tests/app/preheat_configuration_test.py

@@ -21,7 +21,7 @@ def preheat_mode():
 
 @pytest.fixture
 def voila_notebook(notebook_directory):
-    return os.path.join(notebook_directory, 'pre_heat.ipynb')
+    return os.path.join(notebook_directory, 'preheat', 'pre_heat.ipynb')
 
 
 async def send_request(sc, url, wait=0):

tests/app/preheat_multiple_notebooks_test.py

@@ -0,0 +1,49 @@
+import pytest
+import time
+import asyncio
+import os
+
+BASE_DIR = os.path.dirname(__file__)
+NOTEBOOK_EXECUTION_TIME = 2
+NUMBER_PREHEATED_KERNEL = 2
+
+
+@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 preheat_mode():
+    return True
+
+
+@pytest.fixture
+def voila_notebook(notebook_directory):
+    return os.path.join(notebook_directory, 'preheat')
+
+
+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_render_notebook_with_heated_kernel(http_server_client, base_url):
+    await asyncio.sleep(NUMBER_PREHEATED_KERNEL*NOTEBOOK_EXECUTION_TIME + 1)
+    time, text = await send_request(sc=http_server_client, url=f'{base_url}voila/render/pre_heat.ipynb')
+
+    assert 'hello world' in text
+    assert time < 0.5
+
+
+async def test_render_blacklisted_notebook_with_nornal_kernel(http_server_client, base_url):
+    await asyncio.sleep(NUMBER_PREHEATED_KERNEL*NOTEBOOK_EXECUTION_TIME + 1)
+    time, text = await send_request(sc=http_server_client, url=f'{base_url}voila/render/blacklisted.ipynb')
+
+    assert 'hello world' in text
+    assert time > 0.5

tests/configs/preheat/voila.json

@@ -3,6 +3,7 @@
     "preheat_kernel": true
   },
   "VoilaKernelManager": {
+    "preheat_blacklist": ["blacklisted.ipynb"],
     "kernel_pools_size": {
       "default": {
         "kernel": "python3",

tests/notebooks/preheat/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": null,
+   "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
+}

voila/voila_kernel_manager.py

@@ -13,9 +13,9 @@
 from typing import Awaitable, Coroutine, Type, TypeVar, Union
 from typing import Dict as tDict
 from pathlib import Path
-from traitlets.traitlets import Bool, Dict, Float
+from traitlets.traitlets import Bool, Dict, Float, List
 from nbclient.util import ensure_async
-
+import re
 from .notebook_renderer import NotebookRenderer
 
 T = TypeVar('T')
@@ -77,6 +77,10 @@ class VoilaKernelManager(base_class):
                 help='Mapping from notebook name to the number of started kernels to keep on standby.',
             )
 
+            preheat_blacklist = List([],
+                                     config=True,
+                                     help='List of notebooks which do not use pre-heated kernel.')
+
             fill_delay = Float(
                 1,
                 config=True,
@@ -95,16 +99,18 @@ def __init__(self, **kwargs):
                 self.notebook_html: tDict = {}
                 self._pools: tDict[str, Union[str, Coroutine[str]]] = {}
                 self.root_dir = self.parent.root_dir
-                self.notebook_path = os.path.relpath(
+                print(self.parent.notebook_path, self.root_dir)
+                if self.parent.notebook_path is not None:
+                    self.notebook_path = os.path.relpath(
                         self.parent.notebook_path, self.root_dir
-                    )  
-                if self.notebook_path is not None:
+                    )
                     self.fill_if_needed(delay=0, notebook_name=self.notebook_path)
                 else:
+                    self.notebook_path = None
                     all_notebooks = [
                         x.relative_to(self.root_dir)
                         for x in list(Path(self.root_dir).rglob('*.ipynb'))
-                        if '.ipynb_checkpoints' not in str(x)
+                        if self._notebook_filter(x)
                     ]
                     for nb in all_notebooks:
                         self.fill_if_needed(delay=0, notebook_name=str(nb))
@@ -125,7 +131,7 @@ async def start_kernel(
                     and len(self._pools.get(need_refill, ())) > 0
                 ):
                     kernel_id = await self._pop_pooled_kernel(need_refill, **kwargs)
-                    self.log.info('Using pre-heated kernel: %s', kernel_id)
+                    self.log.info('Using pre-heated kernel: %s for %s', kernel_id, need_refill)
                     self.fill_if_needed(delay=None, notebook_name=need_refill, **kwargs)
                 else:
                     kernel_id = await super().start_kernel(
@@ -195,13 +201,24 @@ def fill_if_needed(
                         kernel_env[key] = self.kernel_env_variables[key]
                 kwargs['env'] = kernel_env
 
+                unheated = kernel_size
+
+                def task_counter(tk):
+                    nonlocal unheated
+                    unheated -= 1
+                    if (unheated == 0):
+                        self.log.info(
+                            'Pre-heated %s kernel(s) for notebook %s', kernel_size, notebook_name
+                        )
+
                 for _ in range(kernel_size - len(pool)):
                     fut = super().start_kernel(kernel_name=kernel_name, **kwargs)
                     # Start the work on the loop immediately, so it is ready when needed:
                     task = loop.create_task(
                         wait_before(delay, self._initialize(fut, notebook_name))
                     )
                     pool.append(task)
+                    task.add_done_callback(task_counter)
 
             async def restart_kernel(self, kernel_id: str, **kwargs) -> None:
                 await ensure_async(super().restart_kernel(kernel_id, **kwargs))
@@ -244,7 +261,7 @@ async def shutdown_all(self, *args, **kwargs):
                                 )
 
             async def _initialize(
-                self, kernel_id_future: str, notebook_path: str) -> str:
+                    self, kernel_id_future: str, notebook_path: str) -> str:
                 """Run any configured initialization code in the kernel"""
                 kernel_id = await kernel_id_future
                 gen = self._notebook_renderer_factory(notebook_path)
@@ -272,17 +289,6 @@ async def _initialize(
                     'template': gen.template_name,
                     'theme': gen.theme,
                 }
-                heated_count = len(
-                    [
-                        val
-                        for val in self.notebook_html.values()
-                        if val[0] == notebook_path
-                    ]
-                )
-                self.log.info(
-                    'Pre-heated %s kernel for notebook %s', heated_count, notebook_path
-                )
-
                 return kernel_id
 
             async def cull_kernel_if_idle(self, kernel_id: str):
@@ -319,4 +325,14 @@ def _notebook_renderer_factory(
                     kernel_spec_manager=self.parent.kernel_spec_manager,
                 )
 
+            def _notebook_filter(self, nb_path: Path) -> bool:
+                nb_name = str(nb_path)
+                if '.ipynb_checkpoints' in nb_name:
+                    return False
+                for nb_pattern in self.preheat_blacklist:
+                    pattern = re.compile(nb_pattern)
+                    if (nb_pattern in nb_name) or bool(pattern.match(nb_name)):
+                        return False
+                return True
+
     return VoilaKernelManager