import threading
from typing import Optional
import uvicorn
from fastapi import FastAPI
from fastapi.security import HTTPBearer
from modules.logger import log
from .classes import Model, Config
from .helpers import get_attention_config
from .routes import setup_routes
class OpenAIServer:
def __init__(
self,
model,
tokenizer,
processor=None,
host: str = "127.0.0.1",
port: int = 8888,
server: Optional[uvicorn.Server] = None,
max_context_tokens: Optional[int] = None,
max_new_tokens: Optional[int] = None,
stream: Optional[bool] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
top_k: Optional[int] = None,
repetition_penalty: Optional[float] = None,
api_key: Optional[str] = None,
):
self.model = model
self.tokenizer = tokenizer
self.processor = processor
self.host = host
self.port = port
self.model_info = Model(
name=getattr(model.config, "_name_or_path", "local-transformer"),
cls=model.__class__.__name__,
tokenizer=tokenizer.__class__.__name__,
processor=processor.__class__.__name__ if processor else None,
type=getattr(model.config, "model_type", None)
)
self.config = Config(
max_context_tokens=max_context_tokens if max_context_tokens is not None else 4096,
max_new_tokens=max_new_tokens if max_new_tokens is not None else 512,
stream=stream if stream is not None else False,
temperature=temperature if temperature is not None else 0.2,
top_p=top_p if top_p is not None else 0.9,
top_k=top_k if top_k is not None else 50,
repetition_penalty=repetition_penalty if repetition_penalty is not None else 1.
)
log.info(f"OpenAI: {self.model_info}")
attention = get_attention_config(self)
log.debug(f'OpenAI: {attention}')
if server:
self.use_server = True
self.app = server
self.server = server
self._is_running = True
else:
self.use_server = False
self.app = FastAPI(title="SD.Next OpenAI-compatible LLM Server", version="1.0")
self.server: Optional[uvicorn.Server] = None
self._is_running = False
self._startup_event = threading.Event()
self.api_key = api_key
self._lock = threading.Lock()
self._security = HTTPBearer(auto_error=False)
self.thread: Optional[threading.Thread] = None
setup_routes(self)
def start(self, timeout_seconds: float = 10.0):
"""Spawns the serving interface safely using an active background thread worker."""
if self._is_running:
return
with self._lock:
if self.server is None:
self._startup_event.clear()
config = uvicorn.Config(app=self.app, host=self.host, port=self.port, log_level="info", loop="asyncio", workers=1)
self.server = uvicorn.Server(config)
self.server.install_signal_handlers = lambda *args, **kwargs: None
original_startup = self.server.startup
async def patched_startup(*args, **kwargs):
await original_startup(*args, **kwargs)
self._startup_event.set()
self.server.startup = patched_startup
self.thread = threading.Thread(target=self.server.run, name="TransformersServeWorkerThread", daemon=True)
self.thread.start()
if not self._startup_event.wait(timeout=timeout_seconds):
self.stop()
raise TimeoutError("OpenAI: init timeout")
self._is_running = True
url = f"http://{self.host}:{self.port}/v1"
log.info(f"OpenAI: Server(url={url})")
def stop(self, timeout_seconds: float = 5.0):
"""Safely winds down network sockets and detached worker threads."""
if self.use_server:
return
with self._lock:
if not self._is_running or not self.server:
return
self.server.should_exit = True
if self.thread and self.thread.is_alive():
self.thread.join(timeout=timeout_seconds)
self.server = None
self.thread = None
self.server = None
self.thread = None
self._is_running = False
log.info("OpenAI: Server(None)")