[https://nvbugs/6337235][test] Remove stale MX/GMS model loader skips

[chienchunhung]

Description

Remove stale pytest skips that were added by #15368 for the MX/GMS model-loader coverage gap tracked by https://nvbugs/6337235.

#15471 merged before #15368 and fixed the stale GMS/MX ModelLoader fixtures, so the blanket skip no longer needs to suppress these tests. Unskipping the MX checkpoint-loader tests also exposed stale P2P-path fixtures: current MXCheckpointLoader.load_weights() now requires a matching SourceIdentity before it proceeds with MX transfer. This PR updates those branch-focused tests to provide a fixed matching identity and patch the MX metadata fetch seam, leaving the dedicated SourceIdentity gate coverage in tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py as the source of truth for compatibility behavior.

Test Coverage

• PYTHONPYCACHEPREFIX=/private/tmp/trtllm_pycache python3 -m py_compile tests/unittest/_torch/executor/test_model_loader_gms.py tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py
• git diff --check -- tests/unittest/_torch/executor/test_model_loader_gms.py tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py
• Commit hooks passed for touched files. Skipped waive list check and validate-test-lists locally because this machine uses Python 3.9 for those hooks and fails on Python 3.10 union syntax in scripts/check_test_list.py.
• Focused pytest was attempted with /opt/miniconda3/bin/pytest and PYTHONPATH=$PWD, but this local environment is missing transformers, imported by tests/unittest/conftest.py.
• CI before the fixture repair: /LLM/main/L0_MergeRequest_PR #46496 failed only in A10-PyTorch-1-cbts on the newly unskipped MX tests.

References: #15368, #15471

Summary by CodeRabbit

• Tests
  • Expanded coverage for model loading and checkpoint handling.
  • Previously skipped test cases now run by default, improving validation of MX and GMS-related paths.
  • Updated test scenarios to better cover success, fallback, and timeout cases.

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57799 [ run ] triggered by Bot. Commit: 5fa59ff Link to invocation

[tensorrt-cicd]

PR_Github #57799 [ run ] completed with state SUCCESS. Commit: 5fa59ff
/LLM/main/L0_MergeRequest_PR pipeline #46496 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast --stage-list "A10-PyTorch-1-cbts"

[tensorrt-cicd]

PR_Github #57810 [ run ] triggered by Bot. Commit: 77440c2 Link to invocation

[tensorrt-cicd]

PR_Github #57810 [ run ] completed with state FAILURE. Commit: 77440c2
/LLM/main/L0_MergeRequest_PR pipeline #46507 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast --stage-list "A10-PyTorch-1"

[tensorrt-cicd]

PR_Github #57854 [ run ] triggered by Bot. Commit: 77440c2 Link to invocation

[tensorrt-cicd]

PR_Github #57854 [ run ] completed with state SUCCESS. Commit: 77440c2
/LLM/main/L0_MergeRequest_PR pipeline #46550 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57909 [ run ] triggered by Bot. Commit: 77440c2 Link to invocation

[tensorrt-cicd]

PR_Github #57909 [ run ] completed with state FAILURE. Commit: 77440c2
/LLM/main/L0_MergeRequest_PR pipeline #46602 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58044 [ run ] triggered by Bot. Commit: 88f9cf4 Link to invocation

[tensorrt-cicd]

PR_Github #58044 [ run ] completed with state SUCCESS. Commit: 88f9cf4
/LLM/main/L0_MergeRequest_PR pipeline #46712 completed with status: 'SUCCESS'

CI Report

Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

Removes a skip marker from GMS model loader tests, enabling them to run. Refactors MX checkpoint loader tests to introduce a shared SourceIdentity fixture and helper for stubbing metadata fetch, and removes NVBUG skip decorators from timeout-related tests.

Changes

Test suite re-enablement and MX loader test refactor

Layer / File(s)	Summary
Unskip GMS model loader tests tests/unittest/_torch/executor/test_model_loader_gms.py	Removes the module-level pytestmark = pytest.mark.skip(...) so tests run unconditionally.
SourceIdentity fixture and helper tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py	Adds _SOURCE_IDENTITY and _allow_mx_source_identity(loader) helper stubbing _fetch_source_identity and returning kwargs for load_weights; removes _MX_CHECKPOINT_LOADER_NVBUG marker.
MX P2P and mixed-success tests tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py	Updates success/mixed-success tests to use the new helper and pass source_identity, removing direct _source_identity_compatible mocking; refines fallback tensor mock with numel()/element_size().
MX timeout tests unskipped tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py	Removes NVBUG decorators from default and configured timeout tests, replacing _source_identity_compatible mocking with source_identity_kwargs.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#15387: Directly related to the MXCheckpointLoader test refactor using the SourceIdentity fixture/seam and source-identity compatibility logic.

Suggested reviewers: Funatiq, yuxianq, xxi-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the PR and clearly states the stale MX/GMS model-loader skip removal.
Description check	✅ Passed	The description includes the issue, fix, and test coverage, with enough detail for review.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py (1)

489-494: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add type annotations to the new helper.

_allow_mx_source_identity lacks parameter and return type annotations, per repo guideline requiring functions to always be annotated.

📝 Suggested annotation fix

-def _allow_mx_source_identity(loader):
+def _allow_mx_source_identity(loader: MXCheckpointLoader) -> dict[str, SourceIdentity]:
     """Let MX P2P path tests bypass the metadata fetch seam with a matching identity."""
     loader._fetch_source_identity = MagicMock(return_value=_SOURCE_IDENTITY)
     return {"source_identity": _SOURCE_IDENTITY}

As per coding guidelines, "Always annotate functions with return types (use None if no return); annotate class members and variables when necessary".

🤖 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/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py`
around lines 489 - 494, Add type annotations to the new helper
_allow_mx_source_identity so it follows the repository rule that all functions
must be annotated. Specify the loader parameter type for the checkpoint loader
object being patched, and add an explicit return type for the dictionary it
returns; keep the existing behavior of assigning MagicMock to
loader._fetch_source_identity and returning the source_identity mapping.

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.

Nitpick comments:
In `@tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py`:
- Around line 489-494: Add type annotations to the new helper
_allow_mx_source_identity so it follows the repository rule that all functions
must be annotated. Specify the loader parameter type for the checkpoint loader
object being patched, and add an explicit return type for the dictionary it
returns; keep the existing behavior of assigning MagicMock to
loader._fetch_source_identity and returning the source_identity mapping.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c32ec541-11f6-4295-adff-dc0dea723184

📥 Commits

Reviewing files that changed from the base of the PR and between 148bf8a and 88f9cf4.

📒 Files selected for processing (2)

• tests/unittest/_torch/executor/test_model_loader_gms.py
• tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py

💤 Files with no reviewable changes (1)

• tests/unittest/_torch/executor/test_model_loader_gms.py