API: Fix issues with concurrent requests and queueing

This is the first in many future commits that will overhaul the API
to be more robust and concurrent. The model is admin-first where the
admin can do anything in-case something goes awry.

Previously, calls to long running synchronous background tasks would
block the entire API, making it ignore any terminal signals until
generation is completed.

To fix this, levrage FastAPI's run_in_threadpool to offload the long
running tasks to another thread. However, signals to abort the process
still kept the background thread running and made the terminal hang.

This was due to an issue with Uvicorn not propegating the SIGINT signal
across threads in its event loop. To fix this in a catch-all way, run
the API processes in a separate thread so the main thread can still
kill the process if needed.

In addition, make request error logging more robust and refer to the
console for full error logs rather than creating a long message on the
client-side.

Finally, add state checks to see if a model is fully loaded before
generating a completion.

Signed-off-by: kingbri <bdashore3@proton.me>

common/generators.py

@@ -1,15 +1,16 @@
-"""Generator functions for the tabbyAPI."""
+"""Generator handling"""
+
+import asyncio
 import inspect
-from asyncio import Semaphore
 from functools import partialmethod
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Generator, Union
 
-generate_semaphore = Semaphore(1)
+generate_semaphore = asyncio.Semaphore(1)
 
 
-# Async generation that blocks on a semaphore
-async def generate_with_semaphore(generator: AsyncGenerator):
+async def generate_with_semaphore(generator: Union[AsyncGenerator, Generator]):
     """Generate with a semaphore."""
+
     async with generate_semaphore:
         if inspect.isasyncgenfunction:
             async for result in generator():
@@ -19,9 +20,11 @@ async def generate_with_semaphore(generator: AsyncGenerator):
                 yield result
 
 
-# Block a function with semaphore
 async def call_with_semaphore(callback: partialmethod):
-    if inspect.iscoroutinefunction(callback):
-        return await callback()
+    """Call with a semaphore."""
+
     async with generate_semaphore:
-        return callback()
+        if inspect.iscoroutinefunction(callback):
+            return await callback()
+        else:
+            return callback()

common/utils.py

@@ -1,8 +1,8 @@
-"""Common utilities for the tabbyAPI"""
-import traceback
-from typing import Optional
+"""Common utility functions"""
 
+import traceback
 from pydantic import BaseModel
+from typing import Optional
 
 from common.logger import init_logger
 
@@ -14,30 +14,41 @@ def load_progress(module, modules):
     yield module, modules
 
 
-class TabbyGeneratorErrorMessage(BaseModel):
-    """Common error types."""
+class TabbyRequestErrorMessage(BaseModel):
+    """Common request error type."""
 
     message: str
     trace: Optional[str] = None
 
 
-class TabbyGeneratorError(BaseModel):
-    """Common error types."""
+class TabbyRequestError(BaseModel):
+    """Common request error type."""
 
-    error: TabbyGeneratorErrorMessage
+    error: TabbyRequestErrorMessage
 
 
 def get_generator_error(message: str):
     """Get a generator error."""
-    error_message = TabbyGeneratorErrorMessage(
+
+    generator_error = handle_request_error(message)
+
+    return get_sse_packet(generator_error.model_dump_json())
+
+
+def handle_request_error(message: str):
+    """Log a request error to the console."""
+
+    error_message = TabbyRequestErrorMessage(
         message=message, trace=traceback.format_exc()
     )
 
-    generator_error = TabbyGeneratorError(error=error_message)
+    request_error = TabbyRequestError(error=error_message)
 
-    # Log and send the exception
-    logger.error(generator_error.error.trace)
-    return get_sse_packet(generator_error.model_dump_json())
+    # Log the error and provided message to the console
+    logger.error(error_message.trace)
+    logger.error(message)
+
+    return request_error
 
 
 def get_sse_packet(json_data: str):