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:
2026-03-31 17:06:13 -04:00
parent f2dd25737a
commit 425d307180
3 changed files with 136 additions and 2 deletions
-2
View File
@@ -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)",
+121
View File
@@ -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