diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c84ac640df..d4a33387dc 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1063,6 +1063,7 @@ class LlamaCppBackend: speculative_type: Optional[str] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused + n_parallel: int = 1, ) -> bool: """ Start llama-server with a GGUF model. @@ -1283,7 +1284,7 @@ class LlamaCppBackend: "-c", str(effective_ctx) if effective_ctx > 0 else "0", "--parallel", - "1", # Single-user studio, saves VRAM + str(n_parallel), "--flash-attn", "on", # Force flash attention for speed ] diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 30ff7da49c..97f5358121 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -11,9 +11,10 @@ import time import uuid from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status -from fastapi.responses import StreamingResponse, JSONResponse +from fastapi.responses import StreamingResponse, JSONResponse, Response from typing import Optional import json +import httpx import structlog from loggers import get_logger import asyncio @@ -76,6 +77,7 @@ from models.inference import ( ChatCompletionRequest, ChatCompletionChunk, ChatCompletion, + ChatMessage, ChunkChoice, ChoiceDelta, CompletionChoice, @@ -121,6 +123,7 @@ def get_llama_cpp_backend() -> LlamaCppBackend: @router.post("/load", response_model = LoadResponse) async def load_model( request: LoadRequest, + fastapi_request: Request, current_subject: str = Depends(get_current_subject), ): """ @@ -252,6 +255,8 @@ async def load_model( # Run in a thread so the event loop stays free for progress # polling and other requests during the (potentially long) # GGUF download + llama-server startup. + _n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1) + if config.gguf_hf_repo: # HF mode: download via huggingface_hub then start llama-server success = await asyncio.to_thread( @@ -265,6 +270,7 @@ async def load_model( chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, + n_parallel = _n_parallel, ) else: # Local mode: llama-server loads via -m @@ -278,6 +284,7 @@ async def load_model( chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, speculative_type = request.speculative_type, + n_parallel = _n_parallel, ) if not success: @@ -1816,3 +1823,157 @@ async def openai_list_models( ) return {"object": "list", "data": models} + + +# ===================================================================== +# OpenAI-Compatible Completions Proxy (/completions → /v1/completions) +# ===================================================================== + + +@router.post("/completions") +async def openai_completions( + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + OpenAI-compatible text completions endpoint (non-chat). + + Transparently proxies to the running llama-server's ``/v1/completions``. + Only available when a GGUF model is loaded. + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + body = await request.json() + target_url = f"{llama_backend.base_url}/v1/completions" + is_stream = body.get("stream", False) + + if is_stream: + async def _stream(): + async with httpx.AsyncClient() as client: + async with client.stream("POST", target_url, json = body, timeout = 600) as resp: + async for chunk in resp.aiter_bytes(): + yield chunk + + return StreamingResponse(_stream(), media_type = "text/event-stream") + else: + async with httpx.AsyncClient() as client: + resp = await client.post(target_url, json = body, timeout = 600) + return Response( + content = resp.content, + status_code = resp.status_code, + media_type = "application/json", + ) + + +# ===================================================================== +# OpenAI-Compatible Embeddings Proxy (/embeddings → /v1/embeddings) +# ===================================================================== + + +@router.post("/embeddings") +async def openai_embeddings( + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + OpenAI-compatible embeddings endpoint. + + Transparently proxies to the running llama-server's ``/v1/embeddings``. + Only available when a GGUF model is loaded. + Note: the loaded model must support pooling; otherwise llama-server + will return an error (expected). + """ + llama_backend = get_llama_cpp_backend() + if not llama_backend.is_loaded: + raise HTTPException( + status_code = 503, + detail = "No GGUF model loaded. Load a GGUF model first.", + ) + + body = await request.json() + target_url = f"{llama_backend.base_url}/v1/embeddings" + + async with httpx.AsyncClient() as client: + resp = await client.post(target_url, json = body, timeout = 600) + return Response( + content = resp.content, + status_code = resp.status_code, + media_type = "application/json", + ) + + +# ===================================================================== +# OpenAI Responses API (/responses → /v1/responses) +# ===================================================================== + + +from pydantic import BaseModel +from typing import Union + + +class _ResponsesInputMessage(BaseModel): + role: str + content: str + + +class ResponsesRequest(BaseModel): + """Minimal OpenAI Responses API request.""" + model: str = "default" + input: Union[str, list[_ResponsesInputMessage]] = [] + instructions: Optional[str] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + max_output_tokens: Optional[int] = None + stream: bool = False + + +@router.post("/responses") +async def openai_responses( + payload: ResponsesRequest, + request: Request, + current_subject: str = Depends(get_current_subject), +): + """ + OpenAI Responses API endpoint. + + Converts the Responses-format request into a ChatCompletionRequest + and delegates to the existing chat completions handler. Works with + both GGUF and non-GGUF backends. + """ + # Build messages list from the Responses API input format + messages: list[ChatMessage] = [] + + # System message from instructions + if payload.instructions: + messages.append(ChatMessage(role = "system", content = payload.instructions)) + + # Convert input to messages + if isinstance(payload.input, str): + messages.append(ChatMessage(role = "user", content = payload.input)) + else: + for msg in payload.input: + messages.append(ChatMessage(role = msg.role, content = msg.content)) + + if not messages: + raise HTTPException(status_code = 400, detail = "No input provided.") + + # Build a ChatCompletionRequest and delegate + chat_kwargs = dict( + model = payload.model, + messages = messages, + stream = payload.stream, + ) + if payload.temperature is not None: + chat_kwargs["temperature"] = payload.temperature + if payload.top_p is not None: + chat_kwargs["top_p"] = payload.top_p + if payload.max_output_tokens is not None: + chat_kwargs["max_tokens"] = payload.max_output_tokens + + chat_request = ChatCompletionRequest(**chat_kwargs) + return await openai_chat_completions(chat_request, request) diff --git a/studio/backend/run.py b/studio/backend/run.py index 86c1194661..9675b9ea4c 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -248,6 +248,7 @@ def run_server( port: int = 8888, frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist", silent: bool = False, + llama_parallel_slots: int = 1, ): """ Start the FastAPI server. @@ -257,6 +258,7 @@ def run_server( port: Port to bind to (auto-increments if in use) frontend_path: Path to frontend build directory (optional) silent: Suppress startup messages + llama_parallel_slots: Number of parallel slots for llama-server Note: Signal handlers are NOT registered here so that embedders @@ -331,6 +333,7 @@ def run_server( # binds (port==0) leave it unset and let request handlers fall back # to the ASGI request scope or request.base_url. app.state.server_port = port if port and port > 0 else None + app.state.llama_parallel_slots = llama_parallel_slots # Run server in a daemon thread def _run(): diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 659e5b40d3..3c01a21941 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -319,7 +319,7 @@ def run( # ── 2. Start server (always suppress built-in banner) ───────────── from studio.backend.run import run_server, _resolve_external_ip - run_kwargs = dict(host = host, port = port, silent = True) + run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = 4) if frontend is not None: run_kwargs["frontend_path"] = frontend app = run_server(**run_kwargs)