mineru-rocm/docker/scripts/patch_vllm_platform.py

#!/usr/bin/env python3
"""vllm 平台检测补丁

问题 1：amdsmi 不可用时 platform 回退到 torch.version.hip
问题 2：rocm.py 中 logger.warning_once() 导致循环导入（默认不再改写；仅保留为显式开关）
"""

import os
import subprocess
import sys
import traceback

VLLM_DIR = '/opt/vllm/vllm'
VLLM_SRC_DIR = '/opt/vllm'
VENV_PYTHON = '/opt/mineru_venv/bin/python'
VENV_PIP = '/opt/mineru_venv/bin/pip'


def patch6_init_platform_fallback():
    """补丁 6：platforms/__init__.py —— torch.version.hip 兜底"""
    f = os.path.join(VLLM_DIR, 'platforms', '__init__.py')
    c = open(f).read()
    old = "    return 'vllm.platforms.rocm.RocmPlatform' if is_rocm else None"
    new = ("    # amdsmi fallback: also check torch.version.hip\n"
           "    if not is_rocm:\n"
           "        try:\n"
           "            import torch\n"
           "            if torch.version.hip is not None:\n"
           "                is_rocm = True\n"
           "        except Exception:\n"
           "            pass\n"
           "    return 'vllm.platforms.rocm.RocmPlatform' if is_rocm else None")
    c2 = c.replace(old, new)
    if c2 != c:
        open(f, 'w').write(c2)
        print('Patch 6: __init__.py platform fallback applied.')
    else:
        print('Patch 6: already applied or pattern not found.')


def patch7_rocm_break_import_cycle():
    """补丁 7：platforms/rocm.py —— logger.warning_once → sys.stderr.write

    这个补丁会直接改写 vLLM 源码中的函数调用。部分 vLLM 版本里
    logger.warning_once(...) 参数不一定兼容 sys.stderr.write(...)，改写后可能
    导致 import vllm / import vllm.platforms.rocm 失败。MinerU 会把这种导入失败
    包装成 "Please install vllm"，因此默认禁用，只在显式设置环境变量时启用。
    """
    if os.environ.get('VLLM_PATCH_ROCM_WARNING_ONCE') != '1':
        print('Patch 7: skipped (set VLLM_PATCH_ROCM_WARNING_ONCE=1 to enable).')
        return
    f = os.path.join(VLLM_DIR, 'platforms', 'rocm.py')
    c = open(f).read()
    old = 'logger.warning_once('
    new = 'import sys as _sys\n        _sys.stderr.write('
    c2 = c.replace(old, new)
    if c2 != c:
        open(f, 'w').write(c2)
        print('Patch 7: rocm.py circular import broken.')
    else:
        print('Patch 7: already applied or pattern not found.')


def ensure_vllm_installed():
    """确保 vLLM Python 包在当前虚拟环境中可导入。

    某些 MinerU 依赖安装流程可能会卸载/覆盖 editable vLLM，导致任务运行时报
    "Please install vllm"。如果 /opt/vllm 源码仍在，启动时用 --no-deps 重新注册，
    避免扰动已锁定的 ROCm PyTorch/Triton 依赖。
    """
    try:
        import vllm  # noqa: F401
        return
    except ModuleNotFoundError as exc:
        if exc.name != 'vllm':
            raise

    if not os.path.isdir(VLLM_SRC_DIR):
        raise ModuleNotFoundError(
            'vllm is not installed and /opt/vllm source directory is missing'
        )

    print('vllm package missing; reinstalling editable /opt/vllm with --no-deps...')
    subprocess.check_call([
        VENV_PIP,
        'install',
        '-e', VLLM_SRC_DIR,
        '--no-build-isolation',
        '--no-deps',
    ])
    import vllm  # noqa: F401
    print('vllm editable reinstall OK.')


def main():
    patch6_init_platform_fallback()
    patch7_rocm_break_import_cycle()
    try:
        ensure_vllm_installed()
        import vllm
        from vllm.platforms import current_platform
        print(f'vllm runtime import OK: {vllm.__version__}, platform={type(current_platform).__name__}')
    except Exception:
        print('ERROR: vllm runtime import failed after platform patches:')
        traceback.print_exc()
        raise
    print('vllm platform patches done.')


if __name__ == '__main__':
    main()