truss 0.11.1rc10__py3-none-any.whl → 0.11.1rc12__py3-none-any.whl

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,27 @@ 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:
+            logger = client_ws.app.state.logger
+            logger.info(
+                f"PROXY: Received WebSocketDisconnect from server - Code: {e.code}, Reason: {e.reason}"
+            )
+            await _safe_close_ws(client_ws, logger, e.code or 1000, e.reason)
+            logger.info("PROXY: Client connection closed successfully")
+            return  # Make sure to return!
+        except Exception as e:
+            logger = client_ws.app.state.logger
+            logger.error(
+                f"PROXY: Unexpected error in forward_to_client: {e}", exc_info=True
+            )
+            await _safe_close_ws(client_ws, logger, 1011, "Proxy error")
+            return  # Make sure to return!
 
 
 # NB(nikhil): _handle_websocket_forwarding uses some py311 specific syntax, but in newer

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.1rc12
 uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0

{truss-0.11.1rc10.dist-info → truss-0.11.1rc12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.1rc10
+Version: 0.11.1rc12
 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.dist-info → truss-0.11.1rc12.dist-info}/RECORD

@@ -72,9 +72,9 @@ truss/templates/cache_requirements.txt,sha256=xoPoJ-OVnf1z6oq_RVM3vCr3ionByyqMLj
 truss/templates/copy_cache_files.Dockerfile.jinja,sha256=Os5zFdYLZ_AfCRGq4RcpVTObOTwL7zvmwYcvOzd_Zqo,126
 truss/templates/docker_server_requirements.txt,sha256=PyhOPKAmKW1N2vLvTfLMwsEtuGpoRrbWuNo7tT6v2Mc,18
 truss/templates/server.Dockerfile.jinja,sha256=CUYnF_hgxPGq2re7__0UPWlwzOHMoFkxp6NVKi3U16s,7071
-truss/templates/control/requirements.txt,sha256=KW7_CqkJnJ57giF9Jg8wSOosqn1kEIoV0i3ngQUPy5w,254
+truss/templates/control/requirements.txt,sha256=MdbJzwV3HDCyd16QUMPxwOJ2nmwEYJcYenUN26TEVQ8,254
 truss/templates/control/control/application.py,sha256=jYeta6hWe1SkfLL3W4IDmdYjg3ZuKqI_UagWYs5RB_E,3793
-truss/templates/control/control/endpoints.py,sha256=IarphOa2oVVx9D5ODmS27DqdUQL_epYgcGGxKsfVVEE,12348
+truss/templates/control/control/endpoints.py,sha256=ZAx5MRN67QiUEOiE0QvEGGjqTxBpCC-93rErrsKW8oY,12605
 truss/templates/control/control/server.py,sha256=R4Y219i1dcz0kkksN8obLoX-YXWGo9iW1igindyG50c,3128
 truss/templates/control/control/helpers/context_managers.py,sha256=W6dyFgLBhPa5meqrOb3w_phMtKfaJI-GhwUfpiycDc8,413
 truss/templates/control/control/helpers/custom_types.py,sha256=n_lTudtLTpy4oPV3aDdJ4X2rh3KCV5btYO9UnTeUouQ,5471
@@ -367,8 +367,8 @@ truss_train/definitions.py,sha256=V985HhY4rdXL10DZxpFEpze9ScxzWErMht4WwaPknGU,67
 truss_train/deployment.py,sha256=lWWANSuzBWu2M4oK4qD7n-oVR1JKdmw2Pn5BJQHg-Ck,3074
 truss_train/loader.py,sha256=0o66EjBaHc2YY4syxxHVR4ordJWs13lNXnKjKq2wq0U,1630
 truss_train/public_api.py,sha256=9N_NstiUlmBuLUwH_fNG_1x7OhGCytZLNvqKXBlStrM,1220
-truss-0.11.1rc10.dist-info/METADATA,sha256=VhhRRTySAl9O6RU3G-0hKbvYiZ_eD5xvRJK4vWQGSxs,6673
-truss-0.11.1rc10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.1rc10.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.1rc10.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.1rc10.dist-info/RECORD,,
+truss-0.11.1rc12.dist-info/METADATA,sha256=xHaxci-7BxoXvZUdBJbZMDdRvgP4n7-nMIFjy5lpTJ8,6673
+truss-0.11.1rc12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.1rc12.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.1rc12.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.1rc12.dist-info/RECORD,,