fix custom exceptions (#425)

* always pickle exceptions

* fix

* fix

* fix

Author: aniketmaurya

src/litserve/loops/base.py

@@ -14,6 +14,7 @@
 import asyncio
 import inspect
 import logging
+import pickle
 import signal
 import sys
 import time
@@ -256,6 +257,7 @@ def put_response(
     def put_error_response(
         self, response_queues: List[Queue], response_queue_id: int, uid: str, error: Exception
     ) -> None:
+        error = pickle.dumps(error)
         self.put_response(response_queues, response_queue_id, uid, error, LitAPIStatus.ERROR)
 
     def __del__(self):

src/litserve/loops/simple_loops.py

@@ -109,12 +109,11 @@ def run_single_loop(
                     "Please check the error trace for more details.",
                     uid,
                 )
-                self.put_response(
+                self.put_error_response(
                     response_queues=response_queues,
                     response_queue_id=response_queue_id,
                     uid=uid,
-                    response_data=e,
-                    status=LitAPIStatus.ERROR,
+                    error=e,
                 )
 
     def __call__(
@@ -226,7 +225,7 @@ def run_batched_loop(
                     "Please check the error trace for more details."
                 )
                 for response_queue_id, uid in zip(response_queue_ids, uids):
-                    self.put_response(response_queues, response_queue_id, uid, e, LitAPIStatus.ERROR)
+                    self.put_error_response(response_queues, response_queue_id, uid, e)
 
     def __call__(
         self,

src/litserve/loops/streaming_loops.py

@@ -102,7 +102,7 @@ def run_streaming_loop(
                     "Please check the error trace for more details.",
                     uid,
                 )
-                self.put_response(response_queues, response_queue_id, uid, e, LitAPIStatus.ERROR)
+                self.put_error_response(response_queues, response_queue_id, uid, e)
 
     def __call__(
         self,
@@ -207,7 +207,7 @@ def run_batched_streaming_loop(
                     "Please check the error trace for more details."
                 )
                 for response_queue_id, uid in zip(response_queue_ids, uids):
-                    self.put_response(response_queues, response_queue_id, uid, e, LitAPIStatus.ERROR)
+                    self.put_error_response(response_queues, response_queue_id, uid, e)
 
     def __call__(
         self,

tests/test_loops.py

@@ -132,7 +132,7 @@ def put(self, item, block=False, timeout=None):
         response, status = args
         if status == LitAPIStatus.FINISH_STREAMING:
             raise StopIteration("interrupt iteration")
-        if status == LitAPIStatus.ERROR and isinstance(response, StopIteration):
+        if status == LitAPIStatus.ERROR:
             assert self.count // 2 == self.num_streamed_outputs, (
                 f"Loop count must have incremented for {self.num_streamed_outputs} times."
             )

tests/test_simple.py

@@ -244,3 +244,21 @@ def test_concurrent_requests(lit_server):
         assert response.json() == {"output": i**2}, "Server returns square of the input number"
         count += 1
     assert count == n_requests
+
+
+class CustomError(Exception):
+    def __init__(self, arg1, arg2, arg3):
+        super().__init__("Test exception")
+
+
+class ExceptionAPI(SimpleLitAPI):
+    def predict(self, x):
+        raise CustomError("This", "is", "a test")
+
+
+def test_exception():
+    server = LitServer(ExceptionAPI(), accelerator="cpu", devices=1)
+    with wrap_litserve_start(server) as server, TestClient(server.app) as client:
+        response = client.post("/predict", json={"input": 4.0})
+        assert response.status_code == 500
+        assert response.json() == {"detail": "Internal Server Error"}