mirror of
https://github.com/theroyallab/tabbyAPI.git
synced 2026-03-14 15:57:27 +00:00
56f9b1d1a8
If the generator errors, there's no proper handling to send an error packet and close the connection. This is especially important for unloading models if the load fails at any stage to reclaim a user's VRAM. Raising an exception caused the model_container object to lock and not get freed by the GC. This made sense to propegate SSE errors across all generator functions rather than relying on abort signals. Signed-off-by: kingbri <bdashore3@proton.me>
30 lines
738 B
Python
30 lines
738 B
Python
import traceback
|
|
from pydantic import BaseModel
|
|
from typing import Optional
|
|
|
|
# Wrapper callback for load progress
|
|
def load_progress(module, modules):
|
|
yield module, modules
|
|
|
|
# Common error types
|
|
class TabbyGeneratorErrorMessage(BaseModel):
|
|
message: str
|
|
trace: Optional[str] = None
|
|
|
|
class TabbyGeneratorError(BaseModel):
|
|
error: TabbyGeneratorErrorMessage
|
|
|
|
def get_generator_error(exception: Exception):
|
|
error_message = TabbyGeneratorErrorMessage(
|
|
message = str(exception),
|
|
trace = traceback.format_exc()
|
|
)
|
|
|
|
generator_error = TabbyGeneratorError(
|
|
error = error_message
|
|
)
|
|
|
|
# Log and send the exception
|
|
print(f"\n{generator_error.error.trace}")
|
|
return generator_error.json()
|