truss 0.11.8rc4__tar.gz → 0.11.8rc5__tar.gz

{truss-0.11.8rc4 → truss-0.11.8rc5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truss
-Version: 0.11.8rc4
+Version: 0.11.8rc5
 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.8rc4 → truss-0.11.8rc5}/pyproject.toml

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

{truss-0.11.8rc4 → truss-0.11.8rc5}/truss/templates/control/control/application.py

@@ -3,14 +3,13 @@ import logging
 import logging.config
 import re
 import traceback
-from functools import wraps
 from pathlib import Path
-from typing import Any, Callable, Coroutine, Dict
+from typing import Dict
 
 import httpx
 from endpoints import control_app
 from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse, Response
+from fastapi.responses import JSONResponse
 from helpers.errors import ModelLoadFailed, PatchApplicatonError
 from helpers.inference_server_controller import InferenceServerController
 from helpers.inference_server_process_controller import InferenceServerProcessController
@@ -18,45 +17,54 @@ from helpers.inference_server_starter import async_inference_server_startup_flow
 from helpers.truss_patch.model_container_patch_applier import ModelContainerPatchApplier
 from shared import log_config
 from starlette.datastructures import State
+from starlette.middleware.base import BaseHTTPMiddleware
 
+SANITIZED_EXCEPTION_FRAMES = 2
 
-def sanitized_exception_logging(
-    original_handler: Callable[[Request, Any], Coroutine[Any, Any, Response]],
-) -> Callable[[Request, Any], Coroutine[Any, Any, Response]]:
-    @wraps(original_handler)
-    async def wrapper(request: Request, error: Exception) -> Response:
-        logger = request.app.state.logger
-        traceback_error = traceback.extract_tb(error.__traceback__)
 
-        if traceback_error:
-            last_frame = traceback_error[-1]
-            concise_error = (
-                f'  File "{last_frame.filename}", line {last_frame.lineno}, in {last_frame.name}\n'
-                f"    {last_frame.line}"
-            )
+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:
-            concise_error = f"{type(error).__name__}: {error}"
+            self.concise_msg = f"{type(original_error).__name__}: {original_error}"
+
+    def __str__(self):
+        return self.concise_msg
+
+    def __repr__(self):
+        return self.concise_msg
 
-        logger.error(concise_error)
-        return await original_handler(request, error)
 
-    return wrapper
+# 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):
+        super().__init__(app)
+        self.num_frames = num_frames
+
+    async def dispatch(self, request: Request, call_next):
+        try:
+            return await call_next(request)
+        except Exception as error:
+            raise SanitizedException(error, self.num_frames) from error
 
 
-@sanitized_exception_logging
 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)}})
 
 
-@sanitized_exception_logging
 async def generic_error_handler(_, exc):
     return JSONResponse(
         content={"error": {"type": "unknown", "msg": f"{type(exc)}: {exc}"}}
     )
 
 
-@sanitized_exception_logging
 async def handle_model_load_failed(_, error):
     # Model load failures should result in 503 status
     return JSONResponse({"error": str(error)}, 503)
@@ -118,6 +126,9 @@ def create_app(base_config: Dict):
     )
     app.state = app_state
     app.include_router(control_app)
+    app.add_middleware(
+        SanitizedExceptionMiddleware, num_frames=SANITIZED_EXCEPTION_FRAMES
+    )
 
     @app.on_event("shutdown")
     def on_shutdown():