[https://nvbugs/6337229][fix] Stabilize scaffolding MCP worker test server startup and cleanup (NVIDIA#15958)

Signed-off-by: kleinc <kleinc@nvidia.com>
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>

Author: KleinBlueC

tests/unittest/scaffolding/test_bench.py

@@ -1,8 +1,9 @@
 import asyncio
 import math
-from typing import List
+from typing import Iterator, List
 
 import numpy as np
+import pytest
 
 from tensorrt_llm.executor.request import GenerationRequest
 from tensorrt_llm.executor.result import GenerationResult
@@ -19,6 +20,17 @@
 OUTPUT_STR = "Yes."
 
 
+@pytest.fixture(autouse=True)
+def _restore_native_generation_controller() -> Iterator[None]:
+    original_init = NativeGenerationController.__init__
+    original_process = NativeGenerationController.process
+
+    yield
+
+    NativeGenerationController.__init__ = original_init
+    NativeGenerationController.process = original_process
+
+
 class DummyWorker(Worker):
 
     async def dummy_generation_handler(self, task: GenerationTask):

tests/unittest/scaffolding/test_mcp_worker.py

@@ -1,6 +1,3 @@
-# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-# SPDX-License-Identifier: Apache-2.0
-
 import asyncio
 import json
 import multiprocessing
@@ -15,7 +12,7 @@
 from mcp.server.sse import SseServerTransport
 from starlette.applications import Starlette
 from starlette.requests import Request
-from starlette.responses import StreamingResponse
+from starlette.responses import JSONResponse, StreamingResponse
 from starlette.routing import Mount, Route
 
 from tensorrt_llm.scaffolding.controller import ChatWithMCPController, NativeGenerationController
@@ -67,9 +64,13 @@ async def handle_sse(request: Request) -> StreamingResponse:
         # Return an empty streaming response after SSE connection closes
         return StreamingResponse(iter([]), media_type="text/event-stream")
 
+    async def handle_health(_request: Request) -> JSONResponse:
+        return JSONResponse({"status": "ok"})
+
     return Starlette(
         debug=debug,
         routes=[
+            Route("/health", endpoint=handle_health),
             Route("/sse", endpoint=handle_sse),
             Mount("/messages/", app=sse.handle_post_message),
         ],
@@ -103,9 +104,13 @@ def __init__(self, host: str = "0.0.0.0", port: int = 8080, env: Optional[dict]
         )
         self.proc.start()
 
-        # Wait for server to be ready
-        # self._wait_for_server(url=self.sse_url,
-        #                      timeout=self.MAX_SERVER_START_WAIT_S)
+        try:
+            self._wait_for_server(
+                url=f"{self.url_root}/health", timeout=self.MAX_SERVER_START_WAIT_S
+            )
+        except RuntimeError:
+            self.terminate()
+            raise
 
     def __enter__(self):
         return self
@@ -166,9 +171,9 @@ def get_sse_url(self) -> str:
 
 
 @pytest.fixture(scope="module")
-def mcp_server():
+def mcp_server(unused_tcp_port_factory):
     # Create and start MCP server (similar to RemoteOpenAIServer usage)
-    remote_server = RemoteMCPServer(host="0.0.0.0", port=8080)
+    remote_server = RemoteMCPServer(host="127.0.0.1", port=unused_tcp_port_factory())
 
     # Yield server instance for test use
     yield remote_server
@@ -220,7 +225,6 @@ async def dummy_handler(self, task: ChatTask):
 # ============================================================
 
 
-@pytest.mark.skip(reason="https://nvbugs/6337229")
 @pytest.mark.asyncio
 async def test_mcp_worker_add_numbers(mcp_server):
     # 1. Initialize MCPWorker with mcp_server's SSE URL
@@ -254,10 +258,9 @@ async def test_mcp_worker_add_numbers(mcp_server):
 
     finally:
         # 6. Clean up worker resources
-        worker.shutdown()
+        await worker.async_shutdown()
 
 
-@pytest.mark.skip(reason="https://nvbugs/6337229")
 @pytest.mark.asyncio
 async def test_mcp_worker_echo_message(mcp_server):
     # Initialize MCPWorker
@@ -284,10 +287,9 @@ async def test_mcp_worker_echo_message(mcp_server):
         print(f"✓ Test passed: echo_message returned '{task.result_str}'")
 
     finally:
-        worker.shutdown()
+        await worker.async_shutdown()
 
 
-@pytest.mark.skip(reason="https://nvbugs/6337229")
 @pytest.mark.asyncio
 async def test_mcp_worker_multiple_calls(mcp_server):
     import json
@@ -337,10 +339,9 @@ async def test_mcp_worker_multiple_calls(mcp_server):
         print(f"  - Call 3: echo_message returned '{task3.result_str}'")
 
     finally:
-        worker.shutdown()
+        await worker.async_shutdown()
 
 
-@pytest.mark.skip(reason="https://nvbugs/6337229")
 @pytest.mark.asyncio
 async def test_multiple_calls_to_same_tool(mcp_server):
     sse_url = mcp_server.get_sse_url()
@@ -368,10 +369,9 @@ async def test_multiple_calls_to_same_tool(mcp_server):
     print(f"  - Call 1: add_numbers(10, 20) = {result1}")
     print(f"  - Call 2: add_numbers(100, 200) = {result2}")
 
-    worker.shutdown()
+    await worker.async_shutdown()
 
 
-@pytest.mark.skip(reason="https://nvbugs/6337229")
 @pytest.mark.asyncio
 async def test_scaffolding_with_chat_mcp_controller(mcp_server):
     # Initialize workers
@@ -441,7 +441,7 @@ async def test_scaffolding_with_chat_mcp_controller(mcp_server):
 
     finally:
         # Clean up
-        await mcp_worker.shutdown()
+        await mcp_worker.async_shutdown()
         scaffolding_llm.shutdown(shutdown_workers=False)  # Don't shutdown workers again
 
 

tests/unittest/scaffolding/test_worker.py

@@ -18,9 +18,6 @@
 sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 from llmapi.test_llm import get_model_path
 
-_SCAFFOLDING_MCP_NVBUG = pytest.mark.skip(reason="https://nvbugs/6337229")
-_TRTOAI_PROCESSPOOL_NVBUG = pytest.mark.skip(reason="https://nvbugs/6341072")
-
 
 @pytest.fixture(scope="module")
 def deepseek_distill_7b_path() -> Path:
@@ -37,7 +34,7 @@ def default_prompt():
 @pytest.fixture(scope="module")
 def model_name():
     #return "DeepSeek-R1/DeepSeek-R1-Distill-Qwen-7B"
-    return "gpt-oss-20b"
+    return "gpt_oss/gpt-oss-20b"
 
 
 @pytest.fixture(scope="module", params=['pytorch'])
@@ -68,7 +65,6 @@ def create_trtoai_worker(model_name, async_client):
     )
 
 
-@_TRTOAI_PROCESSPOOL_NVBUG
 @pytest.mark.asyncio(loop_scope="module")
 def test_trtoai_worker_generation(default_prompt, model_name, server):
     worker = create_trtoai_worker(model_name, server.get_async_client())
@@ -82,7 +78,6 @@ def test_trtoai_worker_generation(default_prompt, model_name, server):
         raise e
 
 
-@_SCAFFOLDING_MCP_NVBUG
 @pytest.mark.asyncio(loop_scope="module")
 def test_trtoai_worker_chat(default_prompt, model_name, server):
     worker = create_trtoai_worker(model_name, server.get_async_client())