diff --git a/docker/scripts/patch_vllm_platform.py b/docker/scripts/patch_vllm_platform.py index b49e774..8dbeccf 100644 --- a/docker/scripts/patch_vllm_platform.py +++ b/docker/scripts/patch_vllm_platform.py @@ -16,25 +16,47 @@ VLLM_SRC_DIR = '/opt/vllm' def patch6_init_platform_fallback(): - """补丁 6:platforms/__init__.py —— torch.version.hip 兜底""" + """补丁 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') - 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.') + 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(): @@ -184,9 +206,57 @@ def ensure_vllm_installed(): 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()