fix: HF_ENDPOINT 环境变量操作添加线程锁防竞态

This commit is contained in:
2026-07-10 15:19:43 +08:00
parent 298f13721c
commit 8445af7be0
+10 -5
View File
@@ -18,6 +18,7 @@
vectors = batch_embed(embedder, long_text_list)
"""
import os
import threading
import logging
from typing import Protocol
@@ -29,6 +30,8 @@ logger = logging.getLogger("md-vector-db")
_HF_MIRROR = os.environ.get("HF_MIRROR", "https://hf-mirror.com")
_HF_ENV_LOCK = threading.Lock()
# -- Provider 默认配置 --
_PROVIDER_DEFAULTS: dict[str, dict[str, str | int]] = {
"openai": {
@@ -78,16 +81,18 @@ class LocalEmbedder:
# 通过 HF_ENDPOINT 环境变量设置镜像(sentence-transformers 依赖 huggingface_hub
# 临时设置仅用于模型下载,下载完成后还原
old_endpoint = os.environ.get("HF_ENDPOINT")
os.environ["HF_ENDPOINT"] = _HF_MIRROR
with _HF_ENV_LOCK:
os.environ["HF_ENDPOINT"] = _HF_MIRROR
try:
self._model = SentenceTransformer(
config.local_model, device=device
)
finally:
if old_endpoint is not None:
os.environ["HF_ENDPOINT"] = old_endpoint
else:
os.environ.pop("HF_ENDPOINT", None)
with _HF_ENV_LOCK:
if old_endpoint is not None:
os.environ["HF_ENDPOINT"] = old_endpoint
else:
os.environ.pop("HF_ENDPOINT", None)
@property
def dimension(self) -> int: