truss 0.11.8rc5__py3-none-any.whl → 0.11.8rc7__py3-none-any.whl

truss/templates/control/control/application.py

@@ -4,11 +4,11 @@ import logging.config
 import re
 import traceback
 from pathlib import Path
-from typing import Dict
+from typing import Awaitable, Callable, Dict
 
 import httpx
 from endpoints import control_app
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, Response
 from fastapi.responses import JSONResponse
 from helpers.errors import ModelLoadFailed, PatchApplicatonError
 from helpers.inference_server_controller import InferenceServerController
@@ -22,52 +22,40 @@ from starlette.middleware.base import BaseHTTPMiddleware
 SANITIZED_EXCEPTION_FRAMES = 2
 
 
-class SanitizedException(Exception):
-    def __init__(self, original_error: Exception, num_frames: int):
-        self.original_error = original_error
-        tb_lines = traceback.format_tb(original_error.__traceback__)
-        if tb_lines and num_frames > 0:
-            # NB(nikhil): Take the last num_frames, join them, and strip extra whitespace
-            selected_frames = tb_lines[-num_frames:]
-            self.concise_msg = "".join(selected_frames).rstrip()
-        else:
-            self.concise_msg = f"{type(original_error).__name__}: {original_error}"
-
-    def __str__(self):
-        return self.concise_msg
-
-    def __repr__(self):
-        return self.concise_msg
-
-
-# NB(nikhil): SanitizedExceptionMiddleware runs on every request and cleans up the verbose stack traces
-# that are very noisy for users.
 class SanitizedExceptionMiddleware(BaseHTTPMiddleware):
-    def __init__(self, app, num_frames: int):
+    def __init__(self, app, num_frames: int = SANITIZED_EXCEPTION_FRAMES):
         super().__init__(app)
         self.num_frames = num_frames
 
-    async def dispatch(self, request: Request, call_next):
+    async def dispatch(
+        self, request: Request, call_next: Callable[[Request], Awaitable[Response]]
+    ) -> Response:
         try:
             return await call_next(request)
-        except Exception as error:
-            raise SanitizedException(error, self.num_frames) from error
-
-
-async def handle_patch_error(_, exc):
-    error_type = _camel_to_snake_case(type(exc).__name__)
-    return JSONResponse(content={"error": {"type": error_type, "msg": str(exc)}})
-
-
-async def generic_error_handler(_, exc):
-    return JSONResponse(
-        content={"error": {"type": "unknown", "msg": f"{type(exc)}: {exc}"}}
-    )
-
-
-async def handle_model_load_failed(_, error):
-    # Model load failures should result in 503 status
-    return JSONResponse({"error": str(error)}, 503)
+        except Exception as exc:
+            sanitized_traceback = self._create_sanitized_traceback(exc)
+            if hasattr(request.app.state, "logger"):
+                request.app.state.logger.error(f"Error:\n{sanitized_traceback}")
+
+            if isinstance(exc, ModelLoadFailed):
+                return JSONResponse({"error": str(exc)}, status_code=503)
+            elif isinstance(exc, PatchApplicatonError):
+                error_type = _camel_to_snake_case(type(exc).__name__)
+                return JSONResponse(
+                    {"error": {"type": error_type, "msg": str(exc)}}, status_code=400
+                )
+            else:
+                return JSONResponse(
+                    {"error": {"type": "unknown", "msg": str(exc)}}, status_code=500
+                )
+
+    def _create_sanitized_traceback(self, error: Exception) -> str:
+        tb_lines = traceback.format_tb(error.__traceback__)
+        if tb_lines and self.num_frames > 0:
+            selected_frames = tb_lines[-self.num_frames :]
+            return "".join(selected_frames).rstrip()
+        else:
+            return f"{type(error).__name__}: {error}"
 
 
 def create_app(base_config: Dict):
@@ -118,17 +106,10 @@ def create_app(base_config: Dict):
     app = FastAPI(
         title="Truss Live Reload Server",
         on_startup=[start_background_inference_startup],
-        exception_handlers={
-            PatchApplicatonError: handle_patch_error,
-            ModelLoadFailed: handle_model_load_failed,
-            Exception: generic_error_handler,
-        },
     )
     app.state = app_state
     app.include_router(control_app)
-    app.add_middleware(
-        SanitizedExceptionMiddleware, num_frames=SANITIZED_EXCEPTION_FRAMES
-    )
+    app.add_middleware(SanitizedExceptionMiddleware)
 
     @app.on_event("shutdown")
     def on_shutdown():

truss/templates/control/control/endpoints.py

@@ -5,6 +5,7 @@ from typing import Any, Callable, Dict, Optional, Protocol
 import httpx
 from fastapi import APIRouter, WebSocket
 from fastapi.responses import JSONResponse, StreamingResponse
+from helpers.errors import ModelLoadFailed, ModelNotReady
 from httpx_ws import AsyncWebSocketSession, WebSocketDisconnect, aconnect_ws
 from httpx_ws import _exceptions as httpx_ws_exceptions
 from starlette.requests import ClientDisconnect, Request
@@ -13,11 +14,6 @@ from starlette.websockets import WebSocketDisconnect as StartletteWebSocketDisco
 from tenacity import RetryCallState, Retrying, retry_if_exception_type, wait_fixed
 from wsproto.events import BytesMessage, TextMessage
 
-from truss.templates.control.control.helpers.errors import (
-    ModelLoadFailed,
-    ModelNotReady,
-)
-
 INFERENCE_SERVER_START_WAIT_SECS = 60
 BASE_RETRY_EXCEPTIONS = (
     retry_if_exception_type(httpx.ConnectError)

{truss-0.11.8rc5.dist-info → truss-0.11.8rc7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.8rc5
+Version: 0.11.8rc7
 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.8rc5.dist-info → truss-0.11.8rc7.dist-info}/RECORD

@@ -73,8 +73,8 @@ truss/templates/copy_cache_files.Dockerfile.jinja,sha256=Os5zFdYLZ_AfCRGq4RcpVTO
 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=nqqNmlTwFeV8sV4fqwItwzzd_egADBP_e-cEopXBJ4k,358
-truss/templates/control/control/application.py,sha256=_FmSgIetffVpwUYVeMM6rH9pGjjHhALTdtR2RegVyJY,5226
-truss/templates/control/control/endpoints.py,sha256=VQ1lvZjFvR091yRkiFdvXw1Q7PiNGXT9rJwY7_sX6yg,11828
+truss/templates/control/control/application.py,sha256=XfA5udraulB0z6s4J_S05w0y1TJGcoMuC3jJOjbPVu4,4839
+truss/templates/control/control/endpoints.py,sha256=KzqsLVNJE6r6TCPW8D5FMCtsfHadTwR15A3z_viGxmM,11782
 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
@@ -368,8 +368,8 @@ truss_train/deployment.py,sha256=lWWANSuzBWu2M4oK4qD7n-oVR1JKdmw2Pn5BJQHg-Ck,307
 truss_train/loader.py,sha256=0o66EjBaHc2YY4syxxHVR4ordJWs13lNXnKjKq2wq0U,1630
 truss_train/public_api.py,sha256=9N_NstiUlmBuLUwH_fNG_1x7OhGCytZLNvqKXBlStrM,1220
 truss_train/restore_from_checkpoint.py,sha256=8hdPm-WSgkt74HDPjvCjZMBpvA9MwtoYsxVjOoa7BaM,1176
-truss-0.11.8rc5.dist-info/METADATA,sha256=NIv6jZ6bWqB41Kqdf9SjNWko0fOVcsqc8-FsW9GIPBU,6680
-truss-0.11.8rc5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truss-0.11.8rc5.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
-truss-0.11.8rc5.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
-truss-0.11.8rc5.dist-info/RECORD,,
+truss-0.11.8rc7.dist-info/METADATA,sha256=37bkH41Vl6H2ftd6R4aJZWnxqanKAupqM4uOgFNmYlA,6680
+truss-0.11.8rc7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truss-0.11.8rc7.dist-info/entry_points.txt,sha256=-MwKfHHQHQ6j0HqIgvxrz3CehCmczDLTD-OsRHnjjuU,130
+truss-0.11.8rc7.dist-info/licenses/LICENSE,sha256=FTqGzu85i-uw1Gi8E_o0oD60bH9yQ_XIGtZbA1QUYiw,1064
+truss-0.11.8rc7.dist-info/RECORD,,