truss 0.11.1rc2__py3-none-any.whl → 0.11.1rc4__py3-none-any.whl

truss/cli/chains_commands.py

@@ -359,7 +359,7 @@ def push_chain(
     "--name",
     type=str,
     required=False,
-    help="Name of the chain to be deployed, if not given, the entrypoint name is used.",
+    help="Name of the chain to be watched. If not given, the entrypoint name is used.",
 )
 @click.option(
     "--remote",

truss/templates/control/control/endpoints.py

@@ -1,14 +1,15 @@
 import asyncio
 import logging
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Optional, Protocol
 
 import httpx
 from fastapi import APIRouter, WebSocket
 from fastapi.responses import JSONResponse, StreamingResponse
+from httpx_ws import AsyncWebSocketSession, WebSocketDisconnect, aconnect_ws
 from httpx_ws import _exceptions as httpx_ws_exceptions
-from httpx_ws import aconnect_ws
 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
 
@@ -29,6 +30,15 @@ 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 = WEBSOCKET_NORMAL_CLOSURE_CODE, reason: Optional[str] = None
+    ) -> None: ...
+
 
 @control_app.get("/")
 def index():
@@ -118,13 +128,79 @@ def inference_retries(
         yield attempt
 
 
-async def _safe_close_ws(ws: WebSocket, logger: logging.Logger):
+async def _safe_close_ws(
+    ws: CloseableWebsocket,
+    logger: logging.Logger,
+    code: int,
+    reason: Optional[str] = None,
+):
     try:
-        await ws.close()
+        await ws.close(code, reason)
     except RuntimeError as close_error:
         logger.debug(f"Duplicate close of websocket: `{close_error}`.")
 
 
+async def forward_to_server(
+    client_ws: WebSocket, server_ws: AsyncWebSocketSession
+) -> None:
+    while True:
+        message = await client_ws.receive()
+        if message.get("type") == "websocket.disconnect":
+            raise StartletteWebSocketDisconnect(
+                message.get("code", 1000), message.get("reason")
+            )
+        if "text" in message:
+            await server_ws.send_text(message["text"])
+        elif "bytes" in message:
+            await server_ws.send_bytes(message["bytes"])
+
+
+async def forward_to_client(client_ws: WebSocket, server_ws: AsyncWebSocketSession):
+    while True:
+        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)
+
+
+# NB(nikhil): _handle_websocket_forwarding uses some py311 specific syntax, but in newer
+# versions of truss we're guaranteed to be running the control server with at least that version.
+async def _handle_websocket_forwarding(
+    client_ws: WebSocket, server_ws: AsyncWebSocketSession
+):
+    logger = client_ws.app.state.logger
+    try:
+        async with asyncio.TaskGroup() as tg:  # type: ignore[attr-defined]
+            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
+        # 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):
+            await _safe_close_ws(server_ws, logger, exc.code, exc.reason)
+        else:
+            logger.warning(f"Ungraceful websocket close: {exc}")
+    finally:
+        # 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(
+    client_ws: WebSocket, proxy_client: httpx.AsyncClient, logger
+):
+    async with aconnect_ws("/v1/websocket", proxy_client) as server_ws:  # type: ignore
+        await client_ws.accept()
+        await _handle_websocket_forwarding(client_ws, server_ws)
+
+
 async def proxy_ws(client_ws: WebSocket):
     proxy_client: httpx.AsyncClient = client_ws.app.state.proxy_client
     logger = client_ws.app.state.logger
@@ -132,37 +208,10 @@ async def proxy_ws(client_ws: WebSocket):
     for attempt in inference_retries():
         with attempt:
             try:
-                async with aconnect_ws("/v1/websocket", proxy_client) as server_ws:  # type: ignore
-                    # Unfortunate, but FastAPI and httpx-ws have slightly different abstractions
-                    # for sending data, so it's not easy to create a unified wrapper.
-                    async def forward_to_server():
-                        while True:
-                            message = await client_ws.receive()
-                            if message.get("type") == "websocket.disconnect":
-                                break
-                            if "text" in message:
-                                await server_ws.send_text(message["text"])
-                            elif "bytes" in message:
-                                await server_ws.send_bytes(message["bytes"])
-
-                    async def forward_to_client():
-                        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):
-                                await client_ws.send_bytes(message.data)
-
-                    await client_ws.accept()
-                    try:
-                        await asyncio.gather(forward_to_client(), forward_to_server())
-                    finally:
-                        await _safe_close_ws(client_ws, logger)
+                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/common/tracing.py

@@ -6,6 +6,7 @@ import pathlib
 import time
 from typing import Iterator, List, Optional, Sequence
 
+import opentelemetry.exporter.otlp.proto.http.trace_exporter as oltp_exporter
 import opentelemetry.sdk.resources as resources
 import opentelemetry.sdk.trace as sdk_trace
 import opentelemetry.sdk.trace.export as trace_export
@@ -15,6 +16,7 @@ from shared import secrets_resolver
 logger = logging.getLogger(__name__)
 
 ATTR_NAME_DURATION = "duration_sec"
+OTEL_EXPORTER_OTLP_ENDPOINT = "OTEL_EXPORTER_OTLP_ENDPOINT"
 # Writing trace data to a file is only intended for testing / debugging.
 OTEL_TRACING_NDJSON_FILE = "OTEL_TRACING_NDJSON_FILE"
 # Exporting trace data to a public honeycomb instance (not our cluster collector)
@@ -65,6 +67,13 @@ def get_truss_tracer(secrets: secrets_resolver.Secrets, config) -> trace.Tracer:
         return _truss_tracer
 
     span_processors: List[sdk_trace.SpanProcessor] = []
+    if otlp_endpoint := os.getenv(OTEL_EXPORTER_OTLP_ENDPOINT):
+        if enable_tracing_data:
+            logger.info(f"Exporting trace data to {OTEL_EXPORTER_OTLP_ENDPOINT}.")
+        otlp_exporter = oltp_exporter.OTLPSpanExporter(endpoint=otlp_endpoint)
+        otlp_processor = sdk_trace.export.BatchSpanProcessor(otlp_exporter)
+        span_processors.append(otlp_processor)
+
     if tracing_log_file := os.getenv(OTEL_TRACING_NDJSON_FILE):
         if enable_tracing_data:
             logger.info(f"Exporting trace data to file `{tracing_log_file}`.")
@@ -72,6 +81,21 @@ def get_truss_tracer(secrets: secrets_resolver.Secrets, config) -> trace.Tracer:
         file_processor = sdk_trace.export.SimpleSpanProcessor(json_file_exporter)
         span_processors.append(file_processor)
 
+    if (
+        honeycomb_dataset := os.getenv(HONEYCOMB_DATASET)
+    ) and HONEYCOMB_API_KEY in secrets:
+        honeycomb_api_key = secrets[HONEYCOMB_API_KEY]
+        logger.info("Exporting trace data to honeycomb.")
+        honeycomb_exporter = oltp_exporter.OTLPSpanExporter(
+            endpoint="https://api.honeycomb.io/v1/traces",
+            headers={
+                "x-honeycomb-team": honeycomb_api_key,
+                "x-honeycomb-dataset": honeycomb_dataset,
+            },
+        )
+        honeycomb_processor = sdk_trace.export.BatchSpanProcessor(honeycomb_exporter)
+        span_processors.append(honeycomb_processor)
+
     if span_processors and enable_tracing_data:
         logger.info("Instantiating truss tracer.")
         resource = resources.Resource.create({resources.SERVICE_NAME: "truss-server"})

truss/templates/server/truss_server.py

@@ -76,7 +76,7 @@ async def parse_body(request: Request) -> bytes:
 
 
 async def _safe_close_websocket(
-    ws: WebSocket, reason: Optional[str], status_code: int = 1000
+    ws: WebSocket, status_code: int = 1000, reason: Optional[str] = None
 ) -> None:
     try:
         await ws.close(code=status_code, reason=reason)
@@ -257,14 +257,16 @@ class BasetenEndpoints:
                 try:
                     await ws.accept()
                     await self._model.websocket(ws)
-                    await _safe_close_websocket(ws, None, status_code=1000)
+                    await _safe_close_websocket(ws, status_code=1000, reason=None)
                 except WebSocketDisconnect as ws_error:
                     logging.info(
                         f"Client terminated websocket connection: `{ws_error}`."
                     )
                 except Exception:
                     await _safe_close_websocket(
-                        ws, errors.MODEL_ERROR_MESSAGE, status_code=1011
+                        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

@@ -1,4 +1,5 @@
-from unittest.mock import AsyncMock, MagicMock, patch
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, call, patch
 
 import pytest
 from fastapi import FastAPI, WebSocket
@@ -31,33 +32,38 @@ def client_ws(app):
 
 @pytest.mark.asyncio
 async def test_proxy_ws_bidirectional_messaging(client_ws):
-    """Test that both directions of communication work and clean up properly"""
-    client_ws.receive.side_effect = [
-        {"type": "websocket.receive", "text": "msg1"},
-        {"type": "websocket.receive", "text": "msg2"},
-        {"type": "websocket.disconnect"},
-    ]
+    client_queue = asyncio.Queue()
+    client_ws.receive = client_queue.get
 
+    server_queue = asyncio.Queue()
     mock_server_ws = AsyncMock(spec=AsyncWebSocketSession)
-    mock_server_ws.receive.side_effect = [
-        TextMessage(data="response1"),
-        TextMessage(data="response2"),
-        None,  # server closing connection
-    ]
+    mock_server_ws.receive = server_queue.get
     mock_server_ws.__aenter__.return_value = mock_server_ws
     mock_server_ws.__aexit__.return_value = None
 
+    client_queue.put_nowait({"type": "websocket.receive", "text": "msg1"})
+    client_queue.put_nowait({"type": "websocket.receive", "text": "msg2"})
+    server_queue.put_nowait(TextMessage(data="response1"))
+    server_queue.put_nowait(TextMessage(data="response2"))
+
     with patch(
         "truss.templates.control.control.endpoints.aconnect_ws",
         return_value=mock_server_ws,
     ):
-        await proxy_ws(client_ws)
+        proxy_task = asyncio.create_task(proxy_ws(client_ws))
+        client_queue.put_nowait(
+            {"type": "websocket.disconnect", "code": 1002, "reason": "test-closure"}
+        )
+
+        await proxy_task
 
     assert mock_server_ws.send_text.call_count == 2
     assert mock_server_ws.send_text.call_args_list == [(("msg1",),), (("msg2",),)]
     assert client_ws.send_text.call_count == 2
     assert client_ws.send_text.call_args_list == [(("response1",),), (("response2",),)]
-    client_ws.close.assert_called_once()
+
+    assert mock_server_ws.close.call_args_list[0] == call(1002, "test-closure")
+    client_ws.close.assert_called()
 
 
 @pytest.mark.asyncio

{truss-0.11.1rc2.dist-info → truss-0.11.1rc4.dist-info}/METADATA

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

@@ -8,7 +8,7 @@ truss/base/errors.py,sha256=zDVLEvseTChdPP0oNhBBQCtQUtZJUaof5zeWMIjqz6o,691
 truss/base/trt_llm_config.py,sha256=CRz3AqGDAyv8YpcBWXUrnfjvNAauyo3yf8ZOGVsSt6g,32782
 truss/base/truss_config.py,sha256=7CtiJIwMHtDU8Wzn8UTJUVVunD0pWFl4QUVycK2aIpY,28055
 truss/base/truss_spec.py,sha256=jFVF79CXoEEspl2kXBAPyi-rwISReIGTdobGpaIhwJw,5979
-truss/cli/chains_commands.py,sha256=y6pdIAGCcKOPG9bPuCXPfSA0onQm5x-tT_3blSBfPYg,16971
+truss/cli/chains_commands.py,sha256=bqOXQ-0RPS66vSP_OPQdJ5dvctGiVrsGoSUMbURGdSI,16970
 truss/cli/cli.py,sha256=PaMkuwXZflkU7sa1tEoT_Zmy-iBkEZs1m4IVqcieaeo,30367
 truss/cli/remote_cli.py,sha256=G_xCKRXzgkCmkiZJhUFfsv5YSVgde1jLA5LPQitpZgI,1905
 truss/cli/train_commands.py,sha256=GDye7yXGL_nQvXAlY5MWsdj5x0zYOvcQw0Ubn14TiRU,14365
@@ -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=FM-sgao7I3gMoUTasM3Xq_g2LDoJQe75JxIoaQxzeNo,10031
+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,13 +97,13 @@ 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=ob_nceeGtFPZzKKdk_ZZGLoZrJOGE6hR52xM1sPR97A,19498
+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
-truss/templates/server/common/tracing.py,sha256=TDokphTO0O-b0xZLkkDMU6Z_JIsaZA0aimL6UIQB5eI,4808
+truss/templates/server/common/tracing.py,sha256=XSTXNoRtV8vXwveJoX3H32go0JKnLmznZ2TtrVzIe4M,5967
 truss/templates/server/common/patches/whisper/patch.py,sha256=kDECQ-wmEpeAZFhUTQP457ofueeMsm7DgNy9tqinhJQ,2383
 truss/templates/shared/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss/templates/shared/dynamic_config_resolver.py,sha256=75s42NFhQI5jL7BqlJH_UkuQS7ptbtFh13f2nh6X5Wo,920
@@ -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=tGU3w8zOKC8LfWGdhp-TlV7E603KXg2xGwpqDdf8Pnw,3385
+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
@@ -345,27 +345,27 @@ truss_chains/__init__.py,sha256=QDw1YwdqMaQpz5Oltu2Eq2vzEX9fDrMoqnhtbeh60i4,1278
 truss_chains/framework.py,sha256=CS7tSegPe2Q8UUT6CDkrtSrB3utr_1QN1jTEPjrj5Ug,67519
 truss_chains/private_types.py,sha256=6CaQEPawFLXjEbJ-01lqfexJtUIekF_q61LNENWegFo,8917
 truss_chains/public_api.py,sha256=0AXV6UdZIFAMycUNG_klgo4aLFmBZeKGfrulZEWzR0M,9532
-truss_chains/public_types.py,sha256=q8Oet6MpECW1FhWW25SCExpZhmk4cFmEsqrO30oZIMw,29112
+truss_chains/public_types.py,sha256=8cBstm5m_9aN_U2WoSHQ9s2z6gTdSEl-nVgXyneU3nY,29166
 truss_chains/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/pydantic_numpy.py,sha256=MG8Ji_Inwo_JSfM2n7TPj8B-nbrBlDYsY3SOeBwD8fE,4289
 truss_chains/streaming.py,sha256=DGl2LEAN67YwP7Nn9MK488KmYc4KopWmcHuE6WjyO1Q,12521
 truss_chains/utils.py,sha256=LvpCG2lnN6dqPqyX3PwLH9tyjUzqQN3N4WeEFROMHak,6291
 truss_chains/deployment/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truss_chains/deployment/code_gen.py,sha256=AmAUZ3h1hP3uYkl3J6o096K5RFLuBOP7kOFSnFC_C4U,32568
-truss_chains/deployment/deployment_client.py,sha256=7ckyh71wYVXBQtKfy4gh_MddMQWTNVpOL_FsqnaKFPo,32811
+truss_chains/deployment/deployment_client.py,sha256=haFiVmQek42ewlN_YflBaRDQT4ZYbmT20tvvJOkcUX0,32899
 truss_chains/reference_code/reference_chainlet.py,sha256=5feSeqGtrHDbldkfZCfX2R5YbbW0Uhc35mhaP2pXrHw,1340
 truss_chains/reference_code/reference_model.py,sha256=emH3hb23E_nbP98I37PGp1Xk1hz3g3lQ00tiLo55cSM,322
 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=O_5P-VAUvg0cegEW1uKCOf5EBwD8rEGYVoGMivOmc7k,22374
+truss_chains/remote_chainlet/utils.py,sha256=xX1t3e-BsYkWrxQIqfKRl4PHGuVyW3oleWFQpXSAynI,22949
 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.1rc2.dist-info/METADATA,sha256=-QNAojZwEkUwM3B6Jo9KIbVMpBstsTCKx-qR1S_MFJM,6672
-truss-0.11.1rc2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.1rc2.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.1rc2.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.1rc2.dist-info/RECORD,,
+truss-0.11.1rc4.dist-info/METADATA,sha256=PYD_kydnF-Z7GjTBOB0-JA0lQjQMtiBn7Y-30qyT7wY,6672
+truss-0.11.1rc4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.1rc4.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.1rc4.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.1rc4.dist-info/RECORD,,

truss_chains/deployment/deployment_client.py

@@ -655,7 +655,9 @@ class _Watcher:
         with framework.ChainletImporter.import_target(
             source, entrypoint
         ) as entrypoint_cls:
-            self._deployed_chain_name = name or entrypoint_cls.__name__
+            self._deployed_chain_name = (
+                name or entrypoint_cls.meta_data.chain_name or entrypoint_cls.__name__
+            )
             self._chain_root = _get_chain_root(entrypoint_cls)
             chainlet_names = set(
                 desc.display_name
@@ -677,7 +679,7 @@ class _Watcher:
         )
         if not chain_id:
             raise public_types.ChainsDeploymentError(
-                f"Chain `{chain_id}` was not found."
+                f"Chain `{self._deployed_chain_name}` was not found."
             )
         self._status_page_url = b10_service.URLConfig.status_page_url(
             self._remote_provider.remote_url, b10_service.URLConfig.CHAIN, chain_id

truss_chains/public_types.py

@@ -473,6 +473,7 @@ class WebSocketProtocol(Protocol):
 
     async def close(self, code: int = 1000, reason: Optional[str] = None) -> None: ...
 
+    async def receive(self) -> Union[str, bytes]: ...
     async def receive_text(self) -> str: ...
     async def receive_bytes(self) -> bytes: ...
     async def receive_json(self) -> Any: ...

truss_chains/remote_chainlet/utils.py

@@ -11,6 +11,7 @@ import textwrap
 import threading
 import time
 import traceback
+import typing
 from collections.abc import AsyncIterator
 from typing import (
     TYPE_CHECKING,
@@ -586,6 +587,18 @@ class WebsocketWrapperFastAPI:
     async def close(self, code: int = 1000, reason: Optional[str] = None) -> None:
         await self._websocket.close(code=code, reason=reason)
 
+    async def receive(self) -> Union[str, bytes]:
+        message = await self._websocket.receive()
+
+        if message.get("type") == "websocket.disconnect":
+            # NB(nikhil): Mimics FastAPI `_raise_on_disconnect`, since otherwise the user has no
+            # way of detecting that the client disconnected.
+            raise fastapi.WebSocketDisconnect(message["code"], message.get("reason"))
+        elif message.get("text"):
+            return typing.cast(str, message["text"])
+        else:
+            return typing.cast(bytes, message["bytes"])
+
     async def receive_text(self) -> str:
         return await self._websocket.receive_text()