[TRTLLM-14727][test] Create MX donor-receiver qualification test harness

docs/source/features/model-express.md

@@ -74,6 +74,56 @@ Support for another model family requires a focused qualification change:
    Compare deterministic output token IDs with the standard Hugging Face load
    path before documenting the family as supported.
 
+### Qualification Test
+
+The reusable GPU harness is
+`tests/integration/defs/model_express/test_model_express.py`. It launches an HF
+baseline, a live MX donor, and an MX receiver on disjoint GPU sets. The receiver
+uses a metadata-only view of the donor's canonical snapshot and contains no
+weight shards. A positive result therefore requires direct transfer; disk
+fallback cannot accidentally satisfy the test.
+
+Run the TP=1 smoke test against an isolated ModelExpress 0.4.1 service with
+NIXL enabled:
+
+```bash
+TRTLLM_MX_E2E_REQUIRED=1 \
+MODEL_EXPRESS_URL=http://127.0.0.1:8001 \
+LLM_MODELS_ROOT=/path/to/llm-models \
+pytest -v tests/integration/defs/model_express/test_model_express.py \
+  -k llama-bf16-tp1
+```
+
+Run the TP=2 rank-mapping qualification on four GPUs by selecting
+`llama-bf16-tp2`. `TRTLLM_MX_LLAMA_MODEL` can override the default TinyLlama
+checkpoint path. `TRTLLM_MX_E2E_REQUIRED=1` converts missing service, model,
+or NIXL prerequisites from skips into failures and must be set by a CI
+qualification stage. That stage must also allocate the GPUs declared by the
+selected test row.
+
+The dedicated H100 CI stages own isolated Redis and ModelExpress 0.4.1
+sidecars. The two-GPU TP=1 stage is classified as multi-GPU: it runs
+automatically in post-merge pipelines or when a multi-GPU file changes, while
+direct pre-merge dispatch requires the `ci: full pre-merge approved` label.
+Trigger it directly with:
+
+```text
+/bot run --stage-list "DGX_H100-2_GPUs-PyTorch-ModelExpress-1"
+```
+
+TP=2 is the minimum evidence for adding or changing a parallel profile. Its
+four-GPU stage is intentionally on demand and does not join ordinary
+multi-GPU runs:
+
+```text
+/bot run --stage-list "DGX_H100-4_GPUs-PyTorch-ModelExpress-OnDemand-1"
+```
+
+Both stages set `TRTLLM_MX_E2E_REQUIRED=1`, so missing service, model, client,
+or NIXL prerequisites fail instead of skipping. Do not add every model profile
+to recurring coverage: use the harness for representative rows claimed by the
+support table and keep wider matrices in scheduled qualification.
+
 ### Transform-Layout ABI Rules
 
 An existing transform-layout ABI ID is immutable. Introduce a new ID when a

jenkins/L0_MergeRequest.groovy

@@ -1094,6 +1094,7 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars)
         "tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py",
         "tensorrt_llm/_torch/pyexecutor/model_engine.py",
         "tensorrt_llm/_torch/pyexecutor/py_executor.py",
+        "tensorrt_llm/_torch/weight_sharing/",
         "tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py",
         "tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py",
         "tensorrt_llm/_torch/visual_gen/modules/vae/",
@@ -1126,6 +1127,7 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars)
         "tensorrt_llm/serve/openai_server.py",
         "tensorrt_llm/serve/router.py",
         "tests/integration/defs/cpp/test_multi_gpu.py",
+        "tests/integration/defs/model_express/",
         "tests/integration/test_lists/test-db/l0_b200_multi_gpus_perf_sanity.yml",
         "tests/integration/test_lists/test-db/l0_b200_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node1_gpu8.yml",
         "tests/integration/test_lists/test-db/l0_b200_visual_gen_perf_sanity.yml",
@@ -1155,6 +1157,7 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars)
         "tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node2_gpu8.yml",
         "tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_ctx1_node1_gpu4_gen1_node4_gpu16.yml",
         "tests/integration/test_lists/test-db/l0_gb300_multi_nodes_perf_sanity_node2_gpu8.yml",
+        "tests/integration/test_lists/test-db/l0_model_express.yml",
         "tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml",
         "tests/integration/test_lists/test-db/l0_verl.yml",
         "tests/unittest/auto_deploy/multigpu",

jenkins/L0_Test.groovy

@@ -73,6 +73,11 @@ ARTIFACTORY_CREDENTIALS_ID = "trtllm-artifactory-credentials"
 // DLFW torch image
 DLFW_IMAGE = "urm.nvidia.com/docker/nvidia/pytorch:26.05-py3"
 
+MODEL_EXPRESS_VERSION = "0.4.1"
+MODEL_EXPRESS_NIXL_VERSION = "1.3.1"
+MODEL_EXPRESS_SERVER_IMAGE = "urm.nvidia.com/docker/nvidia/ai-dynamo/modelexpress-server:${MODEL_EXPRESS_VERSION}"
+MODEL_EXPRESS_REDIS_IMAGE = "urm.nvidia.com/docker/redis:7-alpine"
+
 //Ubuntu base image
 UBUNTU_22_04_IMAGE = "urm.nvidia.com/docker/ubuntu:22.04"
 UBUNTU_24_04_IMAGE = "urm.nvidia.com/docker/ubuntu:24.04"
@@ -3137,7 +3142,7 @@ def cacheErrorAndUploadResult(stageName, taskRunner, finallyRunner, noResultIfSu
     }
 }
 
-def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMode = false)
+def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMode = false, modelExpress = false)
 {
     def targetCloud = "kubernetes-cpu"
     def selectors = """
@@ -3148,6 +3153,8 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod
     def nodeLabelPrefix = ""
     def tolerations = ""
     def extraDeviceEnv = ""
+    def serviceInitContainerConfig = ""
+    def serviceContainerConfig = ""
 
     def archSuffix = arch == "arm64" ? "arm" : "amd"
     def jnlpImage = "artifactory.pdx.nvidia.com/sw-ipp-blossom-sre-docker-local/lambda/custom_jnlp_images_${archSuffix}_linux:jdk17"
@@ -3253,9 +3260,14 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod
         if (hasMultipleGPUs)
         {
             // Not a hard requirement, but based on empirical values.
-            memorySize = "${gpuCount * 150}" + "Gi"
-            storageSize = "${gpuCount * 150}" + "Gi"
-            cpuCount = "${gpuCount * 12}"
+            // Keep ModelExpress services inside the existing pod resource envelope;
+            // otherwise their requests can make an otherwise valid GPU pod unschedulable.
+            def serviceCpuReserve = modelExpress ? 6 : 0
+            def serviceMemoryReserveGi = modelExpress ? 16 : 0
+            def serviceStorageReserveGi = modelExpress ? 8 : 0
+            memorySize = "${gpuCount * 150 - serviceMemoryReserveGi}" + "Gi"
+            storageSize = "${gpuCount * 150 - serviceStorageReserveGi}" + "Gi"
+            cpuCount = "${gpuCount * 12 - serviceCpuReserve}"
         }
 
         def gpuType = KubernetesManager.selectGPU(type)
@@ -3346,6 +3358,65 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod
                         - SYS_ADMIN"""
         break
     }
+    if (modelExpress) {
+        if (arch != "amd64") {
+            throw new Exception("ModelExpress CI sidecars currently support amd64 test pods only.")
+        }
+        extraDeviceEnv += """
+                    - name: MODEL_EXPRESS_URL
+                      value: "http://127.0.0.1:8001"
+                    - name: TRTLLM_MX_E2E_REQUIRED
+                      value: "1"
+        """
+        // Mirrors the ModelExpress v0.4.1 Redis deployment and image contract.
+        // The image exposes /app/modelexpress-server and accepts the port/backend settings below.
+        // Use regular containers because the Jenkins Kubernetes launcher does not
+        // reliably attach to pods containing restartable init-container sidecars.
+        // The server waits for Redis, and the E2E preflight waits for port 8001.
+        serviceContainerConfig = """
+                  - name: redis
+                    image: ${MODEL_EXPRESS_REDIS_IMAGE}
+                    args: ["--save", "", "--appendonly", "no"]
+                    ports:
+                    - containerPort: 6379
+                    resources:
+                      requests:
+                        cpu: '1'
+                        memory: 4Gi
+                        ephemeral-storage: 2Gi
+                      limits:
+                        cpu: '1'
+                        memory: 4Gi
+                        ephemeral-storage: 2Gi
+                    imagePullPolicy: Always
+                  - name: model-express-server
+                    image: ${MODEL_EXPRESS_SERVER_IMAGE}
+                    command: ["/bin/bash", "-c"]
+                    args:
+                    - |
+                      until (echo > /dev/tcp/127.0.0.1/6379) >/dev/null 2>&1; do
+                        sleep 1
+                      done
+                      exec /app/modelexpress-server --port 8001
+                    env:
+                    - name: MX_METADATA_BACKEND
+                      value: "redis"
+                    - name: REDIS_URL
+                      value: "redis://127.0.0.1:6379"
+                    ports:
+                    - containerPort: 8001
+                    resources:
+                      requests:
+                        cpu: '4'
+                        memory: 8Gi
+                        ephemeral-storage: 4Gi
+                      limits:
+                        cpu: '4'
+                        memory: 8Gi
+                        ephemeral-storage: 4Gi
+                    imagePullPolicy: Always
+        """
+    }
     // Temporarily avoid an arm64 CPU builder with repeated pod DNS/JNLP failures seen in Build-SBSA #5564.
     def blockedNodeAffinity = targetCloud == "kubernetes-cpu" && arch == "arm64" ? '''
                               - key: "kubernetes.io/hostname"
@@ -3411,6 +3482,7 @@ ${blockedNodeAffinity}
                 nodeSelector: ${selectors}
                 imagePullSecrets:
                   - name: ${ARTIFACTORY_IMAGE_PULL_SECRET}
+${serviceInitContainerConfig}
                 containers:
                   ${containerConfig}
                     env:
@@ -3419,6 +3491,7 @@ ${blockedNodeAffinity}
                         fieldRef:
                           fieldPath: spec.nodeName
                     ${extraDeviceEnv}
+                  ${serviceContainerConfig}
                   - name: jnlp
                     image: ${jnlpImage}
                     args: ['\$(JENKINS_SECRET)', '\$(JENKINS_NAME)']
@@ -4435,6 +4508,13 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
             if (!skipInstallWheel) {
                 trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmPath} && pip3 install --force-reinstall --no-deps TensorRT-LLM/tensorrt_llm-*.whl")
             }
+            if (stageName.contains("-ModelExpress-")) {
+                trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install modelexpress==${MODEL_EXPRESS_VERSION}")
+                // ModelExpress 0.4.1 imports nixl._api, while requirements-dev.txt
+                // installs only the nixl-cu13 backend. Install the matching
+                // namespace shim without pulling the unused CUDA 12 backend.
+                trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install --no-deps nixl==${MODEL_EXPRESS_NIXL_VERSION}")
+            }
         }
 
         trtllm_utils.llmExecStepWithRetry(pipeline, script: "git config --global --add safe.directory \"*\"")
@@ -5468,6 +5548,9 @@ def launchTestJobs(pipeline, testFilter, globalVars)
         "H100_PCIe-PyTorch-Ray-1": ["h100-cr", "l0_h100", 1, 1],
         "H100_PCIe-AutoDeploy-1": ["h100-cr", "l0_h100", 1, 1],
         "H100_PCIe-CPP-1": ["h100-cr", "l0_h100", 1, 1],
+        // platform, test DB, split, splits, GPU count, ModelExpress sidecars
+        "DGX_H100-2_GPUs-PyTorch-ModelExpress-1": ["dgx-h100-x4", "l0_model_express", 1, 1, 2, true],
+        "DGX_H100-4_GPUs-PyTorch-ModelExpress-OnDemand-1": ["dgx-h100-x4", "l0_model_express", 1, 1, 4, true],
         "RTX5090-PyTorch-1": ["rtx-5090", "l0_gb202", 1, 1],
         "RTX5080-PyTorch-1": ["rtx-5080", "l0_gb203", 1, 2],
         "RTX5080-PyTorch-2": ["rtx-5080", "l0_gb203", 2, 2],
@@ -5503,7 +5586,7 @@ def launchTestJobs(pipeline, testFilter, globalVars)
     ]
 
     x86TestConfigs = cbtsResizeSplits(x86TestConfigs)
-    parallelJobs = x86TestConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE, values[0], "amd64", values[4] ?: 1, key.contains("-Perf-")), { attemptTag, isFinalAttempt, retryContext = null ->
+    parallelJobs = x86TestConfigs.collectEntries{key, values -> [key, [createKubernetesPodConfig(LLM_DOCKER_IMAGE, values[0], "amd64", values[4] ?: 1, key.contains("-Perf-"), values.size() > 5 ? values[5] : false), { attemptTag, isFinalAttempt, retryContext = null ->
         def config = VANILLA_CONFIG
         if (key.contains("single-device")) {
             config = SINGLE_DEVICE_CONFIG
@@ -6198,18 +6281,20 @@ def launchTestJobs(pipeline, testFilter, globalVars)
         }, {}, true)
     }]}
 
-    multiGpuJobs = parallelJobs.findAll{(it.key =~ /\d+_GPUs/) && !it.key.contains("Post-Merge")}
+    // OnDemand stages are available through --stage-list/--extra-stage only.
+    multiGpuJobs = parallelJobs.findAll{(it.key =~ /\d+_GPUs/) && !it.key.contains("Post-Merge") && !it.key.contains("-OnDemand-")}
     println multiGpuJobs.keySet()
     multiGpuJobsPostMerge = parallelJobs.findAll{(it.key =~ /\d+_GPUs/) && it.key.contains("Post-Merge")}
 
     parallelJobs += docBuildJobs
     parallelJobs += sanityCheckJobs
     parallelJobs += agentFlowTestJobs
 
+    onDemandJobs = parallelJobs.findAll {it.key.contains("-OnDemand-")}
     postMergeJobs = parallelJobs.findAll {it.key.contains("Post-Merge")}
 
     // Start as a normal pre-merge job
-    parallelJobsFiltered = parallelJobs - multiGpuJobs - postMergeJobs
+    parallelJobsFiltered = parallelJobs - multiGpuJobs - postMergeJobs - onDemandJobs
 
     // Check if the multi GPU related file has changed or not. If changed, add multi GPU test stages.
     if (testFilter[(MULTI_GPU_FILE_CHANGED)]) {
@@ -6220,7 +6305,7 @@ def launchTestJobs(pipeline, testFilter, globalVars)
         echo "AUTO_TRIGGER_TAG_LIST mode is true. Auto trigger tags: ${testFilter[(AUTO_TRIGGER_TAG_LIST)].join(', ')}."
         def autoTriggerTagStages = [:]
         for (tag in testFilter[(AUTO_TRIGGER_TAG_LIST)]) {
-            autoTriggerTagStages += parallelJobs.findAll { it.key.contains(tag) }
+            autoTriggerTagStages += (parallelJobs - onDemandJobs).findAll { it.key.contains(tag) }
         }
         parallelJobsFiltered += autoTriggerTagStages
         if (autoTriggerTagStages.size() > 0) {
@@ -6313,6 +6398,9 @@ def launchTestJobs(pipeline, testFilter, globalVars)
         }
     }
 
+    // Keep manually triggered stages out of every automatic selection path.
+    parallelJobsFiltered -= onDemandJobs
+
     // Check --stage-list, only run the stages in stage-list. Supports wildcard '*'.
     if (testFilter[TEST_STAGE_LIST] != null) {
         echo "Use TEST_STAGE_LIST for filtering. Stages: ${testFilter[(TEST_STAGE_LIST)]}."
@@ -6349,6 +6437,9 @@ def launchTestJobs(pipeline, testFilter, globalVars)
         def needsSanity = cbts.sanity_required
         def needsPerfSanity = cbts.perfsanity_required
         parallelJobsFiltered = parallelJobs.findAll { key, _ ->
+            if (key.contains("-OnDemand-")) {
+                return false
+            }
             if (key =~ /Post-Merge/) return affectedSet.contains(key)
             return affectedSet.contains(key) ||
                    (needsSanity && key =~ /PackageSanityCheck/) ||

tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py

@@ -29,6 +29,7 @@
 
 import inspect
 import json
+import logging
 import os
 import threading
 import traceback
@@ -189,6 +190,16 @@ def _synchronize_cuda_for_mx_publish() -> None:
         torch.cuda.synchronize()
 
 
+def _enable_mx_transfer_logging() -> None:
+    """Enable upstream INFO records when per-rank transfer logs are requested."""
+    if not os.environ.get("MX_TRANSFER_LOG_DIR"):
+        return
+
+    mx_logger = logging.getLogger("modelexpress")
+    if mx_logger.getEffectiveLevel() > logging.INFO:
+        mx_logger.setLevel(logging.INFO)
+
+
 @register_checkpoint_loader("MX")
 class MXCheckpointLoader(HfCheckpointLoader):
     """Checkpoint loader for MX (ModelExpress) P2P weight transfer.
@@ -363,6 +374,12 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict[
                 "`checkpoint_format` to continue without MX."
             ) from exc
 
+        # ModelExpress 0.4.1 installs an INFO-level file handler for
+        # MX_TRANSFER_LOG_DIR, but leaves its logger at Python's WARNING
+        # default outside vLLM. Enable the records in the worker that performs
+        # the transfer so the requested per-rank diagnostics are not empty.
+        _enable_mx_transfer_logging()
+
         try:
             with _MX_TRANSFER_STATE_LOCK:
                 MxClient = mx_transfer.MxClient
@@ -404,6 +421,16 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict[
             )
 
         source_registered = source_metadata is not None
+        if not source_registered and self._query_timeout_s == 0:
+            # A zero timeout explicitly disables source polling. Fall back
+            # before preparing post-transform receiver aliases: that setup
+            # mutates the module graph and is only safe when P2P will proceed.
+            return self._fallback_to_disk(
+                checkpoint_dir,
+                mapping,
+                reason="no MX source is registered and source polling is disabled",
+                **kwargs,
+            )
         if not source_registered and self._local_source_identity is not None:
             # ModelExpress 0.4.1 hashes every SourceIdentity field, including
             # extra_parameters. Proceed to MxLiveWeightLoader.load_weights()