Merge pull request #262 from lsst-sqre/tickets/DM-39325

DM-39325: Make max WebSocket message size configurable

changelog.d/20230522_093320_rra_DM_39325.md

@@ -0,0 +1,3 @@
+### New features
+
+- The maximum allowable size for a WebSocket message from the Jupyter lab is now configurable per business and defaults to 10MB instead of 4MB.

src/mobu/constants.py

@@ -9,6 +9,7 @@
     "NOTEBOOK_REPO_BRANCH",
     "TOKEN_LIFETIME",
     "USERNAME_REGEX",
+    "WEBSOCKET_OPEN_TIMEOUT",
 ]
 
 NOTEBOOK_REPO_URL = "https://github.com/lsst-sqre/notebook-demo.git"
@@ -26,13 +27,6 @@
 forever.
 """
 
-WEBSOCKET_MESSAGE_SIZE_LIMIT = 4 * 1024 * 1024
-"""Largest WebSocket message size allowed from lab (in bytes).
-
-This has to be large enough to hold HTML and image output from executing
-notebook cells, even though we discard that data. Set to `None` for no limit.
-"""
-
 WEBSOCKET_OPEN_TIMEOUT = 60
 """How long to wait for a WebSocket connection to open (in seconds)."""
 

src/mobu/models/business/nublado.py

@@ -184,6 +184,16 @@ class NubladoBusinessOptions(BusinessOptions):
         example=60,
     )
 
+    max_websocket_message_size: int | None = Field(
+        10 * 1024 * 1024,
+        title="Maximum length of WebSocket message (in bytes)",
+        description=(
+            "This has to be large enough to hold HTML and image output from"
+            " executing notebook cells, even though we discard that data."
+            " Set to `null` for no limit."
+        ),
+    )
+
     spawn_settle_time: int = Field(
         10,
         title="How long to wait before polling spawn progress in seconds",

src/mobu/storage/jupyter.py

@@ -27,7 +27,7 @@
 from websockets.exceptions import WebSocketException
 
 from ..config import config
-from ..constants import WEBSOCKET_MESSAGE_SIZE_LIMIT, WEBSOCKET_OPEN_TIMEOUT
+from ..constants import WEBSOCKET_OPEN_TIMEOUT
 from ..exceptions import (
     CodeExecutionError,
     JupyterTimeoutError,
@@ -167,13 +167,15 @@ def __init__(
         jupyter_url: str,
         kernel_name: str = "LSST",
         notebook_name: str | None = None,
+        max_websocket_size: int | None,
         http_client: AsyncClient,
         logger: BoundLogger,
     ) -> None:
         self._username = username
         self._jupyter_url = jupyter_url
         self._kernel_name = kernel_name
         self._notebook_name = notebook_name
+        self._max_websocket_size = max_websocket_size
         self._client = http_client
         self._logger = logger
 
@@ -242,7 +244,7 @@ async def __aenter__(self) -> Self:
                 self._url_for_websocket(url),
                 extra_headers=headers,
                 open_timeout=WEBSOCKET_OPEN_TIMEOUT,
-                max_size=WEBSOCKET_MESSAGE_SIZE_LIMIT,
+                max_size=self._max_websocket_size,
             ).__aenter__()
         except WebSocketException as e:
             user = self._username
@@ -628,7 +630,11 @@ async def is_lab_stopped(self, *, log_running: bool = False) -> bool:
         return result
 
     def open_lab_session(
-        self, notebook_name: str | None = None, *, kernel_name: str = "LSST"
+        self,
+        notebook_name: str | None = None,
+        *,
+        max_websocket_size: int | None = None,
+        kernel_name: str = "LSST",
     ) -> JupyterLabSession:
         """Open a Jupyter lab session.
 
@@ -643,6 +649,8 @@ def open_lab_session(
             session and might influence logging on the lab side. If set, the
             session type will be set to ``notebook``. If not set, the session
             type will be set to ``console``.
+        max_websocket_size
+            Maximum size of a WebSocket message, or `None` for no limit.
         kernel_name
             Name of the kernel to use for the session.
 
@@ -656,6 +664,7 @@ def open_lab_session(
             jupyter_url=self._jupyter_url,
             kernel_name=kernel_name,
             notebook_name=notebook_name,
+            max_websocket_size=max_websocket_size,
             http_client=self._client,
             logger=self._logger,
         )