282 lines
11 KiB
Python
282 lines
11 KiB
Python
#!/usr/bin/env python3
|
||
"""vllm 平台检测补丁
|
||
|
||
问题 1：amdsmi 不可用时 platform 回退到 torch.version.hip
|
||
问题 2：rocm.py 中 logger.warning_once() 导致循环导入（默认不再改写；仅保留为显式开关）
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import sys
|
||
import sysconfig
|
||
import traceback
|
||
|
||
VLLM_DIR = '/opt/vllm/vllm'
|
||
VLLM_SRC_DIR = '/opt/vllm'
|
||
|
||
|
||
def patch6_init_platform_fallback():
|
||
"""补丁 6：platforms/__init__.py —— torch.version.hip 兜底
|
||
|
||
vLLM 的 is_rocm 检测依赖 import amdsmi；amdsmi 未装时 is_rocm=False，
|
||
current_platform 落到 UnspecifiedPlatform，device_type='' →
|
||
RuntimeError: Device string must not be empty。
|
||
|
||
本补丁在 RocmPlatform 的 return 语句前注入 torch.version.hip 兜底。
|
||
采用行结构 + return 定位，不依赖精确字符串匹配
|
||
（vllm main 分支重构频繁，固定字符串匹配已多次失效）。
|
||
"""
|
||
f = os.path.join(VLLM_DIR, 'platforms', '__init__.py')
|
||
lines = open(f).read().splitlines(keepends=True)
|
||
|
||
marker = "'vllm.platforms.rocm.RocmPlatform'"
|
||
return_idx = None
|
||
for i, ln in enumerate(lines):
|
||
if marker in ln and 'return' in ln:
|
||
return_idx = i
|
||
break
|
||
if return_idx is None:
|
||
print('Patch 6: return statement with RocmPlatform not found; skipping.')
|
||
return
|
||
|
||
if 'torch.version.hip is not None' in ''.join(lines):
|
||
print('Patch 6: already applied (torch.version.hip fallback present).')
|
||
return
|
||
|
||
indent = ' ' * (len(lines[return_idx]) - len(lines[return_idx].lstrip()))
|
||
inject = (
|
||
f"{indent}# amdsmi fallback: also check torch.version.hip\n"
|
||
f"{indent}if not is_rocm:\n"
|
||
f"{indent} try:\n"
|
||
f"{indent} import torch as _torch\n"
|
||
f"{indent} if _torch.version.hip is not None:\n"
|
||
f"{indent} is_rocm = True\n"
|
||
f"{indent} except Exception:\n"
|
||
f"{indent} pass\n"
|
||
)
|
||
lines.insert(return_idx, inject)
|
||
open(f, 'w').write(''.join(lines))
|
||
print('Patch 6: __init__.py torch.version.hip fallback applied.')
|
||
|
||
|
||
def patch7_rocm_break_import_cycle():
|
||
"""补丁 7：platforms/rocm.py —— logger.warning_once → sys.stderr.write
|
||
|
||
默认启用安全的“整段 except 块”替换，避免 logger.warning_once() 在平台检测
|
||
期间触发循环导入，导致 current_platform 落到 UnspecifiedPlatform。
|
||
如需禁用，设置 VLLM_PATCH_ROCM_WARNING_ONCE=0。
|
||
"""
|
||
if os.environ.get('VLLM_PATCH_ROCM_WARNING_ONCE') == '0':
|
||
print('Patch 7: skipped (VLLM_PATCH_ROCM_WARNING_ONCE=0).')
|
||
return
|
||
f = os.path.join(VLLM_DIR, 'platforms', 'rocm.py')
|
||
c = open(f).read()
|
||
if 'amdsmi unavailable, using torch.cuda fallback' in c:
|
||
print('Patch 7: already applied.')
|
||
return
|
||
pattern = re.compile(
|
||
r'(\n\s*except Exception as e:\n'
|
||
r'\s*logger\.debug\("Failed to get GCN arch via amdsmi: %s", e\)\n'
|
||
r'\s*logger\.warning_once\(\n'
|
||
r'(?:\s*"[^"]*"\n)+'
|
||
r'\s*\)\n)',
|
||
re.MULTILINE,
|
||
)
|
||
|
||
def repl(match):
|
||
indent = re.search(r'\n(\s*)except Exception as e:', match.group(1)).group(1)
|
||
body_indent = indent + ' '
|
||
return (
|
||
f'\n{indent}except Exception as e:\n'
|
||
f'{body_indent}import sys as _sys\n'
|
||
f'{body_indent}_sys.stderr.write('
|
||
'"vLLM ROCm: amdsmi unavailable, using torch.cuda fallback for GPU detection\\n")\n'
|
||
)
|
||
|
||
c2, n = pattern.subn(repl, c, count=1)
|
||
if c2 != c:
|
||
open(f, 'w').write(c2)
|
||
print('Patch 7: rocm.py logger.warning_once circular import patch applied.')
|
||
else:
|
||
print('Patch 7: already applied or pattern not found.')
|
||
|
||
|
||
def ensure_vllm_dist_info():
|
||
"""为 /opt/vllm 源码导入创建最小 dist-info。
|
||
|
||
vLLM 平台检测会通过 importlib.metadata 查询 vllm 分发元数据/entry points。
|
||
如果只靠 PYTHONPATH 导入源码，没有 dist-info，就会出现：
|
||
"The vLLM package was not found..."，并可能得到 UnspecifiedPlatform。
|
||
|
||
注意：仅当 rocm_platform_plugin 函数实际存在时才注册 entry_point，
|
||
否则 vLLM 会因 AttributeError 回退到 UnspecifiedPlatform。
|
||
"""
|
||
try:
|
||
import vllm
|
||
except Exception:
|
||
return
|
||
|
||
version = getattr(vllm, '__version__', '0.1.dev1') or '0.1.dev1'
|
||
# PEP440 合法性校验：vllm 在 _version.py 缺失时会回退到 'dev'，
|
||
# 这不是合法 PEP440，会触发下游 packaging.version.parse 抛 InvalidVersion
|
||
# （例如 mineru.backend.vlm.utils:set_default_gpu_memory_utilization）。
|
||
try:
|
||
from packaging.version import Version as _PEP440Version
|
||
_PEP440Version(version)
|
||
except Exception:
|
||
print(f'WARNING: vllm.__version__={version!r} is not PEP440-compliant; '
|
||
f'falling back to "0.11.0" for dist-info.')
|
||
version = '0.11.0'
|
||
try:
|
||
vllm.__version__ = version
|
||
except Exception:
|
||
pass
|
||
site_packages = sysconfig.get_paths().get('purelib')
|
||
if not site_packages:
|
||
return
|
||
dist_info = os.path.join(site_packages, f'vllm-{version}.dist-info')
|
||
os.makedirs(dist_info, exist_ok=True)
|
||
|
||
metadata = os.path.join(dist_info, 'METADATA')
|
||
if not os.path.exists(metadata):
|
||
with open(metadata, 'w') as fp:
|
||
fp.write(
|
||
'Metadata-Version: 2.1\n'
|
||
'Name: vllm\n'
|
||
f'Version: {version}\n'
|
||
'Summary: vLLM source tree mounted at /opt/vllm\n'
|
||
)
|
||
|
||
with open(os.path.join(dist_info, 'top_level.txt'), 'w') as fp:
|
||
fp.write('vllm\n')
|
||
with open(os.path.join(dist_info, 'INSTALLER'), 'w') as fp:
|
||
fp.write('mineru-rocm-runtime\n')
|
||
|
||
# 仅当 rocm_platform_plugin 实际存在时才注册，避免 vLLM 加载失败
|
||
entry_points = ''
|
||
try:
|
||
from vllm.platforms.rocm import rocm_platform_plugin # noqa: F401
|
||
entry_points += (
|
||
'[vllm.platform_plugins]\n'
|
||
'rocm = vllm.platforms.rocm:rocm_platform_plugin\n'
|
||
)
|
||
print('rocm_platform_plugin found, registering entry point.')
|
||
except (ImportError, AttributeError, SystemExit, KeyboardInterrupt):
|
||
print('WARNING: rocm_platform_plugin not found in vllm.platforms.rocm; '
|
||
'skipping entry_point registration.')
|
||
except BaseException as e:
|
||
print(f'WARNING: rocm_platform_plugin check failed ({type(e).__name__}: {e}); '
|
||
'skipping entry_point registration.')
|
||
|
||
with open(os.path.join(dist_info, 'entry_points.txt'), 'w') as fp:
|
||
fp.write(entry_points)
|
||
with open(os.path.join(dist_info, 'RECORD'), 'w') as fp:
|
||
fp.write('')
|
||
print(f'vllm dist-info ensured: {dist_info}')
|
||
|
||
|
||
def ensure_vllm_installed():
|
||
"""确保 vLLM Python 包在当前虚拟环境中可导入。
|
||
|
||
某些 MinerU 依赖安装流程后，site-packages 中可能缺少 vLLM 分发元数据，但
|
||
/opt/vllm 源码仍在。此时直接把 /opt/vllm 加入 PYTHONPATH/sys.path 再导入，
|
||
避免 editable 安装触发 pyproject 元数据校验失败。
|
||
"""
|
||
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; trying PYTHONPATH fallback with /opt/vllm...')
|
||
current = os.environ.get('PYTHONPATH', '')
|
||
paths = [p for p in current.split(':') if p]
|
||
if VLLM_SRC_DIR not in paths:
|
||
os.environ['PYTHONPATH'] = f"{VLLM_SRC_DIR}:{current}" if current else VLLM_SRC_DIR
|
||
if VLLM_SRC_DIR not in sys.path:
|
||
sys.path.insert(0, VLLM_SRC_DIR)
|
||
|
||
import vllm # noqa: F401
|
||
print('vllm import recovered via PYTHONPATH fallback.')
|
||
|
||
|
||
def install_sitecustomize_force_rocm():
|
||
"""安装 sitecustomize.py 运行时兜底：每个 Python 进程启动时，
|
||
若 current_platform 落到 UnspecifiedPlatform 但 torch 检测到 ROCm/HIP，
|
||
则强制替换为 RocmPlatform 实例。
|
||
|
||
vllm main 的 is_rocm 检测依赖 amdsmi，ROCm 7.2 容器内 amdsmi 缺失，
|
||
补丁 6 的 torch.version.hip 兜底理论上已够，但子进程加载时序、
|
||
环境差异等可能导致检测仍失败。本 sitecustomize 作为最后一道兜底，
|
||
确保任何 worker 进程都能拿到 device_type='cuda' 的平台。
|
||
"""
|
||
site_packages = sysconfig.get_paths().get('purelib')
|
||
if not site_packages:
|
||
print('Patch 8: cannot locate site-packages; skipping sitecustomize.')
|
||
return
|
||
sc = os.path.join(site_packages, 'sitecustomize.py')
|
||
|
||
marker = '# mineru-rocm: force rocm platform'
|
||
existing = ''
|
||
if os.path.exists(sc):
|
||
existing = open(sc).read()
|
||
if marker in existing:
|
||
print('Patch 8: sitecustomize force-rocm already installed.')
|
||
return
|
||
|
||
snippet = (
|
||
f"\n{marker}\n"
|
||
"try:\n"
|
||
" import torch as _t\n"
|
||
" if getattr(_t.version, 'hip', None) is not None:\n"
|
||
" import vllm.platforms as _vp\n"
|
||
" _cp = getattr(_vp, 'current_platform', None)\n"
|
||
" if _cp is not None and type(_cp).__name__ == 'UnspecifiedPlatform':\n"
|
||
" from vllm.platforms.rocm import RocmPlatform as _RP\n"
|
||
" _rp = _RP()\n"
|
||
" # vllm main 用模块级 __setattr__ 接管 current_platform 赋值\n"
|
||
" try:\n"
|
||
" _vp.current_platform = _rp\n"
|
||
" except Exception:\n"
|
||
" _vp._current_platform = _rp\n"
|
||
"except Exception:\n"
|
||
" pass\n"
|
||
)
|
||
with open(sc, 'a') as fp:
|
||
fp.write(snippet)
|
||
print(f'Patch 8: sitecustomize force-rocm appended to {sc}.')
|
||
|
||
|
||
def main():
|
||
patch6_init_platform_fallback()
|
||
patch7_rocm_break_import_cycle()
|
||
install_sitecustomize_force_rocm()
|
||
try:
|
||
ensure_vllm_installed()
|
||
ensure_vllm_dist_info()
|
||
import vllm
|
||
from vllm.platforms import current_platform
|
||
platform_name = type(current_platform).__name__
|
||
print(f'vllm runtime import OK: {vllm.__version__}, platform={platform_name}')
|
||
if platform_name == 'UnspecifiedPlatform':
|
||
print(
|
||
'WARNING: vLLM platform detection returned UnspecifiedPlatform. '
|
||
'This is expected on containers without GPU access (router/gradio). '
|
||
'On worker containers, check: amdsmi, /dev/kfd, /dev/dri, '
|
||
'and vllm.platform_plugins metadata.'
|
||
)
|
||
except Exception:
|
||
print('WARNING: vllm runtime import failed after platform patches (non-fatal):')
|
||
traceback.print_exc()
|
||
print('vllm platform patches done.')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|