reformat

src/fastserve/base_fastserve.py

@@ -2,7 +2,7 @@
 from contextlib import asynccontextmanager
 from typing import Any, List
 
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 
 from .batching import BatchProcessor
@@ -39,8 +39,12 @@ def _serve(
 
         @self._app.post(path="/endpoint")
         def api(request: INPUT_SCHEMA):
+            print("incoming request")
             wait_obj = self.batch_processing.process(request)
-            return wait_obj.get()
+            result = wait_obj.get()
+            if isinstance(result, Exception):
+                raise HTTPException()
+            return result
 
     def handle(self, batch: List[BaseRequest]):
         n = len(batch)

src/fastserve/batching.py

@@ -119,7 +119,12 @@ def _process_queue(self):
             logger.info(f"Aggregated batch size {len(batch)} in {t1-t0:.2f}s")
             batch_items = [b.item for b in batch]
             logger.debug(batch_items)
-            results = self.func(batch_items)
+            try:
+                results = self.func(batch_items)
+            except Exception as e:
+                logger.error(f"Error while processing batch {batch}")
+                logger.error(e)
+                results = [e] * len(batch)
             if not isinstance(results, list):
                 logger.error(f"returned results must be List but is {type(results)}")
             logger.debug(results)
@@ -134,7 +139,8 @@ def process(self, item: Any):
     def cancel(self):
         logger.info("Terminating Batch Processor...")
         self._cancel_signal.set()
-        self._thread.join()
+        self._thread.terminate()
+        # self._thread.join()
         logger.info("Batch Processor terminated!")
 
     def signal_handler(self, sig, frame):

src/fastserve/models/ssd.py

@@ -4,9 +4,10 @@
 import torch
 from diffusers import StableDiffusionXLPipeline
 from fastapi.responses import StreamingResponse
-from fastserve import BaseRequest, FastServe
 from pydantic import BaseModel
 
+from fastserve import BaseRequest, FastServe
+
 
 class PromptRequest(BaseModel):
     prompt: str  # "An astronaut riding a green horse"