diff --git a/.gitignore b/.gitignore index 11b5578254..a0a068107f 100644 --- a/.gitignore +++ b/.gitignore @@ -221,4 +221,8 @@ inference_testing openspec*/ opsx/ checkpoints/ -development/_private \ No newline at end of file +development/_private +# Local credential material — never commit. Both of these were found untracked +# but NOT ignored, so a single `git add -A` would have committed them. +*.key +modal/deploy_modal_*.sh diff --git a/inference/core/env.py b/inference/core/env.py index 62adfabaf1..4ca544b459 100644 --- a/inference/core/env.py +++ b/inference/core/env.py @@ -945,15 +945,52 @@ WEBEXEC_WS_CONNECT_TIMEOUT_SECONDS = int( os.getenv("WEBEXEC_WS_CONNECT_TIMEOUT_SECONDS", "30") ) -# Set slightly above the server's 700s execution budget (modal_app.py Executor -# timeout) so that when a block hits the server limit, the server's error frame -# arrives before the client read times out. Equal values race and surface an -# ambiguous "connection lost after send" instead of the real server error. +# Set above the Modal ``timeout`` (700s) that bounds a single input in +# modal_app.py. NOTE: there is no server-side per-execution timeout that sends +# an error frame at 700s — when Modal kills the input the connection simply +# dies, so this value only decides how long the client waits before reporting +# that death. Keeping it above 700s avoids racing the (much more common) case +# of a block that finishes just under the Modal budget. WEBEXEC_WS_READ_TIMEOUT_SECONDS = int( os.getenv("WEBEXEC_WS_READ_TIMEOUT_SECONDS", "720") ) +# When a reconnect lands on a container that does not know this executor's +# custom-Python session, runtime state built by earlier frames is gone. +# With this ENABLED the executor fails loudly instead of silently continuing +# with reset globals. +# +# Defaults to False, deliberately. The check cannot tell a stateful block from +# a stateless one: it arms on ANY successful execution, on an executor cached +# per workspace. Container-local Python state also cannot survive a long run by +# construction — a websocket connection is one Modal input capped at 700s, the +# server closes at WEBEXEC_WS_MAX_CONNECTION_SECONDS, namespaces are keyed by +# code hash, and reconnects have no container affinity. So enabling it by +# default imposes a scheduled hard failure on the stateless majority to detect +# a condition the platform guarantees for the stateful minority. Set it True to +# opt into the loud diagnostic where blocks genuinely rely on cross-frame +# globals and a failed run is preferable to a silently reset one. +# str2bool, like every other boolean in this file: it raises on a value that +# is neither "true" nor "false", so a typo ("1", "yes", "True ") fails at boot +# instead of silently resolving to False and disabling this safety net. +WEBEXEC_WS_FAIL_ON_SESSION_LOSS = str2bool( + os.getenv("WEBEXEC_WS_FAIL_ON_SESSION_LOSS", False) +) + WEBEXEC_WS_CONNECTION_POOL_SIZE = int(os.getenv("WEBEXEC_WS_CONNECTION_POOL_SIZE", "1")) +# How long an idle websocket connection keeps heartbeating before the client +# deliberately releases it (closing the socket lets the Modal container scale +# down instead of staying warm — and billed — for the whole connection cap). +# Releasing also discards the custom-Python session: runtime state mutated by +# earlier frames is gone after this much inactivity. +# Set to 0 (or any non-positive value) to disable idle release entirely, matching +# the convention of WEBEXEC_MODAL_EXECUTOR_IDLE_TTL_SECONDS below. +# Must stay well BELOW the server's WEBEXEC_WS_MAX_CONNECTION_SECONDS (600): +# connection age is >= client idle time by construction, so a value at or above +# the cap can never fire and the release path is dead code. +WEBEXEC_WS_IDLE_RELEASE_SECONDS = int( + os.getenv("WEBEXEC_WS_IDLE_RELEASE_SECONDS", "120") +) WEBEXEC_MODAL_EXECUTOR_IDLE_TTL_SECONDS = int( os.getenv("WEBEXEC_MODAL_EXECUTOR_IDLE_TTL_SECONDS", "1800") ) diff --git a/inference/core/interfaces/webrtc_worker/modal.py b/inference/core/interfaces/webrtc_worker/modal.py index eb6146e388..a6030e1ce6 100644 --- a/inference/core/interfaces/webrtc_worker/modal.py +++ b/inference/core/interfaces/webrtc_worker/modal.py @@ -24,7 +24,11 @@ PROJECT, ROBOFLOW_INTERNAL_SERVICE_SECRET, WEBEXEC_TRANSPORT, + WEBEXEC_WS_CONNECT_TIMEOUT_SECONDS, WEBEXEC_WS_CONNECTION_POOL_SIZE, + WEBEXEC_WS_FAIL_ON_SESSION_LOSS, + WEBEXEC_WS_IDLE_RELEASE_SECONDS, + WEBEXEC_WS_READ_TIMEOUT_SECONDS, WEBRTC_DATA_CHANNEL_ACK_WINDOW, WEBRTC_DATA_CHANNEL_BUFFER_SIZE_LIMIT, WEBRTC_GZIP_PREVIEW_FRAME_COMPRESSION, @@ -191,6 +195,12 @@ def check_nvidia_smi_gpu() -> str: "WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE": WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE, "WEBEXEC_TRANSPORT": WEBEXEC_TRANSPORT, "WEBEXEC_WS_CONNECTION_POOL_SIZE": str(WEBEXEC_WS_CONNECTION_POOL_SIZE), + "WEBEXEC_WS_CONNECT_TIMEOUT_SECONDS": str( + WEBEXEC_WS_CONNECT_TIMEOUT_SECONDS + ), + "WEBEXEC_WS_FAIL_ON_SESSION_LOSS": str(WEBEXEC_WS_FAIL_ON_SESSION_LOSS), + "WEBEXEC_WS_IDLE_RELEASE_SECONDS": str(WEBEXEC_WS_IDLE_RELEASE_SECONDS), + "WEBEXEC_WS_READ_TIMEOUT_SECONDS": str(WEBEXEC_WS_READ_TIMEOUT_SECONDS), "TELEMETRY_USE_PERSISTENT_QUEUE": "False", "TELEMETRY_API_PLAN_CACHE_TTL_SECONDS": str( os.getenv("TELEMETRY_API_PLAN_CACHE_TTL_SECONDS", 60) diff --git a/inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py b/inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py index 2f4363e9b1..b08a1c430b 100644 --- a/inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py +++ b/inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py @@ -7,9 +7,12 @@ Two transport modes are available, controlled by ``WEBEXEC_TRANSPORT``: * **http** — JSON POST with gzip compression and persistent ``requests.Session``. -* **websocket** (default) — persistent WebSocket connections with msgpack binary +* **websocket** — persistent WebSocket connections with msgpack binary frames. Eliminates per-request HTTP overhead and base64 image encoding. +``WEBEXEC_TRANSPORT`` defaults to ``http``; deployments opt into the websocket +transport explicitly. + """ import base64 @@ -20,8 +23,9 @@ import sys import threading import time as _time +import uuid from datetime import datetime -from typing import Any, Dict, Optional +from typing import Any, Dict, NamedTuple, Optional from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit import numpy as np @@ -36,6 +40,8 @@ WEBEXEC_MODAL_APP_NAME, WEBEXEC_WS_CONNECT_TIMEOUT_SECONDS, WEBEXEC_WS_CONNECTION_POOL_SIZE, + WEBEXEC_WS_FAIL_ON_SESSION_LOSS, + WEBEXEC_WS_IDLE_RELEASE_SECONDS, WEBEXEC_WS_READ_TIMEOUT_SECONDS, ) from inference.core.logger import logger @@ -78,6 +84,20 @@ # limit are split into a chunk-control frame plus raw chunks. Must match # WEBEXEC_WS_MAX_FRAME_BYTES in modal/modal_app.py. _WS_MAX_FRAME_BYTES = 1024 * 1024 +# Upper bound on the chunk count the server may announce in a ``_chunked`` +# control frame. Must match WEBEXEC_WS_MAX_CHUNKS in modal/modal_app.py. +_WS_MAX_CHUNKS = 1024 +# Ceiling on ONE reassembled response, far below the 1 GiB the chunk count +# alone would allow. Reassembly happens inside the shared inference server +# process, once per concurrent executor, so an oversized result must fail its +# own request rather than exhaust memory for everything else in the process. +# Mirrors WEBEXEC_WS_MAX_REQUEST_BYTES in modal/modal_app.py. +_WS_MAX_RESPONSE_BYTES = 64 * 1024 * 1024 +# Ceiling on ONE outbound request, mirroring WEBEXEC_WS_MAX_REQUEST_BYTES in +# modal/modal_app.py. Enforcing it here turns an oversized payload into a clear +# size error instead of a server-side reassembly abort whose error frame the +# response-id guard discards as a desync. +_WS_MAX_REQUEST_BYTES = 64 * 1024 * 1024 def _split_ws_frames(frame_bytes: bytes, msgpack: Any) -> list: @@ -88,9 +108,122 @@ def _split_ws_frames(frame_bytes: bytes, msgpack: Any) -> list: frame_bytes[i : i + _WS_MAX_FRAME_BYTES] for i in range(0, len(frame_bytes), _WS_MAX_FRAME_BYTES) ] + if len(frame_bytes) > _WS_MAX_REQUEST_BYTES: + # Mirrors the server's WEBEXEC_WS_MAX_REQUEST_BYTES. Without this the + # frame is sent, the server aborts reassembly BEFORE it parses the + # request id, and its error frame comes back unaddressed — which the + # response-id guard reports as a frame desync, hiding the real cause. + raise DynamicBlockError( + public_message=( + f"Custom Python block payload is too large for the websocket " + f"transport: {len(frame_bytes)} bytes, above the " + f"{_WS_MAX_REQUEST_BYTES}-byte limit. Reduce the size of the " + "block's inputs." + ), + context="modal_executor | websocket_payload_too_large", + ) + if len(chunks) > _WS_MAX_CHUNKS: + # Fail here, naming the real cause. Announcing a count the server + # structurally rejects makes it answer with an error frame carrying + # no request_id, which the response-id guard then reports as a frame + # desync — a misleading message on a request that would fail + # identically on every retry. + raise DynamicBlockError( + public_message=( + f"Custom Python block payload is too large for the websocket " + f"transport: {len(frame_bytes)} bytes needs {len(chunks)} " + f"chunks, above the {_WS_MAX_CHUNKS}-chunk limit " + f"({_WS_MAX_FRAME_BYTES * _WS_MAX_CHUNKS} bytes). Reduce the " + "size of the block's inputs." + ), + context="modal_executor | websocket_payload_too_large", + ) return [msgpack.packb({"_chunked": len(chunks)}, use_bin_type=True), *chunks] +# The heartbeat interval is idle_timeout/3, so a server advertising something +# absurd must not be able to brick the transport: too large lets the server +# close before we heartbeat, hence the ceiling. There is deliberately no floor +# on the advertised value — see ``_coerce_idle_timeout``; the floor lives on +# the derived interval instead, where it cannot inflate the client's idea of +# the server's deadline. +_MAX_ADVERTISED_IDLE_TIMEOUT_SECONDS = 300.0 +_DEFAULT_ADVERTISED_IDLE_TIMEOUT_SECONDS = 10.0 +# Floor on the derived interval, so a tiny advertised timeout cannot make the +# keepalive spin with a full RTT every tick. Kept small enough that the +# "worst-case gap is 2 x interval" invariant still holds for any advertised +# timeout at or above 1.5s; below that a deployment is simply misconfigured. +_MIN_HEARTBEAT_INTERVAL_SECONDS = 0.5 + + +def _is_server_originated_response(result: dict) -> bool: + """Whether the server, not the user's block, produced this failure. + + ONLY the ``server_error`` flag decides. ``error_type`` carries an exception + CLASS NAME chosen by untrusted user code, so it can never select a control + path: a block raising ``class UnknownCodeHash(Exception)`` would otherwise + trigger the resend-with-full-code and execute twice — on the HTTP + transport, which is the default and has no request id or dedup to catch it. + + There is deliberately no fallback to matching names. A server that stamps + the flag sets it on its own responses and omits it on user-code failures, + so "absent" means user code, not "old server" — falling back on absence + would make the name list authoritative for exactly the responses it must + not govern. Servers predating the flag are out of scope: this PR's deploy + contract is server-first. + """ + return result.get("server_error") is True + + +def _get_socket_timeout(ws: Any) -> Optional[float]: + """Current socket deadline, or None if this socket cannot report one.""" + if ws is None: + return None + try: + return ws.gettimeout() + except Exception: + return None + + +def _set_socket_timeout(ws: Any, value: Optional[float]) -> None: + """Best-effort deadline change. + + A deadline is a guard against a stalled peer, never a correctness + requirement, so a transport that does not expose ``settimeout`` must not + break the read path that wanted one. + """ + if ws is None or value is None: + return + try: + ws.settimeout(value) + except Exception: + pass + + +def _coerce_idle_timeout(advertised: Any) -> float: + """Clamp the server-advertised idle timeout into a usable range. + + Advisory field: it only sizes a timer, so a missing, non-numeric or + nonsensical value must degrade to the default rather than fail the + connection. + """ + try: + value = float(advertised) + except (TypeError, ValueError): + return _DEFAULT_ADVERTISED_IDLE_TIMEOUT_SECONDS + if value != value or value <= 0: # NaN or non-positive + return _DEFAULT_ADVERTISED_IDLE_TIMEOUT_SECONDS + # Clamp DOWNWARD only. Raising a small advertised value would make the + # client believe it has more headroom than the server actually gives it: + # the keepalive skips a tick when ``idle < interval``, so the real + # worst-case gap between app-level frames is 2 x interval. A server + # advertising 2s clamped up to 3.0 yields a 1.0s interval and a 2.0s + # worst-case gap — at the server's deadline, reinstating the very + # idle-close this protocol exists to fix. The floor belongs on the + # derived interval (see ``_heartbeat_interval``), not on the timeout. + return min(value, _MAX_ADVERTISED_IDLE_TIMEOUT_SECONDS) + + def _build_webexec_endpoint_base(method_label: str) -> str: workspace = MODAL_WORKSPACE_NAME app_name = WEBEXEC_MODAL_APP_NAME @@ -499,6 +632,7 @@ def execute_remote( not send_full_code and not result.get("success", False) and result.get("error_type") == "UnknownCodeHash" + and _is_server_originated_response(result) ): # Server replica doesn't have this hash cached; retry once. self._known_code_hashes.discard(code_hash) @@ -858,11 +992,87 @@ def _deserialize_msgpack_result(result: Any) -> Any: return result +class WebexecSessionLostError(Exception): + """The server no longer holds this session's Python runtime state. + + Raised instead of silently reconnecting: a reconnect that lands on a + fresh container would reset any state the block's code mutated across + frames, producing wrong results with no error. Callers must fail the + job (and replay from a checkpoint) rather than continue. + """ + + +class _ServerInfo(NamedTuple): + """What the handshake learned about the server on the current socket. + + Held as ONE immutable value so readers can never observe a half-updated + combination: the keepalive thread reads this without the io lock, and a + proto=2 seen alongside a stale (None) idle timeout would silently fall + back to the 25s v1 heartbeat interval — long enough for a v2 server to + close the connection under it. Rebinding the whole tuple is atomic; + field-by-field assignment was not. + """ + + proto: int = 1 + idle_timeout: Optional[float] = None + container_id: Optional[str] = None + + +class _ResendUnsafeError(Exception): + """Internal: a resend cannot be proven duplicate-free, so it is refused. + + Never escapes ``_send_recv_with_retry`` — it is translated there into a + ``DynamicBlockError`` explaining that the block may already have run. + """ + + +class _ServerClosingError(Exception): + """Internal: the server announced a graceful close before reading us. + + A v2 server sends ``{"_kind": "closing"}`` immediately before + ``close(1000)`` on its idle and connection-cap paths, both of which are + decided at the TOP of its receive loop — after that point it never reads + again. So a frame we had just written was provably never processed. + + That proof is what makes this different from every other mid-exchange + failure: the request can be retried on ANY container without risking a + duplicate execution, instead of being reported as the ambiguous "the + block may have already executed". Without it, the routine cap-close + (every ``WEBEXEC_WS_MAX_CONNECTION_SECONDS``) surfaces as a spurious + user-visible error on work that never ran. + """ + + class WebSocketModalExecutor: - """Executes Custom Python Blocks via a persistent WebSocket + msgpack.""" + """Executes Custom Python Blocks via a persistent WebSocket + msgpack. + + Protocol v2 (negotiated in ``_connect``): every server message is a + msgpack dict; ``_kind`` marks control frames (``hello``, + ``heartbeat_ack``); execution requests carry a ``request_id`` the + server echoes and dedups, making resend-after-failure safe. Against a + v1 server the executor falls back to the legacy behavior. + """ + # v1 fallback only; with a v2 server the interval derives from the + # idle timeout the server advertises in the hello reply. _KEEPALIVE_IDLE_SECONDS = 25.0 + # Generous enough for a cold container to answer, far below the + # execution-sized read timeout the handshake would otherwise inherit. + _HANDSHAKE_REPLY_TIMEOUT_SECONDS = 60.0 + + # Chunks of one response follow their control frame back-to-back, so this + # bounds a stalled reassembly instead of letting it hold _io_lock for the + # execution-sized read timeout. + _CHUNK_CONTINUATION_TIMEOUT_SECONDS = 30.0 + + # How long ``close()`` waits for _io_lock before giving up on closing the + # socket itself (it still drops its own reference). + _CLOSE_LOCK_TIMEOUT_SECONDS = 1.0 + # Slack on top of the keepalive's worst-case socket hold, to cover + # websocket-client's own close timeout and scheduling. + _CLOSE_JOIN_SLACK_SECONDS = 5.0 + def __init__(self, workspace_id: Optional[str] = None): self.workspace_id = workspace_id or MODAL_ANONYMOUS_WORKSPACE_NAME self._ws: Any = None @@ -872,6 +1082,69 @@ def __init__(self, workspace_id: Optional[str] = None): self._last_activity: float = 0.0 self._keepalive_stop: Optional[threading.Event] = None self._keepalive_thread: Optional[threading.Thread] = None + self._session_id: str = uuid.uuid4().hex + self._server = _ServerInfo() + self._had_success: bool = False + + def _rotate_session(self, reason: str) -> None: + """Start a fresh session after the current one's state is gone. + + Called when session loss is detected (so the loud failure happens + exactly once, after which execution honestly restarts under a new + session id) and when an idle connection is deliberately released. + Without the rotation, the very next reconnect would find the old + session id registered on some container and silently continue with + reset runtime state — the exact bug the session check exists to + prevent. + + KNOWN LIMITS — read these before relying on the check. + + 1. It is per-EXECUTOR, not per-run. Executors are cached per workspace, + so every concurrent workflow run in this process shares one session + id and one ``_had_success`` latch. The error lands on whichever run + reconnects first, not necessarily the one whose state was lost, and + that run may then pass the check silently on the rotated session. + There is no per-run session identity available to fix this: one + ``hello`` binds one session per CONNECTION, and + ``workflow_execution_id`` (already on every frame via + ``workflow_context``) is minted fresh per ``run_workflow`` call — + i.e. per video FRAME — so using it here would report a lost session + on every frame. + + 2. A positive answer proves less than it reads. ``session_known=True`` + means "this container still holds a namespace for this session's + code hash" — NOT "your run's state is intact". Server-side + namespaces are keyed by code hash and never by session, and every + namespace shares one ``_shared_globals`` dict, so two concurrent + runs of the same block on one container already read and overwrite + each other's state. That cross-run corruption predates this check + and is unaffected by it. + + 3. Container-local state cannot survive a long run by construction: a + websocket connection is a single Modal input capped at 700s, the + server closes at ``WEBEXEC_WS_MAX_CONNECTION_SECONDS`` before that, + and reconnects have no container affinity. + + Several paths disarm the latch without an error, by design: a graceful + server close (the scheduled connection cap), idle release + (``WEBEXEC_WS_IDLE_RELEASE_SECONDS``), a v1 -> v2 upgrade, and any + caller that retries after the loud failure. + + Given (1)-(3), ``WEBEXEC_WS_FAIL_ON_SESSION_LOSS`` defaults to False: + the check is an opt-in diagnostic for deployments whose blocks really + do rely on cross-frame globals, not a guarantee the platform can keep. + """ + logger.info( + "[webexec-ws] rotating session %s -> new session (%s)", + self._session_id, + reason, + ) + self._session_id = uuid.uuid4().hex + self._had_success = False + # Not dead: the fail-open path rotates on a LIVE connection, where no + # drop has cleared this. The other call sites drop first, where it is + # redundant but harmless. + self._hashes_sent_on_ws = set() def _get_ws_url(self, workspace_id: str) -> str: if self._ws_url is not None: @@ -906,12 +1179,125 @@ def _connect(self, workspace_id: str) -> None: header=[f"{k}: {v}" for k, v in headers.items()], timeout=WEBEXEC_WS_CONNECT_TIMEOUT_SECONDS, ) - self._ws.settimeout(WEBEXEC_WS_READ_TIMEOUT_SECONDS) # New container -> no compiled namespaces cached yet. self._hashes_sent_on_ws = set() + try: + # The handshake runs while _io_lock is held, so its recv must not + # wait the execution-sized read timeout: a server that accepts + # the upgrade and then stalls (cold boot, starved event loop, + # proxy buffering) would otherwise starve every request thread + # for this workspace, once per connect attempt. + self._ws.settimeout(self._HANDSHAKE_REPLY_TIMEOUT_SECONDS) + self._handshake() + self._ws.settimeout(WEBEXEC_WS_READ_TIMEOUT_SECONDS) + except Exception: + # Never leave a half-negotiated socket cached. + try: + self._ws.close() + except Exception: + pass + self._ws = None + raise self._last_activity = _time.monotonic() self._ensure_keepalive_thread() - logger.info("[webexec-ws] Connected") + logger.info( + "[webexec-ws] Connected (proto=%d idle_timeout=%s)", + self._server.proto, + self._server.idle_timeout, + ) + + def _handshake(self) -> None: + """Negotiate protocol v2; fall back to v1 on a legacy server. + + Sends a hello frame carrying this executor's session id. A v2 + server replies with its idle timeout (so the heartbeat interval is + derived, not guessed) and whether this container still knows the + session. A v1 server treats the hello as an execution request and + replies without ``_kind``; that reply is discarded and the + connection behaves as legacy. + + The session-loss check applies to the v2 path only; see the v1 + branch for why enforcing it there costs far more than it buys. + """ + import msgpack + + # Negotiate into locals and publish one _ServerInfo at the end: the + # surviving keepalive thread reads it unlocked and must never see a + # half-updated view. + hello = msgpack.packb( + {"_kind": "hello", "proto": 2, "session_id": self._session_id}, + use_bin_type=True, + ) + self._ws.send_binary(hello) + reply_raw = self._recv_reassembled(msgpack) + try: + reply = msgpack.unpackb(reply_raw, raw=False) + except Exception as error: + raise ConnectionError( + f"undecodable websocket handshake reply: {error}" + ) from error + if not (isinstance(reply, dict) and reply.get("_kind") == "hello"): + # Legacy server: it executed the hello as an (empty) request and + # returned its error response. Discard it and stay on v1. + # + # The session check is deliberately NOT applied here. A v1 server + # cannot confirm the session, but it also cannot keep a + # connection alive: its idle timeout closes the socket every + # WEBEXEC_WS_IDLE_TIMEOUT_SECONDS (10s by default) and v1's only + # keepalive is a protocol ping, which the ASGI layer answers + # without ever resetting that app-level timer. Reconnects are + # therefore the normal state of a v1 connection, so failing each + # one after a prior success would fail roughly every other + # request — a far worse outcome than the rare silent state reset + # it would prevent, especially as most blocks are stateless. + # Protocol v2 is what makes the guarantee affordable. + logger.info("[webexec-ws] server speaks protocol v1") + proto, idle_timeout, container_id = 1, None, None + else: + session_lost = self._had_success and not reply.get("session_known", False) + proto = 2 + idle_timeout = _coerce_idle_timeout(reply.get("idle_timeout_s")) + container_id = reply.get("container_id") or None + if session_lost and not WEBEXEC_WS_FAIL_ON_SESSION_LOSS: + logger.warning( + "[webexec-ws] session %s is unknown to the container this " + "reconnect landed on; continuing with reset runtime state " + "because WEBEXEC_WS_FAIL_ON_SESSION_LOSS is disabled", + self._session_id, + ) + self._rotate_session("session lost on reconnect (enforcement off)") + session_lost = False + if session_lost: + # Rotate before raising so the failure is loud exactly once: + # the next request starts an honest fresh session instead of + # silently passing the check with reset runtime state. + self._rotate_session("session lost on reconnect") + # Publish the negotiated server info before raising: the + # caller drops the socket, but the keepalive thread reads + # _server unlocked and must not keep a stale v1 interval. + self._server = _ServerInfo( + proto=proto, + idle_timeout=idle_timeout, + container_id=container_id, + ) + raise WebexecSessionLostError( + "The Modal container holding this custom Python session " + "is gone; runtime state mutated by previous frames cannot " + "be restored. Failing instead of silently continuing. " + "Note the session is shared by every concurrent workflow " + "run using this workspace's executor, so this error is " + "raised against whichever run reconnected first — not " + "necessarily the run whose state was lost. Set " + "WEBEXEC_WS_FAIL_ON_SESSION_LOSS=False to downgrade this " + "to a warning." + ) + # Commit together, idle timeout before proto: the keepalive thread + # reads both unlocked, and a proto=2 seen with a stale (None) idle + # timeout falls back to the 25s v1 interval — long enough for a v2 + # server to close the connection under it. + self._server = _ServerInfo( + proto=proto, idle_timeout=idle_timeout, container_id=container_id + ) def _ensure_connection(self, workspace_id: str) -> None: # Hot path: trust the cached socket. A dead connection will surface @@ -926,8 +1312,15 @@ def _ensure_connection(self, workspace_id: str) -> None: self._connect(workspace_id) def _ensure_keepalive_thread(self) -> None: - if self._keepalive_thread is not None and self._keepalive_thread.is_alive(): - return + """Start the keepalive thread for the current connection. + + A surviving thread keeps sleeping on the interval it computed before + the reconnect, so a v1 -> v2 reconnect would leave 25s gaps against a + 10s-idle server. Replace it unconditionally: the old thread is asked + to stop and every exit path clears ``_keepalive_thread`` itself. + """ + if self._keepalive_stop is not None: + self._keepalive_stop.set() self._keepalive_stop = threading.Event() self._keepalive_thread = threading.Thread( target=self._keepalive_loop, @@ -937,20 +1330,41 @@ def _ensure_keepalive_thread(self) -> None: ) self._keepalive_thread.start() + def _clear_keepalive_handle(self) -> None: + """Forget this thread's handle on the way out. + + ``_ensure_keepalive_thread`` no longer guards on ``is_alive()``, but + leaving a dead handle around is still misleading, and the check is + needed so an exiting thread cannot clobber a replacement that a + concurrent reconnect already started. + """ + if self._keepalive_thread is threading.current_thread(): + self._keepalive_thread = None + def _keepalive_loop(self, stop_event: threading.Event) -> None: """Ping the WS when the connection has been idle long enough. Skipped entirely while frames are flowing (``_last_activity`` is updated on every successful RTT). Uses ``acquire(blocking=False)`` so the keepalive never delays a real frame already in flight. + + A connection idle past ``WEBEXEC_WS_IDLE_RELEASE_SECONDS`` is + deliberately released instead of heartbeated forever: heartbeats + keep the Modal container warm (and billed), so an abandoned + connection would otherwise pin it for the full connection cap. + Releasing rotates the session — runtime state is knowingly given + up, so the next request starts fresh instead of failing loudly on + a lost session. """ - interval = self._KEEPALIVE_IDLE_SECONDS - while not stop_event.wait(interval): + while not stop_event.wait(self._heartbeat_interval()): ws = self._ws if ws is None: + self._clear_keepalive_handle() return + # _last_activity tracks real frames only — a heartbeat must not + # count as activity, or the idle-release below could never fire. idle = _time.monotonic() - self._last_activity - if idle < interval: + if idle < self._heartbeat_interval(): continue if not self._io_lock.acquire(blocking=False): # Frame in flight -> that's keepalive enough. @@ -958,39 +1372,217 @@ def _keepalive_loop(self, stop_event: threading.Event) -> None: try: ws = self._ws if ws is None: + self._clear_keepalive_handle() + return + # Recompute under the lock: the value above was read before + # acquiring, so a frame may have completed in between. + # Dropping the session on stale idleness would discard state + # a just-finished frame extended — silently, since rotation + # clears the very latch that would have failed loudly. + idle = _time.monotonic() - self._last_activity + if idle < self._heartbeat_interval(): + continue + if ( + WEBEXEC_WS_IDLE_RELEASE_SECONDS > 0 + and idle >= WEBEXEC_WS_IDLE_RELEASE_SECONDS + ): + logger.info( + "[webexec-ws] connection idle for %.0fs; releasing it " + "(and its custom Python session) so the container " + "can scale down", + idle, + ) + self._drop_ws_connection_locked() + self._rotate_session("idle release") + self._clear_keepalive_handle() return try: - ws.ping() - self._last_activity = _time.monotonic() - logger.debug("[webexec-ws] keepalive ping ok") + self._send_heartbeat(ws) + logger.debug("[webexec-ws] keepalive ok") + except _ServerClosingError: + # The server closed on schedule (connection cap or its own + # idle timeout) with nothing in flight — this thread holds + # _io_lock, so no request can be mid-exchange. Nothing was + # lost that a reconnect will not rebuild, so rotate rather + # than leaving _had_success armed: otherwise the next + # request's handshake lands on some other container, + # reports session_known=False, and raises a session-lost + # error for work that never had state at risk. The cap + # fires every WEBEXEC_WS_MAX_CONNECTION_SECONDS, so + # without this the failure is scheduled, not exceptional. + logger.debug( + "[webexec-ws] server closed gracefully during keepalive; " + "rotating session and dropping conn" + ) + self._drop_ws_connection_locked() + self._rotate_session("server closed gracefully") + self._clear_keepalive_handle() + return except Exception as e: logger.debug( - "[webexec-ws] keepalive ping failed (%s); dropping conn", + "[webexec-ws] keepalive failed (%s); dropping conn", e, ) - try: - ws.close() - except Exception: - pass - self._ws = None - self._hashes_sent_on_ws = set() + self._drop_ws_connection_locked() + self._clear_keepalive_handle() return finally: self._io_lock.release() - def close(self) -> None: - """Best-effort teardown, mainly for tests.""" - if self._keepalive_stop is not None: - self._keepalive_stop.set() - ws = self._ws - self._ws = None - if ws is not None: + def _heartbeat_interval(self) -> float: + """Heartbeat period, derived from the server's advertised idle timeout. + + Must be an application-level frame interval well under the server's + ``receive_bytes`` timeout: protocol-level ws pings are answered by + the ASGI layer and never reset that timer. + + The loop skips a tick when ``idle < interval``, so the real + worst-case gap between two app-level frames is ``2 x interval``. + Dividing by 3 keeps that at 2/3 of the server's deadline. + + The floor is capped by ``idle_timeout / 3`` rather than applied over + it: a bare ``max()`` breaks the invariant for a small advertised + timeout, which is reachable because the server reads its own value + from an int env var. At an advertised 1s a 0.5s floor yields a 1.0s + worst-case gap — exactly the deadline — so the connection would + idle-close every cycle, reintroducing the bug this protocol fixes. + """ + server = self._server + if server.proto == 2 and server.idle_timeout: + third = server.idle_timeout / 3.0 + return min(max(_MIN_HEARTBEAT_INTERVAL_SECONDS, third), third) + return self._KEEPALIVE_IDLE_SECONDS + + # A heartbeat ack is tiny and immediate; waiting the full read timeout + # (sized for user-code execution) would let a half-open connection pin + # _io_lock — and stall every real request — for minutes. + _HEARTBEAT_ACK_TIMEOUT_SECONDS = 10.0 + + def _heartbeat_ack_timeout(self) -> float: + """Ack deadline, never longer than the server's own idle timeout. + + With a low WEBEXEC_WS_IDLE_TIMEOUT_SECONDS the server closes the + connection before the fixed 10s deadline expires, so waiting the full + 10s only holds ``_io_lock`` against a socket already known dead. + """ + server = self._server + if server.proto == 2 and server.idle_timeout: + return max( + 1.0, min(self._HEARTBEAT_ACK_TIMEOUT_SECONDS, server.idle_timeout) + ) + return self._HEARTBEAT_ACK_TIMEOUT_SECONDS + + def _send_heartbeat(self, ws: Any) -> None: + """One heartbeat round-trip. Caller must hold ``_io_lock``. + + The deadline is installed BEFORE the write, not just around the read: + ``settimeout`` maps to ``socket.settimeout``, which bounds sends too. + A black-holed path with a full send buffer would otherwise block the + write for the execution-sized read timeout (720s) while this thread + holds ``_io_lock``, stalling every request for the workspace behind + it — the exact failure ``_HEARTBEAT_ACK_TIMEOUT_SECONDS`` exists to + prevent, reintroduced one line too late. + """ + import msgpack + + ws.settimeout(self._heartbeat_ack_timeout()) + try: + if self._server.proto != 2: + # Legacy server: protocol ping is all v1 offers. It cannot + # reset the server's app-level idle timer, and it only + # writes — the first write after a peer-side close usually + # still succeeds at the TCP layer — so it does not reliably + # detect a dead socket either. Against a v1 server the + # connection is expected to die on idle and be + # re-established by the next request. + ws.ping() + return + ws.send_binary(msgpack.packb({"_kind": "heartbeat"}, use_bin_type=True)) + reply_raw = self._recv_bytes_frame(ws) try: - ws.close() - except Exception: - pass + reply = msgpack.unpackb(reply_raw, raw=False) + except Exception as error: + raise ConnectionError( + f"undecodable heartbeat reply: {error}" + ) from error + finally: + ws.settimeout(WEBEXEC_WS_READ_TIMEOUT_SECONDS) + if isinstance(reply, dict) and reply.get("_kind") == "closing": + # Expected end of a connection's life (idle or the connection + # cap), not a fault: the caller drops the socket either way, but + # this keeps a routine close out of the failure logs. + raise _ServerClosingError( + "server announced a graceful close in response to a heartbeat" + ) + if not (isinstance(reply, dict) and reply.get("_kind") == "heartbeat_ack"): + raise ConnectionError(f"unexpected heartbeat reply: {reply!r}") + + def close(self) -> None: + """Release this executor's connection and keepalive thread. + + A production path, not a test helper: ``block_scaffolding`` calls it + when an executor is evicted from the per-workspace cache, on a request + thread. The keepalive thread must be stopped and joined FIRST — + ``websocket-client``'s ``close()`` performs a ``recv_frame()`` that + bypasses the socket's read lock, so closing while the keepalive is + blocked reading would free the fd underneath it. + """ + keepalive_stop = self._keepalive_stop + if keepalive_stop is not None: + keepalive_stop.set() + keepalive_thread = self._keepalive_thread + if ( + keepalive_thread is not None + and keepalive_thread is not threading.current_thread() + ): + # The ack deadline is installed on the SOCKET, so it bounds the + # heartbeat's send AND its recv — a black-holed peer costs up to + # two of them, plus websocket-client's own close timeout. Joining + # for only one would return with the thread still inside recv(). + keepalive_thread.join( + timeout=2 * self._heartbeat_ack_timeout() + + self._CLOSE_JOIN_SLACK_SECONDS + ) + # Best effort on the lock: an in-flight execution can hold it for the + # whole read timeout, and cache eviction must not block on that. + acquired = self._io_lock.acquire(timeout=self._CLOSE_LOCK_TIMEOUT_SECONDS) + if not acquired: + # Someone still owns the socket — the keepalive if the join above + # timed out, otherwise a request thread. Calling close() here would + # free the fd under their blocked recv(), which is the exact + # use-after-free this method's ordering exists to prevent (and the + # fd could be reused by an unrelated socket before they notice). + # Drop only this executor's reference; their own failure path + # closes the socket. + logger.debug( + "[webexec-ws] close(): io lock still held; releasing the " + "reference without closing the socket" + ) + self._ws = None + self._hashes_sent_on_ws = set() + self._server = self._server._replace(container_id=None) + return + try: + self._drop_ws_connection_locked() + finally: + self._io_lock.release() def _drop_ws_connection(self) -> None: + """Tear down the current connection. Takes ``_io_lock``. + + Without the lock this races ``_connect``/``_handshake``, which publish + ``self._ws`` and ``self._server`` under it: a thread dropping a + connection could otherwise overwrite a socket another thread had just + established, and roll ``_server`` back to a stale snapshot — reverting + ``proto`` to 1 and the heartbeat interval to the v1 25s against a + 10s-idle-timeout v2 server, i.e. reintroducing the bug this protocol + fixes, with no self-heal until the socket dies again. + """ + with self._io_lock: + self._drop_ws_connection_locked() + + def _drop_ws_connection_locked(self) -> None: + """``_drop_ws_connection`` body. Caller must hold ``_io_lock``.""" try: if self._ws is not None: self._ws.close() @@ -998,6 +1590,10 @@ def _drop_ws_connection(self) -> None: pass self._ws = None self._hashes_sent_on_ws = set() + # The id belongs to the connection, not the executor: leaving it set + # would let the resend guard compare against a container this + # executor is no longer talking to. + self._server = self._server._replace(container_id=None) def execute_remote( self, @@ -1045,51 +1641,162 @@ def execute_remote( context="modal_executor | missing_dependency", ) - return self._execute_ws( - block_type_name, - python_code, - inputs, - workspace, - msgpack, - workflow_context or {}, - ) + try: + return self._execute_ws( + block_type_name, + python_code, + inputs, + workspace, + msgpack, + workflow_context or {}, + ) + except WebexecSessionLostError as error: + raise DynamicBlockError( + public_message=( + f"Custom Python session lost for block " + f"`{block_type_name}`: {error} The job must be retried " + "from its last checkpoint (e.g. re-run the chunk) so the " + "block's runtime state is rebuilt from the start." + ), + context="modal_executor | websocket_session_lost", + ) from error def _send_recv_with_retry( self, frame_bytes: bytes, workspace: str, + request_id: Optional[str] = None, ) -> bytes: - """Send frame and receive response, reconnecting once before execution. - - We retry connection/send failures because the frame has not been - accepted by the websocket client. Once ``send_binary`` succeeds, the - remote may already be executing user code; a later ``recv`` failure has - an ambiguous outcome, so we do not resend the frame and risk duplicate - side effects. + """Send frame and receive response, reconnecting on failure. + + Against a v2 server every execution frame carries a ``request_id`` + the server dedups, so resending after a ``recv`` failure is safe — + but only while the reconnect lands on the SAME container: the dedup + registry is per-container, so a resend that reaches a different + container would run the user code a second time. The container is + identified from the handshake and compared under ``_io_lock``, + immediately before the frame goes out; anything else — a different + container, or one that did not identify itself — fails loudly rather + than risking a duplicate execution. + + Against a v1 server (no dedup) the legacy rule holds: once + ``send_binary`` succeeds the outcome is ambiguous and the frame is + not resent. + + ``WebexecSessionLostError`` from the reconnect handshake propagates: + state continuity is gone and the job must fail, not continue. """ import msgpack frames = _split_ws_frames(frame_bytes, msgpack) last_exc: Optional[Exception] = None - for attempt in range(2): + # Set once a frame has been accepted by a container: from then on + # only that same container may answer a resend, from its dedup + # registry. None means "not yet sent", never "any container". + resend_pending = False + sent_to_container: Optional[str] = None + container_at_send: Optional[str] = None + attempts = 3 + for attempt in range(attempts): sent_ok = False try: - self._ensure_connection(workspace) + try: + self._ensure_connection(workspace) + except WebexecSessionLostError: + if resend_pending: + # This frame was already accepted by a container, and + # the reconnect landed somewhere that lost the + # session — so it cannot answer from that container's + # dedup registry. The honest report is the ambiguous + # outcome, not "replay from your last checkpoint", + # which would re-run side effects that may already + # have happened. + raise _ResendUnsafeError() + raise # Hold the lock across send+recv so concurrent callers sharing # this executor's socket can't interleave a request/response. with self._io_lock: + # Compare inside the lock: _connect publishes the socket + # before the handshake commits the new container id, so a + # read taken outside can be stale by the time we send. + if self._ws is None: + # _ensure_connection's hot path runs outside this + # lock, so another thread may have dropped the socket + # in between. Report it as what it is instead of an + # AttributeError on None from send_binary. + raise ConnectionError( + "websocket connection was dropped before the frame " + "could be sent" + ) + container_at_send = self._server.container_id + if resend_pending and ( + container_at_send is None + or sent_to_container is None + or container_at_send != sent_to_container + ): + raise _ResendUnsafeError() for frame in frames: self._ws.send_binary(frame) sent_ok = True resp_bytes = self._recv_reassembled(msgpack) - self._last_activity = _time.monotonic() + self._last_activity = _time.monotonic() return resp_bytes + except _ServerClosingError as e: + # Proof of non-delivery for THIS attempt's write: the server + # decided to close at the top of its receive loop and never + # read again. So when nothing had been delivered before, the + # frame is retriable on ANY container, instead of being + # reported as "may have already executed" — the routine + # connection-cap close would otherwise fail real work that + # never ran. + # + # It proves nothing about an EARLIER attempt. If a previous + # attempt already handed the frame to a container, that + # container may have executed it and lost only the response; + # clearing ``resend_pending`` here would disarm the + # same-container guard below and let the next attempt re-send + # to a container with no dedup record of it — running the + # user's code a second time. So the ambiguity, once incurred, + # must survive a graceful close. + self._drop_ws_connection() + # A graceful close is the server acting on its own schedule + # (connection cap or idle timeout), not evidence that runtime + # state was lost unexpectedly. Rotate so the reconnect starts + # an honest fresh session instead of tripping the session-lost + # check on whichever container it lands on — that check would + # otherwise fire every WEBEXEC_WS_MAX_CONNECTION_SECONDS, for + # stateless blocks that had nothing at risk. The at-most-once + # guard is unaffected: it keys on ``sent_to_container``, not on + # the session id. + self._rotate_session("server closed gracefully") + last_exc = e + logger.debug( + "[webexec-ws] server closed gracefully (attempt %d/%d); " + "frame was not read, retrying", + attempt + 1, + attempts, + ) + continue + except _ResendUnsafeError: + self._drop_ws_connection() + raise DynamicBlockError( + public_message=( + "WebSocket connection to Modal endpoint lost after the " + "request was sent, and the reconnect did not reach the " + "same container. The custom Python block may have " + "already executed, so the frame was not retried." + ), + context="modal_executor | websocket_response", + ) + except WebexecSessionLostError: + raise except Exception as e: self._drop_ws_connection() - if sent_ok: - # recv failed after the frame was sent; the remote may have - # already executed user code, so we don't resend and risk - # duplicate side effects. + resend_is_safe = self._server.proto == 2 and bool(request_id) + if sent_ok and not resend_is_safe: + # v1: recv failed after the frame was sent; the remote may + # have already executed user code, so we don't resend and + # risk duplicate side effects. logger.warning( "[webexec-ws] response receive failed after frame was " "sent; not retrying to avoid duplicate execution: %s", @@ -1103,26 +1810,125 @@ def _send_recv_with_retry( ), context="modal_executor | websocket_response", ) + if sent_ok: + resend_pending = True + sent_to_container = container_at_send last_exc = e logger.warning( - "[webexec-ws] connect/send failed (attempt %d/2): %s", + "[webexec-ws] send/recv failed (attempt %d/%d): %s", attempt + 1, + attempts, e, ) continue + if resend_pending: + # The frame reached a container at least once, so the block may + # already have run. Reported with the same message/context as the + # other delivered-frame exits, so a caller cannot mistake this for + # a request that never left the client and safely re-run it. + raise DynamicBlockError( + public_message=( + "WebSocket connection to Modal endpoint lost after the " + "request was sent, and every retry failed. The custom " + f"Python block may have already executed: {last_exc}" + ), + context="modal_executor | websocket_response", + ) raise DynamicBlockError( public_message=f"WebSocket connection to Modal endpoint failed after retry: {last_exc}", context="modal_executor | websocket_connection", ) + def _recv_bytes_frame(self, ws: Any = None) -> bytes: + """Receive one frame, requiring it to be binary. + + A text frame here is a protocol violation (typically the payload of + a server-side close on an already-dead connection); it must be + treated as connection death, never fed to msgpack. + """ + resp = (ws if ws is not None else self._ws).recv() + if not isinstance(resp, bytes): + raise ConnectionError( + f"non-binary websocket frame ({type(resp).__name__}); " + "treating connection as dead" + ) + return resp + def _recv_reassembled(self, msgpack: Any) -> bytes: - """Receive one logical frame, joining chunked frames if signalled.""" - resp_bytes = self._ws.recv() - if isinstance(resp_bytes, bytes) and len(resp_bytes) < 64: - head = msgpack.unpackb(resp_bytes, raw=False) + """Receive one logical frame, joining chunked frames if signalled. + + A ``closing`` control frame is intercepted here rather than at each + call site: it can land in the handshake slot and in the execution + slot, and both must treat it as "the server never read us" rather + than as a response. + """ + resp_bytes = self._recv_bytes_frame() + if len(resp_bytes) < 64: + try: + head = msgpack.unpackb(resp_bytes, raw=False) + except Exception: + return resp_bytes + if isinstance(head, dict) and head.get("_kind") == "closing": + raise _ServerClosingError( + "server announced a graceful close before reading the frame" + ) if isinstance(head, dict) and "_chunked" in head: - return b"".join(self._ws.recv() for _ in range(head["_chunked"])) + chunk_count = head["_chunked"] + if ( + not isinstance(chunk_count, int) + or isinstance(chunk_count, bool) + or not 1 <= chunk_count <= _WS_MAX_CHUNKS + ): + raise ConnectionError( + f"invalid websocket chunk count from server: " + f"{chunk_count!r}; treating connection as dead" + ) + # Continuation chunks follow the control frame within + # milliseconds, so they must not inherit the execution-sized + # read timeout: a path that black-holes mid-response would + # otherwise block here for WEBEXEC_WS_READ_TIMEOUT_SECONDS + # (720s) while _io_lock is held, stalling every other run on + # this executor. Mirrors the handshake and heartbeat deadlines. + # Restore whatever deadline was in force rather than assuming + # the execution one: this also runs inside _handshake, which + # installs its own shorter reply timeout. + ws = self._ws + previous_timeout = _get_socket_timeout(ws) + _set_socket_timeout(ws, self._CHUNK_CONTINUATION_TIMEOUT_SECONDS) + try: + # Bound the BYTES as well as the chunk count, mirroring the + # server. The count ceiling alone admits 1 GiB, and this + # runs inside the shared inference server process, once per + # concurrent executor — an oversized result must fail the + # request, not the whole process. + parts = [] + reassembled_bytes = 0 + for _ in range(chunk_count): + # Read from the SAME socket the deadline was installed + # on: another thread may have replaced self._ws + # between chunks, and reading that one would both miss + # the continuation deadline and interleave with its + # exchange. + part = self._recv_bytes_frame(ws) + reassembled_bytes += len(part) + if reassembled_bytes > _WS_MAX_RESPONSE_BYTES: + raise ConnectionError( + "chunked websocket response exceeds " + f"{_WS_MAX_RESPONSE_BYTES} bytes; treating " + "connection as dead" + ) + parts.append(part) + return b"".join(parts) + finally: + _set_socket_timeout( + ws, + ( + previous_timeout + if previous_timeout is not None + else WEBEXEC_WS_READ_TIMEOUT_SECONDS + ), + ) return resp_bytes def _execute_ws( @@ -1150,6 +1956,7 @@ def _execute_ws( # compiled namespace by hash. send_full_code = code_hash not in self._hashes_sent_on_ws + request_id = uuid.uuid4().hex frame_bytes = self._build_ws_frame( python_code=python_code, packed_inputs=packed_inputs, @@ -1157,14 +1964,18 @@ def _execute_ws( send_full_code=send_full_code, msgpack=msgpack, workflow_context=workflow_context, + request_id=request_id, ) t_pack = _time.monotonic() - resp_bytes = self._send_recv_with_retry(frame_bytes, workspace) + resp_bytes = self._send_recv_with_retry( + frame_bytes, workspace, request_id=request_id + ) t_rtt = _time.monotonic() - result = msgpack.unpackb(resp_bytes, raw=False) + result = self._unpack_response(resp_bytes, msgpack) + self._check_response_id(result, request_id) # Fresh replica doesn't have this hash cached (can happen after a # reconnect or container restart). Retry once with full code. @@ -1172,12 +1983,16 @@ def _execute_ws( not send_full_code and not result.get("success", False) and result.get("error_type") == "UnknownCodeHash" + and _is_server_originated_response(result) ): self._hashes_sent_on_ws.discard(code_hash) logger.info( "[webexec-ws] server missed cached hash %s, resending full code", code_hash, ) + # A distinct logical request -> a fresh request_id, so the + # server's dedup cache can't answer it with the failed response. + retry_request_id = uuid.uuid4().hex retry_frame = self._build_ws_frame( python_code=python_code, packed_inputs=packed_inputs, @@ -1185,12 +2000,23 @@ def _execute_ws( send_full_code=True, msgpack=msgpack, workflow_context=workflow_context, + request_id=retry_request_id, ) - resp_bytes = self._send_recv_with_retry(retry_frame, workspace) - result = msgpack.unpackb(resp_bytes, raw=False) + resp_bytes = self._send_recv_with_retry( + retry_frame, workspace, request_id=retry_request_id + ) + result = self._unpack_response(resp_bytes, msgpack) + self._check_response_id(result, retry_request_id) if result.get("success", False): self._hashes_sent_on_ws.add(code_hash) + # Only a v2 success establishes continuity worth protecting. A v1 + # connection is closed by the server every idle timeout and its + # state is already reset on every reconnect, so latching here + # would make the first post-upgrade reconnect fail loudly for + # every executor that had ever succeeded against a v1 server. + if self._server.proto == 2: + self._had_success = True t_done = _time.monotonic() @@ -1206,6 +2032,7 @@ def _execute_ws( ) if not result.get("success", False): + self._raise_server_error_if_infrastructure(result, block_type_name) self._raise_code_error(result, block_type_name, python_code) stdout = result.get("stdout") @@ -1219,6 +2046,108 @@ def _execute_ws( return _deserialize_msgpack_result(result.get("result", {})) + def _unpack_response(self, resp_bytes: bytes, msgpack: Any) -> dict: + """Decode a response frame, failing as a transport error. + + A truncated or non-map frame would otherwise escape ``execute_remote`` + as a raw ``OutOfData``/``ExtraData``/``AttributeError`` to the + execution engine, with the (now desynced) socket left in the cache. + """ + try: + result = msgpack.unpackb(resp_bytes, raw=False) + except Exception as error: + self._drop_ws_connection() + raise DynamicBlockError( + public_message=( + "WebSocket response from the Modal endpoint could not be " + f"decoded ({error}); dropping the connection." + ), + context="modal_executor | websocket_response", + ) from error + if not isinstance(result, dict): + self._drop_ws_connection() + raise DynamicBlockError( + public_message=( + "WebSocket response from the Modal endpoint was not a " + f"msgpack map (got {type(result).__name__}); dropping the " + "connection." + ), + context="modal_executor | websocket_response", + ) + return result + + def _raise_server_error_if_infrastructure( + self, + result: dict, + block_type_name: str, + ) -> None: + """Report transport failures as transport failures. + + ``DynamicBlockCodeError`` means "the user's Python raised". Routing a + refused resend, an undecodable frame or a failed response serialization + through it tells the user their block errored when it either never ran + or ran successfully. + """ + error_type = result.get("error_type") or "RuntimeError" + if not _is_server_originated_response(result): + return + raise DynamicBlockError( + public_message=( + f"The Modal webexec server could not complete the request for " + f"block `{block_type_name}` ({error_type}): " + f"{result.get('error', 'Unknown error')}" + ), + context="modal_executor | websocket_server_error", + ) + + def _check_response_id(self, result: Any, request_id: str) -> None: + """Reject a frame that is not this request's response. + + On a v2 connection the echo is mandatory: a late ``heartbeat_ack`` + (tiny, no ``request_id``) landing in the execution recv slot would + otherwise pass as a response, surface as a fabricated + ``RuntimeError: Unknown error`` against the user's block, and leave + the real response queued so the NEXT request desyncs too. + Against a v1 server the field does not exist, so only a mismatch is + enforced there. + + A server-stamped error frame is exempt from the echo requirement: the + server legitimately emits unaddressed ones (a chunk-reassembly abort, a + decode-limit violation, an over-long request id) whose id it never got + to parse. Those carry a precise diagnostic, and reporting them as a + generic desync throws it away — the opposite of what this guard is for. + """ + echoed = result.get("request_id") + if echoed is None and result.get("server_error") is True: + self._drop_ws_connection() + raise DynamicBlockError( + public_message=( + "The Modal webexec server rejected the request " + f"({result.get('error_type') or 'ServerError'}): " + f"{result.get('error', 'Unknown error')}" + ), + context="modal_executor | websocket_server_error", + ) + if self._server.proto == 2 and (echoed is None or "_kind" in result): + self._drop_ws_connection() + raise DynamicBlockError( + public_message=( + "WebSocket response did not carry the in-flight request id " + "(stale or control frame on the connection); dropping the " + "connection." + ), + context="modal_executor | websocket_response_mismatch", + ) + if echoed is not None and echoed != request_id: + self._drop_ws_connection() + raise DynamicBlockError( + public_message=( + "WebSocket response did not match the request in flight " + "(stale frame on the connection); dropping the connection." + ), + context="modal_executor | websocket_response_mismatch", + ) + @staticmethod def _build_ws_frame( python_code: PythonCode, @@ -1227,6 +2156,7 @@ def _build_ws_frame( send_full_code: bool, msgpack: Any, workflow_context: Dict[str, Any], + request_id: Optional[str] = None, ) -> bytes: """Pack a msgpack frame, optionally omitting ``code_str``/``imports``. @@ -1239,6 +2169,8 @@ def _build_ws_frame( "inputs": packed_inputs, "workflow_context": workflow_context, } + if request_id is not None: + payload["request_id"] = request_id if send_full_code: payload["code_str"] = python_code.run_function_code payload["imports"] = python_code.imports or [] @@ -1293,11 +2225,31 @@ class PooledWebSocketModalExecutor: ordered on the connection. The workspace-level executor cache can therefore keep this pool hot without funneling every same-workspace execution through one socket. + + Session continuity ("fail loudly when runtime state is lost") is + per-executor, not per-pool: with a pool size above 1, consecutive frames + of one stateful run can be routed to different executors — and thus + different containers with independent runtime state — whenever + concurrency displaces them from slot 0. Stateful custom Python blocks + should run with the default pool size of 1. """ def __init__(self, workspace_id: Optional[str] = None): self.workspace_id = workspace_id or MODAL_ANONYMOUS_WORKSPACE_NAME pool_size = max(1, WEBEXEC_WS_CONNECTION_POOL_SIZE) + if pool_size > 1 and WEBEXEC_WS_FAIL_ON_SESSION_LOSS: + # Enforced by a docstring only, until now. With more than one + # executor, consecutive frames of one stateful run can be routed + # to different sockets and therefore different containers, so the + # session guarantee the flag asks for cannot hold. + logger.warning( + "[webexec-ws] WEBEXEC_WS_CONNECTION_POOL_SIZE=%d with " + "WEBEXEC_WS_FAIL_ON_SESSION_LOSS enabled: the session " + "continuity guarantee does not hold above pool size 1, since " + "consecutive frames of one run may use different connections. " + "Use pool size 1 for stateful custom Python blocks.", + pool_size, + ) self._executors = [ WebSocketModalExecutor(workspace_id=self.workspace_id) for _ in range(pool_size) diff --git a/modal/modal_app.py b/modal/modal_app.py index 727314bfce..b5b03abc50 100644 --- a/modal/modal_app.py +++ b/modal/modal_app.py @@ -25,7 +25,6 @@ import modal - _thread_local = threading.local() _install_lock = threading.Lock() @@ -35,8 +34,27 @@ WEBEXEC_MODAL_REGION = os.environ.get("WEBEXEC_MODAL_REGION", "us-east-1") WEBEXEC_MODAL_ROUTING_REGION = os.environ.get("WEBEXEC_MODAL_ROUTING_REGION") +# Hard cap on how long one websocket connection is served before this side +# closes it cleanly (code 1000). +# +# It MUST stay below the executor's Modal ``timeout`` (700s, see +# _executor_decorator_kwargs): an ASGI websocket connection is a single Modal +# input, so at 700s Modal kills the input outright — no close frame, in-flight +# executions cancelled mid-run. Closing ourselves first turns that into an +# orderly reconnect. Raise this only together with the Modal timeout. +# +# NOTE: this cap and the protocol v2 session guarantee interact. Sessions are +# container-local and reconnects are not routed with any affinity, so every +# forced close of a STATEFUL run is likely to land on another container and +# surface as a (correct, but scheduled) session-lost failure. Clients that +# cannot tolerate that can disable the check with +# WEBEXEC_WS_FAIL_ON_SESSION_LOSS=False. +# +# Making long stateful runs survive needs session-affine reconnect routing, +# externalized session state, or a drain/handoff before the cap — none of +# which exist yet. WEBEXEC_WS_MAX_CONNECTION_SECONDS = int( - os.getenv("WEBEXEC_WS_MAX_CONNECTION_SECONDS", "3600") + os.getenv("WEBEXEC_WS_MAX_CONNECTION_SECONDS", "600") ) WEBEXEC_WS_IDLE_TIMEOUT_SECONDS = int( os.getenv("WEBEXEC_WS_IDLE_TIMEOUT_SECONDS", "10") @@ -47,6 +65,238 @@ # _WS_MAX_FRAME_BYTES in # inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py. WEBEXEC_WS_MAX_FRAME_BYTES = 1024 * 1024 +# Upper bound on the chunk count a peer may announce in a ``_chunked`` control +# frame. A negative/huge/non-int value would otherwise stall the receive loop +# or allocate without bound. Must match _WS_MAX_CHUNKS in modal_executor.py. +WEBEXEC_WS_MAX_CHUNKS = 1024 +# Ceiling on ONE reassembled request. Deliberately far below +# MAX_CHUNKS * MAX_FRAME_BYTES (1 GiB): the decode happens on a container +# serving up to ``max_inputs`` connections from one event loop, and a gigabyte +# of packed integers expands to tens of GB and minutes of GIL-held work, which +# would push every sibling connection past its idle deadline (or OOM the +# container and destroy every tenant's namespaces on it). +WEBEXEC_WS_MAX_REQUEST_BYTES = int( + os.getenv("WEBEXEC_WS_MAX_REQUEST_BYTES", str(64 * 1024 * 1024)) +) + + +def _safe_unpackb(raw: bytes) -> Any: + """Unpack one frame with explicit collection limits. + + ``msgpack`` derives its default limits from the buffer size, so a large + buffer authorises a proportionally huge object graph. Pinning them keeps a + hostile or malformed frame from expanding into tens of GB of Python + objects inside a container shared by other tenants. + """ + import msgpack + + return msgpack.unpackb( + raw, + raw=False, + max_str_len=WEBEXEC_WS_MAX_REQUEST_BYTES, + max_bin_len=WEBEXEC_WS_MAX_REQUEST_BYTES, + max_array_len=1_000_000, + max_map_len=1_000_000, + ) + + +class _InputsDecodeError(Exception): + """A request's inputs could not be decoded, so user code never ran. + + Raised inside the worker thread (where decoding happens, to keep the + shared event loop responsive) and recognised by the handler so the + failure is still reported as never-executed and the request id stays + resendable — the same contract as when decoding ran on the loop. + """ + + +class _TtlKeySet: + """TTL'd, size-capped set of ids with refreshable timestamps. + + Backs the two container-local registries protocol v2 needs — which + sessions this container holds runtime state for, and which request ids + it has already started executing. Entries are kept in age order so + pruning is an early-exit scan from the front. + + All access happens on the event loop thread, so no locking is needed. + """ + + def __init__( + self, + ttl_seconds: float, + max_entries: int, + min_retention_seconds: float = 0.0, + name: str = "registry", + ): + from collections import OrderedDict + + self._ttl_seconds = ttl_seconds + self._max_entries = max_entries + # Entries younger than this are never evicted to satisfy the size cap. + self._min_retention_seconds = min_retention_seconds + self._name = name + self._seen: "OrderedDict[str, float]" = OrderedDict() + + def _prune(self) -> None: + now = time.monotonic() + # Expire first, so age-based eviction reclaims room before the size cap + # has to. + deadline = now - self._ttl_seconds + while self._seen: + key, last_seen = next(iter(self._seen.items())) + if last_seen >= deadline: + break + del self._seen[key] + # Only THEN enforce the size cap, and never at the cost of an entry the + # peer could still act on. For _ws_executed that entry is the whole + # at-most-once backstop: dropping one inside the client's retry window + # lets a resend fall through to a second execution of user code. Under + # sustained load the count cap would otherwise evict ids seconds old, + # regardless of TTL. Entries are a key plus a float, so retaining them + # for the protected window costs little; exceeding the cap is reported + # rather than silently papered over. + if len(self._seen) <= self._max_entries: + return + protected_deadline = now - self._min_retention_seconds + while len(self._seen) > self._max_entries and self._seen: + key, last_seen = next(iter(self._seen.items())) + if last_seen >= protected_deadline: + print( + f"[webexec] {self._name} is over its {self._max_entries}-entry " + f"cap and every entry is still inside the " + f"{self._min_retention_seconds:.0f}s retention window; keeping " + "them to preserve at-most-once. Raise the cap if this persists." + ) + return + del self._seen[key] + + # Ids the client generates are uuid4 hex (32 chars); this leaves room for + # a prefix without admitting a key large enough to matter. Without it the + # registries bound entry COUNT but not bytes, so a peer sending megabyte + # request ids fills 32768 entries with gigabytes of container-shared + # state — the same argument the chunk byte bound makes below. + _MAX_KEY_LENGTH = 128 + + @classmethod + def _valid(cls, key: object) -> bool: + # Keys arrive from a msgpack frame, so anything can show up here. The + # three accessors below must agree on what counts as a key: an add() + # that stored a non-str while __contains__ rejected it would silently + # void the at-most-once backstop. + return isinstance(key, str) and 0 < len(key) <= cls._MAX_KEY_LENGTH + + def add(self, key: str) -> None: + if not self._valid(key): + return + self._seen.pop(key, None) + self._seen[key] = time.monotonic() + self._prune() + + def discard(self, key: str) -> None: + if not self._valid(key): + return + self._seen.pop(key, None) + + def refresh(self, key: str) -> bool: + """Re-stamp an existing entry. Never adds. Returns whether present. + + Prunes first: without it an entry past its TTL is still reported + present and then re-stamped back to life. + """ + if not self._valid(key): + return False + self._prune() + if key not in self._seen: + return False + self.add(key) + return True + + def __contains__(self, key: object) -> bool: + if not self._valid(key): + return False + self._prune() + return key in self._seen + + +class _WsResponseCache: + """Bounded cache of request_id -> packed response (protocol v2 dedup). + + Lets a client safely resend a request whose response was lost to a + dropped connection: an already-executed request returns its cached + response instead of running the user code a second time. + + The retry window is seconds, so entries carry a short TTL; responses can + embed serialized images, so eviction is byte-capped as well as + entry-capped (LRU order). All access happens on the event loop thread, + so no locking is needed. + + Retention is best effort — a single oversized response or a burst of + concurrent large ones can evict an entry the client still wants. The + at-most-once guarantee therefore does NOT rest on this cache: it rests + on the executed-request registry (``_ws_executed``), which a cache miss + falls back to so the resend gets a loud error instead of a second + execution. + """ + + def __init__( + self, + max_entries: int = 128, + max_bytes: int = 64 * 1024 * 1024, + # The client can take up to its read timeout + # (WEBEXEC_WS_READ_TIMEOUT_SECONDS, 720s) to notice a lost response and + # resend, so a 120s TTL expired before the retry it exists to answer. + ttl_seconds: float = 900.0, + ): + from collections import OrderedDict + + self._max_entries = max_entries + self._max_bytes = max_bytes + self._ttl_seconds = ttl_seconds + self._entries: "OrderedDict[str, Tuple[float, bytes]]" = OrderedDict() + self._total_bytes = 0 + + def _prune_expired(self) -> None: + deadline = time.monotonic() - self._ttl_seconds + while self._entries: + key, (stored_at, payload) = next(iter(self._entries.items())) + if stored_at >= deadline: + break + del self._entries[key] + self._total_bytes -= len(payload) + + def get(self, key: Optional[str]) -> Optional[bytes]: + if not key: + return None + self._prune_expired() + entry = self._entries.get(key) + if entry is None: + return None + # Deliberately no re-stamp / move_to_end: entries stay in age order so + # the early-exit expiry scan above stays correct, and the TTL stays a + # real TTL rather than a sliding window extended by every resend. + return entry[1] + + def put(self, key: str, payload: bytes) -> None: + if not key: + return + if len(payload) * 2 > self._max_bytes: + # Inserting it would evict every other entry and then itself, so a + # single large response would wipe the dedup cache for every + # concurrent connection on this container. Skip the insert; the + # executed-request registry still bars a re-run. + return + old = self._entries.pop(key, None) + if old is not None: + self._total_bytes -= len(old[1]) + self._entries[key] = (time.monotonic(), payload) + self._total_bytes += len(payload) + self._prune_expired() + while self._entries and ( + len(self._entries) > self._max_entries + or self._total_bytes > self._max_bytes + ): + _, (_, evicted) = self._entries.popitem(last=False) + self._total_bytes -= len(evicted) # mirrors inference/core/workflows/execution_engine/v1/dynamic_blocks (avoiding `from inference import ...`) @@ -179,7 +429,10 @@ def get_inference_image(): "ffmpeg", "wget", ) - .pip_install("fastapi[standard]") # Add FastAPI for web endpoints + # FastAPI serves the web endpoints; msgpack is a hard requirement of + # the websocket transport (wsapp crash-loops on boot without it on + # base images that predate msgpack in inference's requirements). + .pip_install("fastapi[standard]", "msgpack") .entrypoint([]) ) return image @@ -197,6 +450,7 @@ def get_inference_image(): "env": { "WEBEXEC_WS_MAX_CONNECTION_SECONDS": str(WEBEXEC_WS_MAX_CONNECTION_SECONDS), "WEBEXEC_WS_IDLE_TIMEOUT_SECONDS": str(WEBEXEC_WS_IDLE_TIMEOUT_SECONDS), + "WEBEXEC_WS_MAX_REQUEST_BYTES": str(WEBEXEC_WS_MAX_REQUEST_BYTES), }, } if WEBEXEC_MODAL_ROUTING_REGION: @@ -220,11 +474,45 @@ class Executor: @modal.enter() def identify(self): + import uuid + print(f"Initializing sandbox for {self.workspace_id}") # Initialize the namespaces dict and shared globals self._code_namespaces = {} self._shared_globals = {} self._namespace_lock = threading.RLock() + # Protocol v2 websocket state (see wsapp): all touched only on the + # event loop thread. + # Prefer Modal's own task id so a client-reported session-lost error + # can be correlated with container logs; fall back to a uuid locally. + # NOTE: this runs in a non-``snap=True`` @modal.enter(), i.e. AFTER a + # memory-snapshot restore, so every restored container gets a fresh + # id. Adding ``snap=True`` here would make restored containers share + # one id and let the client's same-container resend check pass across + # different containers — re-executing user code. Keep it post-restore. + # Always append a random suffix: this id is the sole trust anchor for + # the client's same-container resend guard, so uniqueness must be + # self-guaranteed rather than inherited from MODAL_TASK_ID being set. + # The task id is kept as a prefix for log correlation. + self._container_id = f"{os.environ.get('MODAL_TASK_ID', '')}-{uuid.uuid4().hex}" + self._ws_sessions = _TtlKeySet( + ttl_seconds=self._WS_SESSION_TTL_SECONDS, + max_entries=self._WS_SESSION_MAX_ENTRIES, + min_retention_seconds=self._WS_SESSION_MIN_RETENTION_SECONDS, + name="_ws_sessions", + ) + # Request ids whose user code this container has STARTED running. + # Entries are tiny, so this outlives the response cache by a wide + # margin: it is what makes execution at-most-once even when the + # response payload is gone. + self._ws_executed = _TtlKeySet( + ttl_seconds=self._WS_EXECUTED_TTL_SECONDS, + max_entries=self._WS_EXECUTED_MAX_ENTRIES, + min_retention_seconds=self._WS_EXECUTED_MIN_RETENTION_SECONDS, + name="_ws_executed", + ) + self._ws_response_cache = _WsResponseCache() + self._ws_inflight: Dict[str, "asyncio.Task"] = {} def _get_code_hash(self, code_str: str, imports: list) -> str: """Compute a stable hash for the code to identify unique blocks.""" @@ -239,6 +527,48 @@ def _get_namespace_lock(self) -> threading.RLock: self._namespace_lock = namespace_lock return namespace_lock + # A session id is registered only once user code has SUCCESSFULLY + # executed for it on this container — mirroring the client's + # ``_had_success`` latch. Registering earlier (e.g. at hello time) + # would make a failed "session lost" handshake mark the session as + # known, so the very next reconnect would silently pass the check + # this mechanism exists to fail. An ALREADY known session is refreshed + # on every hello, so a long-lived but currently failing stream cannot + # age out of the registry while its namespaces are still here. + _WS_SESSION_TTL_SECONDS = 7200.0 + _WS_SESSION_MAX_ENTRIES = 4096 + # Request ids live only as long as a client could still resend one: + # 3 attempts with reconnects, bounded by the client's read timeout. + _WS_EXECUTED_TTL_SECONDS = 3600.0 + _WS_EXECUTED_MAX_ENTRIES = 32768 + # No executed marker may be evicted by the SIZE cap while the client could + # still resend the request it guards. The client's window is bounded by its + # read timeout (WEBEXEC_WS_READ_TIMEOUT_SECONDS, 720s) times its retry + # attempts; this covers it with margin. Without it, sustained load evicts + # markers seconds old and a same-container resend re-runs user code. + _WS_EXECUTED_MIN_RETENTION_SECONDS = 900.0 + # A session id must likewise survive the size cap long enough to answer the + # reconnects of a live stream; evicting one early reports a false session + # loss. + _WS_SESSION_MIN_RETENTION_SECONDS = 900.0 + + def _ws_session_seen(self, session_id: str) -> bool: + """Whether user code for this session already ran on this container. + + Answers the client's hello: ``session_known=False`` on a reconnect + tells the client its Python runtime state (mutated globals in the + cached namespaces) lives in some other, gone container — the client + fails loudly instead of silently continuing with reset state. + + A hit also re-stamps the entry: the namespaces this answer is about + are never evicted, so the registry must not expire under a session + that is still actively connecting. + """ + return self._ws_sessions.refresh(session_id) + + def _ws_register_session(self, session_id: str) -> None: + self._ws_sessions.add(session_id) + def _get_cached_namespace(self, code_hash: str) -> Optional[dict]: namespace = self._code_namespaces.get(code_hash) if namespace is not None: @@ -375,6 +705,11 @@ async def execute_block(self, raw_request: Request) -> Dict[str, Any]: f"{code_hash}; client must resend full code." ), "error_type": "UnknownCodeHash", + # This is the server's own control response, not the + # user's code failing. Without the flag the client has to + # match on the error_type NAME, which untrusted user code + # can forge by raising a class of the same name. + "server_error": True, "code_hash": code_hash, } else: @@ -382,6 +717,7 @@ async def execute_block(self, raw_request: Request) -> Dict[str, Any]: "success": False, "error": "Request must include either 'code_str' or 'code_hash'.", "error_type": "InvalidRequest", + "server_error": True, } try: @@ -676,6 +1012,11 @@ def _run_user_code_ws( f"{code_hash}; client must resend full code." ), "error_type": "UnknownCodeHash", + # This is the server's own control response, not the + # user's code failing. Without the flag the client has to + # match on the error_type NAME, which untrusted user code + # can forge by raising a class of the same name. + "server_error": True, "code_hash": code_hash, } else: @@ -683,6 +1024,7 @@ def _run_user_code_ws( "success": False, "error": "Request must include either 'code_str' or 'code_hash'.", "error_type": "InvalidRequest", + "server_error": True, } if run_function_name not in namespace: @@ -824,9 +1166,7 @@ def _decode(obj): return Batch( content=[_decode(v) for v in obj["value"]], indices=( - [tuple(i) for i in indices] - if indices is not None - else None + [tuple(i) for i in indices] if indices is not None else None ), ) if _type == "bytes": @@ -919,17 +1259,328 @@ def wsapp(self): executor_self = self + async def send_payload( + websocket: WebSocket, payload: bytes, request_id: Optional[str] = None + ) -> None: + """Send one logical response, chunking oversized payloads. + + ``request_id`` is echoed on any error frame this helper generates + itself: on protocol v2 the client rejects a response that does not + carry the in-flight id, so an unaddressed error frame reaches the + user as a generic "frame desync" instead of the real cause. + """ + if len(payload) > WEBEXEC_WS_MAX_REQUEST_BYTES: + # The client caps reassembly at the same limit, and it reports + # an over-cap response as a lost connection ("may have already + # executed") rather than as a size problem. Refuse here so the + # real cause is named once and not retried. + await websocket.send_bytes( + _pack_ws_error( + error_type="ResponseTooLarge", + error=( + "The block's result is too large to return over the " + f"websocket transport: {len(payload)} bytes, above " + f"the {WEBEXEC_WS_MAX_REQUEST_BYTES}-byte limit." + ), + request_id=request_id, + ) + ) + return + if len(payload) > WEBEXEC_WS_MAX_FRAME_BYTES: + chunks = [ + payload[i : i + WEBEXEC_WS_MAX_FRAME_BYTES] + for i in range(0, len(payload), WEBEXEC_WS_MAX_FRAME_BYTES) + ] + if len(chunks) > WEBEXEC_WS_MAX_CHUNKS: + # The client enforces the same ceiling on receive, so + # announcing more would strand it waiting for chunks it + # will refuse. Report the size instead of desyncing. + await websocket.send_bytes( + _pack_ws_error( + error_type="ResponseTooLarge", + error=( + f"The block's result is too large to return " + f"over the websocket transport: {len(payload)} " + f"bytes needs {len(chunks)} chunks, above the " + f"{WEBEXEC_WS_MAX_CHUNKS}-chunk limit." + ), + request_id=request_id, + ) + ) + return + await websocket.send_bytes( + msgpack.packb({"_chunked": len(chunks)}, use_bin_type=True) + ) + for chunk in chunks: + await websocket.send_bytes(chunk) + else: + await websocket.send_bytes(payload) + + def _pack_ws_error( + error_type: str, + error: str, + request_id: Optional[str], + ) -> bytes: + """Pack a minimal error response that can never fail to encode.""" + resp: Dict[str, Any] = { + "success": False, + "error": str(error), + "error_type": str(error_type), + # Marks the failure as the transport's, not the user block's: + # every caller of this helper is server infrastructure (a + # frame the server could not decode, a response it could not + # return, a resend it refused). The client uses this to avoid + # reporting a DynamicBlockCodeError against code that either + # never ran or ran fine. + "server_error": True, + } + if request_id: + resp["request_id"] = request_id + return msgpack.packb(resp, use_bin_type=True) + + async def execute_request( + request: dict, request_id: Optional[str], session_id: str + ) -> bytes: + """Run one execution request and return the packed response. + + Registered in ``_ws_inflight`` while running so a resend of the + same ``request_id`` (from this or another connection) awaits the + original execution instead of running the user code a second + time; the completed payload then moves to the response cache for + resends that arrive after completion. + + Only ``asyncio.CancelledError`` escapes (task cancellation must + propagate), and even then a packed error is cached first. Every + other outcome — including ``SystemExit`` from user code calling + ``sys.exit()``, or a result this server cannot serialize — comes + back as a packed response that is also cached. An escaping + exception would kill the connection without a cache entry and + invite the client to resend a request whose side effects already + happened. + """ + payload: Optional[bytes] = None + if not isinstance(request, dict): + # Nothing has executed; report the malformed frame instead of + # raising AttributeError out of the handler and killing the + # connection. + if request_id: + executor_self._ws_inflight.pop(request_id, None) + return _pack_ws_error( + error_type="InvalidRequest", + error=( + "Expected a msgpack map for an execution request, got " + f"{type(request).__name__}; the custom Python block was " + "not run." + ), + request_id=request_id, + ) + code_str = request.get("code_str", "") + imports = request.get("imports", []) + run_function_name = request.get("run_function_name", "") + inputs_raw = request.get("inputs", {}) + client_code_hash = request.get("code_hash", "") + workflow_context = request.get("workflow_context") or {} + + # ``decode_failed`` is ONE of two positive proofs that the user's + # code did not run (the other is ``never_ran``, set for control + # responses returned before the function is looked up). What it is + # not is an inference from ``started_running``: the worker sets it + # before touching anything + # else. ``started_running`` cannot serve that role on its own, + # because ``asyncio.to_thread`` cancellation cancels the awaiting + # future while the pool thread keeps going — so observing it False + # after a CancelledError does not mean the block will not run a + # moment later. Treating that as "never ran" would clear the + # executed marker and invite the client to re-run a block that + # then executes anyway. + started_running = False + decode_failed = False + never_ran = False + + def _run_user_code(): + """Decode, run and pack — all off the event loop. + + This container serves up to ``max_inputs=10`` websocket + connections from ONE event loop. Decoding (``cv2.imdecode`` + plus pydantic validation per image) and encoding + (``ndarray.tolist()`` — a 1920x1080x3 result is 6.2M Python + ints) are both GIL-held CPU work measured in seconds. Doing + either on the loop stalls the other nine connections past the + idle deadline of their ``receive_bytes``, closing sockets + whose heartbeat was already on the wire — i.e. it defeats the + heartbeat this protocol adds. Only packed bytes cross back. + + Decoding stays BEFORE ``started_running`` so a malformed + payload is still reported as never-executed and stays safely + resendable, exactly as when it ran on the loop. + """ + nonlocal started_running, decode_failed, never_ran + try: + inputs = Executor._deserialize_msgpack_inputs(inputs_raw) + except Exception as error: + decode_failed = True + raise _InputsDecodeError(str(error)) from error + started_running = True + resp = Executor._run_user_code_ws( + executor_self, + code_str, + imports, + run_function_name, + inputs, + client_code_hash, + workflow_context, + ) + # Control responses returned BEFORE the user's function is + # looked up: the block provably did not run, so its executed + # marker must not bar a resend. (A namespace-init failure is + # NOT in this set — exec()ing the user's module scope IS + # running their code.) + if not resp.get("success") and resp.get("error_type") in ( + "UnknownCodeHash", + "InvalidRequest", + ): + never_ran = True + succeeded = bool(resp.get("success")) + if succeeded: + resp["result"] = Executor._serialize_msgpack_result(resp["result"]) + if request_id: + resp["request_id"] = request_id + return succeeded, msgpack.packb(resp, use_bin_type=True) + + try: + # Mark BEFORE running: from here on a resend must never + # re-execute, however this call ends. + if request_id: + executor_self._ws_executed.add(request_id) + succeeded, payload = await asyncio.to_thread(_run_user_code) + + if never_ran and request_id: + # Provable non-execution: keep the id resendable. + executor_self._ws_executed.discard(request_id) + + if succeeded: + # Only now has runtime state actually been built here; + # see the note on _ws_register_session. Stays on the loop + # thread, which is where the registry is documented to be + # touched. + executor_self._ws_register_session(session_id) + except BaseException as error: + # BaseException, not Exception: untrusted user code can raise + # SystemExit (``sys.exit()``), which concurrent.futures + # re-raises here, and Modal cancelling the input raises + # CancelledError. Both used to escape uncaught, tearing the + # connection down with no cached payload and leaving the + # request id poisoned for the whole executed-marker TTL. + traceback.print_exc() + if decode_failed: + # The ONLY case we can prove the block did not run: the + # worker positively reported that decoding failed, so it + # never reached the user's code and never will. Safe to + # clear the marker and leave the id resendable. + if request_id: + executor_self._ws_executed.discard(request_id) + payload = _pack_ws_error( + error_type=type( + getattr(error, "__cause__", None) or error + ).__name__, + error=( + "The server could not decode this request's inputs; " + f"the custom Python block was not run: {error}" + ), + request_id=request_id, + ) + elif not started_running: + # The worker had not entered the user code when this was + # observed — but a cancelled ``to_thread`` leaves the pool + # thread running, so the block may still execute. KEEP the + # executed marker: a resend must get a loud + # ResponseNoLongerAvailable rather than a second run, and + # the wording must not invite a checkpoint replay. + payload = _pack_ws_error( + error_type=type(error).__name__, + error=( + "The server failed while dispatching the custom " + "Python block; whether it ran is unknown, so it was " + f"not retried: {error}" + ), + request_id=request_id, + ) + else: + payload = _pack_ws_error( + error_type=type(error).__name__, + error=( + "The custom Python block ran, but the server could " + f"not return its response: {error}" + ), + request_id=request_id, + ) + if isinstance(error, asyncio.CancelledError): + # Cache/clear via ``finally``, then let cancellation + # propagate — swallowing it would desynchronise the task. + raise + finally: + # Cache before clearing in-flight so a resend arriving in + # between finds one or the other, never a gap. + if request_id: + if payload is not None: + executor_self._ws_response_cache.put(request_id, payload) + executor_self._ws_inflight.pop(request_id, None) + return payload + @ws_app.websocket("/ws") async def ws_execute(websocket: WebSocket): await websocket.accept() connected_at = time.monotonic() + conn_session_id = "" + # Last request id seen on this connection, so the terminal error + # handler below can address its frame. On protocol v2 the client + # rejects a response without the in-flight id, so an unaddressed + # error frame is reported to the user as a frame desync and the + # real diagnostic is lost. + conn_last_request_id: Optional[str] = None + # Set once this connection has spoken protocol v2 (i.e. sent a + # hello). A v1 client would try to msgpack-decode the ``closing`` + # frame below as an execution response, so it must never see one. + conn_is_v2 = False + + async def _close_gracefully() -> None: + """Close after telling a v2 client the frame it may have just + sent was never read. + + Both callers close from the TOP of the loop, before any + ``receive_bytes``: once we decide to close we never read + again, so anything the client wrote concurrently is + guaranteed unprocessed. Saying so explicitly is what lets the + client retry it on any container instead of reporting the + ambiguous "may have already executed". The error paths close + with 1011 and their own error frame, so ``closing`` + 1000 is + unambiguously "nothing ran". + """ + if conn_is_v2: + try: + await websocket.send_bytes( + msgpack.packb({"_kind": "closing"}, use_bin_type=True) + ) + except Exception: + # The peer is already gone; the close below is enough. + pass + try: + await websocket.close(code=1000) + except Exception: + pass + try: while True: + # Cleared per iteration so a failure before this frame's id + # is parsed can never echo the PREVIOUS request's id, which + # the client would reject as a stale frame and whose real + # diagnostic would be lost. + conn_last_request_id = None remaining = WEBEXEC_WS_MAX_CONNECTION_SECONDS - ( time.monotonic() - connected_at ) if remaining <= 0: - await websocket.close(code=1000) + await _close_gracefully() return try: raw = await asyncio.wait_for( @@ -937,62 +1588,245 @@ async def ws_execute(websocket: WebSocket): timeout=min(remaining, WEBEXEC_WS_IDLE_TIMEOUT_SECONDS), ) except asyncio.TimeoutError: - await websocket.close(code=1000) + await _close_gracefully() return - request = msgpack.unpackb(raw, raw=False) + request = _safe_unpackb(raw) if isinstance(request, dict) and "_chunked" in request: + chunk_count = request["_chunked"] + if ( + not isinstance(chunk_count, int) + or isinstance(chunk_count, bool) + or not 1 <= chunk_count <= WEBEXEC_WS_MAX_CHUNKS + ): + raise ValueError( + f"invalid websocket chunk count: {chunk_count!r}" + ) parts = [] - for _ in range(request["_chunked"]): - parts.append( - await asyncio.wait_for( - websocket.receive_bytes(), - timeout=WEBEXEC_WS_IDLE_TIMEOUT_SECONDS, + # Bound the BYTES, not just the chunk count: the count + # ceiling alone admits 1024 max-size frames, and the + # join then doubles the peak. The limit is what this + # container can decode INSIDE one idle deadline while + # nine sibling connections wait on the same loop — + # not what the chunk ceiling happens to allow. + reassembled_bytes = 0 + for _ in range(chunk_count): + # Re-derive the budget each chunk: a peer dribbling + # one chunk just inside the idle timeout would + # otherwise hold a max_inputs slot for + # chunk_count * IDLE seconds, far past the + # connection cap, until Modal kills the input with + # no close frame. + chunk_remaining = WEBEXEC_WS_MAX_CONNECTION_SECONDS - ( + time.monotonic() - connected_at + ) + if chunk_remaining <= 0: + raise ValueError( + "connection cap reached during chunked receive" ) + part = await asyncio.wait_for( + websocket.receive_bytes(), + timeout=min( + chunk_remaining, WEBEXEC_WS_IDLE_TIMEOUT_SECONDS + ), ) - request = msgpack.unpackb(b"".join(parts), raw=False) - - code_str = request.get("code_str", "") - imports = request.get("imports", []) - run_function_name = request.get("run_function_name", "") - inputs_raw = request.get("inputs", {}) - client_code_hash = request.get("code_hash", "") - workflow_context = request.get("workflow_context") or {} - - inputs = Executor._deserialize_msgpack_inputs(inputs_raw) - resp = await asyncio.to_thread( - Executor._run_user_code_ws, - executor_self, - code_str, - imports, - run_function_name, - inputs, - client_code_hash, - workflow_context, - ) + reassembled_bytes += len(part) + if reassembled_bytes > WEBEXEC_WS_MAX_REQUEST_BYTES: + raise ValueError( + "chunked websocket request exceeds " + f"{WEBEXEC_WS_MAX_REQUEST_BYTES} bytes" + ) + parts.append(part) + # Off the event loop: unpacking tens of MB is GIL-held + # CPU that would otherwise stall every other + # connection past its idle deadline. + # join INSIDE the worker: argument expressions are + # evaluated on the loop before to_thread is awaited, so + # joining here would copy up to the byte cap on the very + # thread this call exists to keep free. + request = await asyncio.to_thread( + lambda: _safe_unpackb(b"".join(parts)) + ) - if resp.get("success"): - resp["result"] = Executor._serialize_msgpack_result( - resp["result"] + # ---- protocol v2 control frames ---- + kind = request.get("_kind") if isinstance(request, dict) else None + if kind == "hello": + # Coerce at the boundary: everything downstream (the + # session registry, the dedup registries) assumes a + # non-empty str, and a decoded msgpack map here would + # otherwise raise TypeError and kill the connection. + raw_session_id = request.get("session_id") + conn_session_id = ( + raw_session_id if isinstance(raw_session_id, str) else "" + ) + # From here the peer understands v2 control frames, so + # a graceful close may announce itself with ``closing``. + conn_is_v2 = True + await websocket.send_bytes( + msgpack.packb( + { + "_kind": "hello", + "proto": 2, + "idle_timeout_s": WEBEXEC_WS_IDLE_TIMEOUT_SECONDS, + "session_known": executor_self._ws_session_seen( + conn_session_id + ), + # Lets the client tell whether a + # reconnect landed on the same container + # (whose dedup cache it may rely on) or + # a different one. + "container_id": executor_self._container_id, + }, + use_bin_type=True, + ) ) + continue + if kind == "heartbeat": + # An application-level frame is the only thing that + # resets this loop's receive_bytes idle timer; + # protocol ws pings are consumed by the ASGI layer + # and never reach here. + await websocket.send_bytes( + msgpack.packb({"_kind": "heartbeat_ack"}, use_bin_type=True) + ) + continue - payload = msgpack.packb(resp, use_bin_type=True) - if len(payload) > WEBEXEC_WS_MAX_FRAME_BYTES: - chunks = [ - payload[i : i + WEBEXEC_WS_MAX_FRAME_BYTES] - for i in range( - 0, len(payload), WEBEXEC_WS_MAX_FRAME_BYTES + raw_request_id = ( + request.get("request_id") if isinstance(request, dict) else None + ) + if ( + isinstance(raw_request_id, str) + and len(raw_request_id) > _TtlKeySet._MAX_KEY_LENGTH + ): + # Reject rather than coerce to None: the registries + # would silently refuse to store an oversized key, so + # the request would execute with NO dedup record and a + # resend would run the user's code again. Legitimate + # clients send a uuid4 hex. + await websocket.send_bytes( + _pack_ws_error( + error_type="InvalidRequest", + error=( + "request_id exceeds " + f"{_TtlKeySet._MAX_KEY_LENGTH} characters; " + "the block was not run." + ), + request_id=None, ) - ] + ) + continue + if raw_request_id is not None and not isinstance( + raw_request_id, str + ): + # Silently downgrading to "no dedup" would execute the + # block with no at-most-once record, so a resend runs it + # again. An over-long str is already rejected above; + # reject a wrong-typed one the same way. await websocket.send_bytes( - msgpack.packb( - {"_chunked": len(chunks)}, use_bin_type=True + _pack_ws_error( + error_type="InvalidRequest", + error=( + "request_id must be a string; the block was " + "not run." + ), + request_id=None, ) ) - for chunk in chunks: - await websocket.send_bytes(chunk) + continue + request_id = ( + raw_request_id + if isinstance(raw_request_id, str) and raw_request_id + else None + ) + conn_last_request_id = request_id + if request_id: + cached_payload = executor_self._ws_response_cache.get( + request_id + ) + if cached_payload is not None: + # Resend of a request already executed (the + # client lost the response): answer from cache, + # never run user code twice. + await send_payload(websocket, cached_payload, request_id) + continue + # No await between the in-flight lookup and the + # insert below, so two connections can't both start + # the same request. + task = executor_self._ws_inflight.get(request_id) + if task is None: + if request_id in executor_self._ws_executed: + # Already ran here, but its response is no + # longer available (evicted, or the call was + # cancelled). Re-running would duplicate the + # block's side effects, so fail loudly + # instead — execution stays at-most-once. + await send_payload( + websocket, + _pack_ws_error( + error_type="ResponseNoLongerAvailable", + error=( + "This request already executed on " + "this container and its response is " + "no longer available. It was not " + "run again, to avoid duplicating " + "the block's side effects." + ), + request_id=request_id, + ), + ) + continue + task = asyncio.create_task( + execute_request(request, request_id, conn_session_id) + ) + executor_self._ws_inflight[request_id] = task + + # execute_request's ``finally`` clears the entry + # in the normal case, but a task cancelled before + # its body ever starts (loop shutdown, container + # teardown at the input timeout) never reaches it. + # _ws_inflight is the one v2 structure with no cap + # and no TTL, so close that leak here. + # Identity-checked: popping by key alone lets a + # stale callback delete a NEWER task registered + # under the same id, which would re-open the + # duplicate-execution window a resend relies on + # this map to close. + def _clear_inflight( + finished: "asyncio.Task", + rid: str = request_id, + ) -> None: + if executor_self._ws_inflight.get(rid) is finished: + executor_self._ws_inflight.pop(rid, None) + + task.add_done_callback(_clear_inflight) + # shield: this connection dying must not cancel an + # execution another connection may be waiting on + # (or will resend for). + payload = await asyncio.shield(task) else: - await websocket.send_bytes(payload) + payload = await execute_request(request, None, conn_session_id) + await send_payload(websocket, payload, request_id) except WebSocketDisconnect: pass + except Exception as error: + # Anything the per-request paths did not already turn into a + # response frame (a malformed control frame, a broken chunk + # header, a send failure) would otherwise tear the connection + # down silently. Report it, then close with 1011 so the client + # sees a server error rather than an opaque disconnect. + traceback.print_exc() + try: + await websocket.send_bytes( + _pack_ws_error( + error_type=type(error).__name__, + error=f"websocket handler failed: {error}", + request_id=conn_last_request_id, + ) + ) + except Exception: + pass + try: + await websocket.close(code=1011) + except Exception: + pass return ws_app diff --git a/requirements/_requirements.txt b/requirements/_requirements.txt index 0634e41890..eafe316edc 100644 --- a/requirements/_requirements.txt +++ b/requirements/_requirements.txt @@ -45,6 +45,12 @@ slack-sdk~=3.33.4 twilio~=9.3.7 httpx~=0.28.1 msgpack +# Both are hard requirements of WEBEXEC_TRANSPORT=websocket (the custom Python +# block transport in +# inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py). +# websocket-client was previously only present transitively via the docker SDK, +# which made an opt-in transport depend on an unrelated package's dependency. +websocket-client~=1.9.0 pylogix==1.0.5 pymodbus>=3.6.9,<=3.8.3 backoff~=2.2.0 diff --git a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/conftest.py b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/conftest.py index f02aea26f7..d2ddb6d04f 100644 --- a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/conftest.py +++ b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/conftest.py @@ -1,9 +1,22 @@ -"""Shared scaffolding for tests that load ``modal/modal_app.py`` directly. - -``modal_app`` imports the real ``modal`` package at module scope and decorates -``Executor`` with ``@app.cls`` / ``@modal.concurrent`` / ``@modal.fastapi_endpoint``. -Stubbing those out lets the sandbox-side code be imported and driven in-process, -so tests can exercise the shipped path instead of reimplementing it. +"""Shared harness for tests that load the real ``modal/modal_app.py``. + +``modal_app.py`` is deployed separately (``modal deploy``) and is not part of +the ``inference`` package, so it is imported here by path with a stubbed +``modal`` module. Loading the real server module is what makes the websocket +protocol tests genuine behavior tests rather than mock theater. + +Two equivalent harnesses were written independently (one on ``main``, one on +the webexec websocket branch) and met here. This file keeps ONE +implementation and re-exports the older public names, so neither set of tests +had to be rewritten to land the merge: + +* ``load_modal_app(monkeypatch, module_name)`` — the loader. Takes the module + name so each test can hold a FRESH module, which matters because the server + keeps container-local state (compiled namespaces, session and dedup + registries) at module/class scope. +* ``modal_app`` / ``modal_app_with_fake_modal`` — fixtures over that loader. +* ``FakeModalApp`` / ``FakeModalImage`` / ``identity_decorator`` — the stubs, + also exported under their original underscore-free names. """ import importlib.util @@ -14,7 +27,7 @@ import pytest -class FakeModalImage: +class _FakeModalImage: @classmethod def debian_slim(cls, *args, **kwargs): return cls() @@ -33,35 +46,64 @@ def entrypoint(self, *args, **kwargs): return self -class FakeModalApp: - def __init__(self, name: str): - self.name = name +class _FakeModalApp: + def __init__(self, *args, **kwargs): + pass def cls(self, *args, **kwargs): return lambda cls: cls -def identity_decorator(*args, **kwargs): +def _identity_decorator(*args, **kwargs): return lambda obj: obj -@pytest.fixture() -def modal_app_with_fake_modal(monkeypatch): - """Import ``modal/modal_app.py`` with a stubbed ``modal`` package.""" +# Names the tests that arrived via ``main`` import directly. +FakeModalImage = _FakeModalImage +FakeModalApp = _FakeModalApp +identity_decorator = _identity_decorator + + +def load_modal_app(monkeypatch, module_name: str = "modal_app_under_test"): + """Import ``modal/modal_app.py`` fresh, with ``modal`` stubbed out. + + A fresh module per test keeps container-local state (namespaces, session + and dedup registries) from leaking between tests. + """ fake_modal = ModuleType("modal") - fake_modal.App = FakeModalApp - fake_modal.Image = FakeModalImage + fake_modal.App = _FakeModalApp + fake_modal.Image = _FakeModalImage fake_modal.parameter = lambda *args, **kwargs: None - fake_modal.enter = identity_decorator - fake_modal.fastapi_endpoint = identity_decorator - fake_modal.asgi_app = identity_decorator - fake_modal.concurrent = identity_decorator + fake_modal.enter = _identity_decorator + fake_modal.fastapi_endpoint = _identity_decorator + fake_modal.asgi_app = _identity_decorator + fake_modal.concurrent = _identity_decorator monkeypatch.setitem(sys.modules, "modal", fake_modal) modal_app_path = Path(__file__).resolve().parents[5] / "modal" / "modal_app.py" - spec = importlib.util.spec_from_file_location( - "modal_app_under_test", modal_app_path - ) + spec = importlib.util.spec_from_file_location(module_name, modal_app_path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module + + +def build_ws_app(modal_app, run_user_code): + """Build the websocket ASGI app with only user-code execution stubbed.""" + cls = modal_app.Executor + user_cls = cls._get_user_cls() if hasattr(cls, "_get_user_cls") else cls + executor = user_cls.__new__(user_cls) + executor.workspace_id = "test-ws" + user_cls.identify(executor) + user_cls._run_user_code_ws = staticmethod(run_user_code).__func__ + return executor, user_cls.wsapp(executor) + + +@pytest.fixture() +def modal_app(monkeypatch): + return load_modal_app(monkeypatch, "modal_app_under_test") + + +@pytest.fixture() +def modal_app_with_fake_modal(monkeypatch): + """Alias of :func:`modal_app`, kept for the tests that arrived via ``main``.""" + return load_modal_app(monkeypatch, "modal_app_under_test") diff --git a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_ws_contract.py b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_ws_contract.py new file mode 100644 index 0000000000..3ae2d25da5 --- /dev/null +++ b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_ws_contract.py @@ -0,0 +1,555 @@ +"""Client <-> server wire contract for the webexec websocket protocol. + +The client ships inside the ``inference`` package and the server is deployed +separately with ``modal deploy``, so the two halves ride independent release +trains — exactly the situation where a silently drifted constant or wire key +bites. Every value that crosses the boundary is pinned here against BOTH +implementations, loaded as real modules. + +The repo already established this pattern in +``test_modal_code_hash.py::test_client_and_server_code_hashes_stay_in_sync``. +""" + +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import msgpack +import pytest + +from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( + modal_executor as client, +) + +from .conftest import build_ws_app as _ws_app + + +class TestSharedConstants: + def test_max_frame_bytes_match(self, modal_app) -> None: + # Both sides chunk independently; a mismatch means one side splits at + # a size the other never expects. + assert client._WS_MAX_FRAME_BYTES == modal_app.WEBEXEC_WS_MAX_FRAME_BYTES + + def test_max_chunk_count_matches(self, modal_app) -> None: + assert client._WS_MAX_CHUNKS == modal_app.WEBEXEC_WS_MAX_CHUNKS + + def test_connection_cap_stays_under_the_modal_input_timeout( + self, modal_app + ) -> None: + # A websocket connection is one Modal input. If the cap is not below + # the input timeout, Modal kills the connection instead of the server + # closing it cleanly: no close frame, and in-flight executions are + # cancelled mid-run. + assert ( + modal_app.WEBEXEC_WS_MAX_CONNECTION_SECONDS + < modal_app._executor_decorator_kwargs["timeout"] + ) + + def test_connection_cap_is_advertised_to_containers(self, modal_app) -> None: + assert modal_app._executor_decorator_kwargs["env"][ + "WEBEXEC_WS_MAX_CONNECTION_SECONDS" + ] == str(modal_app.WEBEXEC_WS_MAX_CONNECTION_SECONDS) + + +class TestHandshakeContract: + def test_client_handshake_against_the_real_server_reply(self, modal_app) -> None: + from fastapi.testclient import TestClient + + _, app = _ws_app( + modal_app, lambda self, *a, **kw: {"success": True, "result": {}} + ) + + executor = client.WebSocketModalExecutor(workspace_id="test-ws") + + with TestClient(app).websocket_connect("/ws") as ws: + # Drive the REAL client handshake over the REAL server route. + executor._ws = _BridgeWS(ws) + executor._handshake() + + assert executor._server.proto == 2 + # The one constant that cannot drift: learned, not duplicated. + assert executor._server.idle_timeout == float( + modal_app.WEBEXEC_WS_IDLE_TIMEOUT_SECONDS + ) + assert executor._server.container_id is not None + + def test_reconnect_to_the_same_container_reports_session_known( + self, modal_app + ) -> None: + from fastapi.testclient import TestClient + + _, app = _ws_app( + modal_app, lambda self, *a, **kw: {"success": True, "result": {}} + ) + executor = client.WebSocketModalExecutor(workspace_id="test-ws") + test_client = TestClient(app) + + with test_client.websocket_connect("/ws") as ws: + executor._ws = _BridgeWS(ws) + executor._handshake() + first_container = executor._server.container_id + ws.send_bytes( + msgpack.packb( + {"request_id": "r1", "inputs": {}, "code_hash": "h"}, + use_bin_type=True, + ) + ) + assert msgpack.unpackb(ws.receive_bytes(), raw=False)["success"] is True + + executor._had_success = True + with test_client.websocket_connect("/ws") as ws: + executor._ws = _BridgeWS(ws) + # Same container object -> the session is known, no loud failure. + executor._handshake() + + assert executor._server.container_id == first_container + + +class TestRequestFrameContract: + def test_every_key_the_client_sends_is_a_key_the_server_reads( + self, modal_app + ) -> None: + from fastapi.testclient import TestClient + + seen = {} + + def run_user_code( + self, + code_str, + imports, + run_function_name, + inputs, + client_code_hash, + workflow_context, + ): + seen.update( + code_str=code_str, + imports=imports, + run_function_name=run_function_name, + inputs=inputs, + code_hash=client_code_hash, + workflow_context=workflow_context, + ) + return {"success": True, "result": {}} + + _, app = _ws_app(modal_app, run_user_code) + + frame = client.WebSocketModalExecutor._build_ws_frame( + python_code=SimpleNamespace( + run_function_code="def run(x):\n return x\n", + run_function_name="run", + imports=["import os"], + ), + packed_inputs={"x": 1}, + code_hash="hash-1", + send_full_code=True, + msgpack=msgpack, + workflow_context={"workflow_id": "wf-1"}, + request_id="req-1", + ) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(frame) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert resp["success"] is True + assert resp["request_id"] == "req-1" + assert seen == { + "code_str": "def run(x):\n return x\n", + "imports": ["import os"], + "run_function_name": "run", + "inputs": {"x": 1}, + "code_hash": "hash-1", + "workflow_context": {"workflow_id": "wf-1"}, + } + + +class TestChunkingContract: + def test_server_reassembles_a_client_split_frame(self, modal_app) -> None: + from fastapi.testclient import TestClient + + seen = {} + + def run_user_code(self, code_str, imports, name, inputs, code_hash, ctx): + seen["inputs"] = inputs + return {"success": True, "result": {}} + + _, app = _ws_app(modal_app, run_user_code) + big = b"x" * (client._WS_MAX_FRAME_BYTES * 2 + 17) + frame = msgpack.packb( + {"request_id": "req-1", "inputs": {"blob": big}}, use_bin_type=True + ) + frames = client._split_ws_frames(frame, msgpack) + assert len(frames) == 4 # control frame + 3 chunks + + with TestClient(app).websocket_connect("/ws") as ws: + for part in frames: + ws.send_bytes(part) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert resp["success"] is True + assert seen["inputs"] == {"blob": big} + + def test_client_reassembles_a_server_split_payload(self, modal_app) -> None: + from fastapi.testclient import TestClient + + big = b"y" * (modal_app.WEBEXEC_WS_MAX_FRAME_BYTES * 2 + 5) + + def run_user_code(self, *args, **kwargs): + return {"success": True, "result": {"blob": big}} + + _, app = _ws_app(modal_app, run_user_code) + executor = client.WebSocketModalExecutor(workspace_id="test-ws") + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes( + msgpack.packb({"request_id": "req-1", "inputs": {}}, use_bin_type=True) + ) + executor._ws = _BridgeWS(ws) + payload = executor._recv_reassembled(msgpack) + + assert msgpack.unpackb(payload, raw=False)["result"]["blob"] == big + + def test_a_cached_response_is_re_chunked_on_resend(self, modal_app) -> None: + from fastapi.testclient import TestClient + + big = b"z" * (modal_app.WEBEXEC_WS_MAX_FRAME_BYTES * 2 + 5) + calls = [] + + def run_user_code(self, *args, **kwargs): + calls.append(1) + return {"success": True, "result": {"blob": big}} + + _, app = _ws_app(modal_app, run_user_code) + executor = client.WebSocketModalExecutor(workspace_id="test-ws") + frame = msgpack.packb({"request_id": "req-1", "inputs": {}}, use_bin_type=True) + + with TestClient(app).websocket_connect("/ws") as ws: + executor._ws = _BridgeWS(ws) + ws.send_bytes(frame) + first = executor._recv_reassembled(msgpack) + ws.send_bytes(frame) + second = executor._recv_reassembled(msgpack) + + assert len(calls) == 1 + assert first == second + + +class TestErrorTypeContract: + def test_unknown_code_hash_string_matches(self, modal_app) -> None: + from fastapi.testclient import TestClient + + def run_user_code(self, code_str, imports, name, inputs, code_hash, ctx): + # Mirrors the real hash-only path: no code_str and no cached + # namespace for this hash. + return { + "success": False, + "error": "unknown hash", + "error_type": "UnknownCodeHash", + } + + _, app = _ws_app(modal_app, run_user_code) + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes( + msgpack.packb( + {"request_id": "req-1", "inputs": {}, "code_hash": "h"}, + use_bin_type=True, + ) + ) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + # The literal the client keys its resend-with-full-code retry on. + # (This stubs _run_user_code_ws, so it pins the wire literal only; + # the real server's stamping is covered by + # test_unknown_code_hash_is_stamped_as_server_originated below.) + assert resp["error_type"] == "UnknownCodeHash" + + def test_unknown_code_hash_actually_drives_the_full_code_resend(self) -> None: + # Grepping the client source for the literal proves nothing: a broken + # retry condition, or one that resends with a stale hash, passes that + # check. Drive the real path instead — first attempt hash-only, and + # assert the retry carries the full code and succeeds. + executor = client.WebSocketModalExecutor(workspace_id="test-ws") + executor._server = client._ServerInfo( + proto=2, idle_timeout=10.0, container_id="container-1" + ) + code_hash = client._compute_code_hash("def run(): return 1", []) + executor._hashes_sent_on_ws = {code_hash} + sent: list = [] + + def fake_send_recv(frame_bytes, workspace, request_id=None): + sent.append(msgpack.unpackb(frame_bytes, raw=False)) + if len(sent) == 1: + return msgpack.packb( + { + "success": False, + "error": "unknown hash", + "error_type": "UnknownCodeHash", + "server_error": True, + "request_id": request_id, + }, + use_bin_type=True, + ) + return msgpack.packb( + {"success": True, "result": {}, "request_id": request_id}, + use_bin_type=True, + ) + + executor._send_recv_with_retry = fake_send_recv + code = SimpleNamespace( + run_function_code="def run(): return 1", + run_function_name="run", + imports=[], + ) + + executor._execute_ws("MyBlock", code, {}, "test-ws", msgpack, {}) + + assert len(sent) == 2 + assert "code_str" not in sent[0], "first attempt should be hash-only" + assert sent[1]["code_str"] == "def run(): return 1" + # A distinct logical request, or the server's dedup cache would just + # replay the failure it already cached for the first id. + assert sent[0]["request_id"] != sent[1]["request_id"] + + def test_response_no_longer_available_is_recognised_by_the_client( + self, modal_app + ) -> None: + from fastapi.testclient import TestClient + + def run_user_code(self, *args, **kwargs): + return {"success": True, "result": {}} + + executor_obj, app = _ws_app(modal_app, run_user_code) + frame = msgpack.packb({"request_id": "req-1", "inputs": {}}, use_bin_type=True) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(frame) + msgpack.unpackb(ws.receive_bytes(), raw=False) + executor_obj._ws_response_cache = modal_app._WsResponseCache() + ws.send_bytes(frame) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + # The client must classify this as a TRANSPORT failure, not as the + # user's block raising. + from inference.core.workflows.errors import ( + DynamicBlockCodeError, + DynamicBlockError, + ) + + executor = client.WebSocketModalExecutor(workspace_id="test-ws") + with pytest.raises(DynamicBlockError) as excinfo: + executor._raise_server_error_if_infrastructure(resp, "MyBlock") + assert not isinstance(excinfo.value, DynamicBlockCodeError) + + def test_server_stamps_infrastructure_errors(self, modal_app) -> None: + from fastapi.testclient import TestClient + + def run_user_code(self, *args, **kwargs): + raise AssertionError("must not run") + + _, app = _ws_app(modal_app, run_user_code) + with TestClient(app).websocket_connect("/ws") as ws: + # Undecodable inputs: the block is never run. + ws.send_bytes( + msgpack.packb( + {"request_id": "req-1", "inputs": {"i": {"type": "nope"}}}, + use_bin_type=True, + ) + ) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert resp["success"] is False + assert resp.get("server_error") is True + + def test_unknown_code_hash_is_stamped_as_server_originated(self, modal_app) -> None: + # Drives the REAL _run_user_code_ws hash-only path with nothing + # cached, so it pins the stamping the client now relies on instead of + # matching a forgeable error_type name. + from fastapi.testclient import TestClient + + # Deliberately NOT stubbed: this is the one path where the real + # implementation is what we are pinning. + real_run_user_code = modal_app.Executor._run_user_code_ws + _, app = _ws_app(modal_app, real_run_user_code) + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes( + msgpack.packb( + { + "request_id": "req-1", + "inputs": {}, + "code_hash": "never-compiled-here", + "run_function_name": "run", + }, + use_bin_type=True, + ) + ) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert resp["success"] is False + assert resp["error_type"] == "UnknownCodeHash" + assert resp["server_error"] is True + + def test_chunk_count_at_the_limit_round_trips(self, monkeypatch) -> None: + # Only invalid counts were covered, so an off-by-one in either side's + # ``1 <= n <= MAX`` guard passed every test. Pin the valid upper + # boundary from both directions. + # + # The limits are monkeypatched down rather than allocating + # MAX_FRAME_BYTES * MAX_CHUNKS (1 GiB) of real payload: the guard + # under test is the COUNT comparison, and a unit test should not + # allocate a gigabyte to exercise it. + monkeypatch.setattr(client, "_WS_MAX_FRAME_BYTES", 8) + monkeypatch.setattr(client, "_WS_MAX_CHUNKS", 4) + payload = b"z" * (8 * 4) + + frames = client._split_ws_frames(payload, msgpack) + + assert msgpack.unpackb(frames[0], raw=False) == {"_chunked": 4} + assert b"".join(frames[1:]) == payload + + executor = client.WebSocketModalExecutor(workspace_id="test-ws") + executor._ws = SimpleNamespace( + recv=iter(frames).__next__, + gettimeout=lambda: None, + settimeout=lambda _v: None, + ) + assert executor._recv_reassembled(msgpack) == payload + + def test_one_chunk_over_the_limit_is_refused_before_the_wire( + self, monkeypatch + ) -> None: + monkeypatch.setattr(client, "_WS_MAX_FRAME_BYTES", 8) + monkeypatch.setattr(client, "_WS_MAX_CHUNKS", 4) + payload = b"z" * (8 * 4 + 1) + + from inference.core.workflows.errors import DynamicBlockError + + with pytest.raises(DynamicBlockError, match="too large"): + client._split_ws_frames(payload, msgpack) + + def test_oversized_reassembled_response_is_refused(self, monkeypatch) -> None: + # The chunk COUNT ceiling alone still admits 1 GiB, reassembled inside + # the shared inference server process once per concurrent executor. + # An oversized result must fail its own request, not the process. + monkeypatch.setattr(client, "_WS_MAX_RESPONSE_BYTES", 16) + frames = [msgpack.packb({"_chunked": 3}, use_bin_type=True)] + [ + b"z" * 8, + b"z" * 8, + b"z" * 8, + ] + executor = client.WebSocketModalExecutor(workspace_id="test-ws") + executor._ws = SimpleNamespace( + recv=iter(frames).__next__, + gettimeout=lambda: None, + settimeout=lambda _v: None, + ) + + with pytest.raises(ConnectionError, match="exceeds"): + executor._recv_reassembled(msgpack) + + +class _BridgeWS: + """Adapts a starlette TestClient websocket to websocket-client's surface.""" + + def __init__(self, ws: Any): + self._ws = ws + + def send_binary(self, frame: bytes) -> None: + self._ws.send_bytes(frame) + + def recv(self) -> Any: + message = self._ws.receive() + if "bytes" in message and message["bytes"] is not None: + return message["bytes"] + return message.get("text", "") + + def settimeout(self, value: float) -> None: + pass + + def close(self) -> None: + pass + + def ping(self) -> None: + pass + + +class TestGracefulCloseContract: + """The `closing` frame is the client's PROOF that a frame was never read. + + The client resets its delivery state on it and retries on ANY container, + so if the real server could ever emit `closing` after having read a frame, + that would be duplicate execution of user code. These drive the real + server's own close triggers rather than a hand-written frame. + """ + + def test_server_announces_closing_on_the_idle_timeout(self, modal_app) -> None: + from fastapi.testclient import TestClient + + # Idle deadline short enough to fire without slowing the suite. + modal_app.WEBEXEC_WS_IDLE_TIMEOUT_SECONDS = 1 + _, app = _ws_app( + modal_app, lambda self, *a, **kw: {"success": True, "result": {}} + ) + + with TestClient(app).websocket_connect("/ws") as ws: + # Say hello so the connection is v2 — a v1 client must never be + # sent this frame. + ws.send_bytes( + msgpack.packb( + {"_kind": "hello", "proto": 2, "session_id": "s-1"}, + use_bin_type=True, + ) + ) + assert msgpack.unpackb(ws.receive_bytes(), raw=False)["_kind"] == "hello" + # Then go silent and let the server's idle deadline fire. + frame = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert frame == {"_kind": "closing"} + + def test_server_announces_closing_on_the_connection_cap(self, modal_app) -> None: + from fastapi.testclient import TestClient + + # Cap already elapsed => the very first loop iteration closes. + modal_app.WEBEXEC_WS_MAX_CONNECTION_SECONDS = 0 + _, app = _ws_app( + modal_app, lambda self, *a, **kw: {"success": True, "result": {}} + ) + + with TestClient(app).websocket_connect("/ws") as ws: + # No hello: a v1 client must NOT receive a `closing` frame, since + # it would feed the unrecognised map straight to msgpack as a + # response. The server should just close. + with pytest.raises(Exception): + ws.receive_bytes() + + def test_v1_shaped_frame_still_executes(self, modal_app) -> None: + # "Backward compatible against v1 clients" is a wire claim: a frame + # with no _kind, no request_id and no session is what a v1 client + # actually sends. Every other test here sends a request_id. + from fastapi.testclient import TestClient + + seen = {} + + def run_user_code(self, code_str, imports, name, inputs, code_hash, ctx): + seen["inputs"] = inputs + return {"success": True, "result": {"ok": 1}} + + _, app = _ws_app(modal_app, run_user_code) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes( + msgpack.packb( + { + "code_str": "def run():\n return 1\n", + "imports": [], + "run_function_name": "run", + "inputs": {"x": 1}, + }, + use_bin_type=True, + ) + ) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert resp["success"] is True + assert "request_id" not in resp, "must not invent an id the client never sent" + assert seen["inputs"] == {"x": 1} diff --git a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_ws_protocol.py b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_ws_protocol.py new file mode 100644 index 0000000000..55c4f27f90 --- /dev/null +++ b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_ws_protocol.py @@ -0,0 +1,1011 @@ +"""Protocol v2 behavior of WebSocketModalExecutor. + +Covers the failure modes behind the idle-timeout incident: text frames on +dead connections, safe resend with request ids, and loud failure on lost +custom-Python sessions. +""" + +import sys +import threading +import time as _time +from types import SimpleNamespace +from typing import Any, List, Optional + +import msgpack +import pytest + +from inference.core.env import ( + WEBEXEC_WS_IDLE_RELEASE_SECONDS, + WEBEXEC_WS_READ_TIMEOUT_SECONDS, +) +from inference.core.workflows.errors import DynamicBlockCodeError, DynamicBlockError +from inference.core.workflows.execution_engine.v1.dynamic_blocks import modal_executor +from inference.core.workflows.execution_engine.v1.dynamic_blocks.modal_executor import ( + WebexecSessionLostError, + WebSocketModalExecutor, + _ServerInfo, +) + + +def _pack(obj: Any) -> bytes: + return msgpack.packb(obj, use_bin_type=True) + + +class _FakeWS: + """Scripted socket: each recv pops the next canned reply. + + A reply that is an Exception instance is raised instead of returned. + """ + + def __init__(self, replies: Optional[List[Any]] = None): + self.replies = list(replies or []) + self.sent: List[bytes] = [] + self.closed = False + self.timeouts: List[float] = [] + + def send_binary(self, frame: bytes) -> None: + self.sent.append(frame) + + def recv(self) -> Any: + if not self.replies: + raise AssertionError("recv called with no scripted reply") + reply = self.replies.pop(0) + if isinstance(reply, Exception): + raise reply + return reply + + def close(self) -> None: + self.closed = True + + def ping(self) -> None: + pass + + def settimeout(self, value: float) -> None: + self.timeouts.append(value) + + +def _executor_with_ws(ws: _FakeWS, proto: int = 2) -> WebSocketModalExecutor: + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._ws = ws + executor._server = _ServerInfo( + proto=proto, idle_timeout=10.0 if proto == 2 else None + ) + return executor + + +def _hello_reply( + session_known: bool = False, + idle_timeout_s: int = 10, + container_id: Optional[str] = "container-1", +) -> bytes: + reply = { + "_kind": "hello", + "proto": 2, + "idle_timeout_s": idle_timeout_s, + "session_known": session_known, + } + if container_id is not None: + reply["container_id"] = container_id + return _pack(reply) + + +class TestHandshake: + def test_v2_server_sets_proto_and_idle_timeout(self) -> None: + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._ws = _FakeWS([_hello_reply(idle_timeout_s=12)]) + + executor._handshake() + + assert executor._server.proto == 2 + assert executor._server.idle_timeout == 12.0 + sent = msgpack.unpackb(executor._ws.sent[0], raw=False) + assert sent["_kind"] == "hello" + assert sent["session_id"] == executor._session_id + + def test_v1_server_reply_falls_back_to_legacy(self) -> None: + executor = WebSocketModalExecutor(workspace_id="test-ws") + # A v1 server executes the hello as an empty request and answers + # with a plain response dict. + executor._ws = _FakeWS([_pack({"success": False, "error": "no code"})]) + + executor._handshake() + + assert executor._server.proto == 1 + assert executor._server.idle_timeout is None + + def test_v1_fallback_after_prior_success_does_not_fail_the_request(self) -> None: + # A v1 server cannot confirm the session, but it also idle-closes + # every 10s with no way for v1 keepalives to reset that timer, so + # reconnects are its normal state. Failing each one after a prior + # success would fail roughly every other request — worse than the + # rare silent reset it would prevent. The guarantee needs v2. + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._had_success = True + session = executor._session_id + executor._ws = _FakeWS([_pack({"success": False, "error": "no code"})]) + + executor._handshake() + + assert executor._server.proto == 1 + assert executor._had_success is True + assert executor._session_id == session + + def test_handshake_never_exposes_v2_proto_without_its_idle_timeout(self) -> None: + # The keepalive reads server info unlocked, so it must never observe + # proto 2 alongside a stale idle timeout: _heartbeat_interval would + # fall back to 25s, which a v2 server closes under. Publishing one + # immutable value makes the torn read unrepresentable — this test + # samples the state after EVERY bytecode line of the handshake. + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._server = _ServerInfo(proto=1, idle_timeout=None) + executor._ws = _FakeWS([_hello_reply(idle_timeout_s=9)]) + seen = [] + + handshake_code = WebSocketModalExecutor._handshake.__code__ + + def trace(frame, event, _arg): + if frame.f_code is not handshake_code: + return None + if event == "line": + server = executor._server + seen.append((server.proto, server.idle_timeout)) + return trace + + # Restore whatever tracer was active (coverage.py installs one), not + # None: clearing it would silently disable coverage for every test + # that runs after this one in the same worker process. + previous_tracer = sys.gettrace() + sys.settrace(trace) + try: + executor._handshake() + finally: + sys.settrace(previous_tracer) + + assert seen, "tracing did not observe the handshake" + inconsistent = [pair for pair in seen if pair[0] == 2 and pair[1] is None] + assert not inconsistent, f"torn server state observable: {inconsistent}" + assert executor._heartbeat_interval() == 3.0 + + def test_v2_server_container_id_is_recorded(self) -> None: + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._ws = _FakeWS([_hello_reply(container_id="container-7")]) + + executor._handshake() + + assert executor._server.container_id == "container-7" + + def test_session_lost_after_prior_success_raises(self, monkeypatch: Any) -> None: + monkeypatch.setattr(modal_executor, "WEBEXEC_WS_FAIL_ON_SESSION_LOSS", True) + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._had_success = True + executor._ws = _FakeWS([_hello_reply(session_known=False)]) + + with pytest.raises(WebexecSessionLostError): + executor._handshake() + + def test_session_lost_rotates_session_so_it_fails_only_once( + self, monkeypatch: Any + ) -> None: + monkeypatch.setattr(modal_executor, "WEBEXEC_WS_FAIL_ON_SESSION_LOSS", True) + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._had_success = True + executor._hashes_sent_on_ws = {"hash-1"} + old_session = executor._session_id + executor._ws = _FakeWS([_hello_reply(session_known=False)]) + + with pytest.raises(WebexecSessionLostError): + executor._handshake() + + # The failure is loud exactly once: the executor starts an honest + # fresh session instead of finding the old id registered and + # silently continuing with reset runtime state. + assert executor._session_id != old_session + assert executor._had_success is False + assert executor._hashes_sent_on_ws == set() + + executor._ws = _FakeWS([_hello_reply(session_known=False)]) + executor._handshake() + assert executor._server.proto == 2 + + def test_unknown_session_without_prior_state_is_fine(self) -> None: + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._had_success = False + executor._ws = _FakeWS([_hello_reply(session_known=False)]) + + executor._handshake() + + assert executor._server.proto == 2 + + def test_text_frame_during_handshake_is_connection_error(self) -> None: + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._ws = _FakeWS(["connection closed"]) + + with pytest.raises(ConnectionError): + executor._handshake() + + +class TestRecvTyping: + def test_text_frame_raises_instead_of_reaching_msgpack(self) -> None: + executor = _executor_with_ws(_FakeWS(["server says bye"])) + + with pytest.raises(ConnectionError, match="non-binary"): + executor._recv_reassembled(msgpack) + + def test_text_chunk_raises_instead_of_reaching_msgpack(self) -> None: + executor = _executor_with_ws( + _FakeWS([_pack({"_chunked": 2}), b"chunk-1", "bye"]) + ) + + with pytest.raises(ConnectionError, match="non-binary"): + executor._recv_reassembled(msgpack) + + +class TestSendRecvRetry: + def test_v2_does_not_resend_when_container_is_unidentified(self) -> None: + # A server that does not identify its container (pre-container_id + # build) cannot promise the reconnect reaches the same dedup + # registry, so the guard must fail closed rather than assume it did. + first_ws = _FakeWS([ConnectionError("idle close")]) + second_ws = _FakeWS() + executor = _executor_with_ws(first_ws) + executor._server = executor._server._replace(container_id=None) + sockets = [second_ws] + + def fake_ensure(_workspace: str) -> None: + if executor._ws is None: + executor._ws = sockets.pop(0) + + executor._ensure_connection = fake_ensure + + with pytest.raises(DynamicBlockError, match="same container"): + executor._send_recv_with_retry( + _pack({"inputs": {}}), "test-ws", request_id="req-1" + ) + + assert first_ws.closed + assert second_ws.sent == [] + + def test_v2_resends_after_recv_failure_on_same_container(self) -> None: + first_ws = _FakeWS([ConnectionError("idle close")]) + second_ws = _FakeWS([_pack({"success": True, "request_id": "req-1"})]) + executor = _executor_with_ws(first_ws) + executor._server = executor._server._replace(container_id="container-1") + sockets = [second_ws] + + def fake_ensure(_workspace: str) -> None: + if executor._ws is None: + executor._ws = sockets.pop(0) + executor._server = executor._server._replace(container_id="container-1") + + executor._ensure_connection = fake_ensure + + resp = executor._send_recv_with_retry( + _pack({"inputs": {}}), "test-ws", request_id="req-1" + ) + + assert msgpack.unpackb(resp, raw=False)["success"] is True + assert len(second_ws.sent) == 1 + + def test_v2_does_not_resend_when_reconnect_lands_on_other_container( + self, + ) -> None: + # The dedup cache is per-container: a resend that reaches a + # different container would run the user code a second time, so the + # ambiguous outcome must fail loudly instead. + first_ws = _FakeWS([ConnectionError("idle close")]) + second_ws = _FakeWS() + executor = _executor_with_ws(first_ws) + executor._server = executor._server._replace(container_id="container-1") + sockets = [second_ws] + + def fake_ensure(_workspace: str) -> None: + if executor._ws is None: + executor._ws = sockets.pop(0) + executor._server = executor._server._replace(container_id="container-2") + + executor._ensure_connection = fake_ensure + + with pytest.raises(DynamicBlockError, match="same container"): + executor._send_recv_with_retry( + _pack({"inputs": {}}), "test-ws", request_id="req-1" + ) + + # Nothing was sent on the new connection. + assert second_ws.sent == [] + + def test_container_id_is_cleared_when_connection_is_dropped(self) -> None: + # A stale id left behind would let the resend guard match a + # container this executor is no longer talking to. + executor = _executor_with_ws(_FakeWS()) + executor._server = executor._server._replace(container_id="container-1") + + executor._drop_ws_connection() + + assert executor._server.container_id is None + + def test_v1_does_not_resend_after_recv_failure(self) -> None: + ws = _FakeWS([ConnectionError("idle close")]) + executor = _executor_with_ws(ws, proto=1) + executor._ensure_connection = lambda _workspace: None + + with pytest.raises(DynamicBlockError, match="not retried"): + executor._send_recv_with_retry(_pack({"inputs": {}}), "test-ws") + + def test_session_lost_on_resend_reports_ambiguous_outcome(self) -> None: + # The frame was already accepted somewhere. Telling the user to + # replay from a checkpoint would re-run side effects that may + # already have happened, so the ambiguous-outcome error wins over + # the session-lost one. + executor = _executor_with_ws(_FakeWS([ConnectionError("idle close")])) + executor._server = executor._server._replace(container_id="container-1") + + def fake_ensure(_workspace: str) -> None: + if executor._ws is None: + raise WebexecSessionLostError("state gone") + + executor._ensure_connection = fake_ensure + + with pytest.raises(DynamicBlockError, match="may have already executed"): + executor._send_recv_with_retry( + _pack({"inputs": {}}), "test-ws", request_id="req-1" + ) + + def test_session_lost_before_send_still_propagates(self) -> None: + # Nothing was sent yet, so there is no ambiguity to report: the job + # genuinely has to be replayed from its checkpoint. + executor = _executor_with_ws(_FakeWS()) + executor._ws = None + + def fake_ensure(_workspace: str) -> None: + raise WebexecSessionLostError("state gone") + + executor._ensure_connection = fake_ensure + + with pytest.raises(WebexecSessionLostError): + executor._send_recv_with_retry( + _pack({"inputs": {}}), "test-ws", request_id="req-1" + ) + + +def _single_pass_stop_event() -> Any: + """Stop event letting the keepalive body run exactly once.""" + passes = {"n": 0} + + def wait(_timeout: float) -> bool: + passes["n"] += 1 + return passes["n"] > 1 + + return SimpleNamespace(wait=wait, is_set=lambda: passes["n"] > 1) + + +class _LockAcquiredHook: + """Lock that runs a callback once acquired — models a thread winning the + lock only after another has already done its work.""" + + def __init__(self, on_acquire) -> None: + self._lock = threading.Lock() + self._on_acquire = on_acquire + + def acquire(self, blocking: bool = True, timeout: float = -1) -> bool: + acquired = self._lock.acquire(blocking, timeout) + if acquired: + self._on_acquire() + return acquired + + def release(self) -> None: + self._lock.release() + + def __enter__(self) -> "_LockAcquiredHook": + self.acquire() + return self + + def __exit__(self, *_exc: Any) -> None: + self.release() + + +class TestKeepaliveIdleRelease: + def test_releases_connection_and_rotates_session_when_idle(self) -> None: + ws = _FakeWS() + executor = _executor_with_ws(ws) + executor._had_success = True + old_session = executor._session_id + executor._last_activity = _time.monotonic() - ( + WEBEXEC_WS_IDLE_RELEASE_SECONDS + 5 + ) + + executor._keepalive_loop(_single_pass_stop_event()) + + assert executor._ws is None + assert ws.closed + assert executor._session_id != old_session + assert executor._had_success is False + + def test_activity_just_before_the_lock_cancels_the_release(self) -> None: + # The idle value read before acquiring the lock can be stale: a frame + # may complete in between. Releasing on it would discard a live + # session silently, since rotation clears the fail-loudly latch. + ws = _FakeWS() + executor = _executor_with_ws(ws) + executor._had_success = True + old_session = executor._session_id + executor._last_activity = _time.monotonic() - ( + WEBEXEC_WS_IDLE_RELEASE_SECONDS + 5 + ) + executor._io_lock = _LockAcquiredHook( + lambda: setattr(executor, "_last_activity", _time.monotonic()) + ) + + executor._keepalive_loop(_single_pass_stop_event()) + + assert executor._ws is ws + assert not ws.closed + assert executor._session_id == old_session + assert executor._had_success is True + + def test_keepalive_backs_off_while_a_request_holds_the_lock(self) -> None: + # The whole design rests on the request path and the heartbeat never + # touching the socket at the same time: they share one connection and + # both read frames, so an interleave would hand a response to the + # keepalive (or an ack to the request). The keepalive must therefore + # take _io_lock non-blockingly and give up when a real frame holds it. + ws = _FakeWS() + executor = _executor_with_ws(ws) + executor._had_success = True + old_session = executor._session_id + executor._last_activity = _time.monotonic() - ( + WEBEXEC_WS_IDLE_RELEASE_SECONDS + 5 + ) + + # Held for the whole pass by "another thread", never released. + assert executor._io_lock.acquire(blocking=False) + try: + executor._keepalive_loop(_single_pass_stop_event()) + finally: + executor._io_lock.release() + + # Nothing was written, nothing was read, and the idle release did not + # fire — the in-flight frame owns the socket. + assert ws.sent == [] + assert not ws.closed + assert executor._ws is ws + assert executor._session_id == old_session + assert executor._had_success is True + + +class TestGracefulServerClose: + """A close the server announced is proof the frame was never read. + + The server decides to close at the TOP of its receive loop and never + reads again, so a frame written concurrently is unprocessed. Reporting + that as "may have already executed" fails work that never ran — and at + the default connection cap it happens on a schedule. + """ + + def test_closing_frame_in_the_execution_slot_is_not_a_response(self) -> None: + executor = _executor_with_ws(_FakeWS([_pack({"_kind": "closing"})])) + + with pytest.raises(modal_executor._ServerClosingError): + executor._recv_reassembled(msgpack) + + def test_closing_frame_during_a_heartbeat_is_not_an_ack(self) -> None: + ws = _FakeWS([_pack({"_kind": "closing"})]) + executor = _executor_with_ws(ws) + + with pytest.raises(modal_executor._ServerClosingError): + executor._send_heartbeat(ws) + + def test_graceful_close_retries_on_any_container(self, monkeypatch: Any) -> None: + # Attempt 1: the frame goes out, then the server's announced close + # arrives instead of a response. Attempt 2 lands on a DIFFERENT + # container. Without proof of non-delivery the same-container guard + # would refuse the resend and report "may have already executed"; + # with it, the frame is simply retried and succeeds. + executor = WebSocketModalExecutor(workspace_id="test-ws") + containers = iter(["container-1", "container-2"]) + sockets = [ + _FakeWS([_pack({"_kind": "closing"})]), + _FakeWS([_pack({"success": True, "request_id": "req-1", "result": {}})]), + ] + + def fake_ensure_connection(_workspace: str) -> None: + if executor._ws is None: + executor._ws = sockets.pop(0) + executor._server = _ServerInfo( + proto=2, idle_timeout=10.0, container_id=next(containers) + ) + + monkeypatch.setattr(executor, "_ensure_connection", fake_ensure_connection) + + resp = executor._send_recv_with_retry( + _pack({"request_id": "req-1"}), "test-ws", request_id="req-1" + ) + + assert msgpack.unpackb(resp, raw=False)["success"] is True + + def test_graceful_close_does_not_disarm_an_earlier_delivery( + self, monkeypatch: Any + ) -> None: + # A closing frame proves only that THIS attempt's write went unread. + # Attempt 1 hands the frame to container-1, which may have executed it + # and lost only the response. Attempt 2 reaches container-1 again and + # is closed before being read. If that close cleared the delivery + # state, attempt 3 would re-send to container-2 -- which has no dedup + # record of the request -- and run the user's block a SECOND time. + # The ambiguity, once incurred, has to survive. + executor = WebSocketModalExecutor(workspace_id="test-ws") + containers = iter(["container-1", "container-1", "container-2"]) + sockets = [ + # Attempt 1: frame delivered, response lost. + _FakeWS([ConnectionError("response lost after send")]), + # Attempt 2: same container, server closes before reading. + _FakeWS([_pack({"_kind": "closing"})]), + # Attempt 3 must never get here on a different container. + _FakeWS([_pack({"success": True, "request_id": "req-1", "result": {}})]), + ] + + def fake_ensure_connection(_workspace: str) -> None: + if executor._ws is None: + executor._ws = sockets.pop(0) + executor._server = _ServerInfo( + proto=2, idle_timeout=10.0, container_id=next(containers) + ) + + monkeypatch.setattr(executor, "_ensure_connection", fake_ensure_connection) + + third_socket = sockets[2] + + with pytest.raises(DynamicBlockError, match="may have already executed"): + executor._send_recv_with_retry( + _pack({"request_id": "req-1"}), "test-ws", request_id="req-1" + ) + + # The third connection gets established (the guard is checked after + # reconnecting), but NOTHING is written to it: the frame is never + # handed to a container that could not answer it from a dedup record. + assert third_socket.sent == [] + + +class TestResponseIdGuard: + def test_mismatched_request_id_drops_connection(self) -> None: + executor = _executor_with_ws(_FakeWS()) + + with pytest.raises(DynamicBlockError, match="stale"): + executor._check_response_id({"request_id": "other"}, "req-1") + + assert executor._ws is None + + def test_matching_id_passes(self) -> None: + executor = _executor_with_ws(_FakeWS()) + + executor._check_response_id({"request_id": "req-1"}, "req-1") + + def test_v2_requires_the_echoed_id(self) -> None: + # A late heartbeat_ack landing in the execution recv slot has no + # request_id. Accepting it would surface as a fabricated + # "RuntimeError: Unknown error" against the user's block and leave + # the real response queued, desyncing the next request too. + executor = _executor_with_ws(_FakeWS()) + + with pytest.raises(DynamicBlockError, match="in-flight request id"): + executor._check_response_id({"success": True}, "req-1") + + assert executor._ws is None + + def test_v2_rejects_a_control_frame(self) -> None: + executor = _executor_with_ws(_FakeWS()) + + with pytest.raises(DynamicBlockError, match="in-flight request id"): + executor._check_response_id( + {"_kind": "heartbeat_ack", "request_id": "req-1"}, "req-1" + ) + + def test_v1_still_accepts_a_response_without_an_id(self) -> None: + executor = _executor_with_ws(_FakeWS(), proto=1) + + executor._check_response_id({"success": True}, "req-1") + + assert executor._ws is not None + + +class TestHeartbeat: + def test_interval_derives_from_server_idle_timeout(self) -> None: + executor = _executor_with_ws(_FakeWS()) + executor._server = executor._server._replace(idle_timeout=9.0) + + assert executor._heartbeat_interval() == 3.0 + + def test_interval_falls_back_on_v1(self) -> None: + executor = _executor_with_ws(_FakeWS(), proto=1) + + assert ( + executor._heartbeat_interval() + == WebSocketModalExecutor._KEEPALIVE_IDLE_SECONDS + ) + + def test_v2_heartbeat_is_application_frame_with_ack(self) -> None: + ws = _FakeWS([_pack({"_kind": "heartbeat_ack"})]) + executor = _executor_with_ws(ws) + + executor._send_heartbeat(ws) + + sent = msgpack.unpackb(ws.sent[0], raw=False) + assert sent == {"_kind": "heartbeat"} + + def test_v2_heartbeat_text_reply_raises(self) -> None: + ws = _FakeWS(["closing"]) + executor = _executor_with_ws(ws) + + with pytest.raises(ConnectionError): + executor._send_heartbeat(ws) + + def test_v2_heartbeat_uses_short_ack_timeout_and_restores_it(self) -> None: + # The ack must not be awaited under the execution-sized read + # timeout: a half-open connection would pin _io_lock for minutes. + ws = _FakeWS([_pack({"_kind": "heartbeat_ack"})]) + executor = _executor_with_ws(ws) + + executor._send_heartbeat(ws) + + assert ws.timeouts[0] == WebSocketModalExecutor._HEARTBEAT_ACK_TIMEOUT_SECONDS + assert ws.timeouts[-1] == WEBEXEC_WS_READ_TIMEOUT_SECONDS + + def test_v2_heartbeat_restores_read_timeout_on_failure(self) -> None: + ws = _FakeWS([ConnectionError("half-open")]) + executor = _executor_with_ws(ws) + + with pytest.raises(ConnectionError): + executor._send_heartbeat(ws) + + assert ws.timeouts[-1] == WEBEXEC_WS_READ_TIMEOUT_SECONDS + + +class TestSessionLossKillSwitch: + def test_enforcement_can_be_disabled(self, monkeypatch: Any) -> None: + # Prod needs to revert to the pre-v2 (silent continuation) behavior + # without rolling the inference ref back across every consumer. + monkeypatch.setattr(modal_executor, "WEBEXEC_WS_FAIL_ON_SESSION_LOSS", False) + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._had_success = True + session = executor._session_id + executor._ws = _FakeWS([_hello_reply(session_known=False)]) + + executor._handshake() + + assert executor._server.proto == 2 + # The session is still rotated: state really is gone, so the next + # reconnect must not silently pass the check on the old id. + assert executor._session_id != session + assert executor._had_success is False + + def test_enforcement_off_by_default(self) -> None: + # Enforcement is OFF by default: the check cannot distinguish a + # stateful block from a stateless one, and the server's connection cap + # guarantees the trigger on a schedule, so defaulting it on would fail + # the stateless majority roughly every WEBEXEC_WS_MAX_CONNECTION_SECONDS. + assert modal_executor.WEBEXEC_WS_FAIL_ON_SESSION_LOSS is False + + def test_session_lost_publishes_server_info_before_raising( + self, monkeypatch: Any + ) -> None: + monkeypatch.setattr(modal_executor, "WEBEXEC_WS_FAIL_ON_SESSION_LOSS", True) + # The keepalive thread reads _server unlocked; leaving it on the v1 + # default after a v2 handshake would restore the 25s interval + # against a 10s-idle server. + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._had_success = True + executor._ws = _FakeWS([_hello_reply(session_known=False)]) + + with pytest.raises(WebexecSessionLostError): + executor._handshake() + + assert executor._server.proto == 2 + assert executor._server.idle_timeout == 10.0 + + +def _run_one_execution(executor: WebSocketModalExecutor, response: dict) -> Any: + """Drive _execute_ws once with a canned server response.""" + + def _fake_send_recv( + _frame: bytes, _workspace: str, request_id: Any = None + ) -> bytes: + # Echo the id the executor generated, as a real v2 server does. + echoed = dict(response) + if executor._server.proto == 2: + echoed["request_id"] = request_id + return _pack(echoed) + + executor._send_recv_with_retry = _fake_send_recv # type: ignore + return executor._execute_ws( + "MyBlock", + SimpleNamespace( + run_function_code="def run(x):\n return x\n", + run_function_name="run", + imports=[], + ), + {}, + "ws", + msgpack, + {}, + ) + + +class TestHadSuccessLatch: + def test_v1_success_does_not_arm_the_v2_session_check(self) -> None: + # Client rolls out before the Modal app: the executor talks v1 and + # succeeds. If that armed the latch, the first reconnect onto an + # upgraded v2 container would hard-fail every workspace executor + # that had ever succeeded. + executor = _executor_with_ws(_FakeWS(), proto=1) + + _run_one_execution(executor, {"success": True, "result": {}}) + + assert executor._had_success is False + + def test_v2_success_arms_it(self) -> None: + executor = _executor_with_ws(_FakeWS()) + + _run_one_execution(executor, {"success": True, "result": {}}) + + assert executor._had_success is True + + +class TestIdleReleaseDisable: + def test_non_positive_value_disables_idle_release(self, monkeypatch: Any) -> None: + # <= 0 disables, matching WEBEXEC_MODAL_EXECUTOR_IDLE_TTL_SECONDS. + # Treating it as "always release" would drop the connection after one + # heartbeat interval AND bypass the session check (rotation clears the + # latch), silently resetting stateful blocks. + monkeypatch.setattr(modal_executor, "WEBEXEC_WS_IDLE_RELEASE_SECONDS", 0) + ws = _FakeWS([_pack({"_kind": "heartbeat_ack"})]) + executor = _executor_with_ws(ws) + executor._had_success = True + old_session = executor._session_id + executor._last_activity = _time.monotonic() - 100_000 + + executor._keepalive_loop(_single_pass_stop_event()) + + assert executor._ws is ws + assert not ws.closed + assert executor._session_id == old_session + assert executor._had_success is True + + +class TestIdleTimeoutCoercion: + @pytest.mark.parametrize( + "advertised, expected", + [ + (None, 10.0), + ("not-a-number", 10.0), + (0, 10.0), + (-5, 10.0), + (float("nan"), 10.0), + # Never clamped UP: believing the server allows more idle time + # than it does is what reinstates the idle-close this protocol + # fixes. The floor lives on the derived interval instead. + (1, 1.0), + (2, 2.0), + (10_000, 300.0), # clamped down: must stay under the idle release + (30, 30.0), + ], + ) + def test_advisory_field_never_bricks_the_transport( + self, advertised: Any, expected: float + ) -> None: + assert modal_executor._coerce_idle_timeout(advertised) == expected + + @pytest.mark.parametrize("advertised", [1, 1.5, 2, 3, 10, 30, 300]) + def test_heartbeat_gap_stays_under_the_advertised_timeout( + self, advertised: float + ) -> None: + # The keepalive skips a tick when ``idle < interval``, so the real + # worst case between two app-level frames is 2 x interval. If that + # ever reaches the server's deadline the connection dies on idle -- + # exactly the bug protocol v2 exists to fix. + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._server = _ServerInfo( + proto=2, idle_timeout=modal_executor._coerce_idle_timeout(advertised) + ) + + assert 2 * executor._heartbeat_interval() < advertised + + def test_handshake_survives_a_garbage_idle_timeout(self) -> None: + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._ws = _FakeWS([_pack({"_kind": "hello", "idle_timeout_s": "soon"})]) + + executor._handshake() + + assert executor._server.proto == 2 + assert executor._server.idle_timeout == 10.0 + + +class TestPostLoopRetryExhaustion: + def test_delivered_frame_reports_possible_execution(self, monkeypatch: Any) -> None: + # Attempt 1 delivers the frame and loses the response; attempts 2-3 + # reconnect to the same container and fail too. The final error must + # keep the "may have already executed" wording, or a job runner + # re-runs the block and duplicates its side effects. + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._server = _ServerInfo( + proto=2, idle_timeout=10.0, container_id="container-1" + ) + + sends = [] + + class _Sock: + def send_binary(self, frame: bytes) -> None: + sends.append(frame) + + def recv(self) -> Any: + raise ConnectionError("boom") + + def close(self) -> None: + pass + + def _fake_ensure(_workspace: str) -> None: + executor._ws = _Sock() + executor._server = executor._server._replace(container_id="container-1") + + monkeypatch.setattr(executor, "_ensure_connection", _fake_ensure) + + with pytest.raises(DynamicBlockError) as excinfo: + executor._send_recv_with_retry(_pack({"request_id": "r"}), "ws", "r") + + assert "may have already executed" in excinfo.value.public_message + assert "websocket_response" in excinfo.value.context + # Pin the retry count too: without this the test passes unchanged if + # `attempts` regresses to 1, since a single delivered-then-lost frame + # produces the same message. + assert len(sends) == 3, "a delivered frame must be retried twice more" + + def test_never_delivered_frame_reports_a_connect_failure( + self, monkeypatch: Any + ) -> None: + executor = WebSocketModalExecutor(workspace_id="test-ws") + + def _fake_ensure(_workspace: str) -> None: + raise ConnectionError("no route") + + monkeypatch.setattr(executor, "_ensure_connection", _fake_ensure) + + with pytest.raises(DynamicBlockError) as excinfo: + executor._send_recv_with_retry(_pack({"request_id": "r"}), "ws", "r") + + assert "failed after retry" in excinfo.value.public_message + assert "websocket_connection" in excinfo.value.context + + +class TestServerInfrastructureErrors: + @pytest.mark.parametrize( + "result", + [ + { + "success": False, + "error_type": "ResponseNoLongerAvailable", + "error": "x", + "server_error": True, + }, + { + "success": False, + "error_type": "InvalidRequest", + "error": "x", + "server_error": True, + }, + { + "success": False, + "error_type": "ValueError", + "error": "x", + "server_error": True, + }, + ], + ) + def test_transport_failures_are_not_reported_as_user_code_errors( + self, result: dict + ) -> None: + # DynamicBlockCodeError means "the user's Python raised". These + # failures mean the block either never ran or ran fine. + executor = WebSocketModalExecutor(workspace_id="test-ws") + with pytest.raises(DynamicBlockError) as excinfo: + executor._raise_server_error_if_infrastructure(result, "MyBlock") + + assert not isinstance(excinfo.value, DynamicBlockCodeError) + assert "websocket_server_error" in excinfo.value.context + + def test_user_code_errors_still_fall_through(self) -> None: + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._raise_server_error_if_infrastructure( + {"success": False, "error_type": "ZeroDivisionError", "error": "x"}, + "MyBlock", + ) + + @pytest.mark.parametrize( + "error_type", + ["UnknownCodeHash", "InvalidRequest", "ResponseNoLongerAvailable"], + ) + @pytest.mark.parametrize("proto", [1, 2]) + def test_a_forged_infrastructure_name_is_never_trusted( + self, error_type: str, proto: int + ) -> None: + # error_type is an exception CLASS NAME chosen by untrusted user code. + # ONLY the server_error flag may select a control path -- on EITHER + # protocol. A name-based fallback would be authoritative for exactly + # the responses it must not govern, because a stamping server omits + # the flag on user-code failures, so "absent" means user code rather + # than "old server". + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._server = _ServerInfo( + proto=proto, idle_timeout=10.0 if proto == 2 else None + ) + + executor._raise_server_error_if_infrastructure( + {"success": False, "error_type": error_type, "error": "x"}, + "MyBlock", + ) + + +class TestChunkedResponseValidation: + @pytest.mark.parametrize("chunk_count", [0, -1, 10**9, "many", True]) + def test_bogus_chunk_header_is_connection_death(self, chunk_count: Any) -> None: + # A negative count yields an empty join and explodes in unpackb far + # from any handler; a huge one stalls the read timeout per attempt. + ws = _FakeWS([_pack({"_chunked": chunk_count})]) + executor = _executor_with_ws(ws) + + with pytest.raises(ConnectionError, match="chunk count"): + executor._recv_reassembled(msgpack) + + +class TestMalformedResponse: + def test_undecodable_response_becomes_a_transport_error(self) -> None: + executor = _executor_with_ws(_FakeWS()) + + with pytest.raises(DynamicBlockError, match="could not be"): + executor._unpack_response(b"\xc1not-msgpack", msgpack) + + assert executor._ws is None + + def test_non_map_response_becomes_a_transport_error(self) -> None: + executor = _executor_with_ws(_FakeWS()) + + with pytest.raises(DynamicBlockError, match="not a"): + executor._unpack_response(_pack([1, 2, 3]), msgpack) + + assert executor._ws is None + + +class TestCloseTearsDownKeepaliveFirst: + def test_close_stops_and_joins_the_keepalive_thread(self) -> None: + # close() runs on a request thread when the executor cache evicts + # this executor. websocket-client's close() does a recv_frame() that + # bypasses the socket read lock, so the keepalive must be gone + # before the fd is freed. + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._ws = _FakeWS() + executor._server = _ServerInfo( + proto=2, idle_timeout=10.0, container_id="container-1" + ) + executor._ensure_keepalive_thread() + thread = executor._keepalive_thread + assert thread is not None and thread.is_alive() + + executor.close() + + assert not thread.is_alive() + assert executor._ws is None + assert executor._server.container_id is None + + def test_close_does_not_block_on_an_in_flight_execution(self) -> None: + # _io_lock can be held for the whole read timeout by a running block. + # Executor-cache eviction must not wait that long. + executor = WebSocketModalExecutor(workspace_id="test-ws") + executor._ws = _FakeWS() + executor._CLOSE_LOCK_TIMEOUT_SECONDS = 0.05 + executor._io_lock.acquire() + try: + started = _time.monotonic() + executor.close() + elapsed = _time.monotonic() - started + finally: + executor._io_lock.release() + + assert elapsed < 1.0 + assert executor._ws is None diff --git a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_ws_server_dedup.py b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_ws_server_dedup.py new file mode 100644 index 0000000000..cee6ca55e6 --- /dev/null +++ b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_ws_server_dedup.py @@ -0,0 +1,563 @@ +"""Server-side protocol v2 guarantees in modal/modal_app.py. + +Covers the at-most-once execution machinery: the executed-request registry +that survives response-cache eviction, the bounded response cache, and the +session registry that answers the client's hello. +""" + +import time + +import msgpack +import pytest +from types import SimpleNamespace + +from .conftest import build_ws_app as _ws_app + + +class TestTtlKeySet: + def test_membership_add_and_refresh(self, modal_app) -> None: + keys = modal_app._TtlKeySet(ttl_seconds=3600, max_entries=8) + + assert "a" not in keys + keys.add("a") + assert "a" in keys + # refresh never adds, so an unknown key stays unknown + assert keys.refresh("b") is False + assert "b" not in keys + assert keys.refresh("a") is True + + def test_empty_key_is_never_stored(self, modal_app) -> None: + keys = modal_app._TtlKeySet(ttl_seconds=3600, max_entries=8) + + keys.add("") + + assert "" not in keys + + def test_expired_entries_drop_out(self, modal_app) -> None: + keys = modal_app._TtlKeySet(ttl_seconds=-1, max_entries=8) + + keys.add("a") + + assert "a" not in keys + + def test_size_cap_evicts_oldest_first(self, modal_app) -> None: + keys = modal_app._TtlKeySet(ttl_seconds=3600, max_entries=2) + + keys.add("a") + keys.add("b") + keys.add("c") + + assert "a" not in keys + assert "b" in keys and "c" in keys + + def test_refresh_protects_an_entry_from_eviction(self, modal_app) -> None: + keys = modal_app._TtlKeySet(ttl_seconds=3600, max_entries=2) + keys.add("a") + keys.add("b") + + keys.refresh("a") + keys.add("c") + + assert "a" in keys + assert "b" not in keys + + +class TestResponseCache: + def test_round_trip_and_entry_cap(self, modal_app) -> None: + cache = modal_app._WsResponseCache( + max_entries=2, max_bytes=10**9, ttl_seconds=60 + ) + + cache.put("1", b"a") + cache.put("2", b"b") + cache.put("3", b"c") + + assert cache.get("1") is None + assert cache.get("2") == b"b" + assert cache.get("3") == b"c" + + def test_byte_cap_evicts_oldest_first(self, modal_app) -> None: + cache = modal_app._WsResponseCache( + max_entries=16, max_bytes=100, ttl_seconds=60 + ) + cache.put("a", b"x" * 40) + cache.put("b", b"y" * 40) + + # A hit must NOT promote: entries stay in age order so the TTL is a + # real TTL and the early-exit expiry scan stays correct. + cache.get("a") + cache.put("c", b"z" * 40) + + assert cache.get("a") is None + assert cache.get("b") is not None + assert cache.get("c") is not None + + def test_hit_does_not_extend_the_ttl(self, modal_app, monkeypatch) -> None: + # Driven by a fake clock rather than a real busy-wait: the previous + # form spun on wall time for 0.2s, which burns CPU and can flake on a + # loaded CI runner. + now = [1000.0] + monkeypatch.setattr( + modal_app, "time", SimpleNamespace(monotonic=lambda: now[0]) + ) + cache = modal_app._WsResponseCache( + max_entries=8, max_bytes=10**9, ttl_seconds=60.0 + ) + cache.put("a", b"payload") + + # Repeated hits inside the TTL must not slide the expiry forward. + for _ in range(5): + now[0] += 20.0 + if now[0] < 1060.0: + assert cache.get("a") == b"payload" + + now[0] = 1061.0 + assert cache.get("a") is None + + def test_expired_entry_is_not_served(self, modal_app) -> None: + cache = modal_app._WsResponseCache( + max_entries=8, max_bytes=10**9, ttl_seconds=-1 + ) + + cache.put("a", b"payload") + + assert cache.get("a") is None + + def test_oversized_payload_is_refused_without_draining_the_cache( + self, modal_app + ) -> None: + # A response too large to coexist with anything else must not be + # inserted at all: inserting it would evict every other entry and + # then itself, wiping the dedup cache for every concurrent + # connection on this container. + cache = modal_app._WsResponseCache(max_entries=8, max_bytes=100, ttl_seconds=60) + cache.put("keep", b"ok") + + cache.put("huge", b"x" * 500) + + assert cache.get("huge") is None + assert cache.get("keep") == b"ok" + assert cache._total_bytes == 2 + + +class TestExecutionIsAtMostOnce: + def test_resend_of_completed_request_is_answered_from_cache( + self, modal_app, monkeypatch + ) -> None: + from fastapi.testclient import TestClient + + calls = [] + + def run_user_code(self, *args, **kwargs): + calls.append(1) + return {"success": True, "result": {"n": len(calls)}} + + executor, app = _ws_app(modal_app, run_user_code) + frame = msgpack.packb({"request_id": "req-1", "inputs": {}}, use_bin_type=True) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(frame) + first = msgpack.unpackb(ws.receive_bytes(), raw=False) + ws.send_bytes(frame) + second = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert len(calls) == 1, "resend must not run the user code again" + assert first == second + + def test_resend_after_cache_eviction_fails_loudly_instead_of_rerunning( + self, modal_app + ) -> None: + # The response cache is best effort; the executed-request registry + # is what makes execution at-most-once. With the payload gone, a + # resend must get an error, never a second execution. + from fastapi.testclient import TestClient + + calls = [] + + def run_user_code(self, *args, **kwargs): + calls.append(1) + return {"success": True, "result": {}} + + executor, app = _ws_app(modal_app, run_user_code) + frame = msgpack.packb({"request_id": "req-1", "inputs": {}}, use_bin_type=True) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(frame) + msgpack.unpackb(ws.receive_bytes(), raw=False) + # Simulate eviction of the payload while the request id is still + # known to have executed. + executor._ws_response_cache = modal_app._WsResponseCache() + ws.send_bytes(frame) + resent = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert len(calls) == 1 + assert resent["success"] is False + assert resent["error_type"] == "ResponseNoLongerAvailable" + assert resent["request_id"] == "req-1" + + def test_unserializable_result_returns_error_and_is_not_rerun( + self, modal_app + ) -> None: + # A result the server cannot pack must not escape as an exception: + # that would kill the connection with nothing cached and invite a + # resend of a request whose side effects already happened. + from fastapi.testclient import TestClient + + calls = [] + + class _Unpackable: + def __repr__(self): + raise RuntimeError("cannot repr") + + def run_user_code(self, *args, **kwargs): + calls.append(1) + return {"success": True, "result": {"bad": _Unpackable()}} + + executor, app = _ws_app(modal_app, run_user_code) + frame = msgpack.packb({"request_id": "req-1", "inputs": {}}, use_bin_type=True) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(frame) + first = msgpack.unpackb(ws.receive_bytes(), raw=False) + ws.send_bytes(frame) + second = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert len(calls) == 1, "the failed request must not be re-executed" + assert first["success"] is False + assert first["request_id"] == "req-1" + # The resend is answered from the cached error, identically. + assert second == first + + def test_undecodable_inputs_are_reported_as_not_run_and_stay_retryable( + self, modal_app, monkeypatch + ) -> None: + # Nothing executed, so the error must say so — and the request must + # not be marked executed, or a legitimate retry would be refused. + from fastapi.testclient import TestClient + + calls = [] + + def run_user_code(self, *args, **kwargs): + calls.append(1) + return {"success": True, "result": {}} + + executor, app = _ws_app(modal_app, run_user_code) + + def boom(_inputs): + raise ValueError("bad input payload") + + monkeypatch.setattr( + ( + modal_app.Executor._get_user_cls() + if hasattr(modal_app.Executor, "_get_user_cls") + else modal_app.Executor + ), + "_deserialize_msgpack_inputs", + staticmethod(boom), + ) + frame = msgpack.packb({"request_id": "req-1", "inputs": {}}, use_bin_type=True) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(frame) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert calls == [] + assert resp["success"] is False + assert "was not run" in resp["error"] + assert "req-1" not in executor._ws_executed + + def test_session_is_registered_only_after_a_successful_execution( + self, modal_app + ) -> None: + from fastapi.testclient import TestClient + + def run_user_code(self, *args, **kwargs): + return {"success": False, "error": "boom", "error_type": "ValueError"} + + executor, app = _ws_app(modal_app, run_user_code) + hello = msgpack.packb( + {"_kind": "hello", "proto": 2, "session_id": "s1"}, use_bin_type=True + ) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(hello) + first_hello = msgpack.unpackb(ws.receive_bytes(), raw=False) + ws.send_bytes( + msgpack.packb({"request_id": "r1", "inputs": {}}, use_bin_type=True) + ) + msgpack.unpackb(ws.receive_bytes(), raw=False) + ws.send_bytes(hello) + second_hello = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert first_hello["session_known"] is False + # A failed execution builds no runtime state, and asking twice must + # not itself register the session. + assert second_hello["session_known"] is False + assert first_hello["container_id"] == executor._container_id + + +class TestSessionRegistry: + def _executor(self, modal_app): + cls = modal_app.Executor + user_cls = cls._get_user_cls() if hasattr(cls, "_get_user_cls") else cls + + class _Standin: + _WS_SESSION_TTL_SECONDS = user_cls._WS_SESSION_TTL_SECONDS + _WS_SESSION_MAX_ENTRIES = user_cls._WS_SESSION_MAX_ENTRIES + _ws_session_seen = user_cls._ws_session_seen + _ws_register_session = user_cls._ws_register_session + + standin = _Standin() + standin._ws_sessions = modal_app._TtlKeySet( + ttl_seconds=_Standin._WS_SESSION_TTL_SECONDS, + max_entries=_Standin._WS_SESSION_MAX_ENTRIES, + ) + return standin + + def test_unknown_session_is_not_registered_by_asking(self, modal_app) -> None: + # Asking must not create the entry: a hello that reports "unknown" + # and then registers would let the very next reconnect pass the + # check silently. + executor = self._executor(modal_app) + + assert executor._ws_session_seen("s1") is False + assert executor._ws_session_seen("s1") is False + + def test_registered_session_is_known(self, modal_app) -> None: + executor = self._executor(modal_app) + + executor._ws_register_session("s1") + + assert executor._ws_session_seen("s1") is True + assert executor._ws_session_seen("s2") is False + + def test_seen_refreshes_so_an_active_session_cannot_age_out( + self, modal_app + ) -> None: + # The namespaces this answers for are never evicted, so a session + # that keeps reconnecting must not expire out of the registry. + executor = self._executor(modal_app) + executor._ws_sessions = modal_app._TtlKeySet(ttl_seconds=3600, max_entries=2) + executor._ws_register_session("s1") + executor._ws_register_session("s2") + + executor._ws_session_seen("s1") + executor._ws_register_session("s3") + + assert executor._ws_session_seen("s1") is True + assert executor._ws_session_seen("s2") is False + + +class TestBaseExceptionFromUserCode: + def test_system_exit_returns_an_error_frame_instead_of_killing_the_conn( + self, modal_app + ) -> None: + # Untrusted user code can call sys.exit(). concurrent.futures + # re-raises SystemExit in the coroutine, so `except Exception` missed + # it: the exception escaped the shielded task, tore the connection + # down with nothing cached, and poisoned the request id for the whole + # executed-marker TTL. + from fastapi.testclient import TestClient + + def run_user_code(self, *args, **kwargs): + raise SystemExit(2) + + executor, app = _ws_app(modal_app, run_user_code) + frame = msgpack.packb({"request_id": "req-1", "inputs": {}}, use_bin_type=True) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(frame) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + # The connection is still usable afterwards. + ws.send_bytes(frame) + resent = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert resp["success"] is False + assert resp["error_type"] == "SystemExit" + assert resp["server_error"] is True + # The payload was cached, so the resend is answered, not re-run. + assert resent == resp + + def test_a_request_that_never_ran_stays_resendable(self, modal_app) -> None: + # If the worker thread never entered the user code, the executed + # marker must be rolled back — otherwise every resend of that id + # gets ResponseNoLongerAvailable for the whole TTL. + registry = modal_app._TtlKeySet(ttl_seconds=60, max_entries=8) + registry.add("req-1") + assert "req-1" in registry + + registry.discard("req-1") + + assert "req-1" not in registry + + +class TestMalformedFrames: + def test_non_map_execution_frame_is_reported_not_fatal(self, modal_app) -> None: + # request.get(...) on a list used to raise AttributeError outside + # every handler and kill the connection. + from fastapi.testclient import TestClient + + def run_user_code(self, *args, **kwargs): + raise AssertionError("must not run") + + _, app = _ws_app(modal_app, run_user_code) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(msgpack.packb([1, 2, 3], use_bin_type=True)) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert resp["success"] is False + assert resp["error_type"] == "InvalidRequest" + + @pytest.mark.parametrize("chunk_count", [0, -1, 10**9, "many"]) + def test_bogus_chunk_header_closes_with_a_server_error( + self, modal_app, chunk_count + ) -> None: + from fastapi.testclient import TestClient + + def run_user_code(self, *args, **kwargs): + raise AssertionError("must not run") + + _, app = _ws_app(modal_app, run_user_code) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(msgpack.packb({"_chunked": chunk_count}, use_bin_type=True)) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert resp["success"] is False + assert "chunk count" in resp["error"] + + +class TestWireBoundaryTypes: + def test_non_string_request_id_is_rejected_without_running_the_block( + self, modal_app + ) -> None: + # _TtlKeySet.add() used to store any hashable while __contains__ gated + # on isinstance(str), so an int request_id was recorded as executed but + # never found again. Downgrading it to "no request id" is not a fix + # either: the block would run with NO at-most-once record, so a resend + # runs it a second time. It must be refused before execution. + from fastapi.testclient import TestClient + + calls = [] + + def run_user_code(self, *args, **kwargs): + calls.append(1) + return {"success": True, "result": {}} + + executor, app = _ws_app(modal_app, run_user_code) + frame = msgpack.packb({"request_id": 12345, "inputs": {}}, use_bin_type=True) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes(frame) + resp = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert resp["success"] is False + assert resp["error_type"] == "InvalidRequest" + assert resp["server_error"] is True + assert calls == [], "the block must not run without a dedup record" + assert len(executor._ws_executed._seen) == 0 + assert executor._ws_inflight == {} + + def test_registries_agree_on_what_a_key_is(self, modal_app) -> None: + registry = modal_app._TtlKeySet(ttl_seconds=60, max_entries=8) + + for bogus in (12345, None, "", (1, 2)): + registry.add(bogus) + assert bogus not in registry + assert registry.refresh(bogus) is False + assert len(registry._seen) == 0 + + def test_unhashable_session_id_does_not_kill_the_connection( + self, modal_app + ) -> None: + from fastapi.testclient import TestClient + + def run_user_code(self, *args, **kwargs): + return {"success": True, "result": {}} + + _, app = _ws_app(modal_app, run_user_code) + + with TestClient(app).websocket_connect("/ws") as ws: + ws.send_bytes( + msgpack.packb( + {"_kind": "hello", "session_id": {"a": 1}}, use_bin_type=True + ) + ) + reply = msgpack.unpackb(ws.receive_bytes(), raw=False) + + assert reply["_kind"] == "hello" + assert reply["session_known"] is False + + +class TestInflightDedup: + def test_concurrent_resend_awaits_the_running_execution(self, modal_app) -> None: + # The one genuinely new concurrency guarantee: a second connection + # resending the same request_id WHILE the first execution is still + # running must await the in-flight task, not start a second one. + import asyncio + import threading + + from fastapi.testclient import TestClient + + release = threading.Event() + calls = [] + + def run_user_code(self, *args, **kwargs): + calls.append(1) + release.wait(timeout=5) + return {"success": True, "result": {"n": len(calls)}} + + _, app = _ws_app(modal_app, run_user_code) + frame = msgpack.packb({"request_id": "req-1", "inputs": {}}, use_bin_type=True) + # One TestClient context => one portal => one event loop, matching a + # real container where every connection shares the ASGI loop. + client = TestClient(app) + results = {} + + def second_connection(): + with client.websocket_connect("/ws") as ws2: + # Wait until the first execution is actually running. Bounded: + # an unbounded spin in a non-daemon thread turns a genuine + # regression into a CI hang instead of a fast failure. + deadline = time.monotonic() + 10.0 + while not calls: + if time.monotonic() > deadline: + raise AssertionError("first execution never started within 10s") + time.sleep(0.01) + ws2.send_bytes(frame) + release.set() + results["second"] = msgpack.unpackb(ws2.receive_bytes(), raw=False) + + # `with` rather than a bare __enter__/__exit__ pair: an assertion + # raising between them would leak the ASGI portal thread. + with client: + with client.websocket_connect("/ws") as ws1: + ws1.send_bytes(frame) + worker = threading.Thread(target=second_connection, daemon=True) + worker.start() + results["first"] = msgpack.unpackb(ws1.receive_bytes(), raw=False) + worker.join(timeout=10) + + assert len(calls) == 1, "the resend must not start a second execution" + assert results["first"] == results["second"] + + +class TestContainerIdentity: + def test_two_containers_get_distinct_ids(self, modal_app, monkeypatch) -> None: + # The client trusts container_id to decide whether a resend may be + # answered from this container's dedup registry. Two containers + # sharing an id (e.g. if identify() were ever marked snap=True and + # ran pre-snapshot) would let a resend re-execute user code silently. + monkeypatch.delenv("MODAL_TASK_ID", raising=False) + cls = modal_app.Executor + user_cls = cls._get_user_cls() if hasattr(cls, "_get_user_cls") else cls + + ids = set() + for _ in range(2): + executor = user_cls.__new__(user_cls) + executor.workspace_id = "test-ws" + user_cls.identify(executor) + ids.add(executor._container_id) + + assert len(ids) == 2