feat(client): add TensorRT-LLM engine adapter - #565
Conversation
fb5ea64 to
bf9d01c
Compare
bf9d01c to
bc27d6b
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
WalkthroughChangesThe pull request adds a TensorRT-LLM adapter and loader using the shared strategy chain, broadens NIXL and RDMA compatibility, improves UCX device filtering, updates model feature detection, and revises the TensorRT-LLM Kubernetes image, deployment, documentation, and tests. TensorRT-LLM runtime integration
Estimated code review effort: 4 (Complex) | ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (1)
modelexpress_client/python/modelexpress/engines/trtllm/adapter.py (1)
27-27: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueAlign these alias markers with the supported TRT-LLM release.
next_attnandnext_layer_layernormare current TRT-LLM internal wiring names used by several PyTorch backend models, so an upstream rename would make supported checkpoints hit the fail-closedRuntimeErrorpath. Document the TRT-LLM version/commit these names come from, or centralize the marker list with the backend compatibility check.🤖 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 `@modelexpress_client/python/modelexpress/engines/trtllm/adapter.py` at line 27, Update _RUNTIME_ALIAS_COMPONENTS to be explicitly tied to the supported TRT-LLM release: document the TRT-LLM version or commit that defines next_attn and next_layer_layernorm, or reuse a centralized marker list from the backend compatibility check. Preserve these markers for supported checkpoints so they do not incorrectly reach the fail-closed RuntimeError path.
🤖 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.
Inline comments:
In `@examples/p2p_transfer_k8s/client/trtllm/Dockerfile`:
- Around line 4-5: Resolve the incomplete release-image TODO in the Dockerfile
by selecting and validating a TRT-LLM release image that includes the complete
upstream ModelExpress integration before presenting this as a production-style
deployment; alternatively, clearly mark and gate the example as experimental in
the associated deployment configuration and documentation.
- Around line 6-7: Update the Dockerfile after all root-only build steps to
create and select an explicit non-root user and group, then ensure that user can
access /models, /dev/shm, GPU/RDMA devices, and IPC_LOCK at runtime. Preserve
required root privileges during image setup and add the USER declaration only
for the final serving container.
In `@examples/p2p_transfer_k8s/client/trtllm/trtllm-single-node-p2p.yaml`:
- Around line 58-64: Update the workload’s pod anti-affinity from
requiredDuringSchedulingIgnoredDuringExecution to
preferredDuringSchedulingIgnoredDuringExecution so single-node rollouts can
schedule replacement pods; preserve the app label selector and hostname topology
key.
In `@modelexpress_client/python/modelexpress/_nixl.py`:
- Around line 12-18: Update load_nixl_api so it continues only when the
candidate package or its _api submodule is genuinely unavailable. Distinguish
missing-module errors for the current package from ImportError exceptions raised
by dependencies or binary loading, and re-raise the latter instead of trying
subsequent candidates.
In `@modelexpress_client/python/modelexpress/engines/trtllm/loader.py`:
- Around line 115-131: Update cleanup so unpublish_metadata(self._ctx) is
wrapped in best-effort exception handling, logging any failure with
logger.warning while allowing execution to continue. Preserve the existing NIXL
manager shutdown and MX client close cleanup paths so they always run even when
metadata unpublishing fails.
- Around line 99-113: Broaden the exception handling in publish_model to catch
exceptions from synchronization, tensor registration, and metadata publication
so best-effort publishing never propagates failures to the caller. Preserve the
existing warning message, rank context, and continuation behavior while
including the caught exception details.
In `@modelexpress_client/python/modelexpress/transfer_safety.py`:
- Around line 48-53: Update the dtype extraction in the model configuration
feature-building logic to fall back to model_config.torch_dtype when
model_config.dtype is missing or None, not only when the attribute is absent.
Ensure features["dtype"] records the resolved dtype string rather than "None",
while preserving the existing torch. prefix removal.
---
Nitpick comments:
In `@modelexpress_client/python/modelexpress/engines/trtllm/adapter.py`:
- Line 27: Update _RUNTIME_ALIAS_COMPONENTS to be explicitly tied to the
supported TRT-LLM release: document the TRT-LLM version or commit that defines
next_attn and next_layer_layernorm, or reuse a centralized marker list from the
backend compatibility check. Preserve these markers for supported checkpoints so
they do not incorrectly reach the fail-closed RuntimeError path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e26973c3-4999-45b9-809f-20441c112a13
📒 Files selected for processing (29)
ci/TEST_PLAN.mdexamples/p2p_transfer_k8s/README.mdexamples/p2p_transfer_k8s/client/README.mdexamples/p2p_transfer_k8s/client/trtllm/Dockerfileexamples/p2p_transfer_k8s/client/trtllm/README.mdexamples/p2p_transfer_k8s/client/trtllm/trtllm-single-node-p2p.yamlmodelexpress_client/python/README.mdmodelexpress_client/python/modelexpress/__init__.pymodelexpress_client/python/modelexpress/_nixl.pymodelexpress_client/python/modelexpress/adapter.pymodelexpress_client/python/modelexpress/engines/trtllm/__init__.pymodelexpress_client/python/modelexpress/engines/trtllm/adapter.pymodelexpress_client/python/modelexpress/engines/trtllm/loader.pymodelexpress_client/python/modelexpress/gds_transfer.pymodelexpress_client/python/modelexpress/load_strategy/base.pymodelexpress_client/python/modelexpress/load_strategy/context.pymodelexpress_client/python/modelexpress/load_strategy/default_strategy.pymodelexpress_client/python/modelexpress/load_strategy/rdma_strategy.pymodelexpress_client/python/modelexpress/nixl_transfer.pymodelexpress_client/python/modelexpress/transfer_safety.pymodelexpress_client/python/modelexpress/trtllm_live_transfer.pymodelexpress_client/python/modelexpress/ucx_utils.pymodelexpress_client/python/tests/test_nixl_import.pymodelexpress_client/python/tests/test_transfer_safety.pymodelexpress_client/python/tests/test_trtllm_adapter.pymodelexpress_client/python/tests/test_trtllm_live_transfer.pymodelexpress_client/python/tests/test_trtllm_loader.pymodelexpress_client/python/tests/test_trtllm_native_example.pymodelexpress_client/python/tests/test_ucx_utils.py
💤 Files with no reviewable changes (2)
- modelexpress_client/python/tests/test_trtllm_live_transfer.py
- modelexpress_client/python/modelexpress/trtllm_live_transfer.py
Signed-off-by: Zheng Luo <zheluo@nvidia.com>
bc27d6b to
b7b42d5
Compare
Signed-off-by: Zheng Luo <zheluo@nvidia.com> (cherry picked from commit 3542dda)
Summary
modelexpress.engines.trtllm.TrtllmAdapterandMxModelLoadersoTensorRT-LLM uses the same
LoadStrategyChainas vLLM and SGLang.SourceIdentityand include the completeTRT compatibility fingerprint and post-transform protocol in MX discovery.
publication, republish, and cleanup to shared ModelExpress components.
TensorRT-LLM's native loader return contract on fallback.
trtllm_live_transferintegration path.installing a second NIXL stack.
Kubernetes RDMA device plugin.
current Llama-only qualification scope.
Why
TensorRT-LLM already owns model construction, compatibility identity, Llama
post-transform qualification, and post-load lifecycle. ModelExpress should own
the transport and source-strategy behavior shared across engines. Keeping two
independent implementations made fallback, cleanup, observability, and
multi-engine behavior drift over time.
This PR establishes one MX-owned strategy path while retaining TensorRT-LLM as
the authority for model-specific correctness.
The RDMA device-plugin fix addresses a production container failure found
during E2E testing:
/sys/class/infinibandexposed every host NIC, but the podwas allocated only one
/dev/infiniband/uverbsNdevice. Auto-pinning couldselect an inaccessible host NIC and make NIXL fail with
Destination is unreachable.Cross-repository dependency
The TensorRT-LLM bridge is NVIDIA/TensorRT-LLM#17029. That PR imports
modelexpress.engines.trtllm.MxModelLoaderand therefore depends on this PR.The initial supported runtime scope remains:
NOT STARTEDThis draft does not claim PP/EP, speculative decoding, multi-node TP, or
PD-disaggregated production readiness.
Validation
87 passed4 passed30 passedgit diff --checkBefore marking ready, rerun the cross-node E2E from the final PR commits and
complete the planned TP=4 CI qualification.
Summary by CodeRabbit
New Features
Documentation