"""HF Dataset repo 持久化同步. 为什么需要: HF Spaces 免费版磁盘是临时的 (容器重启后 /data 内的非持久卷数据会丢). 唯一免费的持久化方案是把 /data 同步到 HF Dataset repo (Git LFS). 调用模式 (A 改良版): - lifespan startup: restore_from_hf() (拉 snapshot) - 重要写操作后: await push_to_hf() (同步阻塞 + verify, upload/delete 用) - 低频写操作: schedule_push() (5s debounce, chat 等用) - lifespan shutdown: await push_to_hf() (兜底) 健壮性 (A 改良版): - ✅ push 后 verify: list_repo_files 确认数据真的上去了 - ✅ push 失败/超时状态写入 _state, /readyz 暴露给前端 - ✅ debounce 30s → 5s (缩小"重启即丢"窗口) - ✅ 重要操作 (upload/delete) 改同步 push, 不再吃 debounce 黑洞 - ✅ 独立 ThreadPoolExecutor, 不跟 Chroma / BGE-M3 业务抢线程 - ✅ _background_tasks 持有 task 引用, 防 asyncio GC - ✅ push 失败也重置 pending_push, 不卡死后续写 - 首次部署 (repo 不存在) → RepositoryNotFoundError → fresh_start - HF_TOKEN 缺失 → disabled 降级 """ from __future__ import annotations import asyncio import concurrent.futures import logging import shutil import time from pathlib import Path from typing import Literal from huggingface_hub import ( create_repo, list_repo_files, snapshot_download, upload_folder, upload_file, ) from huggingface_hub.errors import RepositoryNotFoundError from app.config import settings from app.core.paths import data_dir, sqlite_dir, chroma_dir, upload_dir logger = logging.getLogger(__name__) # ✅ 独立 ThreadPoolExecutor, 不跟业务 (Chroma / BGE-M3 / run_in_executor) 抢线程 # 2 个 worker 够用: push 是单飞, 第二个留给 restore (启动期偶发重入) _persist_executor = concurrent.futures.ThreadPoolExecutor( max_workers=2, thread_name_prefix="persist", ) # ✅ 持有后台 task 引用, 避免 asyncio GC 回收未完成的 debounce 任务 # (曾因为 _state["pending_push"] = True 但 _delayed_push 被 GC, 导致永远卡死) _background_tasks: set[asyncio.Task] = set() # debounce 间隔 (秒): 5s 足够让同一秒内的多次写合并, 又不会留太久黑洞窗口 DEBOUNCE_SECONDS = 5 # 状态机: 持久化是否启用 / 启动模式 / 最后一次 push 结果 _state: dict[str, object] = { "mode": "disabled", # disabled | cold_restore | fresh_start "last_push_at": 0.0, # unix timestamp "pending_push": False, # debounce 窗口期 "last_push_status": "idle", # idle | uploading | verifying | ok | failed "last_error": None, # str | None "last_verify": None, # dict | None (verify 步骤结果) } def persist_mode() -> Literal["disabled", "cold_restore", "fresh_start"]: return _state["mode"] # type: ignore[return-value] async def restore_from_hf() -> None: """从 HF Dataset repo 拉取数据到本地. 调用时机: FastAPI lifespan 启动. """ if not settings.is_persist_enabled(): logger.info("Persistence disabled (HF_PERSIST_REPO or HF_TOKEN not set)") _state["mode"] = "disabled" return repo_id = settings.hf_persist_repo token = settings.hf_token.get_secret_value() local_root = data_dir() # ⚠️ 必须含所有可能存在的子目录, 否则 restore 漏掉 → 数据看似在 Dataset, # 但 Space 重启后 /data 是空的. # 历史: 早期只用 chroma/, 后改 ENALBE_COLBERT=true 用 colbert/, # 但 restore 的 allow_patterns 没改, → colbert/ 永远没被 restore 回来. target_subdirs = ["chroma", "colbert", "sqlite", "uploads"] loop = asyncio.get_running_loop() try: await loop.run_in_executor( _persist_executor, lambda: snapshot_download( repo_id=repo_id, repo_type="dataset", local_dir=str(local_root), token=token, allow_patterns=[f"{d}/*" for d in target_subdirs] + target_subdirs, ), ) _state["mode"] = "cold_restore" logger.info("Persisted data restored from %s", repo_id) except RepositoryNotFoundError: _state["mode"] = "fresh_start" logger.info( "Persist repo %s not found (first deploy?). " "Will create on first push.", repo_id, ) except Exception as e: # noqa: BLE001 logger.error("Persist restore failed (will start fresh): %s", e, exc_info=True) _state["mode"] = "fresh_start" _state["last_error"] = f"restore: {type(e).__name__}: {e}"[:500] async def push_to_hf() -> bool: """同步推送本地数据到 HF Dataset repo. 阻塞, 含 verify. 返回 bool: True=upload + verify 都成功; False=失败 (失败原因在 _state['last_error']). A 改良版要点: 1. upload_folder 完不立即返回, list_repo_files 确认数据真的上去了 2. verify 结果写入 _state['last_verify'], /readyz 暴露 3. push 状态写入 _state['last_push_status'], 前端能实时看 4. ⚠️ chroma/ 大文件 (ChromaDB HNSW index) upload_folder (走 git 协议) 经常被静默丢弃 (单文件超过 HF 限制). push 完后如果本地有 chroma 数据但 dataset 上 chroma/ 缺失, 用 hf_hub 低层 API 一个一个 upload_file 补上. """ if not settings.is_persist_enabled(): _state["last_push_status"] = "idle" return False if _state["mode"] == "fresh_start": # 首次需要先 create_repo await _ensure_repo_exists() repo_id = settings.hf_persist_repo token = settings.hf_token.get_secret_value() local_root = data_dir() loop = asyncio.get_running_loop() _state["last_push_status"] = "uploading" _state["pending_push"] = False started = time.time() try: await loop.run_in_executor( _persist_executor, lambda: upload_folder( folder_path=str(local_root), repo_id=repo_id, repo_type="dataset", token=token, commit_message=f"sync {started:.0f}", ignore_patterns=[".cache/*", "*.tmp", "*.lock"], ), ) upload_ms = int((time.time() - started) * 1000) logger.info("upload_folder done in %dms, verifying...", upload_ms) # ============ A 改良版核心: verify ============ _state["last_push_status"] = "verifying" files = await loop.run_in_executor( _persist_executor, lambda: list_repo_files( repo_id=repo_id, repo_type="dataset", token=token, ), ) has_chroma = any(f.startswith("chroma/") for f in files) has_colbert = any(f.startswith("colbert/") for f in files) # ENABLE_COLBERT=true 时用这个 has_sqlite = any(f.startswith("sqlite/") for f in files) has_uploads = any(f.startswith("uploads/") for f in files) verify_result = { "total_files": len(files), "has_chroma": has_chroma, "has_colbert": has_colbert, "has_sqlite": has_sqlite, "has_uploads": has_uploads, "checked_at": time.time(), "upload_ms": upload_ms, } _state["last_verify"] = verify_result # 即使 upload 报成功但 repo 啥都没, 也标记 ok (cold start 后还没数据是正常的) # 但要把 verify 结果暴露出去 # ============ ChromaDB HNSW 大文件补传 ============ # upload_folder 走 git 协议, ChromaDB 的 chroma.sqlite3 + data_level0.bin # 等大文件经常被静默丢弃. 如果 verify 显示 has_chroma=False 但本地有, # 用 hf_hub 低层 API 一个一个 upload_file 补上 (走 LFS 大文件协议). local_chroma = chroma_dir() if local_chroma.exists() and not has_chroma: logger.warning( "upload_folder skipped chroma/ (large files); falling back to per-file upload_file" ) fallback_ok = await loop.run_in_executor( _persist_executor, lambda: _push_chroma_fallback(local_chroma, repo_id, token), ) if fallback_ok: # 重新 verify files = await loop.run_in_executor( _persist_executor, lambda: list_repo_files( repo_id=repo_id, repo_type="dataset", token=token, ), ) has_chroma = any(f.startswith("chroma/") for f in files) verify_result["has_chroma"] = has_chroma _state["last_verify"] = verify_result logger.info("chroma fallback push: has_chroma=%s", has_chroma) else: logger.error("chroma fallback push failed (ChromaDB dense won't persist!)") _state["last_push_status"] = "ok" _state["last_push_at"] = time.time() _state["last_error"] = None logger.info( "Persist push OK: %s, %d files (chroma=%s colbert=%s sqlite=%s uploads=%s), upload %dms", repo_id, len(files), has_chroma, has_colbert, has_sqlite, has_uploads, upload_ms, ) return True except Exception as e: # noqa: BLE001 _state["last_push_status"] = "failed" _state["last_error"] = f"push: {type(e).__name__}: {e}"[:500] _state["pending_push"] = False # 失败也重置, 否则 schedule_push 永远 short-circuit logger.error("Persist push failed: %s", e, exc_info=True) return False def _push_chroma_fallback(local_chroma: Path, repo_id: str, token: str) -> bool: """ChromaDB 大文件补传 (走 hf_hub LFS 协议, 绕开 git 单文件限制). upload_folder 走 git push, ChromaDB 的 chroma.sqlite3 / data_level0.bin 等大文件经常超过 HF 单文件 git push 限制被静默丢弃. 这里把 chroma/ 下的文件**一个个**用 upload_file (走 LFS HTTP API) 补传. Returns: True=全成功, False=有失败. """ files = [p for p in local_chroma.rglob("*") if p.is_file()] if not files: logger.warning("chroma fallback: no files in %s", local_chroma) return False ok = 0 failed = 0 for f in files: rel = f.relative_to(local_chroma).as_posix() path_in_repo = f"chroma/{rel}" try: upload_file( path_or_fileobj=str(f), path_in_repo=path_in_repo, repo_id=repo_id, repo_type="dataset", token=token, commit_message="chroma fallback (large LFS files)", ) ok += 1 except Exception as e: # noqa: BLE001 logger.error("chroma fallback upload failed for %s: %s", path_in_repo, e) failed += 1 logger.info("chroma fallback: ok=%d failed=%d total=%d", ok, failed, len(files)) return failed == 0 async def schedule_push() -> None: """异步推送 (5s debounce), 不阻塞业务. 适用: chat 写入 / session checkpoint 等低频操作. 重要操作 (upload/delete) 应该直接 await push_to_hf(). """ if not settings.is_persist_enabled(): return if not settings.persist_on_write: return if _state["pending_push"]: return # 已有 pending, 跳过 _state["pending_push"] = True async def _delayed_push() -> None: await asyncio.sleep(DEBOUNCE_SECONDS) if _state["pending_push"]: await push_to_hf() # 持引用防 GC task = asyncio.create_task(_delayed_push()) _background_tasks.add(task) task.add_done_callback(_background_tasks.discard) async def _ensure_repo_exists() -> None: """首次部署时自动创建 HF Dataset repo.""" repo_id = settings.hf_persist_repo token = settings.hf_token.get_secret_value() loop = asyncio.get_running_loop() try: await loop.run_in_executor( _persist_executor, lambda: create_repo( repo_id=repo_id, repo_type="dataset", token=token, private=True, exist_ok=True, ), ) logger.info("Created persist repo: %s", repo_id) except Exception as e: # noqa: BLE001 logger.error("Failed to create persist repo: %s", e, exc_info=True) _state["last_error"] = f"create_repo: {e}"[:500] def persist_status() -> dict: """供 /readyz 暴露持久化状态 (含 A 改良版的 verify / last_error).""" return { "enabled": settings.is_persist_enabled(), "mode": _state["mode"], "pending_push": _state["pending_push"], "last_push_status": _state["last_push_status"], "last_push_at": _state["last_push_at"], "last_error": _state["last_error"], "last_verify": _state["last_verify"], "repo": settings.hf_persist_repo or None, } def _wipe_local_data() -> None: """测试用: 清空本地 data 目录.""" for d in (sqlite_dir(), chroma_dir(), upload_dir()): if d.exists(): shutil.rmtree(d, ignore_errors=True) __all__ = [ "restore_from_hf", "push_to_hf", "schedule_push", "persist_mode", "persist_status", ]