Port voila-dashboards#969 to this PR

trungleduc · Sep 27, 2021 · 5c488cd · 5c488cd
1 parent c69c70b
commit 5c488cd
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/voila/notebook_renderer.py b/voila/notebook_renderer.py
@@ -214,6 +214,8 @@ async def _jinja_notebook_execute(self, nb, kernel_id):
         #  (it seems to be local to our block)
         nb.cells = result.cells
 
+        await self._cleanup_resources()
+
     async def _jinja_cell_generator(self, nb, kernel_id, timeout_callback):
         """Generator that will execute a single notebook cell at a time"""
         nb, _ = ClearOutputPreprocessor().preprocess(
@@ -279,6 +281,12 @@ async def _jinja_cell_generator(self, nb, kernel_id, timeout_callback):
             finally:
                 yield output_cell
 
+        await self._cleanup_resources()
+
+    async def _cleanup_resources(self):
+        await ensure_async(self.executor.km.cleanup_resources())
+        await ensure_async(self.executor.kc.stop_channels())
+
     async def load_notebook(self, path):
 
         model = await ensure_async(self.contents_manager.get(path=path))

diff --git a/voila/voila_kernel_manager.py b/voila/voila_kernel_manager.py
@@ -186,7 +186,6 @@ def fill_if_needed(
                     default_config = self.kernel_pools_size.get('default', {})
                     kernel_name = default_config.get('kernel', 'python3')
                     kernel_size = default_config.get('pool_size', 1)
-                self.log.info(f'Pre-heat {kernel_name} kernel for {notebook_name}')
                 pool = self._pools.get(notebook_name, [])
                 self._pools[notebook_name] = pool
                 if 'path' not in kwargs:
@@ -201,14 +200,14 @@ def fill_if_needed(
                         kernel_env[key] = self.kernel_env_variables[key]
                 kwargs['env'] = kernel_env
 
-                unheated = kernel_size
+                heated = len(pool)
 
                 def task_counter(tk):
-                    nonlocal unheated
-                    unheated -= 1
-                    if (unheated == 0):
+                    nonlocal heated
+                    heated += 1
+                    if (heated == kernel_size):
                         self.log.info(
-                            'Pre-heated %s kernel(s) for notebook %s', kernel_size, notebook_name
+                            'Kernel pool of %s is filled with %s kernel(s)', notebook_name, kernel_size
                         )
 
                 for _ in range(kernel_size - len(pool)):