truss 0.11.8rc6__tar.gz → 0.11.8rc8__tar.gz

{truss-0.11.8rc6 → truss-0.11.8rc8}/PKG-INFO

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

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

{truss-0.11.8rc6 → truss-0.11.8rc8}/truss/base/trt_llm_config.py

@@ -329,9 +329,16 @@ pip install truss==0.10.8
             raise ValueError("Using fp8 context fmha requires paged context fmha")
         if (
             self.plugin_configuration.use_fp8_context_fmha
-            and not self.quantization_type == TrussTRTLLMQuantizationType.FP8_KV
+            and self.quantization_type
+            not in (
+                TrussTRTLLMQuantizationType.FP8_KV,
+                TrussTRTLLMQuantizationType.FP4_KV,
+            )
         ):
-            raise ValueError("Using fp8 context fmha requires fp8 kv cache dtype")
+            raise ValueError(
+                "Using fp8 context fmha requires fp8 kv, or fp4 with kv cache dtype"
+            )
+
         return self
 
     def _validate_speculator_config(self):

{truss-0.11.8rc6 → truss-0.11.8rc8}/truss/templates/control/control/application.py

@@ -10,61 +10,63 @@ import httpx
 from endpoints import control_app
 from fastapi import FastAPI, Request, 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
-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
+
+from truss.templates.control.control.helpers.errors import (
+    ModelLoadFailed,
+    PatchApplicatonError,
+)
+from truss.templates.control.control.helpers.inference_server_controller import (
+    InferenceServerController,
+)
+from truss.templates.control.control.helpers.inference_server_process_controller import (
+    InferenceServerProcessController,
+)
+from truss.templates.control.control.helpers.inference_server_starter import (
+    async_inference_server_startup_flow,
+)
+from truss.templates.control.control.helpers.truss_patch.model_container_patch_applier import (
+    ModelContainerPatchApplier,
+)
 
 SANITIZED_EXCEPTION_FRAMES = 2
 
 
-def create_sanitized_traceback(error: Exception, num_frames: int) -> str:
-    tb_lines = traceback.format_tb(error.__traceback__)
-    if tb_lines and num_frames > 0:
-        selected_frames = tb_lines[-num_frames:]
-        return "".join(selected_frames).rstrip()
-    else:
-        return f"{type(error).__name__}: {error}"
-
-
-def sanitize_exception(num_frames=SANITIZED_EXCEPTION_FRAMES):
-    def decorator(
-        handler_func: Callable[[Request, Exception], Awaitable[Response]],
-    ) -> Callable[[Request, Exception], Awaitable[Response]]:
-        async def wrapper(request: Request, exc: Exception) -> Response:
-            sanitized_traceback = create_sanitized_traceback(exc, num_frames)
-            if hasattr(request.app.state, "logger"):
-                request.app.state.logger.error(sanitized_traceback)
-
-            return await handler_func(request, exc)
-
-        return wrapper
-
-    return decorator
-
-
-@sanitize_exception()
-async def handle_patch_error(request, exc):
-    error_type = _camel_to_snake_case(type(exc).__name__)
-    return JSONResponse(content={"error": {"type": error_type, "msg": str(exc)}})
-
-
-@sanitize_exception()
-async def generic_error_handler(request, exc):
-    print("CALLED GENERIC ========== ")
-    return JSONResponse(
-        content={"error": {"type": "unknown", "msg": f"{type(exc)}: {exc}"}}
-    )
-
-
-@sanitize_exception()
-async def handle_model_load_failed(request, error):
-    # Model load failures should result in 503 status
-    print("CALLED MODEL LOAD ========== ")
-    return JSONResponse({"error": str(error)}, 503)
+class SanitizedExceptionMiddleware(BaseHTTPMiddleware):
+    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: Callable[[Request], Awaitable[Response]]
+    ) -> Response:
+        try:
+            return await call_next(request)
+        except Exception as exc:
+            sanitized_traceback = self._create_sanitized_traceback(exc)
+            request.app.state.logger.error(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):
@@ -115,14 +117,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)
 
     @app.on_event("shutdown")
     def on_shutdown():