Return 500 error codes on user error. (#545)

* Return 500 error codes on user error.

* Fix issue w/ missing values.

* Remove stacktrace.

* Fix pr.

* Fix tests.

* Fix tests.

* Fix all tests.

* Fix test.

* Fix broked TGI model_id swap

* Bump rc version

* Fix create dir that already exists

* bump version to test

* Drop old make from in-memory model code

* Drop trainable config from cli comand

* Fix training failures

* Fix tests

* Try again to make test pass

* Add flask depedency for integration tests

* Skip flaky test

* Bump pyproject.

* Update for PR feedback.

---------

Co-authored-by: Bola Malek <bola@baseten.co>

Author: squidarth

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.6.5rc7"
+version = "0.6.5rc11"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/templates/server/common/errors.py

@@ -94,3 +94,7 @@ async def not_implemented_error_handler(_, exc):
     return JSONResponse(
         status_code=HTTPStatus.NOT_IMPLEMENTED, content={"error": str(exc)}
     )
+
+
+async def http_exception_handler(_, exc):
+    return JSONResponse(status_code=exc.status_code, content={"error": exc.detail})

truss/templates/server/common/truss_server.py

@@ -234,6 +234,7 @@ def create_application(self):
                 errors.ModelNotFound: errors.model_not_found_handler,
                 errors.ModelNotReady: errors.model_not_ready_handler,
                 NotImplementedError: errors.not_implemented_error_handler,
+                HTTPException: errors.http_exception_handler,
                 Exception: errors.generic_exception_handler,
             },
         )

truss/templates/server/model_wrapper.py

@@ -5,7 +5,6 @@
 import os
 import sys
 import time
-import traceback
 from collections.abc import Generator
 from contextlib import asynccontextmanager
 from enum import Enum
@@ -16,6 +15,7 @@
 from anyio import Semaphore, to_thread
 from common.patches import apply_patches
 from common.retry import retry
+from fastapi import HTTPException
 from shared.secrets_resolver import SecretsResolver
 
 MODEL_BASENAME = "model"
@@ -177,23 +177,11 @@ async def preprocess(
             return payload
 
         if inspect.iscoroutinefunction(self._model.preprocess):
-            return await self._model.preprocess(payload)
+            return await _intercept_exceptions_async(self._model.preprocess)(payload)
         else:
-            return await to_thread.run_sync(self._model.preprocess, payload)
-
-    def _predict_sync_with_error_handling(self, payload):
-        try:
-            return self._model.predict(payload)
-        except Exception:
-            logging.exception("Exception while running predict")
-            return {"error": {"traceback": traceback.format_exc()}}
-
-    async def _predict_async_with_error_handling(self, payload):
-        try:
-            return await self._model.predict(payload)
-        except Exception:
-            logging.exception("Exception while running predict")
-            return {"error": {"traceback": traceback.format_exc()}}
+            return await to_thread.run_sync(
+                _intercept_exceptions_sync(self._model.preprocess), payload
+            )
 
     async def predict(
         self,
@@ -214,9 +202,11 @@ async def predict(
             return self._model.predict(payload)
 
         if inspect.iscoroutinefunction(self._model.predict):
-            return await self._predict_async_with_error_handling(payload)
+            return await _intercept_exceptions_async(self._model.predict)(payload)
 
-        return await to_thread.run_sync(self._predict_sync_with_error_handling, payload)
+        return await to_thread.run_sync(
+            _intercept_exceptions_sync(self._model.predict), payload
+        )
 
     async def postprocess(
         self,
@@ -238,9 +228,11 @@ async def postprocess(
             return self._model.postprocess(response, headers)
 
         if inspect.iscoroutinefunction(self._model.postprocess):
-            return await self._model.postprocess(response)
+            return await _intercept_exceptions_async(self._model.postprocess)(response)
 
-        return await to_thread.run_sync(self._model.postprocess, response)
+        return await to_thread.run_sync(
+            _intercept_exceptions_sync(self._model.postprocess), response
+        )
 
     async def write_response_to_queue(
         self, queue: asyncio.Queue, generator: AsyncGenerator
@@ -368,3 +360,30 @@ def _signature_accepts_kwargs(signature: inspect.Signature) -> bool:
 
 def _elapsed_ms(since_micro_seconds: float) -> int:
     return int((time.perf_counter() - since_micro_seconds) * 1000)
+
+
+def _handle_exception():
+    # Note that logger.exception logs the stacktrace, such that the user can
+    # debug this error from the logs.
+    logging.exception("Internal Server Error")
+    raise HTTPException(status_code=500, detail={"message": "Internal Server Error"})
+
+
+def _intercept_exceptions_sync(func):
+    def inner(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception:
+            _handle_exception()
+
+    return inner
+
+
+def _intercept_exceptions_async(func):
+    async def inner(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except Exception:
+            _handle_exception()
+
+    return inner

truss/tests/test_model_inference.py

@@ -1,5 +1,6 @@
 import concurrent
 import inspect
+import json
 import logging
 import tempfile
 import textwrap
@@ -32,6 +33,19 @@ def _create_truss(truss_dir: Path, config_contents: str, model_contents: str):
         file.write(model_contents)
 
 
+def _log_contains_error(line: dict, error: str):
+    return (
+        line["levelname"] == "ERROR"
+        and line["message"] == "Exception while running predict"
+        and error in line["exc_info"]
+    )
+
+
+def assert_logs_contain_error(logs: str, error: str):
+    loglines = logs.splitlines()
+    assert any(_log_contains_error(json.loads(line), error) for line in loglines)
+
+
 class PropagatingThread(Thread):
     """
     PropagatingThread allows us to run threads and keep track of exceptions
@@ -317,20 +331,11 @@ def predict(self, request):
             return self._secrets["secret"]
 
     config = """model_name: secrets-truss
-cpu: "3"
-memory: 14Gi
-use_gpu: true
-accelerator: A10G
 secrets:
     secret: null
     """
 
-    config_with_no_secret = """model_name: secrets-truss
-cpu: "3"
-memory: 14Gi
-use_gpu: true
-accelerator: A10G
-    """
+    config_with_no_secret = "model_name: secrets-truss"
 
     with ensure_kill_all(), tempfile.TemporaryDirectory(dir=".") as tmp_work_dir:
         truss_dir = Path(tmp_work_dir, "truss")
@@ -344,9 +349,8 @@ def predict(self, request):
         full_url = f"{truss_server_addr}/v1/models/model:predict"
 
         response = requests.post(full_url, json={})
-        assert response.json() == "secret_value"
 
-        _create_truss(truss_dir, config, textwrap.dedent(inspect.getsource(Model)))
+        assert response.json() == "secret_value"
 
     with ensure_kill_all(), tempfile.TemporaryDirectory(dir=".") as tmp_work_dir:
         # Case where the secret is not specified in the config
@@ -357,33 +361,149 @@ def predict(self, request):
         )
         tr = TrussHandle(truss_dir)
         LocalConfigHandler.set_secret("secret", "secret_value")
-        _ = tr.docker_run(local_port=8090, detach=True, wait_for_server_ready=True)
+        container = tr.docker_run(
+            local_port=8090, detach=True, wait_for_server_ready=True
+        )
         truss_server_addr = "http://localhost:8090"
         full_url = f"{truss_server_addr}/v1/models/model:predict"
 
         response = requests.post(full_url, json={})
 
         assert "error" in response.json()
-        assert "not specified in the config" in response.json()["error"]["traceback"]
+        assert_logs_contain_error(container.logs(), "not specified in the config")
+        assert "Error while running predict" in response.json()["error"]["message"]
 
     with ensure_kill_all(), tempfile.TemporaryDirectory(dir=".") as tmp_work_dir:
-        # Case where the secret is not specified in the config
+        # Case where the secret is not mounted
         truss_dir = Path(tmp_work_dir, "truss")
 
         _create_truss(truss_dir, config, textwrap.dedent(inspect.getsource(Model)))
         tr = TrussHandle(truss_dir)
         LocalConfigHandler.remove_secret("secret")
-        _ = tr.docker_run(local_port=8090, detach=True, wait_for_server_ready=True)
+        container = tr.docker_run(
+            local_port=8090, detach=True, wait_for_server_ready=True
+        )
+        truss_server_addr = "http://localhost:8090"
+        full_url = f"{truss_server_addr}/v1/models/model:predict"
+
+        response = requests.post(full_url, json={})
+        assert response.status_code == 500
+        assert_logs_contain_error(
+            container.logs(), "'secret' not found. Please check available secrets."
+        )
+        assert "Error while running predict" in response.json()["error"]["message"]
+
+
+@pytest.mark.integration
+def test_truss_with_errors():
+    model = """
+    class Model:
+        def predict(self, request):
+            raise ValueError("error")
+    """
+
+    config = "model_name: error-truss"
+
+    with ensure_kill_all(), tempfile.TemporaryDirectory(dir=".") as tmp_work_dir:
+        truss_dir = Path(tmp_work_dir, "truss")
+
+        _create_truss(truss_dir, config, textwrap.dedent(model))
+
+        tr = TrussHandle(truss_dir)
+        container = tr.docker_run(
+            local_port=8090, detach=True, wait_for_server_ready=True
+        )
+        truss_server_addr = "http://localhost:8090"
+        full_url = f"{truss_server_addr}/v1/models/model:predict"
+
+        response = requests.post(full_url, json={})
+        assert response.status_code == 500
+        assert "error" in response.json()
+
+        assert_logs_contain_error(container.logs(), "ValueError: error")
+
+        assert "Error while running predict" in response.json()["error"]["message"]
+
+    model_preprocess_error = """
+    class Model:
+        def preprocess(self, request):
+            raise ValueError("error")
+
+        def predict(self, request):
+            return {"a": "b"}
+    """
+
+    with ensure_kill_all(), tempfile.TemporaryDirectory(dir=".") as tmp_work_dir:
+        truss_dir = Path(tmp_work_dir, "truss")
+
+        _create_truss(truss_dir, config, textwrap.dedent(model_preprocess_error))
+
+        tr = TrussHandle(truss_dir)
+        container = tr.docker_run(
+            local_port=8090, detach=True, wait_for_server_ready=True
+        )
         truss_server_addr = "http://localhost:8090"
         full_url = f"{truss_server_addr}/v1/models/model:predict"
 
         response = requests.post(full_url, json={})
+        assert response.status_code == 500
+        assert "error" in response.json()
+
+        assert_logs_contain_error(container.logs(), "ValueError: error")
+        assert "Error while running predict" in response.json()["error"]["message"]
+
+    model_postprocess_error = """
+    class Model:
+        def predict(self, request):
+            return {"a": "b"}
+
+        def postprocess(self, response):
+            raise ValueError("error")
+    """
+
+    with ensure_kill_all(), tempfile.TemporaryDirectory(dir=".") as tmp_work_dir:
+        truss_dir = Path(tmp_work_dir, "truss")
+
+        _create_truss(truss_dir, config, textwrap.dedent(model_postprocess_error))
+
+        tr = TrussHandle(truss_dir)
+        container = tr.docker_run(
+            local_port=8090, detach=True, wait_for_server_ready=True
+        )
+        truss_server_addr = "http://localhost:8090"
+        full_url = f"{truss_server_addr}/v1/models/model:predict"
 
+        response = requests.post(full_url, json={})
+        assert response.status_code == 500
         assert "error" in response.json()
-        assert (
-            "not found. Please check available secrets."
-            in response.json()["error"]["traceback"]
+        assert_logs_contain_error(container.logs(), "ValueError: error")
+        assert "Error while running predict" in response.json()["error"]["message"]
+
+    model_async = """
+    class Model:
+        async def predict(self, request):
+            raise ValueError("error")
+    """
+
+    with ensure_kill_all(), tempfile.TemporaryDirectory(dir=".") as tmp_work_dir:
+        truss_dir = Path(tmp_work_dir, "truss")
+
+        _create_truss(truss_dir, config, textwrap.dedent(model_async))
+
+        tr = TrussHandle(truss_dir)
+        container = tr.docker_run(
+            local_port=8090, detach=True, wait_for_server_ready=True
         )
+        truss_server_addr = "http://localhost:8090"
+        full_url = f"{truss_server_addr}/v1/models/model:predict"
+
+        response = requests.post(full_url, json={})
+        assert response.status_code == 500
+        assert "error" in response.json()
+
+        assert_logs_contain_error(container.logs(), "ValueError: error")
+
+        assert "Error while running predict" in response.json()["error"]["message"]
 
 
 @pytest.mark.integration