truss 0.11.1rc1__tar.gz → 0.11.1rc3__tar.gz

{truss-0.11.1rc1 → truss-0.11.1rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.1rc1
+Version: 0.11.1rc3
 Summary: A seamless bridge from model development to model delivery
 Project-URL: Repository, https://github.com/basetenlabs/truss
 Project-URL: Homepage, https://truss.baseten.co

{truss-0.11.1rc1 → truss-0.11.1rc3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truss"
-version = "0.11.1rc1"
+version = "0.11.1rc3"
 description = "A seamless bridge from model development to model delivery"
 authors = [
     { name = "Pankaj Gupta", email = "no-reply@baseten.co" },

{truss-0.11.1rc1 → truss-0.11.1rc3}/truss/templates/control/control/endpoints.py

@@ -30,9 +30,14 @@ BASE_RETRY_EXCEPTIONS = (
 
 control_app = APIRouter()
 
+WEBSOCKET_NORMAL_CLOSURE_CODE = 1000
+WEBSOCKET_SERVER_ERROR_CODE = 1011
+
 
 class CloseableWebsocket(Protocol):
-    async def close(self, code: int = 1000, reason: Optional[str] = None) -> None: ...
+    async def close(
+        self, code: int = WEBSOCKET_NORMAL_CLOSURE_CODE, reason: Optional[str] = None
+    ) -> None: ...
 
 
 @control_app.get("/")
@@ -126,7 +131,7 @@ def inference_retries(
 async def _safe_close_ws(
     ws: CloseableWebsocket,
     logger: logging.Logger,
-    code: int = 1000,
+    code: int,
     reason: Optional[str] = None,
 ):
     try:
@@ -153,8 +158,6 @@ async def forward_to_server(
 async def forward_to_client(client_ws: WebSocket, server_ws: AsyncWebSocketSession):
     while True:
         message = await server_ws.receive()
-        if message is None:
-            break
         if isinstance(message, TextMessage):
             await client_ws.send_text(message.data)
         elif isinstance(message, BytesMessage):
@@ -172,7 +175,11 @@ async def _handle_websocket_forwarding(
             tg.create_task(forward_to_client(client_ws, server_ws))
             tg.create_task(forward_to_server(client_ws, server_ws))
     except ExceptionGroup as eg:  # type: ignore[name-defined] # noqa: F821
-        exc = eg.exceptions[0]  # NB(nikhil): Only care about the first one.
+        # NB(nikhil): The first websocket proxy method to raise an error will
+        # be surfaced here, and that contains the information we want to forward to the
+        # other websocket. Further errors might raise as a result of cancellation, but we
+        # can safely ignore those.
+        exc = eg.exceptions[0]
         if isinstance(exc, WebSocketDisconnect):
             await _safe_close_ws(client_ws, logger, exc.code, exc.reason)
         elif isinstance(exc, StartletteWebSocketDisconnect):
@@ -180,8 +187,10 @@ async def _handle_websocket_forwarding(
         else:
             logger.warning(f"Ungraceful websocket close: {exc}")
     finally:
-        await _safe_close_ws(client_ws, logger)
-        await _safe_close_ws(server_ws, logger)
+        # NB(nikhil): In most common cases, both websockets would have been successfully
+        # closed with applicable codes above, these lines are just a failsafe.
+        await _safe_close_ws(client_ws, logger, code=WEBSOCKET_SERVER_ERROR_CODE)
+        await _safe_close_ws(server_ws, logger, code=WEBSOCKET_SERVER_ERROR_CODE)
 
 
 async def _attempt_websocket_proxy(
@@ -202,7 +211,7 @@ async def proxy_ws(client_ws: WebSocket):
                 await _attempt_websocket_proxy(client_ws, proxy_client, logger)
             except httpx_ws_exceptions.HTTPXWSException as e:
                 logger.warning(f"WebSocket connection rejected: {e}")
-                await _safe_close_ws(client_ws, logger)
+                await _safe_close_ws(client_ws, logger, WEBSOCKET_SERVER_ERROR_CODE)
                 break
 
 

{truss-0.11.1rc1 → truss-0.11.1rc3}/truss/templates/control/requirements.txt

@@ -6,7 +6,7 @@ loguru>=0.7.2
 python-json-logger>=2.0.2
 tenacity>=8.1.0
  # To avoid divergence, this should follow the latest release.
-truss==0.9.100
+truss==0.11.1rc3
 uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0

{truss-0.11.1rc1 → truss-0.11.1rc3}/truss/templates/server/common/errors.py

@@ -18,6 +18,7 @@ _BASETEN_DOWNSTREAM_ERROR_CODE = 600
 _BASETEN_CLIENT_ERROR_CODE = 700
 
 MODEL_ERROR_MESSAGE = "Internal Server Error (in model/chainlet)."
+WEBSOCKET_SERVER_ERROR_CODE = 1011
 
 
 class ModelMissingError(Exception):

{truss-0.11.1rc1 → truss-0.11.1rc3}/truss/templates/server/truss_server.py

@@ -264,7 +264,9 @@ class BasetenEndpoints:
                     )
                 except Exception:
                     await _safe_close_websocket(
-                        ws, status_code=1011, reason=errors.MODEL_ERROR_MESSAGE
+                        ws,
+                        status_code=errors.WEBSOCKET_SERVER_ERROR_CODE,
+                        reason=errors.MODEL_ERROR_MESSAGE,
                     )
                     raise  # Re raise to let `intercept_exceptions` deal with it.
 

{truss-0.11.1rc1 → truss-0.11.1rc3}/truss/tests/templates/control/control/test_endpoints.py

@@ -51,8 +51,6 @@ async def test_proxy_ws_bidirectional_messaging(client_ws):
         return_value=mock_server_ws,
     ):
         proxy_task = asyncio.create_task(proxy_ws(client_ws))
-        await asyncio.sleep(0.5)
-
         client_queue.put_nowait(
             {"type": "websocket.disconnect", "code": 1002, "reason": "test-closure"}
         )