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_PRIVATE_KEY: "${VAPID_PRIVATE_KEY:-}"
|
||||||
VAPID_PUBLIC_KEY: "${VAPID_PUBLIC_KEY:-}"
|
VAPID_PUBLIC_KEY: "${VAPID_PUBLIC_KEY:-}"
|
||||||
VAPID_CLAIMS_SUB: "${VAPID_CLAIMS_SUB:-mailto:admin@fabledassistant.local}"
|
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:
|
healthcheck:
|
||||||
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:5000/api/health')"]
|
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:5000/api/health')"]
|
||||||
interval: 10s
|
interval: 10s
|
||||||
|
|||||||
@@ -343,6 +343,21 @@ async def start_briefing_scheduler(loop: asyncio.AbstractEventLoop) -> None:
|
|||||||
replace_existing=True,
|
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()
|
_scheduler.start()
|
||||||
logger.info(
|
logger.info(
|
||||||
"Briefing scheduler started with %d user(s) across %d job(s)",
|
"Briefing scheduler started with %d user(s) across %d job(s)",
|
||||||
|
|||||||
@@ -2,7 +2,9 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import io
|
import io
|
||||||
import logging
|
import logging
|
||||||
|
import os
|
||||||
import time
|
import time
|
||||||
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -14,6 +16,12 @@ _pipeline: "KPipeline | None" = None
|
|||||||
_pipeline_lock = asyncio.Lock()
|
_pipeline_lock = asyncio.Lock()
|
||||||
_load_error: str | None = None
|
_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
|
# Static list of supported Kokoro voice IDs and display labels
|
||||||
_VOICES: list[dict] = [
|
_VOICES: list[dict] = [
|
||||||
{"id": "af_heart", "label": "Heart (American Female, warm)"},
|
{"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:
|
async def load_tts_model() -> None:
|
||||||
"""Load the Kokoro pipeline. Called once at startup when voice is enabled."""
|
"""Load the Kokoro pipeline. Called once at startup when voice is enabled."""
|
||||||
global _pipeline, _load_error
|
global _pipeline, _load_error
|
||||||
@@ -44,8 +77,16 @@ async def load_tts_model() -> None:
|
|||||||
try:
|
try:
|
||||||
from kokoro import KPipeline
|
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...")
|
logger.info("Loading Kokoro TTS pipeline...")
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
|
|
||||||
def _load_and_warm():
|
def _load_and_warm():
|
||||||
p = KPipeline(lang_code="a") # "a" = American English
|
p = KPipeline(lang_code="a") # "a" = American English
|
||||||
# Pre-load all voice tensors so synthesis never hits HuggingFace at request time
|
# 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
|
return p
|
||||||
|
|
||||||
_pipeline = await loop.run_in_executor(None, _load_and_warm)
|
_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")
|
logger.info("Kokoro TTS pipeline loaded and voices pre-warmed")
|
||||||
except Exception:
|
except Exception:
|
||||||
_load_error = "Failed to load Kokoro TTS pipeline"
|
_load_error = "Failed to load Kokoro TTS pipeline"
|
||||||
logger.exception(_load_error)
|
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:
|
async def reload_tts_model() -> None:
|
||||||
"""Unload and reload the TTS pipeline. Safe to call at runtime."""
|
"""Unload and reload the TTS pipeline. Safe to call at runtime."""
|
||||||
global _pipeline, _load_error
|
global _pipeline, _load_error
|
||||||
|
|||||||
Reference in New Issue
Block a user