601 lines · 22.0 KB · Python
RawDownload
"""vLLM Studio Controller - FastAPI Application.""" from __future__ import annotations import asyncioimport loggingimport osfrom contextlib import asynccontextmanagerfrom pathlib import Pathfrom typing import Optional from fastapi import FastAPI, Requestfrom fastapi.middleware.cors import CORSMiddleware from . import __version__from .config import settingsfrom .events import event_managerfrom .gpu import get_gpu_infofrom .metrics import (    update_active_model, update_gpu_metrics, update_sse_metrics,    get_metrics_content, get_metrics_content_type)from .process import find_inference_processfrom .store import RecipeStore, ChatStore, PeakMetricsStore, LifetimeMetricsStore, MCPStore import httpximport re # Route modulesfrom .routes import (    system_router,    models_router,    lifecycle_router,    chats_router,    logs_router,    monitoring_router,    mcp_router,    proxy_router,    usage_router,) logger = logging.getLogger(__name__) # Global stores (singleton pattern)_store: Optional[RecipeStore] = None_chat_store: Optional[ChatStore] = None_peak_metrics_store: Optional[PeakMetricsStore] = None_lifetime_metrics_store: Optional[LifetimeMetricsStore] = None_mcp_store: Optional[MCPStore] = None # Background task handle for metrics collection_metrics_task: Optional[asyncio.Task] = None # Cache for vLLM metrics (for rate calculations)_last_vllm_metrics: dict = {}_last_metrics_time: float = 0 # Session tracking for average TPS calculations_session_start_tokens: dict = {}  # {'prompt': 0, 'generation': 0}_session_start_time: float = 0_session_peak_prefill: float = 0_session_peak_generation: float = 0  async def _scrape_vllm_metrics(port: int) -> dict:    """Scrape metrics from vLLM Prometheus endpoint."""    try:        async with httpx.AsyncClient(timeout=5.0) as client:            response = await client.get(f"http://localhost:{port}/metrics")            if response.status_code != 200:                return {}             metrics = {}            for line in response.text.split('\n'):                if line.startswith('#') or not line.strip():                    continue                 # Parse Prometheus format: metric_name{labels} value                match = re.match(r'^([a-zA-Z_:][a-zA-Z0-9_:]*)\{?[^}]*\}?\s+([\d.eE+-]+)$', line)                if match:                    name, value = match.groups()                    try:                        metrics[name] = float(value)                    except ValueError:                        pass             return metrics    except Exception:        return {}  def get_store() -> RecipeStore:    """Get the recipe store singleton."""    global _store    if _store is None:        _store = RecipeStore(settings.db_path)    return _store  def get_chat_store() -> ChatStore:    """Get the chat store singleton."""    global _chat_store    if _chat_store is None:        _chat_store = ChatStore(settings.data_dir / "chats.db")    return _chat_store  def get_peak_metrics_store() -> PeakMetricsStore:    """Get the peak metrics store singleton."""    global _peak_metrics_store    if _peak_metrics_store is None:        _peak_metrics_store = PeakMetricsStore(settings.db_path)    return _peak_metrics_store  def get_lifetime_metrics_store() -> LifetimeMetricsStore:    """Get the lifetime metrics store singleton."""    global _lifetime_metrics_store    if _lifetime_metrics_store is None:        _lifetime_metrics_store = LifetimeMetricsStore(settings.db_path)    return _lifetime_metrics_store  def get_mcp_store() -> MCPStore:    """Get the MCP store singleton."""    global _mcp_store    if _mcp_store is None:        _mcp_store = MCPStore(settings.db_path)    return _mcp_store  async def _collect_and_broadcast_metrics():    """Background task to collect metrics and broadcast updates."""    global _last_vllm_metrics, _last_metrics_time    global _session_start_tokens, _session_start_time, _session_peak_prefill, _session_peak_generation    import time     while True:        try:            # Collect current state            current = find_inference_process(settings.inference_port)            gpu_list = get_gpu_info()             # Update Prometheus metrics            if current:                update_active_model(                    model_path=current.model_path,                    backend=current.backend,                    served_name=current.served_model_name                )            else:                update_active_model()             update_gpu_metrics([gpu.model_dump() for gpu in gpu_list])             sse_stats = event_manager.get_stats()            update_sse_metrics(sse_stats)             # Update lifetime metrics (energy tracking)            lifetime_store = get_lifetime_metrics_store()            total_power_watts = sum(gpu.power_draw for gpu in gpu_list)            # Convert watts to watt-hours for 5 second interval: W * (5/3600) = Wh            energy_wh = total_power_watts * (5 / 3600)            lifetime_store.increment('energy_wh', energy_wh)            lifetime_store.increment('uptime_seconds', 5)             # Broadcast status update via SSE (matching /status endpoint format)            await event_manager.publish_status({                "running": current is not None,                "process": current.model_dump() if current else None,                "inference_port": settings.inference_port,            })             # Also broadcast GPU info            await event_manager.publish_gpu([gpu.model_dump() for gpu in gpu_list])             # Collect and broadcast vLLM metrics            if current:                vllm_metrics = await _scrape_vllm_metrics(settings.inference_port)                now = time.time()                elapsed = now - _last_metrics_time if _last_metrics_time > 0 else 5                 # Calculate real-time throughput rates                prompt_throughput = 0.0                generation_throughput = 0.0                curr_prompt = int(vllm_metrics.get('vllm:prompt_tokens_total', 0))                curr_gen = int(vllm_metrics.get('vllm:generation_tokens_total', 0))                 if vllm_metrics and _last_vllm_metrics and elapsed > 0:                    prev_prompt = _last_vllm_metrics.get('vllm:prompt_tokens_total', 0)                    prev_gen = _last_vllm_metrics.get('vllm:generation_tokens_total', 0)                     if curr_prompt > prev_prompt:                        prompt_throughput = (curr_prompt - prev_prompt) / elapsed                    if curr_gen > prev_gen:                        generation_throughput = (curr_gen - prev_gen) / elapsed                 # Session tracking - initialize on first tokens                if curr_gen > 0 and not _session_start_tokens:                    _session_start_tokens = {'prompt': curr_prompt, 'generation': curr_gen}                    _session_start_time = now                    _session_peak_prefill = 0                    _session_peak_generation = 0                 # Calculate session averages                session_avg_prefill = 0.0                session_avg_generation = 0.0                session_elapsed = now - _session_start_time if _session_start_time > 0 else 0                 if _session_start_tokens and session_elapsed > 0:                    prompt_delta = curr_prompt - _session_start_tokens.get('prompt', 0)                    gen_delta = curr_gen - _session_start_tokens.get('generation', 0)                    if prompt_delta > 0:                        session_avg_prefill = prompt_delta / session_elapsed                    if gen_delta > 0:                        session_avg_generation = gen_delta / session_elapsed                 # Track session peaks (best during this session)                if prompt_throughput > _session_peak_prefill:                    _session_peak_prefill = prompt_throughput                if generation_throughput > _session_peak_generation:                    _session_peak_generation = generation_throughput                 _last_vllm_metrics = vllm_metrics                _last_metrics_time = now                 # Get and update peak metrics                peak_store = get_peak_metrics_store()                model_id = current.served_model_name or current.model_path.split('/')[-1]                 # Auto-update peaks if current throughput exceeds stored peaks                if prompt_throughput > 0 or generation_throughput > 0:                    peak_store.update_if_better(                        model_id=model_id,                        prefill_tps=prompt_throughput if prompt_throughput > 0 else None,                        generation_tps=generation_throughput if generation_throughput > 0 else None                    )                 peak_data = peak_store.get(model_id)                 # Get lifetime metrics                lifetime_data = lifetime_store.get_all()                 # Broadcast combined metrics                await event_manager.publish_metrics({                    # Real-time from vLLM                    "running_requests": int(vllm_metrics.get('vllm:num_requests_running', 0)),                    "pending_requests": int(vllm_metrics.get('vllm:num_requests_waiting', 0)),                    "kv_cache_usage": vllm_metrics.get('vllm:kv_cache_usage_perc', 0),                    "prompt_tokens_total": curr_prompt,                    "generation_tokens_total": curr_gen,                     # Real-time throughput (current 5s window)                    "prompt_throughput": round(prompt_throughput, 1),                    "generation_throughput": round(generation_throughput, 1),                     # Session averages (since first token)                    "session_avg_prefill": round(session_avg_prefill, 1),                    "session_avg_generation": round(session_avg_generation, 1),                     # Session peaks (best this session)                    "session_peak_prefill": round(_session_peak_prefill, 1),                    "session_peak_generation": round(_session_peak_generation, 1),                     # All-time peak metrics (persisted)                    "peak_prefill_tps": peak_data.get('prefill_tps') if peak_data else None,                    "peak_generation_tps": peak_data.get('generation_tps') if peak_data else None,                    "peak_ttft_ms": peak_data.get('ttft_ms') if peak_data else None,                     # Lifetime metrics                    "lifetime_prompt_tokens": lifetime_data.get('prompt_tokens_total', 0),                    "lifetime_completion_tokens": lifetime_data.get('completion_tokens_total', 0),                    "lifetime_requests": lifetime_data.get('requests_total', 0),                    "lifetime_energy_kwh": lifetime_data.get('energy_wh', 0) / 1000,                    "lifetime_uptime_hours": lifetime_data.get('uptime_seconds', 0) / 3600,                     # Current power                    "current_power_watts": total_power_watts,                     # Cost efficiency metrics                    "kwh_per_million_input": (lifetime_data.get('energy_wh', 0) / 1000) / (lifetime_data.get('prompt_tokens_total', 1) / 1_000_000) if lifetime_data.get('prompt_tokens_total', 0) > 0 else None,                    "kwh_per_million_output": (lifetime_data.get('energy_wh', 0) / 1000) / (lifetime_data.get('completion_tokens_total', 1) / 1_000_000) if lifetime_data.get('completion_tokens_total', 0) > 0 else None,                })            else:                # Reset session tracking when no model is running                _session_start_tokens = {}                _session_start_time = 0                _session_peak_prefill = 0                _session_peak_generation = 0         except asyncio.CancelledError:            break        except Exception as e:            logger.error(f"Metrics collection error: {e}")         await asyncio.sleep(5)  @asynccontextmanagerasync def lifespan(app: FastAPI):    """Application lifespan manager."""    global _metrics_task     logger.info(f"Starting vLLM Studio Controller v{__version__}")    logger.info(f"Data directory: {settings.data_dir}")    logger.info(f"Database: {settings.db_path}")     # Ensure data directory exists    settings.data_dir.mkdir(parents=True, exist_ok=True)     # Initialize stores    get_store()    get_chat_store()    get_peak_metrics_store()    get_lifetime_metrics_store()    get_mcp_store()     # Start background metrics collection    _metrics_task = asyncio.create_task(_collect_and_broadcast_metrics())     yield     # Cleanup    if _metrics_task:        _metrics_task.cancel()        try:            await _metrics_task        except asyncio.CancelledError:            pass     logger.info("Controller shutdown complete")  # Create FastAPI applicationapp = FastAPI(    title="vLLM Studio Controller",    description="Model lifecycle management for vLLM and SGLang inference servers",    version=__version__,    lifespan=lifespan,) # CORS middlewareapp.add_middleware(    CORSMiddleware,    allow_origins=["*"],    allow_credentials=True,    allow_methods=["*"],    allow_headers=["*"],)  # Access logging middleware@app.middleware("http")async def log_requests(request: Request, call_next):    """Log all incoming requests."""    # Skip noisy endpoints    skip_paths = {"/health", "/metrics", "/events", "/status"}    if request.url.path not in skip_paths:        logger.debug(f"{request.method} {request.url.path}")     response = await call_next(request)    return response  # Include route modulesapp.include_router(system_router)app.include_router(models_router)app.include_router(lifecycle_router)app.include_router(chats_router)app.include_router(logs_router)app.include_router(monitoring_router)app.include_router(mcp_router)app.include_router(proxy_router)app.include_router(usage_router)  # Tokenization endpoints (kept in main app for simplicity)@app.post("/v1/tokenize", tags=["Tokenization"])async def tokenize_text(request: Request):    """Tokenize text using the running model's tokenizer (vLLM-compatible)."""    import httpx     current = find_inference_process(settings.inference_port)    if not current:        return {"error": "No model running", "num_tokens": 0}     try:        body = await request.json()        async with httpx.AsyncClient(timeout=30) as client:            r = await client.post(                f"http://localhost:{settings.inference_port}/tokenize",                json=body            )            if r.status_code == 200:                return r.json()            return {"error": f"Tokenization failed: {r.status_code}", "num_tokens": 0}    except Exception as e:        return {"error": str(e), "num_tokens": 0}  @app.post("/v1/detokenize", tags=["Tokenization"])async def detokenize_tokens(request: Request):    """Detokenize tokens using the running model's tokenizer (vLLM-compatible)."""    import httpx     current = find_inference_process(settings.inference_port)    if not current:        return {"error": "No model running", "text": ""}     try:        body = await request.json()        async with httpx.AsyncClient(timeout=30) as client:            r = await client.post(                f"http://localhost:{settings.inference_port}/detokenize",                json=body            )            if r.status_code == 200:                return r.json()            return {"error": f"Detokenization failed: {r.status_code}", "text": ""}    except Exception as e:        return {"error": str(e), "text": ""}  @app.post("/v1/count-tokens", tags=["Tokenization"])async def count_text_tokens(request: Request):    """Count tokens in text using the running model's tokenizer."""    import httpx     current = find_inference_process(settings.inference_port)    if not current:        return {"error": "No model running", "num_tokens": 0}     try:        body = await request.json()        text = body.get("text", "")        model = body.get("model", current.served_model_name or "default")         async with httpx.AsyncClient(timeout=30) as client:            r = await client.post(                f"http://localhost:{settings.inference_port}/tokenize",                json={"model": model, "prompt": text}            )            if r.status_code == 200:                data = r.json()                tokens = data.get("tokens", [])                return {"num_tokens": len(tokens), "model": model}            return {"error": f"Token count failed: {r.status_code}", "num_tokens": 0}    except Exception as e:        return {"error": str(e), "num_tokens": 0}  @app.post("/v1/tokenize-chat-completions", tags=["Tokenization"])async def tokenize_chat_completions(request: Request):    """Estimate token count for a chat completion request.     Applies the model's chat template and counts tokens for messages and tools.    Returns total input tokens and per-component breakdown.    """    import httpx     current = find_inference_process(settings.inference_port)    if not current:        return {"error": "No model running", "input_tokens": 0}     try:        body = await request.json()        messages = body.get("messages", [])        tools = body.get("tools", [])        model = body.get("model", current.served_model_name or "default")         # Count message tokens by applying chat template        messages_tokens = 0        async with httpx.AsyncClient(timeout=30) as client:            # Try to get templated prompt            try:                # Most vLLM setups support /v1/chat/completions with max_tokens=1 and stream=false                # to get the prompt tokens in the response                test_request = {                    "model": model,                    "messages": messages,                    "max_tokens": 1,                    "stream": False,                }                if tools:                    test_request["tools"] = tools                 r = await client.post(                    f"http://localhost:{settings.inference_port}/v1/chat/completions",                    json=test_request                )                if r.status_code == 200:                    data = r.json()                    usage = data.get("usage", {})                    prompt_tokens = usage.get("prompt_tokens", 0)                    return {                        "input_tokens": prompt_tokens,                        "breakdown": {                            "messages": prompt_tokens,                            "tools": 0,  # Included in messages for templated approach                        },

Lines 1–500 of 601.