Merge branch 'main' into fix/server_llama_call_thread_starvation

abetlen · web-flow · commit 9e0728bb4782 · 2024-12-06T05:03:15.000-05:00
diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py
@@ -312,10 +312,14 @@ async def create_completion(
         else:
             kwargs["logits_processor"].extend(_min_tokens_logits_processor)
 
-    iterator_or_completion: Union[
-        llama_cpp.CreateCompletionResponse,
-        Iterator[llama_cpp.CreateCompletionStreamResponse],
-    ] = await run_in_threadpool(llama, **kwargs)
+    try:
+        iterator_or_completion: Union[
+            llama_cpp.CreateCompletionResponse,
+            Iterator[llama_cpp.CreateCompletionStreamResponse],
+        ] = await run_in_threadpool(llama, **kwargs)
+    except Exception as err:
+        exit_stack.close()
+        raise err
 
     if isinstance(iterator_or_completion, Iterator):
         # EAFP: It's easier to ask for forgiveness than permission
@@ -504,9 +508,13 @@ async def create_chat_completion(
         else:
             kwargs["logits_processor"].extend(_min_tokens_logits_processor)
 
-    iterator_or_completion: Union[
-        llama_cpp.ChatCompletion, Iterator[llama_cpp.ChatCompletionChunk]
-    ] = await run_in_threadpool(llama.create_chat_completion, **kwargs)
+    try:
+        iterator_or_completion: Union[
+            llama_cpp.ChatCompletion, Iterator[llama_cpp.ChatCompletionChunk]
+        ] = await run_in_threadpool(llama.create_chat_completion, **kwargs)
+    except Exception as err:
+        exit_stack.close()
+        raise err
 
     if isinstance(iterator_or_completion, Iterator):
         # EAFP: It's easier to ask for forgiveness than permission
diff --git a/llama_cpp/server/errors.py b/llama_cpp/server/errors.py
@@ -134,8 +134,6 @@ def error_message_wrapper(
         ] = None,
     ) -> Tuple[int, ErrorResponse]:
         """Wraps error message in OpenAI style error response"""
-        print(f"Exception: {str(error)}", file=sys.stderr)
-        traceback.print_exc(file=sys.stderr)
         if body is not None and isinstance(
             body,
             (
@@ -149,6 +147,10 @@ def error_message_wrapper(
                 if match is not None:
                     return callback(body, match)
 
+        # Only print the trace on unexpected exceptions
+        print(f"Exception: {str(error)}", file=sys.stderr)
+        traceback.print_exc(file=sys.stderr)
+
         # Wrap other errors as internal server error
         return 500, ErrorResponse(
             message=str(error),
