From 7861ebdad745773db278a348f539a781535a61a3 Mon Sep 17 00:00:00 2001
From: Simo Lin <25425177+slin1237@users.noreply.github.com>
Date: Thu, 11 Jun 2026 12:47:20 -0700
Subject: [PATCH] fix(e2e): unbreak trtllm lanes broken by mcp 2.0.0a1
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Every trtllm e2e job dies at the model-download step since the
TensorRT-LLM 1.3.0rc18 bump (#1663): rc18 newly depends on mcp
(unconstrained), the trtllm install runs pip with --pre (required for
NVIDIA pre-release wheels), and pip resolves mcp to 2.0.0a1 — which
removed the legacy mcp.server.fastmcp module. Our e2e dependency floor
mcp>=1.0 treats the alpha as already satisfied, and e2e_test/infra
eagerly imports the FastMCP-based mock at package import, so even
ci_download_model.sh crashes with ModuleNotFoundError before any test
runs. rc14 never depended on mcp, which is why this only surfaced now.

Two fixes, both on our side of the contract:

- Declare the real requirement: mcp>=1.2,<2 (fastmcp was added in 1.2
  and removed in 2.0). PEP 440 excludes 2.0 pre-releases from <2, so
  the e2e deps install now replaces the alpha with stable 1.x instead
  of keeping it. TensorRT-LLM's own mcp requirement is unconstrained,
  so the downgrade stays consistent.

- Stop importing the mock MCP server eagerly in infra/__init__: only
  the sglang-only responses lane uses it, yet the eager import made
  model download and every chat/router lane depend on the mcp SDK.
  The symbols now resolve lazily via module __getattr__.

Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
---
 e2e_test/infra/__init__.py | 17 ++++++++++++++++-
 e2e_test/pyproject.toml    |  3 ++-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/e2e_test/infra/__init__.py b/e2e_test/infra/__init__.py
index e1734cab09..89044d746f 100644
--- a/e2e_test/infra/__init__.py
+++ b/e2e_test/infra/__init__.py
@@ -41,7 +41,6 @@
 from .gateway import Gateway, WorkerInfo, launch_cloud_gateway
 from .gpu_monitor import GPUMonitor
 from .gpu_monitor import should_monitor as should_monitor_gpu
-from .mock_mcp import IMAGE_GENERATION_PNG_BASE64, MockMcpServer, mock_mcp_server
 from .model_specs import (  # Default model paths; Model groups
     CHAT_MODELS,
     DEFAULT_EMBEDDING_MODEL_PATH,
@@ -161,3 +160,19 @@
     # Evaluation
     "run_eval",
 ]
+
+# The mock MCP server is only used by the agentic (responses) lane and drags
+# in the `mcp` SDK at import time. Resolve its symbols lazily so infra
+# consumers that never touch it — model download, chat/router lanes — keep
+# working even when the venv's `mcp` lacks `mcp.server.fastmcp` (e.g.
+# TensorRT-LLM's `pip install --pre` resolving mcp to a 2.x pre-release,
+# which removed the module).
+_LAZY_MOCK_MCP = ("IMAGE_GENERATION_PNG_BASE64", "MockMcpServer", "mock_mcp_server")
+
+
+def __getattr__(name: str) -> object:
+    if name in _LAZY_MOCK_MCP:
+        from . import mock_mcp
+
+        return getattr(mock_mcp, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/e2e_test/pyproject.toml b/e2e_test/pyproject.toml
index 0d3e04e24f..f6de88dd0d 100644
--- a/e2e_test/pyproject.toml
+++ b/e2e_test/pyproject.toml
@@ -11,7 +11,8 @@ dependencies = [
   "grpcio-health-checking",
   "httpx",
   "jinja2",
-  "mcp>=1.0",
+  # infra.mock_mcp_server needs mcp.server.fastmcp: added in 1.2, removed in 2.0
+  "mcp>=1.2,<2",
   "numpy",
   "openai",
   "pandas",