Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

DM-39325: Make max WebSocket message size configurable #262

Merged
merged 1 commit into from
May 22, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions changelog.d/20230522_093320_rra_DM_39325.md
Original file line number Diff line number Diff line change
@@ -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.
8 changes: 1 addition & 7 deletions src/mobu/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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)."""

Expand Down
10 changes: 10 additions & 0 deletions src/mobu/models/business/nublado.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
15 changes: 12 additions & 3 deletions src/mobu/storage/jupyter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand All @@ -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.

Expand All @@ -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,
)
Expand Down