truss 0.11.1rc1__py3-none-any.whl → 0.11.1rc3__py3-none-any.whl

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/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/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/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/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"}
         )

{truss-0.11.1rc1.dist-info → truss-0.11.1rc3.dist-info}/METADATA

@@ -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.dist-info → truss-0.11.1rc3.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=Kk0tYID7trPk5gwX38Wrt2-YGWZAXFJCJRcqJ8ZzCjc,251
+truss/templates/control/requirements.txt,sha256=D2kIrXfCKlWl8LO7quTUlCFYuT3Dn_MVAlCG_0YjHQY,253
 truss/templates/control/control/application.py,sha256=jYeta6hWe1SkfLL3W4IDmdYjg3ZuKqI_UagWYs5RB_E,3793
-truss/templates/control/control/endpoints.py,sha256=z8OJVBdAlJRl4mAdACVgYGGXaSx8Z8CR-_IVu70IidA,11259
+truss/templates/control/control/endpoints.py,sha256=VQ1lvZjFvR091yRkiFdvXw1Q7PiNGXT9rJwY7_sX6yg,11828
 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
@@ -97,9 +97,9 @@ truss/templates/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3
 truss/templates/server/main.py,sha256=kWXrdD8z8IpamyWxc8qcvd5ck9gM1Kz2QH5qHJCnmOQ,222
 truss/templates/server/model_wrapper.py,sha256=k75VVISwwlsx5EGb82UZsu8kCM_i6Yi3-Hd0-Kpm1yo,42055
 truss/templates/server/requirements.txt,sha256=XblmpfxAmRo3X1V_9oMj8yjdpZ5Wk-C2oa3z6nq4OGw,672
-truss/templates/server/truss_server.py,sha256=m1DdCBqT80p8Sfft60vpSeH1ZNoUVpTdYPIcqkIE8CA,19519
+truss/templates/server/truss_server.py,sha256=noXfGJMsKIhgF4oI_8LC1UHkcx8Vg8nGSITZJ_bkRFQ,19598
 truss/templates/server/common/__init__.py,sha256=qHIqr68L5Tn4mV6S-PbORpcuJ4jmtBR8aCuRTIWDvNo,85
-truss/templates/server/common/errors.py,sha256=qWeZlmNI8ZGbZbOIp_mtS6IKvUFIzhj3QH8zp-xTp9o,8554
+truss/templates/server/common/errors.py,sha256=My0P6-Y7imVTICIhazHT0vlSu3XJDH7As06OyVzu4Do,8589
 truss/templates/server/common/patches.py,sha256=uEOzvDnXsHOkTSa8zygGYuR4GHhrFNVHNQc5peJcwvo,1393
 truss/templates/server/common/retry.py,sha256=dtz6yvwLoY0i55FnxECz57zEOKjAhGMYvvM-k9jiR9c,624
 truss/templates/server/common/schema.py,sha256=WLFtVyEKmk4whg5_gk6Gt1vOD6wM5fWKLb4zNuD0bkw,6042
@@ -161,7 +161,7 @@ truss/tests/remote/baseten/test_auth.py,sha256=ttu4bDnmwGfo3oiNut4HVGnh-QnjAefwZ
 truss/tests/remote/baseten/test_core.py,sha256=6NzJTDmoSUv6Muy1LFEYIUg10-cqw-hbLyeTSWcdNjY,26117
 truss/tests/remote/baseten/test_remote.py,sha256=y1qSPL1t7dBeYI3xMFn436fttG7wkYdAoENTz7qKObg,23634
 truss/tests/remote/baseten/test_service.py,sha256=ufZbtQlBNIzFCxRt_iE-APLpWbVw_3ViUpSh6H9W5nU,1945
-truss/tests/templates/control/control/test_endpoints.py,sha256=wUC24DvQPgYbYmIxUIFvRtNjlATAoO-1r9XY38iidLI,3678
+truss/tests/templates/control/control/test_endpoints.py,sha256=fxTiiCR0ltaHCL_-v-22Ie1qVgnch1lqcj3w0U3R-fk,3644
 truss/tests/templates/control/control/test_server.py,sha256=r1O3VEK9eoIL2-cg8nYLXYct_H3jf5rGp1wLT1KBdeA,9488
 truss/tests/templates/control/control/test_server_integration.py,sha256=EdDY3nLzjrRCJ5LI5yZsNCEImSRkxTL7Rn9mGnK67zA,11837
 truss/tests/templates/control/control/helpers/test_context_managers.py,sha256=3LoonRaKu_UvhaWs1eNmEQCZq-iJ3aIjI0Mn4amC8Bw,283
@@ -364,8 +364,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.1rc1.dist-info/METADATA,sha256=Wwt0gl5KmgZp_yyELNdWsVKlR_KO7HOPEKwFzHZeDFM,6672
-truss-0.11.1rc1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.1rc1.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.1rc1.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.1rc1.dist-info/RECORD,,
+truss-0.11.1rc3.dist-info/METADATA,sha256=g7LI89s5I17UGwJVLTnS3SM3sIDp7zMTW383slMfgXo,6672
+truss-0.11.1rc3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.1rc3.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.1rc3.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.1rc3.dist-info/RECORD,,