workaround_env_fix.py

"""进程启动时的环境变量修复（仅依赖 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}")