x

docker/scripts/patch_vllm_platform.py

@@ -6,9 +6,14 @@
 """
 
 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():
@@ -56,10 +61,42 @@ def patch7_rocm_break_import_cycle():
         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__}')