| """进程启动时的环境变量修复（仅依赖 stdlib，不引入 OpenMP 等）。""" |
|
|
| import os |
| import multiprocessing |
|
|
|
|
| def diagnose_and_fix_thread_env_vars() -> None: |
| """ |
| 诊断并修复 OMP_NUM_THREADS 和 MKL_NUM_THREADS 环境变量。 |
| |
| 在 HF Space 的 CUDA 容器中，可能预设了无效的环境变量值， |
| 这会导致 bitsandbytes 库初始化时 libgomp 报错。 |
| """ |
| actual_cores = multiprocessing.cpu_count() |
| env_vars = ['OMP_NUM_THREADS', 'MKL_NUM_THREADS'] |
| is_first_fix = True |
|
|
| for env_var in env_vars: |
| value = os.environ.get(env_var) |
| if value is None: |
| continue |
|
|
| stripped = value.strip() |
| is_valid = False |
| reason = "" |
|
|
| if not stripped: |
| reason = "值为空字符串" |
| elif not stripped.isdigit(): |
| reason = f"包含非数字字符: {repr(stripped)}" |
| else: |
| try: |
| int_value = int(stripped) |
| if int_value <= 0: |
| reason = f"值 <= 0: {int_value}" |
| else: |
| is_valid = True |
| except ValueError: |
| reason = f"无法转换为整数: {repr(stripped)}" |
|
|
| if not is_valid: |
| if is_first_fix: |
| print(f"🔍 检测到无效的线程环境变量（实际 CPU 核数: {actual_cores}）:") |
| is_first_fix = False |
| os.environ[env_var] = str(actual_cores) |
| print(f" {env_var}:") |
| print(f" - 原始值: {repr(value)}") |
| print(f" - 问题: {reason}") |
| print(f" - 🔧 已自动修复: {env_var}={actual_cores}") |
|
|
