mirror of https://github.com/razor-ai/soup.git
504 lines
15 KiB
Python
504 lines
15 KiB
Python
"""soup serve — local inference server with OpenAI-compatible API."""
|
|
|
|
import json
|
|
import logging
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import typer
|
|
from rich.console import Console
|
|
from rich.panel import Panel
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
console = Console()
|
|
|
|
|
|
def serve(
|
|
model: str = typer.Option(
|
|
...,
|
|
"--model",
|
|
"-m",
|
|
help="Path to LoRA adapter directory or full model",
|
|
),
|
|
base_model: Optional[str] = typer.Option(
|
|
None,
|
|
"--base",
|
|
"-b",
|
|
help="Base model ID. Auto-detected from adapter_config.json if not set.",
|
|
),
|
|
port: int = typer.Option(
|
|
8000,
|
|
"--port",
|
|
"-p",
|
|
help="Port to serve on",
|
|
),
|
|
host: str = typer.Option(
|
|
"0.0.0.0",
|
|
"--host",
|
|
help="Host to bind to",
|
|
),
|
|
device: Optional[str] = typer.Option(
|
|
None,
|
|
"--device",
|
|
help="Device: cuda, mps, cpu. Auto-detected if not set.",
|
|
),
|
|
max_tokens_default: int = typer.Option(
|
|
512,
|
|
"--max-tokens",
|
|
help="Default max tokens for generation",
|
|
),
|
|
backend: str = typer.Option(
|
|
"transformers",
|
|
"--backend",
|
|
help="Inference backend: transformers (default) or vllm",
|
|
),
|
|
tensor_parallel: int = typer.Option(
|
|
1,
|
|
"--tensor-parallel",
|
|
"--tp",
|
|
help="Number of GPUs for tensor parallelism (vLLM only)",
|
|
),
|
|
gpu_memory_utilization: float = typer.Option(
|
|
0.9,
|
|
"--gpu-memory",
|
|
help="Fraction of GPU memory to use (vLLM only, 0.0-1.0)",
|
|
),
|
|
):
|
|
"""Start a local inference server with OpenAI-compatible API."""
|
|
# Lazy imports for fast CLI startup
|
|
try:
|
|
import uvicorn # noqa: F401
|
|
from fastapi import FastAPI # noqa: F401
|
|
from fastapi.responses import StreamingResponse # noqa: F401
|
|
except ImportError:
|
|
console.print(
|
|
"[red]FastAPI/uvicorn not installed.[/]\n"
|
|
"Install with: [bold]pip install 'soup-cli[serve]'[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
# Validate backend
|
|
backend = backend.lower()
|
|
if backend not in ("transformers", "vllm"):
|
|
console.print(
|
|
f"[red]Unknown backend: {backend}[/]\n"
|
|
"Supported backends: [bold]transformers[/], [bold]vllm[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
# Auto-detect vLLM: if installed but not selected, show hint
|
|
if backend == "transformers":
|
|
from soup_cli.utils.vllm import is_vllm_available
|
|
|
|
if is_vllm_available():
|
|
console.print(
|
|
"[dim]Hint: vLLM is installed. Use [bold]--backend vllm[/] "
|
|
"for 2-4x better throughput.[/]"
|
|
)
|
|
|
|
# Validate vLLM availability
|
|
if backend == "vllm":
|
|
from soup_cli.utils.vllm import is_vllm_available
|
|
|
|
if not is_vllm_available():
|
|
console.print(
|
|
"[red]vLLM not installed.[/]\n"
|
|
"Install with: [bold]pip install 'soup-cli[serve-fast]'[/]"
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
model_path = Path(model)
|
|
if not model_path.exists():
|
|
console.print(f"[red]Model path not found: {model_path}[/]")
|
|
raise typer.Exit(1)
|
|
|
|
# Detect adapter
|
|
adapter_config_path = model_path / "adapter_config.json"
|
|
is_adapter = adapter_config_path.exists()
|
|
|
|
# Resolve base model
|
|
if is_adapter and not base_model:
|
|
base_model = _detect_base_model(adapter_config_path)
|
|
if not base_model:
|
|
console.print(
|
|
"[red]Cannot detect base model from adapter_config.json.[/]\n"
|
|
"Please specify with [bold]--base[/] flag."
|
|
)
|
|
raise typer.Exit(1)
|
|
|
|
# Detect device (only for transformers backend)
|
|
if not device and backend == "transformers":
|
|
from soup_cli.utils.gpu import detect_device
|
|
|
|
device, _ = detect_device()
|
|
elif not device:
|
|
device = "cuda"
|
|
|
|
backend_label = "vLLM" if backend == "vllm" else "transformers"
|
|
console.print(
|
|
Panel(
|
|
f"Model: [bold]{model_path}[/]\n"
|
|
+ (f"Base: [bold]{base_model}[/]\n" if is_adapter else "")
|
|
+ f"Device: [bold]{device}[/]\n"
|
|
f"Type: [bold]{'LoRA adapter' if is_adapter else 'Full model'}[/]\n"
|
|
f"Backend: [bold]{backend_label}[/]"
|
|
+ (f"\nTP: [bold]{tensor_parallel}[/]" if backend == "vllm" else ""),
|
|
title="Loading model",
|
|
)
|
|
)
|
|
|
|
if backend == "vllm":
|
|
app = _serve_vllm(
|
|
model_path=model_path,
|
|
base_model=base_model,
|
|
is_adapter=is_adapter,
|
|
max_tokens_default=max_tokens_default,
|
|
tensor_parallel=tensor_parallel,
|
|
gpu_memory_utilization=gpu_memory_utilization,
|
|
)
|
|
else:
|
|
# Transformers backend (original)
|
|
model_obj, tokenizer = _load_model(
|
|
model_path=str(model_path),
|
|
base_model=base_model,
|
|
is_adapter=is_adapter,
|
|
device=device,
|
|
)
|
|
console.print("[bold green]Model loaded![/]")
|
|
|
|
app = _create_app(
|
|
model_obj=model_obj,
|
|
tokenizer=tokenizer,
|
|
device=device,
|
|
model_name=str(model_path.name),
|
|
max_tokens_default=max_tokens_default,
|
|
)
|
|
|
|
console.print(
|
|
Panel(
|
|
f"URL: [bold]http://{host}:{port}[/]\n"
|
|
f"Backend: [bold]{backend_label}[/]\n"
|
|
f"Endpoints: [bold]/v1/chat/completions[/], [bold]/v1/models[/], [bold]/health[/]\n\n"
|
|
f"Example:\n"
|
|
f" curl http://localhost:{port}/v1/chat/completions \\\n"
|
|
f' -H "Content-Type: application/json" \\\n'
|
|
f" -d '{{"
|
|
f'"model": "{model_path.name}", '
|
|
f'"messages": [{{"role": "user", "content": "Hello!"}}]'
|
|
f"}}'\n\n"
|
|
f"Press [bold]Ctrl+C[/] to stop.",
|
|
title="[bold green]Server Ready[/]",
|
|
)
|
|
)
|
|
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host=host, port=port, log_level="warning")
|
|
|
|
|
|
def _serve_vllm(
|
|
model_path: Path,
|
|
base_model: Optional[str],
|
|
is_adapter: bool,
|
|
max_tokens_default: int,
|
|
tensor_parallel: int,
|
|
gpu_memory_utilization: float,
|
|
):
|
|
"""Set up vLLM engine and create FastAPI app."""
|
|
from soup_cli.utils.vllm import create_vllm_app, create_vllm_engine
|
|
|
|
console.print("[dim]Initializing vLLM engine...[/]")
|
|
engine, engine_model_name = create_vllm_engine(
|
|
model_path=str(model_path),
|
|
base_model=base_model,
|
|
is_adapter=is_adapter,
|
|
tensor_parallel_size=tensor_parallel,
|
|
gpu_memory_utilization=gpu_memory_utilization,
|
|
)
|
|
console.print("[bold green]vLLM engine ready![/]")
|
|
|
|
adapter_path = str(model_path) if is_adapter else None
|
|
|
|
app = create_vllm_app(
|
|
engine=engine,
|
|
engine_model_name=engine_model_name,
|
|
model_name=str(model_path.name),
|
|
adapter_path=adapter_path,
|
|
max_tokens_default=max_tokens_default,
|
|
)
|
|
|
|
return app
|
|
|
|
|
|
def _detect_base_model(adapter_config_path: Path) -> Optional[str]:
|
|
"""Read base_model_name_or_path from adapter_config.json."""
|
|
try:
|
|
with open(adapter_config_path, encoding="utf-8") as f:
|
|
config = json.load(f)
|
|
return config.get("base_model_name_or_path")
|
|
except (json.JSONDecodeError, OSError):
|
|
return None
|
|
|
|
|
|
def _load_model(
|
|
model_path: str,
|
|
base_model: Optional[str],
|
|
is_adapter: bool,
|
|
device: str,
|
|
):
|
|
"""Load model and tokenizer."""
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
console.print("[dim]Loading tokenizer...[/]")
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
|
if tokenizer.pad_token is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
|
|
if is_adapter:
|
|
from peft import PeftModel
|
|
|
|
console.print(f"[dim]Loading base model: {base_model}...[/]")
|
|
base = AutoModelForCausalLM.from_pretrained(
|
|
base_model,
|
|
trust_remote_code=True,
|
|
device_map="auto",
|
|
dtype=torch.float16,
|
|
)
|
|
console.print(f"[dim]Loading LoRA adapter: {model_path}...[/]")
|
|
model_obj = PeftModel.from_pretrained(base, model_path)
|
|
else:
|
|
console.print(f"[dim]Loading model: {model_path}...[/]")
|
|
model_obj = AutoModelForCausalLM.from_pretrained(
|
|
model_path,
|
|
trust_remote_code=True,
|
|
device_map="auto",
|
|
dtype=torch.float16,
|
|
)
|
|
|
|
model_obj.eval()
|
|
return model_obj, tokenizer
|
|
|
|
|
|
def _generate_response(
|
|
model,
|
|
tokenizer,
|
|
messages: list[dict],
|
|
max_tokens: int = 512,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
stream: bool = False,
|
|
):
|
|
"""Generate a response from the model."""
|
|
import torch
|
|
|
|
# Apply chat template
|
|
if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template:
|
|
text = tokenizer.apply_chat_template(
|
|
messages, tokenize=False, add_generation_prompt=True
|
|
)
|
|
else:
|
|
parts = []
|
|
for msg in messages:
|
|
role = msg["role"]
|
|
content = msg["content"]
|
|
if role == "system":
|
|
parts.append(f"System: {content}")
|
|
elif role == "user":
|
|
parts.append(f"User: {content}")
|
|
elif role == "assistant":
|
|
parts.append(f"Assistant: {content}")
|
|
parts.append("Assistant:")
|
|
text = "\n".join(parts)
|
|
|
|
inputs = tokenizer(text, return_tensors="pt")
|
|
input_ids = inputs["input_ids"].to(model.device)
|
|
attention_mask = inputs["attention_mask"].to(model.device)
|
|
|
|
with torch.no_grad():
|
|
gen_kwargs = {
|
|
"input_ids": input_ids,
|
|
"attention_mask": attention_mask,
|
|
"max_new_tokens": max_tokens,
|
|
"do_sample": temperature > 0,
|
|
"pad_token_id": tokenizer.pad_token_id,
|
|
}
|
|
if temperature > 0:
|
|
gen_kwargs["temperature"] = temperature
|
|
gen_kwargs["top_p"] = top_p
|
|
|
|
outputs = model.generate(**gen_kwargs)
|
|
|
|
new_tokens = outputs[0][input_ids.shape[1]:]
|
|
response = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
|
|
|
|
prompt_tokens = input_ids.shape[1]
|
|
completion_tokens = len(new_tokens)
|
|
|
|
return response, prompt_tokens, completion_tokens
|
|
|
|
|
|
def _create_app(
|
|
model_obj,
|
|
tokenizer,
|
|
device: str,
|
|
model_name: str,
|
|
max_tokens_default: int,
|
|
):
|
|
"""Create the FastAPI application with OpenAI-compatible endpoints."""
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel as PydanticBaseModel
|
|
from pydantic import Field
|
|
|
|
app = FastAPI(title="Soup Inference Server", version="1.0.0")
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# --- Request/Response models ---
|
|
|
|
class ChatMessage(PydanticBaseModel):
|
|
role: str
|
|
content: str
|
|
|
|
class ChatCompletionRequest(PydanticBaseModel):
|
|
model: str = model_name
|
|
messages: list[ChatMessage]
|
|
temperature: float = Field(default=0.7, ge=0.0, le=2.0)
|
|
top_p: float = Field(default=0.9, ge=0.0, le=1.0)
|
|
max_tokens: Optional[int] = Field(default=None, ge=1, le=16384)
|
|
stream: bool = False
|
|
|
|
# --- Endpoints ---
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "model": model_name, "device": device}
|
|
|
|
@app.get("/v1/models")
|
|
def list_models():
|
|
return {
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"id": model_name,
|
|
"object": "model",
|
|
"owned_by": "soup",
|
|
}
|
|
],
|
|
}
|
|
|
|
@app.post("/v1/chat/completions")
|
|
def chat_completions(request: ChatCompletionRequest):
|
|
messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
|
|
max_tokens = request.max_tokens or max_tokens_default
|
|
|
|
if request.stream:
|
|
return StreamingResponse(
|
|
_stream_response(
|
|
model_obj, tokenizer, messages,
|
|
max_tokens=max_tokens,
|
|
temperature=request.temperature,
|
|
top_p=request.top_p,
|
|
model_name=model_name,
|
|
),
|
|
media_type="text/event-stream",
|
|
)
|
|
|
|
try:
|
|
response_text, prompt_tokens, completion_tokens = _generate_response(
|
|
model_obj, tokenizer, messages,
|
|
max_tokens=max_tokens,
|
|
temperature=request.temperature,
|
|
top_p=request.top_p,
|
|
)
|
|
except Exception:
|
|
logger.exception("Generation error")
|
|
raise HTTPException(status_code=500, detail="Internal server error")
|
|
|
|
return {
|
|
"id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
|
|
"object": "chat.completion",
|
|
"created": int(time.time()),
|
|
"model": model_name,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": response_text,
|
|
},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": prompt_tokens + completion_tokens,
|
|
},
|
|
}
|
|
|
|
return app
|
|
|
|
|
|
def _stream_response(
|
|
model, tokenizer, messages,
|
|
max_tokens, temperature, top_p, model_name,
|
|
):
|
|
"""Generator that yields SSE chunks for streaming responses."""
|
|
chat_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
|
|
created = int(time.time())
|
|
|
|
# Generate full response (true token-by-token streaming requires TextIteratorStreamer)
|
|
response_text, _, _ = _generate_response(
|
|
model, tokenizer, messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
)
|
|
|
|
# Simulate streaming by sending word-by-word
|
|
words = response_text.split(" ")
|
|
for idx, word in enumerate(words):
|
|
chunk_text = word if idx == 0 else f" {word}"
|
|
chunk = {
|
|
"id": chat_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": model_name,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {"content": chunk_text},
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
yield f"data: {json.dumps(chunk)}\n\n"
|
|
|
|
# Final chunk
|
|
final_chunk = {
|
|
"id": chat_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created,
|
|
"model": model_name,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
}
|
|
yield f"data: {json.dumps(final_chunk)}\n\n"
|
|
yield "data: [DONE]\n\n"
|