truss 0.11.1rc2__tar.gz → 0.11.1rc4__tar.gz

{truss-0.11.1rc2 → truss-0.11.1rc4}/PKG-INFO

@@ -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 → truss-0.11.1rc4}/pyproject.toml

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

{truss-0.11.1rc2 → truss-0.11.1rc4}/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-0.11.1rc2 → truss-0.11.1rc4}/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-0.11.1rc2 → truss-0.11.1rc4}/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-0.11.1rc2 → truss-0.11.1rc4}/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-0.11.1rc2 → truss-0.11.1rc4}/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-0.11.1rc2 → truss-0.11.1rc4}/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-0.11.1rc2 → truss-0.11.1rc4}/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 → truss-0.11.1rc4}/truss-chains/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-0.11.1rc2 → truss-0.11.1rc4}/truss-chains/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-0.11.1rc2 → truss-0.11.1rc4}/truss-chains/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()