Update Dockerfile.ci to use a new vllm wheel URL and refactor MooncakeConnector import logic in monkey_patch.py for improved compatibility. Adjust OpenAIServingChat to call the harmony request method from the correct context.

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

docker/Dockerfile.ci

@@ -17,7 +17,7 @@ RUN apt-get update && \
 # Note: --prerelease=allow is required because nightly/commit wheels use
 # pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
 # causing it to fall back to the older stable release from PyPI.
-RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/e9163b536e721c431500f6f43ace22fcb3532e7e
+RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2754231ba3a72f41e62922d1552c33e8f3f6a9d1
 
 RUN uv pip install --system --no-cache-dir ".[dev]"
 

vllm_omni/distributed/kv_transfer/monkey_patch.py

@@ -8,6 +8,7 @@ so the decode side references the correct KV entry.
 
 from __future__ import annotations
 
+import importlib
 import logging
 import sys
 from dataclasses import dataclass
@@ -30,30 +31,23 @@ class PatchedRecvReqMeta:
 
 def _import_mooncake_module():
     """Import MooncakeConnector module, supporting both vLLM >=0.16 and older."""
-    try:
-        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import mooncake_connector
-
-        return mooncake_connector
-    except ImportError:
-        pass
-    try:
-        from vllm.distributed.kv_transfer.kv_connector.v1 import mooncake_connector
-
-        return mooncake_connector
-    except ImportError:
-        return None
+    for mod_path in (
+        "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
+        "vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector",
+    ):
+        try:
+            return importlib.import_module(mod_path)
+        except (ImportError, ModuleNotFoundError):
+            continue
+    return None
 
 
 def _create_patched_mooncake_connector():
     """Return a subclass of MooncakeConnector with remote_request_id support."""
-    try:
-        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
-            MooncakeConnector as _OriginalMooncakeConnector,
-        )
-    except (ImportError, AttributeError):
-        from vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector import (
-            MooncakeConnector as _OriginalMooncakeConnector,
-        )
+    _mc_mod = _import_mooncake_module()
+    if _mc_mod is None:
+        raise ImportError("Cannot import MooncakeConnector from upstream vLLM")
+    _OriginalMooncakeConnector = _mc_mod.MooncakeConnector
 
     class PatchedMooncakeConnector(_OriginalMooncakeConnector):
         """Fixes request-ID mismatch in PD disaggregation by injecting

vllm_omni/entrypoints/openai/serving_chat.py

@@ -256,7 +256,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 )
             else:
                 should_include_tools = tool_dicts is not None
-                conversation, engine_prompts = self._make_request_with_harmony(request, should_include_tools)
+                conversation, engine_prompts = self.openai_serving_render._make_request_with_harmony(
+                    request, should_include_tools
+                )
 
         except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e:
             logger.exception("Error in preprocessing prompt inputs")