""" 模型 Provider 状态探测模块 为 CPU-only HuggingFace Space 环境提供模型可用性探测能力。 不强制加载完整大模型，仅报告 provider 在当前环境的状态。 """ import logging import time import os from datetime import datetime, timezone from pathlib import Path from typing import Optional, Dict, Any, List from enum import Enum from pydantic import BaseModel, Field from app.config.settings import settings logger = logging.getLogger(__name__) class ProviderAvailability(str, Enum): """Provider 可用状态""" AVAILABLE = "available" # 模型依赖就绪，预期可运行 UNAVAILABLE = "unavailable" # 模型依赖缺失或环境不支持 NOT_VERIFIED = "not_verified" # 无法验证（如缺少凭证、资源限制等） class ProviderStatus(BaseModel): """单个模型 Provider 的状态信息 供前端和 API 展示模型可用性，不强制加载完整模型。 """ provider: str = Field(..., description="Provider 名称，如 whisper_small") model_id: str = Field(..., description="HuggingFace 模型 ID") status: ProviderAvailability = Field( default=ProviderAvailability.NOT_VERIFIED, description="可用状态：available/unavailable/not_verified" ) cache_path: Optional[str] = Field(None, description="模型缓存路径") reason: Optional[str] = Field(None, description="状态原因说明") last_probe_at: Optional[datetime] = Field(None, description="最后探测时间") probe_duration_ms: Optional[float] = Field(None, description="探测耗时（毫秒）") extra: Dict[str, Any] = Field(default_factory=dict, description="额外信息") class ModelProbeResult(BaseModel): """模型探测汇总结果 包含所有 provider 的状态列表和探测摘要。 """ providers: List[ProviderStatus] = Field(default_factory=list) probe_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) cpu_only: bool = True # 当前环境是否为 CPU-only hf_home: str = "" hf_cache_dir: str = "" def probe_cpu_only_environment() -> Dict[str, Any]: """探测 CPU-only 环境的基本信息 不加载任何模型，仅检查环境变量、缓存目录和 GPU 可用性。 Returns: Dict 包含 cpu_only 标志、hf_home、缓存路径等信息 """ result = { "cpu_only": True, "hf_home": str(settings.HF_HOME), "hf_cache_dir": str(settings.HF_CACHE_DIR), "hf_home_exists": settings.HF_HOME.exists(), "hf_cache_dir_exists": settings.HF_CACHE_DIR.exists(), } # 检查 GPU 可用性（不导入 torch，只检查环境变量） cuda_visible = os.getenv("CUDA_VISIBLE_DEVICES", "") if cuda_visible and cuda_visible != "-1": # 有 GPU 环境变量设置，但本项目目标是 CPU-only result["cuda_visible_devices"] = cuda_visible result["note"] = "检测到 CUDA 环境变量，但项目以 CPU-only 模式运行" # 检查 HuggingFace 环境变量 hf_home_env = os.getenv("HF_HOME", "") if hf_home_env: result["hf_home_env"] = hf_home_env hf_token = os.getenv("HF_TOKEN", "") or os.getenv("HUGGINGFACE_HUB_TOKEN", "") result["hf_token_configured"] = bool(hf_token) # 检查 HuggingFace Hub 是否可用（尝试 import，不调用网络） try: import huggingface_hub # noqa: F401 result["huggingface_hub_available"] = True result["huggingface_hub_version"] = huggingface_hub.__version__ except ImportError: result["huggingface_hub_available"] = False result["huggingface_hub_version"] = None return result def probe_provider( provider: str, model_id: str, check_dependencies: List[str], check_cache: bool = True, ) -> ProviderStatus: """探测单个 Provider 的可用性 执行轻量检查：依赖包是否可导入、模型缓存是否存在。 不加载完整模型权重。 Args: provider: Provider 名称 model_id: HuggingFace 模型 ID check_dependencies: 需要检查的 Python 依赖包列表 check_cache: 是否检查模型缓存 Returns: ProviderStatus 包含探测结果 """ start_time = time.perf_counter() reasons = [] status = ProviderAvailability.AVAILABLE # 检查依赖包 missing_deps = [] for dep in check_dependencies: try: __import__(dep) except ImportError: missing_deps.append(dep) if missing_deps: status = ProviderAvailability.UNAVAILABLE reasons.append(f"缺失依赖: {', '.join(missing_deps)}") # 检查模型缓存 cache_dir_name = "models--" + model_id.replace("/", "--") cache_path = settings.HF_CACHE_DIR / cache_dir_name if check_cache: if cache_path.exists(): # 检查缓存中是否有实际模型文件（blobs 或 snapshots） has_blobs = (cache_path / "blobs").exists() and any((cache_path / "blobs").iterdir()) has_snapshots = (cache_path / "snapshots").exists() and any((cache_path / "snapshots").iterdir()) if not has_blobs and not has_snapshots: # 缓存目录存在但无模型文件，标记为需要下载 reasons.append("模型缓存目录存在但无模型文件，首次使用需下载") else: reasons.append(f"模型未缓存，首次使用将从 HuggingFace Hub 下载") # 检查 CPU-only 环境的额外限制 cpu_env = probe_cpu_only_environment() if not cpu_env.get("huggingface_hub_available"): status = ProviderAvailability.UNAVAILABLE reasons.append("huggingface_hub 包不可用，无法访问 HuggingFace Hub") # 如果缺少依赖且无法下载，标记为 unavailable if missing_deps and not cpu_env.get("huggingface_hub_available"): status = ProviderAvailability.UNAVAILABLE duration_ms = (time.perf_counter() - start_time) * 1000 return ProviderStatus( provider=provider, model_id=model_id, status=status, cache_path=str(cache_path) if check_cache else None, reason="; ".join(reasons) if reasons else "依赖和缓存检查通过", last_probe_at=datetime.now(timezone.utc), probe_duration_ms=round(duration_ms, 2), extra={ "missing_dependencies": missing_deps, "cpu_only": cpu_env["cpu_only"], }, ) def probe_all_providers() -> ModelProbeResult: """探测所有已知 Provider 的可用性 汇总 ASR、OCR 等所有模型 provider 的状态。 Returns: ModelProbeResult 包含所有 provider 状态 """ cpu_env = probe_cpu_only_environment() providers: List[ProviderStatus] = [] # ASR providers（均基于 transformers pipeline） providers.append(probe_provider( provider="whisper_small", model_id="openai/whisper-small", check_dependencies=["transformers", "torch"], )) providers.append(probe_provider( provider="glm_asr_nano", model_id="zai-org/GLM-ASR-Nano-2512", check_dependencies=["transformers", "torch"], )) # OCR provider providers.append(probe_provider( provider="ppocrv6", model_id="PaddlePaddle/PP-OCRv6_medium_det_safetensors", check_dependencies=["paddleocr", "paddle"], )) # 补充 OCR rec 模型 providers.append(probe_provider( provider="ppocrv6_rec", model_id="PaddlePaddle/PP-OCRv6_medium_rec_safetensors", check_dependencies=["paddleocr", "paddle"], )) return ModelProbeResult( providers=providers, cpu_only=cpu_env["cpu_only"], hf_home=cpu_env["hf_home"], hf_cache_dir=cpu_env["hf_cache_dir"], )