feat: automatic Kokoro offline mode + daily update check
- On first load: model runs online (downloads .pt files), then stores the current HF commit SHA to /data/kokoro_commit_hash.txt and switches the process to offline mode (HF_HUB_OFFLINE) for all future requests - On subsequent restarts: presence of the commit file triggers offline mode before the pipeline loads, skipping all HuggingFace network validation - Daily at 03:00 UTC: scheduler temporarily lifts offline mode, fetches the latest commit SHA from HF, and only reloads the pipeline if the model has actually changed — then restores offline mode - Removed HF_HUB_OFFLINE from docker-compose.yml; behaviour is now automatic and not a hoster/user concern Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -25,8 +25,6 @@ services:
|
||||
VAPID_PRIVATE_KEY: "${VAPID_PRIVATE_KEY:-}"
|
||||
VAPID_PUBLIC_KEY: "${VAPID_PUBLIC_KEY:-}"
|
||||
VAPID_CLAIMS_SUB: "${VAPID_CLAIMS_SUB:-mailto:admin@fabledassistant.local}"
|
||||
# Prevent HuggingFace from making network calls at startup once models are cached:
|
||||
HF_HUB_OFFLINE: "${HF_HUB_OFFLINE:-0}"
|
||||
healthcheck:
|
||||
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:5000/api/health')"]
|
||||
interval: 10s
|
||||
|
||||
@@ -343,6 +343,21 @@ async def start_briefing_scheduler(loop: asyncio.AbstractEventLoop) -> None:
|
||||
replace_existing=True,
|
||||
)
|
||||
|
||||
def _run_kokoro_update_check() -> None:
|
||||
from fabledassistant.services.tts import check_for_kokoro_updates
|
||||
future = asyncio.run_coroutine_threadsafe(check_for_kokoro_updates(), _loop)
|
||||
try:
|
||||
future.result(timeout=300)
|
||||
except Exception as exc:
|
||||
logger.error("Kokoro update check failed: %s", exc)
|
||||
|
||||
_scheduler.add_job(
|
||||
_run_kokoro_update_check,
|
||||
CronTrigger(hour=3, minute=0, timezone="UTC"),
|
||||
id="kokoro_update_check_daily",
|
||||
replace_existing=True,
|
||||
)
|
||||
|
||||
_scheduler.start()
|
||||
logger.info(
|
||||
"Briefing scheduler started with %d user(s) across %d job(s)",
|
||||
|
||||
@@ -2,7 +2,9 @@
|
||||
import asyncio
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -14,6 +16,12 @@ _pipeline: "KPipeline | None" = None
|
||||
_pipeline_lock = asyncio.Lock()
|
||||
_load_error: str | None = None
|
||||
|
||||
# Repo identifier used for HuggingFace update checks
|
||||
KOKORO_REPO = "hexgrad/Kokoro-82M"
|
||||
|
||||
# Persisted across restarts so we can detect model updates and run offline otherwise
|
||||
_COMMIT_HASH_FILE = Path("/data/kokoro_commit_hash.txt")
|
||||
|
||||
# Static list of supported Kokoro voice IDs and display labels
|
||||
_VOICES: list[dict] = [
|
||||
{"id": "af_heart", "label": "Heart (American Female, warm)"},
|
||||
@@ -30,6 +38,31 @@ _VOICES: list[dict] = [
|
||||
]
|
||||
|
||||
|
||||
def _read_stored_commit() -> str | None:
|
||||
try:
|
||||
if _COMMIT_HASH_FILE.exists():
|
||||
return _COMMIT_HASH_FILE.read_text().strip() or None
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _write_stored_commit(sha: str) -> None:
|
||||
try:
|
||||
_COMMIT_HASH_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
_COMMIT_HASH_FILE.write_text(sha)
|
||||
except Exception:
|
||||
logger.warning("Failed to write Kokoro commit hash to %s", _COMMIT_HASH_FILE)
|
||||
|
||||
|
||||
def _set_hf_offline(offline: bool) -> None:
|
||||
"""Toggle HuggingFace offline mode in the current process environment."""
|
||||
if offline:
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
else:
|
||||
os.environ.pop("HF_HUB_OFFLINE", None)
|
||||
|
||||
|
||||
async def load_tts_model() -> None:
|
||||
"""Load the Kokoro pipeline. Called once at startup when voice is enabled."""
|
||||
global _pipeline, _load_error
|
||||
@@ -44,8 +77,16 @@ async def load_tts_model() -> None:
|
||||
try:
|
||||
from kokoro import KPipeline
|
||||
|
||||
# If we have a stored commit hash the model files are already cached
|
||||
# locally — run offline to skip HuggingFace network validation entirely.
|
||||
already_cached = _read_stored_commit() is not None
|
||||
if already_cached:
|
||||
_set_hf_offline(True)
|
||||
logger.info("Kokoro model previously cached — loading in offline mode")
|
||||
|
||||
logger.info("Loading Kokoro TTS pipeline...")
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _load_and_warm():
|
||||
p = KPipeline(lang_code="a") # "a" = American English
|
||||
# Pre-load all voice tensors so synthesis never hits HuggingFace at request time
|
||||
@@ -57,12 +98,92 @@ async def load_tts_model() -> None:
|
||||
return p
|
||||
|
||||
_pipeline = await loop.run_in_executor(None, _load_and_warm)
|
||||
|
||||
# Record the current commit hash on first successful load so future
|
||||
# restarts know the model is cached and can skip HF network checks.
|
||||
if not already_cached:
|
||||
asyncio.create_task(_record_initial_commit())
|
||||
|
||||
logger.info("Kokoro TTS pipeline loaded and voices pre-warmed")
|
||||
except Exception:
|
||||
_load_error = "Failed to load Kokoro TTS pipeline"
|
||||
logger.exception(_load_error)
|
||||
|
||||
|
||||
async def _record_initial_commit() -> None:
|
||||
"""Fetch and store the current Kokoro commit hash after first download."""
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _fetch():
|
||||
from huggingface_hub import model_info
|
||||
return model_info(KOKORO_REPO).sha
|
||||
|
||||
sha = await loop.run_in_executor(None, _fetch)
|
||||
if sha:
|
||||
_write_stored_commit(sha)
|
||||
# Now that files are cached, switch to offline mode for subsequent runs
|
||||
_set_hf_offline(True)
|
||||
logger.info("Kokoro commit hash stored (%s) — future restarts will use offline mode", sha[:8])
|
||||
except Exception:
|
||||
logger.warning("Could not record Kokoro commit hash", exc_info=True)
|
||||
|
||||
|
||||
async def check_for_kokoro_updates() -> None:
|
||||
"""Check HuggingFace for Kokoro model updates.
|
||||
|
||||
Intended to be called on a daily schedule. Temporarily lifts offline mode
|
||||
to query the HF API, compares the latest commit SHA against the locally
|
||||
stored one, and reloads the pipeline only if the model has changed.
|
||||
"""
|
||||
from fabledassistant.services.voice_config import is_voice_enabled
|
||||
if not await is_voice_enabled():
|
||||
return
|
||||
|
||||
stored_sha = _read_stored_commit()
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _fetch_latest_sha():
|
||||
# Temporarily lift offline mode to reach the HF API
|
||||
_set_hf_offline(False)
|
||||
try:
|
||||
from huggingface_hub import model_info
|
||||
return model_info(KOKORO_REPO).sha
|
||||
finally:
|
||||
# Restore offline mode regardless of outcome
|
||||
_set_hf_offline(True)
|
||||
|
||||
latest_sha = await loop.run_in_executor(None, _fetch_latest_sha)
|
||||
except Exception:
|
||||
logger.warning("Kokoro update check failed — will retry tomorrow", exc_info=True)
|
||||
# Ensure offline mode is restored if the executor raised before the finally
|
||||
_set_hf_offline(bool(stored_sha))
|
||||
return
|
||||
|
||||
if latest_sha == stored_sha:
|
||||
logger.debug("Kokoro model is up to date (sha: %s)", (latest_sha or "")[:8])
|
||||
return
|
||||
|
||||
logger.info(
|
||||
"Kokoro model update detected (%s → %s), reloading pipeline",
|
||||
(stored_sha or "none")[:8],
|
||||
(latest_sha or "")[:8],
|
||||
)
|
||||
# Go online for the reload so Kokoro can download the updated files
|
||||
_set_hf_offline(False)
|
||||
try:
|
||||
await reload_tts_model()
|
||||
if latest_sha:
|
||||
_write_stored_commit(latest_sha)
|
||||
logger.info("Kokoro model updated and reloaded successfully")
|
||||
except Exception:
|
||||
logger.exception("Kokoro pipeline reload after update failed")
|
||||
finally:
|
||||
_set_hf_offline(True)
|
||||
|
||||
|
||||
async def reload_tts_model() -> None:
|
||||
"""Unload and reload the TTS pipeline. Safe to call at runtime."""
|
||||
global _pipeline, _load_error
|
||||
|
||||
Reference in New Issue
Block a user