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:
2026-02-26 22:49:06 -05:00
parent e119331645
commit 5e83c8a56d
2 changed files with 40 additions and 4 deletions
@@ -22,7 +22,7 @@ from fabledassistant.services.generation_buffer import (
GenerationBuffer, GenerationBuffer,
GenerationState, GenerationState,
) )
from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, stream_chat, stream_chat_with_tools, summarize_history_for_context from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, stream_chat, stream_chat_with_tools, summarize_history_for_context, wait_for_model_loaded
from fabledassistant.services.chat import update_conversation_title from fabledassistant.services.chat import update_conversation_title
from fabledassistant.services.intent import IntentResult, classify_intent from fabledassistant.services.intent import IntentResult, classify_intent
from fabledassistant.services.logging import log_generation from fabledassistant.services.logging import log_generation
@@ -210,9 +210,12 @@ async def run_generation(
buf.append_event("status", {"status": "Summarizing conversation history..."}) buf.append_event("status", {"status": "Summarizing conversation history..."})
history_to_use, history_summary = await summarize_history_for_context(history, intent_model) history_to_use, history_summary = await summarize_history_for_context(history, intent_model)
# Phase 3: Build context and start intent classification in parallel. # Phase 3: Build context, classify intent, and wait for model — all in parallel.
# We block on context (need messages to stream) — intent is consumed # build_context is fast DB/search ops that don't need the main model.
# after context is ready, at the start of round 0. # classify_intent uses the small intent model, not the main model.
# wait_for_model_loaded polls /api/ps so the main stream starts without 500 errors.
model_load_task = asyncio.create_task(wait_for_model_loaded(model, timeout=90.0))
context_task = asyncio.create_task(build_context( context_task = asyncio.create_task(build_context(
user_id, history_to_use, context_note_id, user_content, user_id, history_to_use, context_note_id, user_content,
history_summary=history_summary, history_summary=history_summary,
@@ -235,6 +238,14 @@ async def run_generation(
# Emit context event # Emit context event
buf.append_event("context", {"context": context_meta}) buf.append_event("context", {"context": context_meta})
# Wait for main model to be loaded before starting any generation.
# If it's already loaded (common case), this returns immediately.
if not model_load_task.done():
buf.append_event("status", {"status": "Loading model..."})
loaded = await model_load_task
if not loaded:
logger.warning("Model %s did not load within 90s — proceeding anyway", model)
t_start = time.monotonic() t_start = time.monotonic()
timing: dict = { timing: dict = {
"intent_ms": None, "intent_ms": None,
+25
View File
@@ -2,6 +2,7 @@ import asyncio
import json import json
import logging import logging
import re import re
import time
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Literal 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) 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( async def stream_chat(
messages: list[dict], messages: list[dict],
model: str, model: str,