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) vectors = batch_embed(embedder, long_text_list)
""" """
import os import os
import threading
import logging import logging
from typing import Protocol 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_MIRROR = os.environ.get("HF_MIRROR", "https://hf-mirror.com")
_HF_ENV_LOCK = threading.Lock()
# -- Provider 默认配置 -- # -- Provider 默认配置 --
_PROVIDER_DEFAULTS: dict[str, dict[str, str | int]] = { _PROVIDER_DEFAULTS: dict[str, dict[str, str | int]] = {
"openai": { "openai": {
@@ -78,16 +81,18 @@ class LocalEmbedder:
# 通过 HF_ENDPOINT 环境变量设置镜像(sentence-transformers 依赖 huggingface_hub # 通过 HF_ENDPOINT 环境变量设置镜像(sentence-transformers 依赖 huggingface_hub
# 临时设置仅用于模型下载,下载完成后还原 # 临时设置仅用于模型下载,下载完成后还原
old_endpoint = os.environ.get("HF_ENDPOINT") old_endpoint = os.environ.get("HF_ENDPOINT")
os.environ["HF_ENDPOINT"] = _HF_MIRROR with _HF_ENV_LOCK:
os.environ["HF_ENDPOINT"] = _HF_MIRROR
try: try:
self._model = SentenceTransformer( self._model = SentenceTransformer(
config.local_model, device=device config.local_model, device=device
) )
finally: finally:
if old_endpoint is not None: with _HF_ENV_LOCK:
os.environ["HF_ENDPOINT"] = old_endpoint if old_endpoint is not None:
else: os.environ["HF_ENDPOINT"] = old_endpoint
os.environ.pop("HF_ENDPOINT", None) else:
os.environ.pop("HF_ENDPOINT", None)
@property @property
def dimension(self) -> int: def dimension(self) -> int: