raise HTTPException (#145)

_requirements/test.txt

@@ -1,12 +1,11 @@
-coverage[toml] >=5.0
+coverage[toml] >=7.5.3
 pytest >=8.0
 pytest-cov
 mypy ==1.9.0
 pytest-asyncio
 asgi-lifespan
 python-multipart
 psutil
-
 requests
 lightning >2.0.0
 torch >2.0.0

src/litserve/utils.py

@@ -49,13 +49,14 @@ async def wait_for_queue_timeout(coro: Coroutine, timeout: Optional[float], uid:
 
 def load_and_raise(response):
     try:
-        pickle.loads(response)
-        raise HTTPException(500, "Internal Server Error")
+        exception = pickle.loads(response)
+        raise exception
     except pickle.PickleError:
         logger.exception(
             f"main process failed to load the exception from the parallel worker process. "
             f"{response} couldn't be unpickled."
         )
+        raise
 
 
 async def azip(*async_iterables):

tests/test_lit_server.py

@@ -18,7 +18,7 @@
 from multiprocessing import Pipe, Manager
 from asgi_lifespan import LifespanManager
 from litserve import LitAPI
-from fastapi import Request, Response
+from fastapi import Request, Response, HTTPException
 
 import torch
 import torch.nn as nn
@@ -416,3 +416,16 @@ def test_custom_api_path():
     with TestClient(server.app) as client:
         response = client.post(url, json={"input": 4.0})
         assert response.status_code == 200, "Server response should be 200 (OK)"
+
+
+class TestHTTPExceptionAPI(ls.examples.SimpleLitAPI):
+    def decode_request(self, request):
+        raise HTTPException(501, "decode request is bad")
+
+
+def test_http_exception():
+    server = LitServer(TestHTTPExceptionAPI())
+    with TestClient(server.app) as client:
+        response = client.post("/predict", json={"input": 4.0})
+        assert response.status_code == 501, "Server raises 501 error"
+        assert response.text == '{"detail":"decode request is bad"}', "decode request is bad"

tests/test_specs.py

@@ -14,6 +14,7 @@
 
 import pytest
 from asgi_lifespan import LifespanManager
+from fastapi import HTTPException
 from httpx import AsyncClient
 from litserve.examples.openai_spec_example import TestAPI, TestAPIWithCustomEncode, TestAPIWithToolCalls
 from litserve.specs.openai import OpenAISpec, ChatMessage
@@ -137,12 +138,13 @@ def setup(self, device):
 
     def predict(self, prompt):
         yield "This is a sample generated text"
-        raise Exception("random error")
+        raise HTTPException(501, "test LitAPI.predict error")
 
 
 @pytest.mark.asyncio()
 async def test_fail_http(openai_request_data):
     server = ls.LitServer(WrongLitAPI(), spec=ls.OpenAISpec())
     async with LifespanManager(server.app) as manager, AsyncClient(app=manager.app, base_url="http://test") as ac:
-        resp = await ac.post("/v1/chat/completions", json=openai_request_data, timeout=10)
-        assert resp.status_code == 500, "Server raises an exception so client should fail"
+        res = await ac.post("/v1/chat/completions", json=openai_request_data, timeout=10)
+        assert res.status_code == 501, "Server raises 501 error"
+        assert res.text == '{"detail":"test LitAPI.predict error"}'