Fix wait_for_request (#1125)

Author: trungleduc

docs/source/customize.rst

@@ -410,7 +410,7 @@ In preheating kernel mode, users can prepend with ``wait_for_request`` from ``vo
 
 ``wait_for_request`` will pause the execution of the notebook in the preheated kernel at this cell and wait for an actual user to connect to Voilà, set the request info environment variables and then continue the execution of the remaining cells.
 
-If the Voilà websocket handler is not started with the default protocol (`ws`), the default IP address (`127.0.0.1`) or the default port (`8866`), users need to provide these values through the environment variables ``VOILA_APP_PROTOCOL``, ``VOILA_APP_IP`` and ``VOILA_APP_PORT``. The easiest way is to set these variables in the `voila.json` configuration file, for example:
+If the Voilà websocket handler is not started with the default protocol (`ws`), the default IP address (`127.0.0.1`) the default port (`8866`) or with url suffix, users need to provide these values through the environment variables ``VOILA_WS_PROTOCOL``, ``VOILA_APP_IP``, ``VOILA_APP_PORT`` and ``VOILA_WS_BASE_URL``. The easiest way is to set these variables in the `voila.json` configuration file, for example:
 
 .. code-block:: python
 
@@ -423,7 +423,7 @@ If the Voilà websocket handler is not started with the default protocol (`ws`),
                "kernel_env_variables": { 
                   "VOILA_APP_IP": "192.168.1.1",
                   "VOILA_APP_PORT": "6789",
-                  "VOILA_APP_PROTOCOL": "wss"
+                  "VOILA_WS_PROTOCOL": "wss"
                }
             }
          },

voila/utils.py

@@ -11,6 +11,7 @@
 import json
 import os
 import threading
+import warnings
 from enum import Enum
 from functools import partial
 from typing import Awaitable, Dict
@@ -30,7 +31,8 @@ class ENV_VARIABLE(str, Enum):
     VOILA_BASE_URL = 'VOILA_BASE_URL'
     VOILA_APP_IP = 'VOILA_APP_IP'
     VOILA_APP_PORT = 'VOILA_APP_PORT'
-    VOILA_APP_PROTOCOL = 'VOILA_APP_PROTOCOL'
+    VOILA_WS_PROTOCOL = 'VOILA_WS_PROTOCOL'
+    VOILA_WS_BASE_URL = 'VOILA_WS_BASE_URL'
     SERVER_NAME = 'SERVER_NAME'
     SERVER_PORT = 'SERVER_PORT'
     SCRIPT_NAME = 'SCRIPT_NAME'
@@ -57,9 +59,14 @@ def get_server_root_dir(settings):
 
 
 async def _get_request_info(ws_url: str) -> Awaitable:
-    async with websockets.connect(ws_url) as websocket:
-        ri = await websocket.recv()
-    return ri
+    try:
+        async with websockets.connect(ws_url, open_timeout=5) as websocket:
+            ri = await websocket.recv()
+    except (TimeoutError, ConnectionRefusedError):
+        warnings.warn(f'Failed to connect to {ws_url}')
+        return None
+    else:
+        return ri
 
 
 def wait_for_request(url: str = None) -> str:
@@ -73,17 +80,16 @@ def wait_for_request(url: str = None) -> str:
         Defaults to None.
 
     """
-
     preheat_mode = os.getenv(ENV_VARIABLE.VOILA_PREHEAT, 'False')
     if preheat_mode == 'False':
         return
 
     request_info = None
     if url is None:
-        protocol = os.getenv(ENV_VARIABLE.VOILA_APP_PROTOCOL, 'ws')
+        protocol = os.getenv(ENV_VARIABLE.VOILA_WS_PROTOCOL, 'ws')
         server_ip = os.getenv(ENV_VARIABLE.VOILA_APP_IP, '127.0.0.1')
         server_port = os.getenv(ENV_VARIABLE.VOILA_APP_PORT, '8866')
-        base_url = os.getenv(ENV_VARIABLE.VOILA_BASE_URL, '/')
+        base_url = os.getenv(ENV_VARIABLE.VOILA_WS_BASE_URL, '/')
         url = f'{protocol}://{server_ip}:{server_port}{base_url}voila/query'
 
     kernel_id = os.getenv(ENV_VARIABLE.VOILA_KERNEL_ID)
@@ -101,8 +107,9 @@ def inner():
     except (KeyboardInterrupt, SystemExit):
         asyncio.get_event_loop().stop()
 
-    for k, v in json.loads(request_info).items():
-        os.environ[k] = v
+    if request_info is not None:
+        for k, v in json.loads(request_info).items():
+            os.environ[k] = v
 
 
 def get_query_string(url: str = None) -> str:

voila/voila_kernel_manager.py

@@ -175,13 +175,16 @@ def fill_if_needed(
                 except RuntimeError:
                     loop = asyncio.new_event_loop()
                     asyncio.set_event_loop(loop)
-                if notebook_name in self.kernel_pools_config:
-                    kernel_size = self.kernel_pools_config[notebook_name].get(
-                        'pool_size', 1
-                    )
-                else:
-                    default_config = self.kernel_pools_config.get('default', {})
-                    kernel_size = default_config.get('pool_size', 1)
+                default_config: dict = self.kernel_pools_config.get('default', {})
+                notebook_config: dict = self.kernel_pools_config.get(
+                    notebook_name, default_config
+                )
+                kernel_env_variables: dict = notebook_config.get(
+                    'kernel_env_variables', default_config.get('kernel_env_variables', {})
+                )
+                kernel_size: int = notebook_config.get(
+                    'pool_size', default_config.get('pool_size', 1)
+                )
                 pool = self._pools.get(notebook_name, [])
                 self._pools[notebook_name] = pool
                 if 'path' not in kwargs:
@@ -193,7 +196,7 @@ def fill_if_needed(
                 kernel_env = os.environ.copy()
                 kernel_env_arg = kwargs.get('env', {})
                 kernel_env.update(kernel_env_arg)
-                kernel_env_variables = self.kernel_pools_config.get(notebook_name, {}).get('kernel_env_variables', {})
+
                 for key in kernel_env_variables:
                     if key not in kernel_env:
                         kernel_env[key] = kernel_env_variables[key]