Add explicit warm-wait before generation starts
Instead of relying solely on retry-on-500, poll /api/ps before starting any LLM stream so the main model has time to fully load into VRAM. - llm.py: add wait_for_model_loaded(model, timeout=90s) — polls /api/ps every 2s, returns True when model appears in loaded list - generation_task.py: launch model_load_task in parallel with build_context and classify_intent (both use fast/small-model ops that don't need the main model); after context is built, await the load task — shows "Loading model..." status only if the user actually has to wait; logs a warning and proceeds if 90s timeout elapses Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -2,6 +2,7 @@ import asyncio
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from collections.abc import AsyncGenerator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Literal
|
||||
@@ -75,6 +76,30 @@ async def ensure_model(model: str) -> None:
|
||||
logger.warning("Failed to pull model '%s' — chat may not work", model, exc_info=True)
|
||||
|
||||
|
||||
async def wait_for_model_loaded(model: str, timeout: float = 90.0) -> bool:
|
||||
"""Poll /api/ps every 2s until the model appears in Ollama's loaded-model list.
|
||||
|
||||
Returns True when the model is loaded, False if timeout elapses first.
|
||||
Used before generation to avoid streaming 500s during cold model loads.
|
||||
"""
|
||||
base = model.removesuffix(":latest")
|
||||
deadline = time.monotonic() + timeout
|
||||
while True:
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||||
resp = await client.get(f"{Config.OLLAMA_URL}/api/ps")
|
||||
resp.raise_for_status()
|
||||
loaded = {m["name"] for m in resp.json().get("models", [])}
|
||||
if model in loaded or f"{base}:latest" in loaded or base in loaded:
|
||||
return True
|
||||
except Exception:
|
||||
pass # Ollama may still be starting up
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
return False
|
||||
await asyncio.sleep(min(2.0, remaining))
|
||||
|
||||
|
||||
async def stream_chat(
|
||||
messages: list[dict],
|
||||
model: str,
|
||||
|
||||
Reference in New Issue
Block a user