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:
@@ -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,
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user