diff --git a/changelog.d/20230522_093320_rra_DM_39325.md b/changelog.d/20230522_093320_rra_DM_39325.md new file mode 100644 index 00000000..ddfa7efd --- /dev/null +++ b/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. diff --git a/src/mobu/constants.py b/src/mobu/constants.py index e3a9fa0d..349e6eb8 100644 --- a/src/mobu/constants.py +++ b/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).""" diff --git a/src/mobu/models/business/nublado.py b/src/mobu/models/business/nublado.py index d0688ab3..d7849e55 100644 --- a/src/mobu/models/business/nublado.py +++ b/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", diff --git a/src/mobu/storage/jupyter.py b/src/mobu/storage/jupyter.py index 5f4d6bfd..537080ab 100644 --- a/src/mobu/storage/jupyter.py +++ b/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,6 +167,7 @@ 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: @@ -174,6 +175,7 @@ def __init__( 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, )