truss 0.11.1rc6__py3-none-any.whl → 0.11.1rc8__py3-none-any.whl

truss/templates/control/control/endpoints.py

@@ -11,7 +11,8 @@ 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.events import BytesMessage, TextMessage
+from wsproto.connection import ConnectionState
+from wsproto.events import BytesMessage, CloseConnection, TextMessage
 
 from truss.templates.control.control.helpers.errors import (
     ModelLoadFailed,
@@ -158,7 +159,13 @@ 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, TextMessage):
+        if isinstance(message, CloseConnection):
+            # 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):
             await client_ws.send_text(message.data)
         elif isinstance(message, BytesMessage):
             await client_ws.send_bytes(message.data)

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.1rc6
+truss==0.11.1rc7
 uvicorn>=0.24.0
 uvloop>=0.19.0
 websockets>=10.0

{truss-0.11.1rc6.dist-info → truss-0.11.1rc8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.1rc6
+Version: 0.11.1rc8
 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.1rc6.dist-info → truss-0.11.1rc8.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=MiVoU5n8GTj5ygN-iL6neTD1AIqhVYIaJAHkcnsHGvA,253
+truss/templates/control/requirements.txt,sha256=Mgqjo6vKoDS-mJNKx0V3UUDspYnB-Y6ti0VTee7n99k,253
 truss/templates/control/control/application.py,sha256=jYeta6hWe1SkfLL3W4IDmdYjg3ZuKqI_UagWYs5RB_E,3793
-truss/templates/control/control/endpoints.py,sha256=VQ1lvZjFvR091yRkiFdvXw1Q7PiNGXT9rJwY7_sX6yg,11828
+truss/templates/control/control/endpoints.py,sha256=hW1fw19tPLNPV9cW1Jio4tt-uvCbJDUxDdNe7vcvNe4,12267
 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
@@ -361,14 +361,14 @@ truss_chains/reference_code/reference_model.py,sha256=emH3hb23E_nbP98I37PGp1Xk1h
 truss_chains/remote_chainlet/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/remote_chainlet/model_skeleton.py,sha256=8ZReLOO2MLcdg7bNZ61C-6j-e68i2Z-fFlyV3sz0qH8,2376
 truss_chains/remote_chainlet/stub.py,sha256=Y2gDUzMY9WRaQNHIz-o4dfLUfFyYV9dUhIRQcfgrY8g,17209
-truss_chains/remote_chainlet/utils.py,sha256=RJ74JeB_jzq0wjzxkkVrcnoh_fdWhiq5-FtZTYQdgyQ,23260
+truss_chains/remote_chainlet/utils.py,sha256=nnJQG3hXIoGghWEGPIAqKfRaueCLgSdwNI2x9B7BSaY,23080
 truss_train/__init__.py,sha256=7hE6j6-u6UGzCGaNp3CsCN0kAVjBus1Ekups-Bk0fi4,837
 truss_train/definitions.py,sha256=V985HhY4rdXL10DZxpFEpze9ScxzWErMht4WwaPknGU,6789
 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.1rc6.dist-info/METADATA,sha256=FpPBQ0BmRG8IPSDsYnA7f9Su_I3L0dqTycCIsC7tKv4,6672
-truss-0.11.1rc6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.1rc6.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.1rc6.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.1rc6.dist-info/RECORD,,
+truss-0.11.1rc8.dist-info/METADATA,sha256=dd9uyZ92UZODbZliBCxANNIqni-Mnap4kazJA2elp40,6672
+truss-0.11.1rc8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.1rc8.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.1rc8.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.1rc8.dist-info/RECORD,,

truss_chains/remote_chainlet/utils.py

@@ -382,20 +382,13 @@ def pydantic_set_field_dict(obj: pydantic.BaseModel) -> dict[str, pydantic.BaseM
 # Error Propagation Utils. #############################################################
 
 
-# NB(nikhil): Deployed chainlets have access to FastAPI, but local testing doesn't necessarily
-# have that dependency. We have a helpful error message via `utils.make_optional_import_error`
-# for those cases.
-def _safe_import_fastapi():
+def _handle_exception(exception: Exception) -> NoReturn:
+    """Raises `HTTPException` with `RemoteErrorDetail`."""
     try:
-        import fastapi  # noqa: F401
+        import fastapi
     except ImportError:
         raise utils.make_optional_import_error("fastapi")
 
-
-def _handle_exception(exception: Exception) -> NoReturn:
-    """Raises `HTTPException` with `RemoteErrorDetail`."""
-    _safe_import_fastapi()
-
     if hasattr(exception, "__module__"):
         exception_module_name = exception.__module__
     else:
@@ -595,7 +588,11 @@ class WebsocketWrapperFastAPI:
         await self._websocket.close(code=code, reason=reason)
 
     async def receive(self) -> Union[str, bytes]:
-        _safe_import_fastapi()
+        try:
+            import fastapi
+        except ImportError:
+            raise utils.make_optional_import_error("fastapi")
+
         message = await self._websocket.receive()
 
         if message.get("type") == "websocket.disconnect":