truss 0.11.1rc10__tar.gz → 0.11.1rc11__tar.gz

{truss-0.11.1rc10 → truss-0.11.1rc11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.1rc10
+Version: 0.11.1rc11
 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.1rc10 → truss-0.11.1rc11}/pyproject.toml

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

{truss-0.11.1rc10 → truss-0.11.1rc11}/truss/templates/control/control/endpoints.py

@@ -11,8 +11,7 @@ from starlette.requests import ClientDisconnect, Request
 from starlette.responses import Response
 from starlette.websockets import WebSocketDisconnect as StartletteWebSocketDisconnect
 from tenacity import RetryCallState, Retrying, retry_if_exception_type, wait_fixed
-from wsproto.connection import ConnectionState
-from wsproto.events import BytesMessage, CloseConnection, TextMessage
+from wsproto.events import BytesMessage, TextMessage
 
 from truss.templates.control.control.helpers.errors import (
     ModelLoadFailed,
@@ -158,19 +157,16 @@ async def forward_to_server(
 
 async def forward_to_client(client_ws: WebSocket, server_ws: AsyncWebSocketSession):
     while True:
-        message = await server_ws.receive()
-        if isinstance(message, CloseConnection):
-            print("RECEIVED CLOSE === ")
-            # NB(nikhil): httpx requires explicit confirmation of the close frame.
-            if server_ws.connection.state == ConnectionState.REMOTE_CLOSING:
-                await server_ws.send(message.response())  # type: ignore[arg-type]
-
-            raise WebSocketDisconnect(code=message.code, reason=message.reason)
-        elif isinstance(message, TextMessage):
-            print("TEXT PROXYING === ")
-            await client_ws.send_text(message.data)
-        elif isinstance(message, BytesMessage):
-            await client_ws.send_bytes(message.data)
+        try:
+            message = await server_ws.receive()
+            if isinstance(message, TextMessage):
+                await client_ws.send_text(message.data)
+            elif isinstance(message, BytesMessage):
+                await client_ws.send_bytes(message.data)
+        except WebSocketDisconnect as e:
+            await _safe_close_ws(
+                client_ws, client_ws.app.state.logger, e.code, e.reason
+            )
 
 
 # NB(nikhil): _handle_websocket_forwarding uses some py311 specific syntax, but in newer

{truss-0.11.1rc10 → truss-0.11.1rc11}/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.11.1rc10
+truss==0.11.1rc11
 uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0