Improve startup, error handling and shutdown of Jupyter kernels (#4364)

panel/io/jupyter_executor.py

@@ -23,7 +23,6 @@
 from bokeh.server.views.ws import WSHandler
 from bokeh.util.token import get_session_id, get_token_payload
 from ipykernel.comm import Comm
-from IPython.display import HTML, publish_display_data
 
 from ..util import edit_readonly
 from .resources import Resources
@@ -38,6 +37,14 @@ class _RequestProxy:
     cookies: Dict[str, str]
     headers: Dict[str, str | List[str]]
 
+class Mimebundle:
+
+    def __init__(self, mimebundle):
+        self._mimebundle = mimebundle
+
+    def _repr_mimebundle_(self, include=None, exclude=None):
+        return self._mimebundle, {}
+
 
 class PanelExecutor(WSHandler):
     """
@@ -66,8 +73,11 @@ def __init__(self, path, token, root_url):
             path_versioner=StaticHandler.append_version
         )
         self._set_state()
-        self.session = self._create_server_session()
-        self.connection = ServerConnection(self.protocol, self, None, self.session)
+        try:
+            self.session = self._create_server_session()
+            self.connection = ServerConnection(self.protocol, self, None, self.session)
+        except Exception:
+            self.session = None
 
     def _get_payload(self, token: str) -> Dict[str, Any]:
         payload = get_token_payload(token)
@@ -145,7 +155,6 @@ def _create_server_session(self) -> ServerSession:
         session_context._set_session(session)
         return session
 
-
     async def write_message(
         self, message: Union[bytes, str, Dict[str, Any]],
         binary: bool = False, locked: bool = True
@@ -156,11 +165,13 @@ async def write_message(
         else:
             self.comm.send(message, metadata=metadata)
 
-    def render(self) -> HTML:
+    def render(self) -> Mimebundle:
         """
         Renders the application to an IPython.display.HTML object to
         be served by the `PanelJupyterHandler`.
         """
+        if self.session is None:
+            return Mimebundle({'text/error': 'Session did not start correctly'})
         with set_curdoc(self.session.document):
             html = server_html_page_for_session(
                 self.session,
@@ -169,5 +180,4 @@ def render(self) -> HTML:
                 template=self.session.document.template,
                 template_variables=self.session.document.template_variables
             )
-        publish_display_data({'application/bokeh-extensions': extension_dirs})
-        return HTML(html)
+        return Mimebundle({'text/html': html, 'application/bokeh-extensions': extension_dirs})

panel/io/jupyter_server_extension.py

@@ -169,17 +169,21 @@ async def _get_info(self, msg_id, timeout=KERNEL_TIMEOUT):
                 msg = await ensure_async(self.kernel.iopub_channel.get_msg(timeout=None))
             except Empty:
                 if not await ensure_async(self.kernel.is_alive()):
-                    raise RuntimeError("Kernel died before establishing Comm connection to Panel app.")
+                    raise RuntimeError("Kernel died before establishing Comm connection to Panel application.")
                 continue
             if msg['parent_header'].get('msg_id') != msg_id:
                 continue
             msg_type = msg['header']['msg_type']
-            if msg_type == 'display_data' and 'application/bokeh-extensions' in msg['content']['data']:
-                extension_dirs = msg['content']['data']['application/bokeh-extensions']
-            elif msg_type == 'execute_result':
-                result = msg
+            if msg_type == 'execute_result':
+                data = msg['content']['data']
+                if 'text/error' in data:
+                    raise RuntimeError("Panel application errored during startup.")
+                extension_dirs = data['application/bokeh-extensions']
+                result = data['text/html']
             elif msg_type == 'comm_open' and msg['content']['target_name'] == self.session_id:
                 comm_id = msg['content']['comm_id']
+            elif msg_type == 'stream' and msg['content']['name'] == 'stderr':
+                logger.info(msg['content']['text'])
         return result, comm_id, extension_dirs
 
     @tornado.web.authenticated
@@ -262,7 +266,7 @@ async def get(self, path=None):
 
         # Wait for comm to open and rendered HTML to be returned by the kernel
         try:
-            msg, comm_id, ext_dirs = await self._get_info(msg_id)
+            html, comm_id, ext_dirs = await self._get_info(msg_id)
         except (TimeoutError, RuntimeError) as e:
             await self.kernel_manager.shutdown_kernel(kernel_id, now=True)
             html = ERROR_TEMPLATE.render(
@@ -280,7 +284,6 @@ async def get(self, path=None):
             state._kernels[self.session_id] = (self.kernel, comm_id, kernel_id, False)
             loop = tornado.ioloop.IOLoop.current()
             loop.call_later(CONNECTION_TIMEOUT, self._check_connected)
-            html = msg['content']['data']['text/html']
             self.finish(html)
 
     async def _check_connected(self):
@@ -291,7 +294,6 @@ async def _check_connected(self):
             await self.kernel_manager.shutdown_kernel(kernel_id, now=True)
 
 
-
 class PanelWSProxy(WSHandler, JupyterHandler):
     """
     The PanelWSProxy serves as a proxy between the frontend and the
@@ -414,7 +416,9 @@ def on_close(self) -> None:
         if self.session_id in state._kernels:
             del state._kernels[self.session_id]
         self._ping_job.stop()
+        reply = self.kernel.shutdown(reply=True)
         future = self.kernel_manager.shutdown_kernel(self.kernel_id, now=True)
+        asyncio.ensure_future(reply)
         asyncio.ensure_future(future)
         self.kernel = None