From 2cd6cd78fb0ebb2f3d3592747bd66ed5ea21e14f Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:17:21 -0700 Subject: [PATCH 01/16] [TRTLLM-14727][test] add MX donor receiver qualification harness Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- docs/source/features/model-express.md | 47 ++ jenkins/L0_Test.groovy | 102 +++- .../defs/model_express/mx_e2e_worker.py | 109 ++++ .../defs/model_express/test_model_express.py | 486 ++++++++++++++++++ .../test_lists/test-db/l0_model_express.yml | 37 ++ 5 files changed, 776 insertions(+), 5 deletions(-) create mode 100644 tests/integration/defs/model_express/mx_e2e_worker.py create mode 100644 tests/integration/defs/model_express/test_model_express.py create mode 100644 tests/integration/test_lists/test-db/l0_model_express.yml diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index e573d2f8cb71..8c8db8ae678a 100644 --- a/docs/source/features/model-express.md +++ b/docs/source/features/model-express.md @@ -74,6 +74,53 @@ 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. Run the recurring two-GPU TP=1 smoke test 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 diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index b1a4f8a756a6..1e1cf931ffc5 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -73,6 +73,9 @@ ARTIFACTORY_CREDENTIALS_ID = "trtllm-artifactory-credentials" // DLFW torch image DLFW_IMAGE = "urm.nvidia.com/docker/nvidia/pytorch:26.05-py3" +MODEL_EXPRESS_SERVER_IMAGE = "urm.nvidia.com/docker/nvidia/ai-dynamo/modelexpress-server:0.4.1" +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 +3140,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 +3151,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" @@ -3346,6 +3351,80 @@ 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" + """ + serviceInitContainerConfig = """ + initContainers: + - name: redis + image: ${MODEL_EXPRESS_REDIS_IMAGE} + args: ["--save", "", "--appendonly", "no"] + restartPolicy: Always + ports: + - containerPort: 6379 + startupProbe: + tcpSocket: + port: 6379 + initialDelaySeconds: 1 + periodSeconds: 1 + failureThreshold: 30 + livenessProbe: + tcpSocket: + port: 6379 + periodSeconds: 10 + resources: + requests: + cpu: 500m + memory: 1Gi + ephemeral-storage: 2Gi + limits: + cpu: 500m + memory: 1Gi + ephemeral-storage: 2Gi + imagePullPolicy: Always + """ + serviceContainerConfig = """ + - name: model-express-server + image: ${MODEL_EXPRESS_SERVER_IMAGE} + command: ["/app/modelexpress-server"] + args: ["--port", "8001"] + env: + - name: MX_METADATA_BACKEND + value: "redis" + - name: REDIS_URL + value: "redis://127.0.0.1:6379" + ports: + - containerPort: 8001 + readinessProbe: + tcpSocket: + port: 8001 + initialDelaySeconds: 2 + periodSeconds: 2 + failureThreshold: 30 + livenessProbe: + tcpSocket: + port: 8001 + initialDelaySeconds: 10 + periodSeconds: 10 + resources: + requests: + cpu: '1' + memory: 4Gi + ephemeral-storage: 4Gi + limits: + cpu: '1' + memory: 4Gi + 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 +3490,7 @@ ${blockedNodeAffinity} nodeSelector: ${selectors} imagePullSecrets: - name: ${ARTIFACTORY_IMAGE_PULL_SECRET} +${serviceInitContainerConfig} containers: ${containerConfig} env: @@ -3419,6 +3499,7 @@ ${blockedNodeAffinity} fieldRef: fieldPath: spec.nodeName ${extraDeviceEnv} + ${serviceContainerConfig} - name: jnlp image: ${jnlpImage} args: ['\$(JENKINS_SECRET)', '\$(JENKINS_NAME)'] @@ -5468,6 +5549,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 +5587,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,7 +6282,8 @@ 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")} @@ -6206,10 +6291,11 @@ def launchTestJobs(pipeline, testFilter, globalVars) 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 +6306,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 +6399,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 +6438,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/) || diff --git a/tests/integration/defs/model_express/mx_e2e_worker.py b/tests/integration/defs/model_express/mx_e2e_worker.py new file mode 100644 index 000000000000..269810fcb672 --- /dev/null +++ b/tests/integration/defs/model_express/mx_e2e_worker.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Run one process role for the ModelExpress donor/receiver E2E test.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import KvCacheConfig + +_PROMPT_TOKEN_IDS = ( + (1, 42, 7, 9), + (1, 17, 23, 5, 11), +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--role", choices=("baseline", "donor", "receiver"), required=True) + parser.add_argument("--model", required=True) + parser.add_argument("--tp-size", type=int, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--mx-url") + parser.add_argument("--ready-file", type=Path) + parser.add_argument("--stop-file", type=Path) + return parser.parse_args() + + +def _llm_kwargs(args: argparse.Namespace) -> dict[str, object]: + kwargs: dict[str, object] = { + "model": args.model, + "backend": "pytorch", + "checkpoint_format": "HF" if args.role == "baseline" else "MX", + "tensor_parallel_size": args.tp_size, + "dtype": "bfloat16", + "attn_backend": "TRTLLM", + "skip_tokenizer_init": True, + "max_batch_size": len(_PROMPT_TOKEN_IDS), + "max_num_tokens": 64, + "max_seq_len": 64, + "kv_cache_config": KvCacheConfig(free_gpu_memory_fraction=0.15), + } + if args.role != "baseline": + if not args.mx_url: + raise ValueError("MX donor and receiver roles require --mx-url") + kwargs["mx_config"] = { + "server_url": args.mx_url, + # ModelExpress 0.4.1 treats zero as no query. The donor should + # fall back immediately; the receiver starts only after donor readiness. + "server_query_timeout_s": 0 if args.role == "donor" else 30, + } + return kwargs + + +def main() -> None: + args = _parse_args() + if args.role == "donor" and (args.ready_file is None or args.stop_file is None): + raise ValueError("The donor role requires --ready-file and --stop-file") + + started = time.perf_counter() + with LLM(**_llm_kwargs(args)) as llm: + load_seconds = time.perf_counter() - started + sampling_params = SamplingParams( + max_tokens=8, + temperature=0.0, + top_k=1, + end_id=-1, + pad_id=0, + ) + results = list( + llm.generate( + [list(prompt) for prompt in _PROMPT_TOKEN_IDS], + sampling_params=sampling_params, + ) + ) + payload = { + "role": args.role, + "tp_size": args.tp_size, + "load_seconds": load_seconds, + "token_ids": [list(result.outputs[0].token_ids) for result in results], + } + args.output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + + if args.role == "donor": + assert args.ready_file is not None + assert args.stop_file is not None + args.ready_file.write_text("ready\n", encoding="utf-8") + while not args.stop_file.exists(): + time.sleep(0.2) + + +if __name__ == "__main__": + main() diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py new file mode 100644 index 000000000000..a4c026477dd1 --- /dev/null +++ b/tests/integration/defs/model_express/test_model_express.py @@ -0,0 +1,486 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Real-GPU ModelExpress donor/receiver qualification tests.""" + +from __future__ import annotations + +import json +import os +import re +import shutil +import signal +import socket +import subprocess +import sys +import time +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import NoReturn +from urllib.parse import urlparse + +import pytest +import torch + +from tensorrt_llm._torch.weight_sharing import ArtifactIdentity + +_WORKER = Path(__file__).with_name("mx_e2e_worker.py") +_WEIGHT_SUFFIXES = frozenset({".bin", ".ckpt", ".gguf", ".pt", ".pth", ".safetensors"}) +_RECEIVER_FAILURE_MARKERS = ( + "falling back to disk", + "partial fallback", + "size mismatch", + "still missing", + "mx p2p transfer failed", + "source sourceidentity incompatible", + "sourceidentity mismatch", + "invalid sourceidentity", +) +_MATCHED_PARAMS_PATTERN = re.compile(r"Matched\s+(\d+)/(\d+)\s+params", re.IGNORECASE) +_TRANSFERRED_PARAMS_PATTERN = re.compile( + r"Rank\s+(\d+):\s+transferred\s+(\d+)\s+params", + re.IGNORECASE, +) +_MX_PREFLIGHT_SCRIPT = """ +import importlib.metadata as metadata +import importlib.util + + +def module_exists(name): + try: + return importlib.util.find_spec(name) is not None + except ModuleNotFoundError: + return False + + +assert module_exists("modelexpress.trtllm_live_transfer") or module_exists( + "modelexpress.engines.trtllm" +), "no TRT-LLM adapter is available" +from modelexpress.nixl_transfer import is_nixl_available + +assert is_nixl_available(), "NIXL is unavailable" +print(f"MODELEXPRESS_VERSION={metadata.version('modelexpress')}") +""" + + +@dataclass(frozen=True) +class MxE2ECase: + """One real-model qualification row for the shared MX harness.""" + + model_env: str + default_model_subdir: str + repository_cache_prefix: str + tp_size: int + + +_MX_CASES = ( + pytest.param( + MxE2ECase( + model_env="TRTLLM_MX_LLAMA_MODEL", + default_model_subdir="llama-models-v2/TinyLlama-1.1B-Chat-v1.0", + repository_cache_prefix="models--trtllm-mx-e2e--llama-tp1", + tp_size=1, + ), + id="llama-bf16-tp1", + marks=pytest.mark.skip_less_device(2), + ), + pytest.param( + MxE2ECase( + model_env="TRTLLM_MX_LLAMA_MODEL", + default_model_subdir="llama-models-v2/TinyLlama-1.1B-Chat-v1.0", + repository_cache_prefix="models--trtllm-mx-e2e--llama-tp2", + tp_size=2, + ), + id="llama-bf16-tp2", + marks=pytest.mark.skip_less_device(4), + ), +) + + +def _qualification_required() -> bool: + return os.environ.get("TRTLLM_MX_E2E_REQUIRED") == "1" + + +def _skip_or_fail(message: str) -> NoReturn: + if _qualification_required(): + pytest.fail(message) + pytest.skip(message) + + +def _resolve_model_path(case: MxE2ECase) -> Path: + configured_path = os.environ.get(case.model_env) + if configured_path: + model_path = Path(configured_path).expanduser() + else: + configured_root = os.environ.get("LLM_MODELS_ROOT") + if configured_root: + models_root = Path(configured_root).expanduser() + else: + models_root = Path("/home/scratch.trt_llm_data_ci/llm-models") + if not models_root.exists(): + models_root = Path("/scratch.trt_llm_data/llm-models") + model_path = models_root / case.default_model_subdir + + if not model_path.is_dir(): + _skip_or_fail(f"MX E2E model directory does not exist: {model_path}. Set {case.model_env}.") + return model_path.absolute() + + +def _require_mx_environment(required_gpus: int) -> tuple[str, tuple[str, ...]]: + mx_url = os.environ.get("MODEL_EXPRESS_URL") + if not mx_url: + _skip_or_fail("MODEL_EXPRESS_URL must point to an isolated ModelExpress test service") + assert mx_url is not None + + parsed_url = urlparse(mx_url) + if parsed_url.scheme not in ("http", "https") or not parsed_url.hostname: + _skip_or_fail(f"MODEL_EXPRESS_URL is invalid: {mx_url!r}") + try: + port = parsed_url.port or (443 if parsed_url.scheme == "https" else 80) + except ValueError as error: + _skip_or_fail(f"MODEL_EXPRESS_URL is invalid: {mx_url!r}: {error}") + deadline = time.monotonic() + 30 + while True: + try: + with socket.create_connection((parsed_url.hostname, port), timeout=1): + break + except OSError as error: + if time.monotonic() >= deadline: + _skip_or_fail(f"ModelExpress service {mx_url!r} is unreachable: {error}") + time.sleep(1) + + preflight = subprocess.run( + [sys.executable, "-c", _MX_PREFLIGHT_SCRIPT], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + if preflight.returncode != 0: + detail = (preflight.stdout + preflight.stderr).strip() + _skip_or_fail(f"ModelExpress/NIXL preflight failed: {detail}") + version_lines = [ + line for line in preflight.stdout.splitlines() if line.startswith("MODELEXPRESS_VERSION=") + ] + if len(version_lines) != 1: + _skip_or_fail(f"ModelExpress preflight did not report its version: {preflight.stdout!r}") + print(f"MX E2E client preflight passed: {version_lines[0]}") + + configured_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + if configured_devices and configured_devices.strip().lower() != "all": + gpu_ids = tuple( + device.strip() for device in configured_devices.split(",") if device.strip() + ) + else: + gpu_ids = tuple(str(index) for index in range(torch.cuda.device_count())) + if len(gpu_ids) < required_gpus: + _skip_or_fail(f"MX E2E requires {required_gpus} GPUs, but only {len(gpu_ids)} are visible") + return mx_url, gpu_ids[:required_gpus] + + +def _checkpoint_files(model_path: Path) -> tuple[Path, ...]: + return tuple(sorted(path for path in model_path.rglob("*") if path.is_file())) + + +def _build_canonical_snapshot(case: MxE2ECase, model_path: Path, run_dir: Path) -> Path: + files = _checkpoint_files(model_path) + if not files: + pytest.fail(f"MX E2E model directory is empty: {model_path}") + if not any(path.suffix.lower() in _WEIGHT_SUFFIXES for path in files): + pytest.fail(f"MX E2E donor checkpoint contains no recognized weight files: {model_path}") + + source_identity = ArtifactIdentity.from_checkpoint(model_path) + run_id = uuid.uuid4().hex[:12] + repository_cache_name = f"{case.repository_cache_prefix}-{run_id}" + snapshot = ( + run_dir / "donor-hf-cache" / repository_cache_name / "snapshots" / source_identity.digest + ) + for path in files: + destination = snapshot / path.relative_to(model_path) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.symlink_to(path.resolve()) + return snapshot + + +def _build_metadata_only_snapshot(donor_snapshot: Path, run_dir: Path) -> Path: + repository_cache_name = donor_snapshot.parents[1].name + revision = donor_snapshot.name + receiver_snapshot = ( + run_dir / "receiver-hf-cache" / repository_cache_name / "snapshots" / revision + ) + + for path in _checkpoint_files(donor_snapshot): + if path.suffix.lower() in _WEIGHT_SUFFIXES: + continue + destination = receiver_snapshot / path.relative_to(donor_snapshot) + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(path, destination, follow_symlinks=True) + + if not (receiver_snapshot / "config.json").is_file(): + pytest.fail("The metadata-only MX receiver snapshot has no config.json") + receiver_weights = tuple( + path + for path in _checkpoint_files(receiver_snapshot) + if path.suffix.lower() in _WEIGHT_SUFFIXES + ) + if receiver_weights: + pytest.fail(f"The MX receiver snapshot unexpectedly contains weights: {receiver_weights}") + return receiver_snapshot + + +def _worker_command( + *, + role: str, + model_path: Path, + tp_size: int, + output_path: Path, + mx_url: str | None = None, + ready_file: Path | None = None, + stop_file: Path | None = None, +) -> list[str]: + command = [ + sys.executable, + str(_WORKER), + "--role", + role, + "--model", + str(model_path), + "--tp-size", + str(tp_size), + "--output", + str(output_path), + ] + if mx_url is not None: + command.extend(("--mx-url", mx_url)) + if ready_file is not None: + command.extend(("--ready-file", str(ready_file))) + if stop_file is not None: + command.extend(("--stop-file", str(stop_file))) + return command + + +def _worker_environment(gpu_ids: tuple[str, ...], transfer_log_dir: Path) -> dict[str, str]: + transfer_log_dir.mkdir(parents=True, exist_ok=True) + environment = os.environ.copy() + environment.update( + { + "CUDA_VISIBLE_DEVICES": ",".join(gpu_ids), + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + "MX_TRANSFER_LOG_DIR": str(transfer_log_dir), + "PYTHONUNBUFFERED": "1", + "TLLM_LOG_LEVEL": "INFO", + } + ) + return environment + + +def _log_tail(log_path: Path, lines: int = 120) -> str: + if not log_path.exists(): + return "" + return "\n".join(log_path.read_text(encoding="utf-8", errors="replace").splitlines()[-lines:]) + + +def _run_worker( + command: list[str], environment: dict[str, str], log_path: Path, timeout_s: int +) -> None: + with log_path.open("w", encoding="utf-8") as log_file: + process = subprocess.Popen( + command, + env=environment, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + try: + returncode = process.wait(timeout=timeout_s) + except subprocess.TimeoutExpired: + _signal_process_group(process, signal.SIGTERM) + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + _signal_process_group(process, signal.SIGKILL) + process.wait(timeout=30) + pytest.fail(f"MX E2E worker timed out: {' '.join(command)}\n{_log_tail(log_path)}") + if returncode != 0: + pytest.fail( + f"MX E2E worker exited with status {returncode}: " + f"{' '.join(command)}\n{_log_tail(log_path)}" + ) + + +def _wait_for_donor( + process: subprocess.Popen[bytes], ready_file: Path, log_path: Path, timeout_s: int +) -> None: + deadline = time.monotonic() + timeout_s + while not ready_file.exists(): + returncode = process.poll() + if returncode is not None: + pytest.fail( + f"MX donor exited before publication with status {returncode}\n{_log_tail(log_path)}" + ) + if time.monotonic() >= deadline: + pytest.fail(f"MX donor did not become ready within {timeout_s}s\n{_log_tail(log_path)}") + time.sleep(0.5) + + +def _signal_process_group(process: subprocess.Popen[bytes], sig: signal.Signals) -> None: + if process.poll() is not None: + return + try: + os.killpg(process.pid, sig) + except ProcessLookupError: + pass + + +def _stop_donor(process: subprocess.Popen[bytes], stop_file: Path) -> int: + if process.poll() is not None: + assert process.returncode is not None + return process.returncode + stop_file.write_text("stop\n", encoding="utf-8") + try: + return process.wait(timeout=120) + except subprocess.TimeoutExpired: + _signal_process_group(process, signal.SIGTERM) + try: + return process.wait(timeout=30) + except subprocess.TimeoutExpired: + _signal_process_group(process, signal.SIGKILL) + return process.wait(timeout=30) + + +def _load_tokens(output_path: Path) -> list[list[int]]: + payload = json.loads(output_path.read_text(encoding="utf-8")) + token_ids = payload.get("token_ids") + if not isinstance(token_ids, list) or not token_ids or not all(token_ids): + pytest.fail(f"MX E2E worker produced invalid token IDs: {payload}") + return token_ids + + +def _transfer_log_text(transfer_log_dir: Path) -> str: + logs = tuple( + path for path in transfer_log_dir.rglob("*") if path.is_file() and path.stat().st_size > 0 + ) + if not logs: + pytest.fail(f"ModelExpress created no receiver transfer logs in {transfer_log_dir}") + return "\n".join(path.read_text(encoding="utf-8", errors="replace") for path in sorted(logs)) + + +def _assert_transfer_evidence( + case: MxE2ECase, + receiver_log_path: Path, + receiver_transfer_log_dir: Path, +) -> None: + receiver_log = receiver_log_path.read_text(encoding="utf-8", errors="replace") + all_receiver_logs = receiver_log + "\n" + _transfer_log_text(receiver_transfer_log_dir) + all_receiver_logs_lower = all_receiver_logs.lower() + + for marker in _RECEIVER_FAILURE_MARKERS: + assert marker not in all_receiver_logs_lower + + matched_params = _MATCHED_PARAMS_PATTERN.findall(all_receiver_logs) + assert len(matched_params) >= case.tp_size + assert all(int(matched) == int(total) > 0 for matched, total in matched_params) + + transferred_params = _TRANSFERRED_PARAMS_PATTERN.findall(all_receiver_logs) + transferred_ranks = {int(rank) for rank, count in transferred_params if int(count) > 0} + assert transferred_ranks.issuperset(range(case.tp_size)) + + +@pytest.mark.parametrize("case", _MX_CASES) +def test_mx_donor_receiver(case: MxE2ECase, tmp_path: Path) -> None: + """Compare HF, MX donor, and no-weight-shards MX receiver outputs.""" + required_gpus = case.tp_size * 2 + mx_url, gpu_ids = _require_mx_environment(required_gpus) + model_path = _resolve_model_path(case) + timeout_s = int(os.environ.get("TRTLLM_MX_E2E_TIMEOUT_S", "1200")) + + donor_snapshot = _build_canonical_snapshot(case, model_path, tmp_path) + receiver_snapshot = _build_metadata_only_snapshot(donor_snapshot, tmp_path) + donor_identity = ArtifactIdentity.from_checkpoint(donor_snapshot) + receiver_identity = ArtifactIdentity.from_checkpoint(receiver_snapshot) + assert donor_identity.scheme == "hf_snapshot_revision" + assert receiver_identity == donor_identity + + donor_gpu_ids = gpu_ids[: case.tp_size] + receiver_gpu_ids = gpu_ids[case.tp_size :] + baseline_output = tmp_path / "baseline.json" + donor_output = tmp_path / "donor.json" + receiver_output = tmp_path / "receiver.json" + baseline_log = tmp_path / "baseline.log" + donor_log = tmp_path / "donor.log" + receiver_log = tmp_path / "receiver.log" + donor_ready = tmp_path / "donor.ready" + donor_stop = tmp_path / "donor.stop" + + # Leave donor GPUs untouched in case MPI worker teardown trails the baseline process. + _run_worker( + _worker_command( + role="baseline", + model_path=donor_snapshot, + tp_size=case.tp_size, + output_path=baseline_output, + ), + _worker_environment(receiver_gpu_ids, tmp_path / "baseline-transfer-logs"), + baseline_log, + timeout_s, + ) + + donor_environment = _worker_environment(donor_gpu_ids, tmp_path / "donor-transfer-logs") + donor_returncode = None + with donor_log.open("w", encoding="utf-8") as donor_log_file: + donor_process = subprocess.Popen( + _worker_command( + role="donor", + model_path=donor_snapshot, + tp_size=case.tp_size, + output_path=donor_output, + mx_url=mx_url, + ready_file=donor_ready, + stop_file=donor_stop, + ), + env=donor_environment, + stdout=donor_log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + try: + _wait_for_donor(donor_process, donor_ready, donor_log, timeout_s) + _run_worker( + _worker_command( + role="receiver", + model_path=receiver_snapshot, + tp_size=case.tp_size, + output_path=receiver_output, + mx_url=mx_url, + ), + _worker_environment(receiver_gpu_ids, tmp_path / "receiver-transfer-logs"), + receiver_log, + timeout_s, + ) + finally: + donor_returncode = _stop_donor(donor_process, donor_stop) + + assert donor_returncode == 0, ( + f"MX donor exited with status {donor_returncode}\n{_log_tail(donor_log)}" + ) + baseline_tokens = _load_tokens(baseline_output) + assert _load_tokens(donor_output) == baseline_tokens + assert _load_tokens(receiver_output) == baseline_tokens + _assert_transfer_evidence( + case, + receiver_log, + tmp_path / "receiver-transfer-logs", + ) diff --git a/tests/integration/test_lists/test-db/l0_model_express.yml b/tests/integration/test_lists/test-db/l0_model_express.yml new file mode 100644 index 000000000000..69adb45bda1b --- /dev/null +++ b/tests/integration/test_lists/test-db/l0_model_express.yml @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +version: 0.0.1 +l0_model_express: +- condition: + ranges: + system_gpu_count: + gte: 2 + lte: 2 + wildcards: + gpu: + - '*h100*' + linux_distribution_name: ubuntu* + cpu: x86_64 + terms: + stage: pre_merge + backend: pytorch + orchestrator: mpi + tests: + - model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp1] +- condition: + ranges: + system_gpu_count: + gte: 4 + lte: 4 + wildcards: + gpu: + - '*h100*' + linux_distribution_name: ubuntu* + cpu: x86_64 + terms: + stage: pre_merge + backend: pytorch + orchestrator: mpi + tests: + - model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp2] From 52f786f7bf5ed6ee73c08bd14f383ff4ed139d26 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 10 Aug 2026 14:40:14 -0700 Subject: [PATCH 02/16] [TRTLLM-14727][fix] address MX harness review feedback Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- docs/source/features/model-express.md | 5 +- .../defs/model_express/mx_e2e_worker.py | 12 +++- .../defs/model_express/test_model_express.py | 59 ++++++++----------- 3 files changed, 36 insertions(+), 40 deletions(-) diff --git a/docs/source/features/model-express.md b/docs/source/features/model-express.md index 8c8db8ae678a..4b5d809081ab 100644 --- a/docs/source/features/model-express.md +++ b/docs/source/features/model-express.md @@ -102,7 +102,10 @@ 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. Run the recurring two-GPU TP=1 smoke test with: +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" diff --git a/tests/integration/defs/model_express/mx_e2e_worker.py b/tests/integration/defs/model_express/mx_e2e_worker.py index 269810fcb672..ce5605bd0a3b 100644 --- a/tests/integration/defs/model_express/mx_e2e_worker.py +++ b/tests/integration/defs/model_express/mx_e2e_worker.py @@ -61,8 +61,9 @@ def _llm_kwargs(args: argparse.Namespace) -> dict[str, object]: raise ValueError("MX donor and receiver roles require --mx-url") kwargs["mx_config"] = { "server_url": args.mx_url, - # ModelExpress 0.4.1 treats zero as no query. The donor should - # fall back immediately; the receiver starts only after donor readiness. + # ModelExpress 0.4.1 skips polling at zero but still sleeps once + # for five seconds before disk fallback. The receiver starts only + # after donor readiness, so it can use a bounded discovery window. "server_query_timeout_s": 0 if args.role == "donor" else 30, } return kwargs @@ -73,8 +74,9 @@ def main() -> None: if args.role == "donor" and (args.ready_file is None or args.stop_file is None): raise ValueError("The donor role requires --ready-file and --stop-file") + llm_kwargs = _llm_kwargs(args) started = time.perf_counter() - with LLM(**_llm_kwargs(args)) as llm: + with LLM(**llm_kwargs) as llm: load_seconds = time.perf_counter() - started sampling_params = SamplingParams( max_tokens=8, @@ -89,10 +91,14 @@ def main() -> None: sampling_params=sampling_params, ) ) + mx_config = llm_kwargs.get("mx_config") payload = { "role": args.role, "tp_size": args.tp_size, "load_seconds": load_seconds, + "server_query_timeout_s": ( + mx_config.get("server_query_timeout_s") if isinstance(mx_config, dict) else None + ), "token_ids": [list(result.outputs[0].token_ids) for result in results], } args.output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index a4c026477dd1..10a3bb16a859 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -20,7 +20,6 @@ import os import re import shutil -import signal import socket import subprocess import sys @@ -33,6 +32,7 @@ import pytest import torch +from defs.trt_test_alternative import cleanup_process_tree, popen from tensorrt_llm._torch.weight_sharing import ArtifactIdentity @@ -296,24 +296,20 @@ def _log_tail(log_path: Path, lines: int = 120) -> str: def _run_worker( command: list[str], environment: dict[str, str], log_path: Path, timeout_s: int ) -> None: - with log_path.open("w", encoding="utf-8") as log_file: - process = subprocess.Popen( - command, - env=environment, - stdout=log_file, - stderr=subprocess.STDOUT, - start_new_session=True, - ) - try: + try: + with ( + log_path.open("w", encoding="utf-8") as log_file, + popen( + command, + env=environment, + stdout=log_file, + stderr=subprocess.STDOUT, + suppress_output_info=True, + ) as process, + ): returncode = process.wait(timeout=timeout_s) - except subprocess.TimeoutExpired: - _signal_process_group(process, signal.SIGTERM) - try: - process.wait(timeout=30) - except subprocess.TimeoutExpired: - _signal_process_group(process, signal.SIGKILL) - process.wait(timeout=30) - pytest.fail(f"MX E2E worker timed out: {' '.join(command)}\n{_log_tail(log_path)}") + except subprocess.TimeoutExpired: + pytest.fail(f"MX E2E worker timed out: {' '.join(command)}\n{_log_tail(log_path)}") if returncode != 0: pytest.fail( f"MX E2E worker exited with status {returncode}: " @@ -336,15 +332,6 @@ def _wait_for_donor( time.sleep(0.5) -def _signal_process_group(process: subprocess.Popen[bytes], sig: signal.Signals) -> None: - if process.poll() is not None: - return - try: - os.killpg(process.pid, sig) - except ProcessLookupError: - pass - - def _stop_donor(process: subprocess.Popen[bytes], stop_file: Path) -> int: if process.poll() is not None: assert process.returncode is not None @@ -353,11 +340,7 @@ def _stop_donor(process: subprocess.Popen[bytes], stop_file: Path) -> int: try: return process.wait(timeout=120) except subprocess.TimeoutExpired: - _signal_process_group(process, signal.SIGTERM) - try: - return process.wait(timeout=30) - except subprocess.TimeoutExpired: - _signal_process_group(process, signal.SIGKILL) + cleanup_process_tree(process, has_session=True) return process.wait(timeout=30) @@ -440,8 +423,9 @@ def test_mx_donor_receiver(case: MxE2ECase, tmp_path: Path) -> None: donor_environment = _worker_environment(donor_gpu_ids, tmp_path / "donor-transfer-logs") donor_returncode = None - with donor_log.open("w", encoding="utf-8") as donor_log_file: - donor_process = subprocess.Popen( + with ( + donor_log.open("w", encoding="utf-8") as donor_log_file, + popen( _worker_command( role="donor", model_path=donor_snapshot, @@ -454,8 +438,9 @@ def test_mx_donor_receiver(case: MxE2ECase, tmp_path: Path) -> None: env=donor_environment, stdout=donor_log_file, stderr=subprocess.STDOUT, - start_new_session=True, - ) + suppress_output_info=True, + ) as donor_process, + ): try: _wait_for_donor(donor_process, donor_ready, donor_log, timeout_s) _run_worker( @@ -476,6 +461,8 @@ def test_mx_donor_receiver(case: MxE2ECase, tmp_path: Path) -> None: assert donor_returncode == 0, ( f"MX donor exited with status {donor_returncode}\n{_log_tail(donor_log)}" ) + donor_payload = json.loads(donor_output.read_text(encoding="utf-8")) + assert donor_payload["server_query_timeout_s"] == 0 baseline_tokens = _load_tokens(baseline_output) assert _load_tokens(donor_output) == baseline_tokens assert _load_tokens(receiver_output) == baseline_tokens From 77e580654649891c67c8513982e028d1474e327b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 10 Aug 2026 14:45:43 -0700 Subject: [PATCH 03/16] [TRTLLM-14727][fix] harden MX qualification harness Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 3 +++ .../defs/model_express/mx_e2e_worker.py | 9 ++++++++ .../defs/model_express/test_model_express.py | 22 +++++++++++++++---- 3 files changed, 30 insertions(+), 4 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 1e1cf931ffc5..22b8b8514ab7 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -3361,6 +3361,9 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod - 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. + // Redis is a native sidecar: Kubernetes < 1.33 must enable SidecarContainers. serviceInitContainerConfig = """ initContainers: - name: redis diff --git a/tests/integration/defs/model_express/mx_e2e_worker.py b/tests/integration/defs/model_express/mx_e2e_worker.py index ce5605bd0a3b..7eca6e8b4b88 100644 --- a/tests/integration/defs/model_express/mx_e2e_worker.py +++ b/tests/integration/defs/model_express/mx_e2e_worker.py @@ -39,6 +39,7 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--mx-url") parser.add_argument("--ready-file", type=Path) parser.add_argument("--stop-file", type=Path) + parser.add_argument("--max-serve-seconds", type=float, default=1800.0) return parser.parse_args() @@ -73,6 +74,8 @@ def main() -> None: args = _parse_args() if args.role == "donor" and (args.ready_file is None or args.stop_file is None): raise ValueError("The donor role requires --ready-file and --stop-file") + if args.role == "donor" and args.max_serve_seconds <= 0: + raise ValueError("The donor role requires --max-serve-seconds > 0") llm_kwargs = _llm_kwargs(args) started = time.perf_counter() @@ -107,7 +110,13 @@ def main() -> None: assert args.ready_file is not None assert args.stop_file is not None args.ready_file.write_text("ready\n", encoding="utf-8") + deadline = time.monotonic() + args.max_serve_seconds while not args.stop_file.exists(): + if time.monotonic() >= deadline: + raise TimeoutError( + f"The stop file {args.stop_file} did not appear within " + f"{args.max_serve_seconds}s" + ) time.sleep(0.2) diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index 10a3bb16a859..0fda407ec710 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -249,6 +249,7 @@ def _worker_command( mx_url: str | None = None, ready_file: Path | None = None, stop_file: Path | None = None, + max_serve_seconds: int | None = None, ) -> list[str]: command = [ sys.executable, @@ -268,6 +269,8 @@ def _worker_command( command.extend(("--ready-file", str(ready_file))) if stop_file is not None: command.extend(("--stop-file", str(stop_file))) + if max_serve_seconds is not None: + command.extend(("--max-serve-seconds", str(max_serve_seconds))) return command @@ -371,15 +374,25 @@ def _assert_transfer_evidence( all_receiver_logs_lower = all_receiver_logs.lower() for marker in _RECEIVER_FAILURE_MARKERS: - assert marker not in all_receiver_logs_lower + assert marker not in all_receiver_logs_lower, ( + f"MX receiver logs contain failure marker {marker!r}" + ) matched_params = _MATCHED_PARAMS_PATTERN.findall(all_receiver_logs) - assert len(matched_params) >= case.tp_size - assert all(int(matched) == int(total) > 0 for matched, total in matched_params) + assert len(matched_params) >= case.tp_size, ( + f"Expected at least {case.tp_size} matched-parameter summaries, got {matched_params}" + ) + assert all(int(matched) == int(total) > 0 for matched, total in matched_params), ( + f"MX receiver reported incomplete parameter matches: {matched_params}" + ) transferred_params = _TRANSFERRED_PARAMS_PATTERN.findall(all_receiver_logs) transferred_ranks = {int(rank) for rank, count in transferred_params if int(count) > 0} - assert transferred_ranks.issuperset(range(case.tp_size)) + expected_ranks = set(range(case.tp_size)) + assert transferred_ranks.issuperset(expected_ranks), ( + f"Expected transferred ranks {expected_ranks}, observed {transferred_ranks}; " + f"transfer summaries: {transferred_params}" + ) @pytest.mark.parametrize("case", _MX_CASES) @@ -434,6 +447,7 @@ def test_mx_donor_receiver(case: MxE2ECase, tmp_path: Path) -> None: mx_url=mx_url, ready_file=donor_ready, stop_file=donor_stop, + max_serve_seconds=timeout_s + 120, ), env=donor_environment, stdout=donor_log_file, From 93b0cc8d28c49047889e59c0f0c7ac63cce4f0f5 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:51:13 -0700 Subject: [PATCH 04/16] [TRTLLM-14727][fix] bind MX evidence to ranks Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../defs/model_express/test_model_express.py | 55 +++++++++++++------ .../test_lists/qa/llm_function_core.txt | 4 ++ 2 files changed, 42 insertions(+), 17 deletions(-) diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index 0fda407ec710..bf2366fc17d9 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -49,6 +49,7 @@ "invalid sourceidentity", ) _MATCHED_PARAMS_PATTERN = re.compile(r"Matched\s+(\d+)/(\d+)\s+params", re.IGNORECASE) +_RANK_LOG_PATTERN = re.compile(r"rank(\d+)\.log", re.IGNORECASE) _TRANSFERRED_PARAMS_PATTERN = re.compile( r"Rank\s+(\d+):\s+transferred\s+(\d+)\s+params", re.IGNORECASE, @@ -355,13 +356,23 @@ def _load_tokens(output_path: Path) -> list[list[int]]: return token_ids -def _transfer_log_text(transfer_log_dir: Path) -> str: +def _transfer_logs_by_rank(transfer_log_dir: Path) -> dict[int, str]: logs = tuple( path for path in transfer_log_dir.rglob("*") if path.is_file() and path.stat().st_size > 0 ) if not logs: pytest.fail(f"ModelExpress created no receiver transfer logs in {transfer_log_dir}") - return "\n".join(path.read_text(encoding="utf-8", errors="replace") for path in sorted(logs)) + + logs_by_rank = {} + for path in sorted(logs): + match = _RANK_LOG_PATTERN.fullmatch(path.name) + if match is None: + pytest.fail(f"Unexpected ModelExpress receiver transfer log: {path}") + rank = int(match.group(1)) + if rank in logs_by_rank: + pytest.fail(f"ModelExpress created multiple receiver transfer logs for rank {rank}") + logs_by_rank[rank] = path.read_text(encoding="utf-8", errors="replace") + return logs_by_rank def _assert_transfer_evidence( @@ -370,7 +381,12 @@ def _assert_transfer_evidence( receiver_transfer_log_dir: Path, ) -> None: receiver_log = receiver_log_path.read_text(encoding="utf-8", errors="replace") - all_receiver_logs = receiver_log + "\n" + _transfer_log_text(receiver_transfer_log_dir) + transfer_logs = _transfer_logs_by_rank(receiver_transfer_log_dir) + expected_ranks = set(range(case.tp_size)) + assert set(transfer_logs) == expected_ranks, ( + f"Expected receiver transfer logs for ranks {expected_ranks}, got {set(transfer_logs)}" + ) + all_receiver_logs = receiver_log + "\n" + "\n".join(transfer_logs.values()) all_receiver_logs_lower = all_receiver_logs.lower() for marker in _RECEIVER_FAILURE_MARKERS: @@ -378,21 +394,26 @@ def _assert_transfer_evidence( f"MX receiver logs contain failure marker {marker!r}" ) - matched_params = _MATCHED_PARAMS_PATTERN.findall(all_receiver_logs) - assert len(matched_params) >= case.tp_size, ( - f"Expected at least {case.tp_size} matched-parameter summaries, got {matched_params}" - ) - assert all(int(matched) == int(total) > 0 for matched, total in matched_params), ( - f"MX receiver reported incomplete parameter matches: {matched_params}" - ) + for rank in sorted(expected_ranks): + rank_log = transfer_logs[rank] + matched_params = _MATCHED_PARAMS_PATTERN.findall(rank_log) + assert len(matched_params) == 1, ( + f"Expected one matched-parameter summary for rank {rank}, got {matched_params}" + ) + matched, total = (int(value) for value in matched_params[0]) + assert matched == total > 0, ( + f"MX receiver rank {rank} reported incomplete parameter match {matched}/{total}" + ) - transferred_params = _TRANSFERRED_PARAMS_PATTERN.findall(all_receiver_logs) - transferred_ranks = {int(rank) for rank, count in transferred_params if int(count) > 0} - expected_ranks = set(range(case.tp_size)) - assert transferred_ranks.issuperset(expected_ranks), ( - f"Expected transferred ranks {expected_ranks}, observed {transferred_ranks}; " - f"transfer summaries: {transferred_params}" - ) + transferred_params = _TRANSFERRED_PARAMS_PATTERN.findall(rank_log) + assert len(transferred_params) == 1, ( + f"Expected one transfer summary for rank {rank}, got {transferred_params}" + ) + transferred_rank, transferred_count = (int(value) for value in transferred_params[0]) + assert transferred_rank == rank and transferred_count == matched, ( + f"MX receiver rank {rank} matched {matched} params but reported transfer summary " + f"{transferred_params[0]}" + ) @pytest.mark.parametrize("case", _MX_CASES) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 46497f84b6cb..115d95ca2c46 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -950,3 +950,7 @@ test_e2e.py::test_relaxed_acceptance_quickstart_advanced_deepseek_r1_8gpus[DeepS test_e2e.py::test_trtllm_benchmark_serving[gpt_oss/gpt-oss-20b] test_e2e.py::test_trtllm_multimodal_benchmark_serving unittest/llmapi/apps/_test_openai_embeddings.py + +# ModelExpress donor/receiver qualification +model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp1] +model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp2] From d6cacf2098e756a6971bc49f3da07d6f3a1b6dd1 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 11 Aug 2026 13:00:34 -0700 Subject: [PATCH 05/16] [TRTLLM-14727][fix] make MX CI sidecar cluster compatible Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 36 +++++++++++++++++++++++++----------- 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 22b8b8514ab7..eafd46ff21e2 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -3363,7 +3363,8 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod """ // Mirrors the ModelExpress v0.4.1 Redis deployment and image contract. // The image exposes /app/modelexpress-server and accepts the port/backend settings below. - // Redis is a native sidecar: Kubernetes < 1.33 must enable SidecarContainers. + // Redis is a native sidecar. A one-shot init container waits for Redis because + // the CI cluster rejects probe fields on restartable init containers. serviceInitContainerConfig = """ initContainers: - name: redis @@ -3372,16 +3373,6 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod restartPolicy: Always ports: - containerPort: 6379 - startupProbe: - tcpSocket: - port: 6379 - initialDelaySeconds: 1 - periodSeconds: 1 - failureThreshold: 30 - livenessProbe: - tcpSocket: - port: 6379 - periodSeconds: 10 resources: requests: cpu: 500m @@ -3392,6 +3383,29 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod memory: 1Gi ephemeral-storage: 2Gi imagePullPolicy: Always + - name: wait-for-redis + image: ${MODEL_EXPRESS_REDIS_IMAGE} + command: ["/bin/sh", "-c"] + args: + - | + attempt=0 + until redis-cli -h 127.0.0.1 ping | grep -q PONG; do + attempt=\$((attempt + 1)) + if [ "\$attempt" -ge 30 ]; then + exit 1 + fi + sleep 1 + done + resources: + requests: + cpu: 100m + memory: 64Mi + ephemeral-storage: 1Gi + limits: + cpu: 100m + memory: 64Mi + ephemeral-storage: 1Gi + imagePullPolicy: Always """ serviceContainerConfig = """ - name: model-express-server From 58be424e72052017fff60dbe8f0e5c05af725671 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 11 Aug 2026 14:15:32 -0700 Subject: [PATCH 06/16] [TRTLLM-14727][fix] reserve CI resources for MX sidecars Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index eafd46ff21e2..dacc73ec6cc9 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -3258,9 +3258,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 ? 2 : 0 + def serviceMemoryReserveGi = modelExpress ? 8 : 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) From 43f3808b07d409cbe9dbfe9c13bdd0e11a274187 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:17:09 -0700 Subject: [PATCH 07/16] [TRTLLM-14727][fix] surface MX service startup failures Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 28 ++++++++++++---------------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index dacc73ec6cc9..6e98f5ae995e 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -3260,8 +3260,8 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod // Not a hard requirement, but based on empirical values. // Keep ModelExpress services inside the existing pod resource envelope; // otherwise their requests can make an otherwise valid GPU pod unschedulable. - def serviceCpuReserve = modelExpress ? 2 : 0 - def serviceMemoryReserveGi = modelExpress ? 8 : 0 + 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" @@ -3380,12 +3380,12 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod - containerPort: 6379 resources: requests: - cpu: 500m - memory: 1Gi + cpu: '1' + memory: 4Gi ephemeral-storage: 2Gi limits: - cpu: 500m - memory: 1Gi + cpu: '1' + memory: 4Gi ephemeral-storage: 2Gi imagePullPolicy: Always - name: wait-for-redis @@ -3412,6 +3412,8 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod ephemeral-storage: 1Gi imagePullPolicy: Always """ + // The E2E preflight waits for port 8001. A readiness probe here would keep + // Jenkins from attaching and hide startup errors behind a pod timeout. serviceContainerConfig = """ - name: model-express-server image: ${MODEL_EXPRESS_SERVER_IMAGE} @@ -3424,12 +3426,6 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod value: "redis://127.0.0.1:6379" ports: - containerPort: 8001 - readinessProbe: - tcpSocket: - port: 8001 - initialDelaySeconds: 2 - periodSeconds: 2 - failureThreshold: 30 livenessProbe: tcpSocket: port: 8001 @@ -3437,12 +3433,12 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod periodSeconds: 10 resources: requests: - cpu: '1' - memory: 4Gi + cpu: '4' + memory: 8Gi ephemeral-storage: 4Gi limits: - cpu: '1' - memory: 4Gi + cpu: '4' + memory: 8Gi ephemeral-storage: 4Gi imagePullPolicy: Always """ From 418c01f8ef088f234f7252adf4c3b8ab5bb154d2 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:43:22 -0700 Subject: [PATCH 08/16] [TRTLLM-14727][test] install MX client in E2E stages Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 6e98f5ae995e..6348f17e393e 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -73,7 +73,8 @@ ARTIFACTORY_CREDENTIALS_ID = "trtllm-artifactory-credentials" // DLFW torch image DLFW_IMAGE = "urm.nvidia.com/docker/nvidia/pytorch:26.05-py3" -MODEL_EXPRESS_SERVER_IMAGE = "urm.nvidia.com/docker/nvidia/ai-dynamo/modelexpress-server:0.4.1" +MODEL_EXPRESS_VERSION = "0.4.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 @@ -4534,6 +4535,9 @@ 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}") + } } trtllm_utils.llmExecStepWithRetry(pipeline, script: "git config --global --add safe.directory \"*\"") From 991ac1acaa1d5cfeeb4eb3a58136c8643d76bbac Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 11 Aug 2026 17:04:33 -0700 Subject: [PATCH 09/16] [TRTLLM-14727][fix] use regular MX CI sidecars Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 46 +++++++++--------------------------------- 1 file changed, 9 insertions(+), 37 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 6348f17e393e..876b73287ff2 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -3369,14 +3369,13 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod """ // Mirrors the ModelExpress v0.4.1 Redis deployment and image contract. // The image exposes /app/modelexpress-server and accepts the port/backend settings below. - // Redis is a native sidecar. A one-shot init container waits for Redis because - // the CI cluster rejects probe fields on restartable init containers. - serviceInitContainerConfig = """ - initContainers: + // 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"] - restartPolicy: Always ports: - containerPort: 6379 resources: @@ -3389,37 +3388,15 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod memory: 4Gi ephemeral-storage: 2Gi imagePullPolicy: Always - - name: wait-for-redis - image: ${MODEL_EXPRESS_REDIS_IMAGE} - command: ["/bin/sh", "-c"] + - name: model-express-server + image: ${MODEL_EXPRESS_SERVER_IMAGE} + command: ["/bin/bash", "-c"] args: - | - attempt=0 - until redis-cli -h 127.0.0.1 ping | grep -q PONG; do - attempt=\$((attempt + 1)) - if [ "\$attempt" -ge 30 ]; then - exit 1 - fi + until (echo > /dev/tcp/127.0.0.1/6379) >/dev/null 2>&1; do sleep 1 done - resources: - requests: - cpu: 100m - memory: 64Mi - ephemeral-storage: 1Gi - limits: - cpu: 100m - memory: 64Mi - ephemeral-storage: 1Gi - imagePullPolicy: Always - """ - // The E2E preflight waits for port 8001. A readiness probe here would keep - // Jenkins from attaching and hide startup errors behind a pod timeout. - serviceContainerConfig = """ - - name: model-express-server - image: ${MODEL_EXPRESS_SERVER_IMAGE} - command: ["/app/modelexpress-server"] - args: ["--port", "8001"] + exec /app/modelexpress-server --port 8001 env: - name: MX_METADATA_BACKEND value: "redis" @@ -3427,11 +3404,6 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod value: "redis://127.0.0.1:6379" ports: - containerPort: 8001 - livenessProbe: - tcpSocket: - port: 8001 - initialDelaySeconds: 10 - periodSeconds: 10 resources: requests: cpu: '4' From 5e0182196a740bd6d0fbf4fb763692133c28d17c Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 11 Aug 2026 18:15:01 -0700 Subject: [PATCH 10/16] [TRTLLM-14727][fix] install NIXL namespace for MX CI Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 876b73287ff2..e4bdc74c4b5b 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -74,6 +74,7 @@ ARTIFACTORY_CREDENTIALS_ID = "trtllm-artifactory-credentials" 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" @@ -4509,6 +4510,10 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO } 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}") } } From 00e0afcb1cdba864bcaf75230c096c153ec77d4f Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 11 Aug 2026 20:40:31 -0700 Subject: [PATCH 11/16] [TRTLLM-14727][fix] preserve cold MX fallback graph Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../checkpoints/mx/checkpoint_loader.py | 10 +++++++ .../mx/test_mx_checkpoint_loader.py | 29 +++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py index 05db808f1d5e..cc53838841a1 100644 --- a/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py @@ -404,6 +404,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() diff --git a/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py b/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py index 77c0f3bd6f1a..967f357cb8f3 100644 --- a/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py +++ b/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py @@ -1015,6 +1015,35 @@ def _assert_timeout(*args, **kwargs): mx_loader.load_weights.assert_called_once() assert "MX_SOURCE_QUERY_TIMEOUT" not in os.environ + def test_zero_timeout_falls_back_before_receiver_preparation(self): + identity = _identity() + disk_weights = {"disk.weight": MagicMock()} + prepare_receiver = MagicMock() + loader = MXCheckpointLoader(mx_server_url="http://mx:8001", query_timeout_s=0) + fake_mx = _build_fake_modelexpress() + + with ( + _install_fake_modelexpress(fake_mx), + patch.object( + HfCheckpointLoader, "load_weights", return_value=disk_weights + ) as mock_super_load, + ): + result = loader.load_weights( + "/nonexistent", + mapping=MagicMock(), + model=MagicMock(), + source_identity=identity, + allow_post_transform_weights=True, + prepare_post_transform_receiver=prepare_receiver, + ) + + assert result is disk_weights + assert loader.is_weights_preloaded() is False + assert loader.is_post_transform_weights_preloaded() is False + prepare_receiver.assert_not_called() + fake_mx.trtllm_live_transfer.MxLiveWeightLoader.assert_not_called() + mock_super_load.assert_called_once() + def test_existing_source_keeps_upstream_default_when_unset(self): identity = _identity() From 38dd1c6b3a0667de64c9258873119eeeab88fb61 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 12 Aug 2026 11:51:05 -0700 Subject: [PATCH 12/16] [TRTLLM-14727][fix] isolate MX workers from MPI UCX Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../defs/model_express/test_model_express.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index bf2366fc17d9..f8b57f8de4cd 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -54,6 +54,11 @@ r"Rank\s+(\d+):\s+transferred\s+(\d+)\s+params", re.IGNORECASE, ) +_DONOR_PROCESS_FAILURE_MARKERS = ( + b"Segfault encountered", + b"Primary job terminated normally, but", + b"process returned a non-zero exit code", +) _MX_PREFLIGHT_SCRIPT = """ import importlib.metadata as metadata import importlib.util @@ -284,6 +289,9 @@ def _worker_environment(gpu_ids: tuple[str, ...], transfer_log_dir: Path) -> dic "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", "MX_TRANSFER_LOG_DIR": str(transfer_log_dir), + # ModelExpress initializes its own NIXL UCX context. Keep OpenMPI + # on ob1 so two independently packaged UCX stacks do not coexist. + "OMPI_MCA_pml": "ob1", "PYTHONUNBUFFERED": "1", "TLLM_LOG_LEVEL": "INFO", } @@ -325,12 +333,21 @@ def _wait_for_donor( process: subprocess.Popen[bytes], ready_file: Path, log_path: Path, timeout_s: int ) -> None: deadline = time.monotonic() + timeout_s + log_offset = 0 + failure_window = b"" while not ready_file.exists(): returncode = process.poll() if returncode is not None: pytest.fail( f"MX donor exited before publication with status {returncode}\n{_log_tail(log_path)}" ) + if log_path.exists(): + with log_path.open("rb") as log_file: + log_file.seek(log_offset) + failure_window = (failure_window + log_file.read())[-4096:] + log_offset = log_file.tell() + if any(marker in failure_window for marker in _DONOR_PROCESS_FAILURE_MARKERS): + pytest.fail(f"MX donor worker failed before publication\n{_log_tail(log_path)}") if time.monotonic() >= deadline: pytest.fail(f"MX donor did not become ready within {timeout_s}s\n{_log_tail(log_path)}") time.sleep(0.5) From 2491db5e7f6c2ed657b893096ac31235b7a49e4b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:06:34 -0700 Subject: [PATCH 13/16] [TRTLLM-14727][fix] isolate MX workers from all MPI UCX paths Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../defs/model_express/test_model_express.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index f8b57f8de4cd..fa4b7b221d16 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -59,6 +59,7 @@ b"Primary job terminated normally, but", b"process returned a non-zero exit code", ) +_DONOR_PROCESS_FAILURE_OVERLAP = max(len(marker) for marker in _DONOR_PROCESS_FAILURE_MARKERS) - 1 _MX_PREFLIGHT_SCRIPT = """ import importlib.metadata as metadata import importlib.util @@ -289,9 +290,12 @@ def _worker_environment(gpu_ids: tuple[str, ...], transfer_log_dir: Path) -> dic "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", "MX_TRANSFER_LOG_DIR": str(transfer_log_dir), - # ModelExpress initializes its own NIXL UCX context. Keep OpenMPI - # on ob1 so two independently packaged UCX stacks do not coexist. + # NIXL's wheel bundles UCX. Keep OpenMPI's UCX-capable components + # from loading the container UCX stack into the same process. "OMPI_MCA_pml": "ob1", + "OMPI_MCA_osc": "pt2pt", + "OMPI_MCA_btl": "self,vader,tcp", + "OMPI_MCA_coll": "^hcoll,ucc", "PYTHONUNBUFFERED": "1", "TLLM_LOG_LEVEL": "INFO", } @@ -344,10 +348,11 @@ def _wait_for_donor( if log_path.exists(): with log_path.open("rb") as log_file: log_file.seek(log_offset) - failure_window = (failure_window + log_file.read())[-4096:] + failure_window += log_file.read() log_offset = log_file.tell() if any(marker in failure_window for marker in _DONOR_PROCESS_FAILURE_MARKERS): pytest.fail(f"MX donor worker failed before publication\n{_log_tail(log_path)}") + failure_window = failure_window[-_DONOR_PROCESS_FAILURE_OVERLAP:] if time.monotonic() >= deadline: pytest.fail(f"MX donor did not become ready within {timeout_s}s\n{_log_tail(log_path)}") time.sleep(0.5) From 35b19406e77be5586f078d96394c39914ba9a654 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 13 Aug 2026 11:18:07 -0700 Subject: [PATCH 14/16] [TRTLLM-14727][fix] propagate MX transfer logs to workers Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tests/integration/defs/model_express/mx_e2e_worker.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/integration/defs/model_express/mx_e2e_worker.py b/tests/integration/defs/model_express/mx_e2e_worker.py index 7eca6e8b4b88..4b7e48a59c15 100644 --- a/tests/integration/defs/model_express/mx_e2e_worker.py +++ b/tests/integration/defs/model_express/mx_e2e_worker.py @@ -18,6 +18,7 @@ import argparse import json +import os import time from pathlib import Path @@ -60,6 +61,9 @@ def _llm_kwargs(args: argparse.Namespace) -> dict[str, object]: if args.role != "baseline": if not args.mx_url: raise ValueError("MX donor and receiver roles require --mx-url") + transfer_log_dir = os.environ.get("MX_TRANSFER_LOG_DIR") + if transfer_log_dir: + kwargs["env_overrides"] = {"MX_TRANSFER_LOG_DIR": transfer_log_dir} kwargs["mx_config"] = { "server_url": args.mx_url, # ModelExpress 0.4.1 skips polling at zero but still sleeps once From 9e83bd815d8979a99a0a6eef2eb6fbfd3d61225b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 13 Aug 2026 12:36:21 -0700 Subject: [PATCH 15/16] [TRTLLM-14727][fix] enable MX transfer evidence logs Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../checkpoints/mx/checkpoint_loader.py | 17 ++++++++++++++++ .../defs/model_express/test_model_express.py | 19 ++++++++++++------ .../mx/test_mx_checkpoint_loader.py | 20 +++++++++++++++++++ 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py index cc53838841a1..714396bd9142 100644 --- a/tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py +++ b/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 diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index fa4b7b221d16..db92eb4099d7 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -378,12 +378,19 @@ def _load_tokens(output_path: Path) -> list[list[int]]: return token_ids -def _transfer_logs_by_rank(transfer_log_dir: Path) -> dict[int, str]: - logs = tuple( - path for path in transfer_log_dir.rglob("*") if path.is_file() and path.stat().st_size > 0 - ) +def _transfer_logs_by_rank(transfer_log_dir: Path, receiver_log: str) -> dict[int, str]: + log_files = tuple(path for path in transfer_log_dir.rglob("*") if path.is_file()) + logs = tuple(path for path in log_files if path.stat().st_size > 0) if not logs: - pytest.fail(f"ModelExpress created no receiver transfer logs in {transfer_log_dir}") + entries = ", ".join( + f"{path.relative_to(transfer_log_dir)} ({path.stat().st_size} bytes)" + for path in log_files + ) + pytest.fail( + f"ModelExpress created no non-empty receiver transfer logs in " + f"{transfer_log_dir}; entries: {entries or ''}\n" + f"Receiver log:\n{receiver_log}" + ) logs_by_rank = {} for path in sorted(logs): @@ -403,7 +410,7 @@ def _assert_transfer_evidence( receiver_transfer_log_dir: Path, ) -> None: receiver_log = receiver_log_path.read_text(encoding="utf-8", errors="replace") - transfer_logs = _transfer_logs_by_rank(receiver_transfer_log_dir) + transfer_logs = _transfer_logs_by_rank(receiver_transfer_log_dir, receiver_log) expected_ranks = set(range(case.tp_size)) assert set(transfer_logs) == expected_ranks, ( f"Expected receiver transfer logs for ranks {expected_ranks}, got {set(transfer_logs)}" diff --git a/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py b/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py index 967f357cb8f3..7432d70019ea 100644 --- a/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py +++ b/tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py @@ -14,6 +14,7 @@ """ import json +import logging import os import sys from contextlib import ExitStack @@ -154,6 +155,25 @@ def test_post_transform_signal_requires_p2p_and_identity_match(self): loader._source_identity_compatible_for_last_load = True assert loader.is_post_transform_weights_preloaded() is True + @pytest.mark.parametrize( + ("effective_level", "expected_level"), + ((logging.WARNING, logging.INFO), (logging.DEBUG, None)), + ) + def test_transfer_log_dir_enables_info_records( + self, monkeypatch, effective_level, expected_level + ): + monkeypatch.setenv("MX_TRANSFER_LOG_DIR", "/tmp/mx-transfer-logs") + mx_logger = MagicMock() + mx_logger.getEffectiveLevel.return_value = effective_level + + with patch.object(mx_checkpoint_loader.logging, "getLogger", return_value=mx_logger): + mx_checkpoint_loader._enable_mx_transfer_logging() + + if expected_level is None: + mx_logger.setLevel.assert_not_called() + else: + mx_logger.setLevel.assert_called_once_with(expected_level) + # --------------------------------------------------------------------------- # Registry From 8e4a7e7bf8eb55c9c37640423c17b567dcbcc564 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 13 Aug 2026 19:47:55 -0700 Subject: [PATCH 16/16] [TRTLLM-14727][test] Address MX harness review feedback Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_MergeRequest.groovy | 3 ++ .../defs/model_express/test_model_express.py | 36 ++++++++++++++----- .../test_lists/qa/llm_function_core.txt | 4 --- 3 files changed, 30 insertions(+), 13 deletions(-) diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index b386b1aca564..996ac079314e 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/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", diff --git a/tests/integration/defs/model_express/test_model_express.py b/tests/integration/defs/model_express/test_model_express.py index db92eb4099d7..f4edd3a1d4fb 100644 --- a/tests/integration/defs/model_express/test_model_express.py +++ b/tests/integration/defs/model_express/test_model_express.py @@ -33,6 +33,7 @@ import pytest import torch from defs.trt_test_alternative import cleanup_process_tree, popen +from packaging.version import InvalidVersion, Version from tensorrt_llm._torch.weight_sharing import ArtifactIdentity @@ -60,6 +61,8 @@ b"process returned a non-zero exit code", ) _DONOR_PROCESS_FAILURE_OVERLAP = max(len(marker) for marker in _DONOR_PROCESS_FAILURE_MARKERS) - 1 +_MINIMUM_MODELEXPRESS_VERSION = Version("0.4.1") +_MODELEXPRESS_VERSION_PREFIX = "MODELEXPRESS_VERSION=" _MX_PREFLIGHT_SCRIPT = """ import importlib.metadata as metadata import importlib.util @@ -168,22 +171,37 @@ def _require_mx_environment(required_gpus: int) -> tuple[str, tuple[str, ...]]: _skip_or_fail(f"ModelExpress service {mx_url!r} is unreachable: {error}") time.sleep(1) - preflight = subprocess.run( - [sys.executable, "-c", _MX_PREFLIGHT_SCRIPT], - capture_output=True, - text=True, - timeout=30, - check=False, - ) + try: + preflight = subprocess.run( + [sys.executable, "-c", _MX_PREFLIGHT_SCRIPT], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + except subprocess.TimeoutExpired as error: + _skip_or_fail(f"ModelExpress/NIXL preflight timed out after {error.timeout} seconds") if preflight.returncode != 0: detail = (preflight.stdout + preflight.stderr).strip() _skip_or_fail(f"ModelExpress/NIXL preflight failed: {detail}") version_lines = [ - line for line in preflight.stdout.splitlines() if line.startswith("MODELEXPRESS_VERSION=") + line + for line in preflight.stdout.splitlines() + if line.startswith(_MODELEXPRESS_VERSION_PREFIX) ] if len(version_lines) != 1: _skip_or_fail(f"ModelExpress preflight did not report its version: {preflight.stdout!r}") - print(f"MX E2E client preflight passed: {version_lines[0]}") + modelexpress_version = version_lines[0].removeprefix(_MODELEXPRESS_VERSION_PREFIX) + try: + parsed_modelexpress_version = Version(modelexpress_version) + except InvalidVersion: + _skip_or_fail(f"ModelExpress client reported invalid version {modelexpress_version!r}") + if parsed_modelexpress_version < _MINIMUM_MODELEXPRESS_VERSION: + _skip_or_fail( + f"Unsupported ModelExpress client version {modelexpress_version!r}; " + f"requires >= {_MINIMUM_MODELEXPRESS_VERSION}" + ) + print(f"MX E2E client preflight passed: ModelExpress {modelexpress_version}") configured_devices = os.environ.get("CUDA_VISIBLE_DEVICES") if configured_devices and configured_devices.strip().lower() != "all": diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 115d95ca2c46..46497f84b6cb 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -950,7 +950,3 @@ test_e2e.py::test_relaxed_acceptance_quickstart_advanced_deepseek_r1_8gpus[DeepS test_e2e.py::test_trtllm_benchmark_serving[gpt_oss/gpt-oss-20b] test_e2e.py::test_trtllm_multimodal_benchmark_serving unittest/llmapi/apps/_test_openai_embeddings.py - -# ModelExpress donor/receiver qualification -model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp1] -model_express/test_model_express.py::test_mx_donor_receiver[llama-bf16-tp2]