x

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()