[https://nvbugs/6337229, https://nvbugs/6341072][fix] Stabilize scaffolding MCP worker test server startup and cleanup#15958
Conversation
GenerationTask no longer has a `result` field; its output is stored in `output_str`/`output_tokens`. Update the majority-vote result copy to use `output_str`, fixing test_majority_vote. Signed-off-by: kleinc <kleinc@nvidia.com>
…eController fix Remove the _SCAFFOLDING_EXECUTOR_NVBUG skip marker now that the underlying AttributeError in MajorityVoteController.process is fixed. Verified passing: test_unbatched_scaffolding_sync, test_batched_scaffolding_sync, test_async_scaffolding_generation, test_majority_vote, and test_trtllm_worker_generation. Signed-off-by: kleinc <kleinc@nvidia.com>
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
…erver startup and cleanup Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughChanges the majority-vote controller to write output_str instead of result on the winning task. Adds a /health endpoint to the test MCP server and readiness-wait logic in RemoteMCPServer with dynamic port allocation. Updates test cleanup to async_shutdown and removes NVBUG skip decorators, re-enabling previously skipped tests. ChangesScaffolding controller and test infrastructure updates
Estimated code review effort: 2 (Simple) | ~12 minutes Sequence Diagram(s)sequenceDiagram
participant Fixture as mcp_server fixture
participant RemoteMCPServer
participant StarletteApp as Starlette App
Fixture->>RemoteMCPServer: start server on dynamic port
RemoteMCPServer->>StarletteApp: GET /health
alt server ready
StarletteApp-->>RemoteMCPServer: JSON OK response
RemoteMCPServer-->>Fixture: initialization complete
else server fails to start
RemoteMCPServer->>RemoteMCPServer: terminate process
RemoteMCPServer-->>Fixture: raise RuntimeError
end
Related Issues: None referenced. Related PRs: None referenced. Suggested labels: test, scaffolding Suggested reviewers: None specified. 🐰 A hop through skipped tests, now free to run, 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/scaffolding/test_mcp_worker.py (1)
346-372: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMissing try/finally leaves
async_shutdown()unreached on assertion failure.Every other test in this file wraps worker usage in
try/finallyso cleanup runs even if an assertion fails (see lines 237-261, 273-290, 306-342). This test does not, so a failed assertion (e.g. Lines 361-366) will skipawait worker.async_shutdown()entirely, leaking the worker's background tasks — directly undermining this PR's goal of stabilizing worker cleanup.🔧 Proposed fix
async def test_multiple_calls_to_same_tool(mcp_server): sse_url = mcp_server.get_sse_url() worker = MCPWorker(urls=[sse_url]) await worker.init_in_asyncio_event_loop() - task1 = MCPCallTask() - task1.tool_name = "add_numbers" - task1.args = json.dumps({"a": 10, "b": 20}) - - task2 = MCPCallTask() - task2.tool_name = "add_numbers" - task2.args = json.dumps({"a": 100, "b": 200}) - - status1 = await worker.run_task(task1) - status2 = await worker.run_task(task2) - assert status1 == TaskStatus.SUCCESS - assert status2 == TaskStatus.SUCCESS - result1 = int(task1.result_str) - result2 = int(task2.result_str) - assert result1 == 30, f"Expected 30, but got {result1}" - assert result2 == 300, f"Expected 300, but got {result2}" - - print("✓ Test passed: Multiple calls to same tool successful") - print(f" - Call 1: add_numbers(10, 20) = {result1}") - print(f" - Call 2: add_numbers(100, 200) = {result2}") - - await worker.async_shutdown() + try: + task1 = MCPCallTask() + task1.tool_name = "add_numbers" + task1.args = json.dumps({"a": 10, "b": 20}) + + task2 = MCPCallTask() + task2.tool_name = "add_numbers" + task2.args = json.dumps({"a": 100, "b": 200}) + + status1 = await worker.run_task(task1) + status2 = await worker.run_task(task2) + assert status1 == TaskStatus.SUCCESS + assert status2 == TaskStatus.SUCCESS + result1 = int(task1.result_str) + result2 = int(task2.result_str) + assert result1 == 30, f"Expected 30, but got {result1}" + assert result2 == 300, f"Expected 300, but got {result2}" + + print("✓ Test passed: Multiple calls to same tool successful") + print(f" - Call 1: add_numbers(10, 20) = {result1}") + print(f" - Call 2: add_numbers(100, 200) = {result2}") + finally: + await worker.async_shutdown()As per path instructions, "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM" — this gap in cleanup coverage should be closed before merge.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/scaffolding/test_mcp_worker.py` around lines 346 - 372, The test in MCPWorker cleanup is missing a try/finally, so an assertion failure can skip async_shutdown() and leave background tasks running. Wrap the body of test_multiple_calls_to_same_tool with the same try/finally pattern used by the other tests in test_mcp_worker.py, keeping worker.run_task, the assertions, and result checks in the try block and moving worker.async_shutdown() into the finally block so cleanup always runs.Source: Path instructions
🧹 Nitpick comments (1)
tests/unittest/scaffolding/test_worker.py (1)
36-36: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove stale commented-out code.
Leftover commented line for the old DeepSeek path adds no value now that the fixture returns a different model. As per coding guidelines, comments should be reserved for things needing explanation rather than disabled code.
🧹 Proposed cleanup
def model_name(): - `#return` "DeepSeek-R1/DeepSeek-R1-Distill-Qwen-7B" return "gpt_oss/gpt-oss-20b"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/scaffolding/test_worker.py` at line 36, Remove the stale commented-out DeepSeek return line from the test fixture in test_worker.py. Update the fixture that returns the model path so it only contains the active return value and no disabled code; keep the cleanup focused around the fixture or helper that provides the model string.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tests/unittest/scaffolding/test_mcp_worker.py`:
- Around line 346-372: The test in MCPWorker cleanup is missing a try/finally,
so an assertion failure can skip async_shutdown() and leave background tasks
running. Wrap the body of test_multiple_calls_to_same_tool with the same
try/finally pattern used by the other tests in test_mcp_worker.py, keeping
worker.run_task, the assertions, and result checks in the try block and moving
worker.async_shutdown() into the finally block so cleanup always runs.
---
Nitpick comments:
In `@tests/unittest/scaffolding/test_worker.py`:
- Line 36: Remove the stale commented-out DeepSeek return line from the test
fixture in test_worker.py. Update the fixture that returns the model path so it
only contains the active return value and no disabled code; keep the cleanup
focused around the fixture or helper that provides the model string.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 88b3970d-ea66-4489-a7a7-918266b0a6a1
📒 Files selected for processing (4)
tensorrt_llm/scaffolding/controller.pytests/unittest/scaffolding/test_mcp_worker.pytests/unittest/scaffolding/test_scaffolding.pytests/unittest/scaffolding/test_worker.py
💤 Files with no reviewable changes (1)
- tests/unittest/scaffolding/test_scaffolding.py
|
PR_Github #57687 [ run ] triggered by Bot. Commit: |
|
PR_Github #57687 [ run ] completed with state
|
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
|
/bot run --disable-fail-fast |
Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
KleinBlueC
force-pushed
the
fix/scaffolding_mcp_worker_test
branch
from
43b35ff to
681a1ce
Compare
|
PR_Github #57729 [ run ] triggered by Bot. Commit: |
|
PR_Github #57729 [ run ] completed with state |
KleinBlueC
changed the title
…erver startup and cleanup (NVIDIA#15958) Signed-off-by: kleinc <kleinc@nvidia.com> Signed-off-by: KleinBlueC <102355066+KleinBlueC@users.noreply.github.com>
3 participants
Bug Fixes
Tests
Fix
Description
Rewrite tested MCP worker server in scaffolding unit tests for stability and revise model path resolution.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.