From 3d004e1cf8fc495f77e73d3028a4db4d4aa0b569 Mon Sep 17 00:00:00 2001 From: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> Date: Wed, 6 May 2026 13:40:52 -0700 Subject: [PATCH 1/4] move test list check to top-level pipeline Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> --- jenkins/L0_MergeRequest.groovy | 57 +++++++++++++++++-- jenkins/L0_Test.groovy | 46 +-------------- jenkins/ci_versions.properties | 7 +++ scripts/check_test_list.py | 37 +++++++++--- .../integration/test_lists/test-db/README.md | 4 +- 5 files changed, 92 insertions(+), 59 deletions(-) create mode 100644 jenkins/ci_versions.properties diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index de9da034765c..a44721eec9a5 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -571,6 +571,23 @@ def launchReleaseCheck(pipeline, globalVars) }) } +def launchTestListCheck(pipeline, globalVars) +{ + def key = "Check Test List" + def image = globalVars["LLM_DOCKER_IMAGE"] + trtllm_utils.launchKubernetesPod(pipeline, createKubernetesPodConfig(image, "package"), "trt-llm", { + stage("[${key}] Run") { + echoNodeAndGpuInfo(pipeline, key) + sh "git config --global --add safe.directory \"*\"" + trtllm_utils.checkoutSource(LLM_REPO, env.gitlabCommit, LLM_ROOT, false, true) + + def llmPath = sh(script: "realpath ${LLM_ROOT}", returnStdout: true).trim() + sh "NVIDIA_TRITON_SERVER_VERSION=26.05 LLM_ROOT=${llmPath} LLM_BACKEND_ROOT=${llmPath}/triton_backend " + + "python3 ${llmPath}/scripts/check_test_list.py --l0 --qa --waive --validate --check-duplicate-waives" + } + }) +} + def getGitlabMRChangedFile(pipeline, function, filePath="") { def result = null def pageId = 0 @@ -1613,6 +1630,15 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) launchReleaseCheck(this, globalVars) } }, + "Check Test List": { + script { + if (GEN_POST_MERGE_BUILDS_ONLY) { + echo "Skipping Check Test List (GenPostMergeBuilds mode: builds only)" + return + } + launchTestListCheck(this, globalVars) + } + }, "x86_64-Linux": { script { // CBTS deliberately does NOT short-circuit at the arch / Build @@ -2048,15 +2074,25 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) echo "Will run job to build ngc containers and running in-pipeline scanning for them" } + def alwaysFailFastStages = ["Release-Check", "Check Test List"] as Set parallelJobs = stages.collectEntries{key, value -> [key, { script { stage(key) { - value() + if (enableFailFast || key in alwaysFailFastStages) { + value() + } else { + // Avoid interrupting other stages on failure. + catchError(catchInterruptions: false) { + value() + } + } } } }]} - parallelJobs.failFast = enableFailFast + // With --disable-fail-fast, ordinary build/test failures are suppressed above after + // marking their stage and build failed, so those branches do not trigger this fail-fast. + parallelJobs.failFast = true pipeline.parallel parallelJobs } @@ -2148,11 +2184,20 @@ pipeline { steps { script { if (isReleaseCheckMode) { - stage("Release-Check") { - script { - launchReleaseCheck(this, globalVars) + def releaseCheckStages = [ + "Release-Check": { + stage("Release-Check") { + launchReleaseCheck(this, globalVars) + } + }, + "Check Test List": { + stage("Check Test List") { + launchTestListCheck(this, globalVars) + } } - } + ] + releaseCheckStages.failFast = true + parallel releaseCheckStages } else { // globalVars[CACHED_CHANGED_FILE_LIST] is only used in setupPipelineEnvironment // Remove it to workaround the "Argument list too long" error diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index adcc4c847e97..ddf39b527977 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -3539,30 +3539,6 @@ def runLLMAgentFlowTest(pipeline, stageName) sh "cd ${WORKSPACE}/${stageName} && sed -i 's/testsuite name=\"pytest\"/testsuite name=\"${stageName}\"/g' results.xml || true" } -def launchTestListCheck(pipeline) -{ - stageName = "Test List Check" - trtllm_utils.launchKubernetesPod(pipeline, createKubernetesPodConfig(LLM_DOCKER_IMAGE, "a10"), "trt-llm", { - try { - echoNodeAndGpuInfo(pipeline, stageName) - sh "nvidia-smi && nvidia-smi -q && nvidia-smi topo -m" - // download TRT-LLM tarfile - def tarName = BUILD_CONFIGS[VANILLA_CONFIG][TARNAME] - def llmTarfile = "https://urm.nvidia.com/artifactory/${ARTIFACT_PATH}/${tarName}" - trtllm_utils.llmExecStepWithRetry(pipeline, script: "pwd && wget -nv ${llmTarfile} && ls -alh") - sh "tar -zxf ${tarName}" - def llmPath = sh (script: "realpath .", returnStdout: true).trim() - def llmSrc = "${llmPath}/TensorRT-LLM/src" - trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install -r ${llmSrc}/requirements-dev.txt") - sh "NVIDIA_TRITON_SERVER_VERSION=26.05 LLM_ROOT=${llmSrc} LLM_BACKEND_ROOT=${llmSrc}/triton_backend python3 ${llmSrc}/scripts/check_test_list.py --l0 --qa --waive" - } catch (InterruptedException e) { - throw e - } catch (Exception e) { - throw e - } - }) -} - def generateTimeoutTestResultXml(pipeline, stageName) { def scriptPath = sh( script: "find . -name generate_timeout_xml.py | head -n 1 | xargs realpath", @@ -3771,7 +3747,8 @@ def renderTestDB(pipeline, testContext, llmSrc, stageName, preDefinedMakoOpts=nu } } - sh "pip3 install --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/sw-tensorrt-pypi/simple --ignore-installed trt-test-db==1.8.5+bc6df7" + def ciVersions = readProperties file: "${llmSrc}/jenkins/ci_versions.properties" + sh "pip3 install --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/sw-tensorrt-pypi/simple --ignore-installed trt-test-db==${ciVersions.TRT_TEST_DB_VERSION}" // CBTS Layer 3: download the pre-built cbts_test_db/ tarball that the // orchestrator uploaded to Artifactory (see getCbtsResult in // L0_MergeRequest.groovy). This avoids re-running main.py locally and @@ -6458,25 +6435,6 @@ pipeline { } } } - stage("Check Test List") - { - when { - expression { - // Only run the test list validation when necessary - globalVars[RUN_MODE] != "nightly_release" && - env.targetArch == X86_64_TRIPLE && - testFilter[ONLY_ONE_GROUP_CHANGED] != "Docs" && - !(env.JOB_NAME ==~ /.*Multi-GPU.*/) && - !(env.JOB_NAME ==~ /.*BuildDockerImageSanityTest.*/) - } - } - steps - { - script { - launchTestListCheck(this) - } - } - } stage("Test") { steps { script { diff --git a/jenkins/ci_versions.properties b/jenkins/ci_versions.properties new file mode 100644 index 000000000000..d2d9d4cc0e48 --- /dev/null +++ b/jenkins/ci_versions.properties @@ -0,0 +1,7 @@ +# CI tool versions shared across Jenkins pipelines and Python scripts. +# Format: KEY=VALUE (no quotes, no spaces around '='). +# +# Consumed by: +# - jenkins/L0_Test.groovy (readProperties) +# - scripts/check_test_list.py (key=value parse) +TRT_TEST_DB_VERSION=1.8.5+bc6df7 diff --git a/scripts/check_test_list.py b/scripts/check_test_list.py index 97bb4332c1d9..40c5038d2e28 100755 --- a/scripts/check_test_list.py +++ b/scripts/check_test_list.py @@ -527,17 +527,41 @@ def validate_test_lists(test_lists_dir: str, test_base_dir: str): # ============================================================================= +def _get_trt_test_db_version(): + """Read TRT_TEST_DB_VERSION from jenkins/ci_versions.properties.""" + props_file = Path( + __file__).resolve().parent.parent / "jenkins" / "ci_versions.properties" + with open(props_file) as f: + for line in f: + line = line.strip() + if line.startswith("TRT_TEST_DB_VERSION="): + return line.split("=", 1)[1] + raise RuntimeError(f"TRT_TEST_DB_VERSION not found in {props_file}") + + def install_python_dependencies(llm_src): subprocess.run(f"cd {llm_src} && pip3 install -r requirements-dev.txt", shell=True, check=True) + + whl = glob.glob(f"{llm_src}/../tensorrt_llm-*.whl") + if whl: + subprocess.run(f"pip3 install --force-reinstall --no-deps {whl[0]}", + shell=True, + check=True) + else: + # No pre-built wheel available — editable install with precompiled + # bindings downloaded from PyPI (avoids C++ compilation). + env = {**os.environ, "TRTLLM_USE_PRECOMPILED": "1"} + subprocess.run(f"cd {llm_src} && pip3 install --no-deps -e .", + shell=True, + check=True, + env=env) + + trt_test_db_ver = _get_trt_test_db_version() subprocess.run( - f"pip3 install --force-reinstall --no-deps {llm_src}/../tensorrt_llm-*.whl", - shell=True, - check=True) - subprocess.run( - "pip3 install --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/sw-tensorrt-pypi/simple " - "--ignore-installed trt-test-db==1.8.5+bc6df7", + f"pip3 install --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/sw-tensorrt-pypi/simple " + f"--ignore-installed trt-test-db=={trt_test_db_ver}", shell=True, check=True) @@ -765,7 +789,6 @@ def main(): script_dir = os.path.dirname(os.path.realpath(__file__)) llm_src = os.path.abspath(os.path.join(script_dir, "../")) - # Only skip installing dependencies if ONLY --check-duplicates or --validate is used if args.l0 or args.qa or args.waive: install_python_dependencies(llm_src) diff --git a/tests/integration/test_lists/test-db/README.md b/tests/integration/test_lists/test-db/README.md index 74e8a137cceb..658b10f01577 100644 --- a/tests/integration/test_lists/test-db/README.md +++ b/tests/integration/test_lists/test-db/README.md @@ -4,10 +4,10 @@ This folder contains test definition which is consumed by `trt-test-db` tool bas ## Installation -Install `trt-test-db` using the following command: +Install `trt-test-db` using the following command (substitute `TRT_TEST_DB_VERSION` with the version from `jenkins/ci_versions.properties`): ```bash -pip3 install --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/sw-tensorrt-pypi/simple --ignore-installed trt-test-db==1.8.5+bc6df7 +pip3 install --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/sw-tensorrt-pypi/simple --ignore-installed trt-test-db== ``` ## Test Definition From 8bc6a8beccf4c5192368f929b455e0221b18ff63 Mon Sep 17 00:00:00 2001 From: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> Date: Mon, 10 Aug 2026 12:50:30 -0700 Subject: [PATCH 2/4] stubify instead of precompile Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> --- scripts/check_test_list.py | 73 +- .../defs/accuracy/test_cli_flow.py | 870 ------------------ tests/integration/defs/conftest.py | 32 +- tests/integration/defs/stubify_bindings.py | 344 +++++++ .../integration/test_lists/test-db/README.md | 9 + tests/unittest/utils/util.py | 12 +- 6 files changed, 433 insertions(+), 907 deletions(-) delete mode 100644 tests/integration/defs/accuracy/test_cli_flow.py create mode 100644 tests/integration/defs/stubify_bindings.py diff --git a/scripts/check_test_list.py b/scripts/check_test_list.py index 40c5038d2e28..703bce83a803 100755 --- a/scripts/check_test_list.py +++ b/scripts/check_test_list.py @@ -17,6 +17,15 @@ --waive: Check only the tests in $LLM_ROOT/tests/integration/test_list/waives.txt. --validate: Run AST-based validation of test list entries against source files. +Collection stub (``--l0`` / ``--qa`` / ``--waive``): + These modes run ``pytest --co`` with ``tests/integration/defs/stubify_bindings.py`` + (loaded only via ``-p stubify_bindings``, not by default) so TensorRT-LLM need + not be compiled and no ``tensorrt_llm`` wheel is downloaded. The stub fabricates + the compiled modules on demand; its ``_EXPLICIT`` table is only for symbols + whose real *value* is consumed at import time. Full local builds will not catch + stub gaps — watch Jenkins Check Test List. Pre-commit still runs only + ``--validate`` / waive duplicate checks (no stubbed ``--co``). + Note: All the perf tests will be excluded since they are generated dynamically. """ @@ -523,7 +532,7 @@ def validate_test_lists(test_lists_dir: str, test_base_dir: str): # ============================================================================= -# L0 / QA / Waive verification (runtime, requires pytest + model weights) +# L0 / QA / Waive verification (runtime pytest --co with bindings collection stub) # ============================================================================= @@ -540,24 +549,11 @@ def _get_trt_test_db_version(): def install_python_dependencies(llm_src): + """Install Python deps for collection — no TRT-LLM wheel or compile.""" subprocess.run(f"cd {llm_src} && pip3 install -r requirements-dev.txt", shell=True, check=True) - whl = glob.glob(f"{llm_src}/../tensorrt_llm-*.whl") - if whl: - subprocess.run(f"pip3 install --force-reinstall --no-deps {whl[0]}", - shell=True, - check=True) - else: - # No pre-built wheel available — editable install with precompiled - # bindings downloaded from PyPI (avoids C++ compilation). - env = {**os.environ, "TRTLLM_USE_PRECOMPILED": "1"} - subprocess.run(f"cd {llm_src} && pip3 install --no-deps -e .", - shell=True, - check=True, - env=env) - trt_test_db_ver = _get_trt_test_db_version() subprocess.run( f"pip3 install --extra-index-url https://urm.nvidia.com/artifactory/api/pypi/sw-tensorrt-pypi/simple " @@ -566,6 +562,35 @@ def install_python_dependencies(llm_src): check=True) +def _collection_pytest_env(llm_src): + """Env for stubbed ``pytest --co``: PYTHONPATH + bindings stub flags. + + LLM_MODELS_ROOT is deliberately left alone: helpers degrade gracefully when + no models root exists, but an empty one makes model lookups fail. + """ + existing = os.environ.get("PYTHONPATH", "") + pythonpath = os.pathsep.join(p for p in (llm_src, existing) if p) + return { + **os.environ, + "PYTHONPATH": pythonpath, + "TRTLLM_BINDINGS_STUB": "1", + "TRT_LLM_NO_LIB_INIT": "1", + } + + +def _run_collection_pytest(llm_src, test_list): + """Run pytest --co with the collection bindings stub plugin.""" + env = _collection_pytest_env(llm_src) + subprocess.run( + f"cd {llm_src}/tests/integration/defs && " + f"pytest -p stubify_bindings --test-list={test_list} " + f"--output-dir={llm_src} -s --co -q", + shell=True, + check=True, + env=env, + ) + + def verify_l0_test_lists(llm_src): test_db_path = f"{llm_src}/tests/integration/test_lists/test-db" test_list = f"{llm_src}/l0_test.txt" @@ -615,11 +640,7 @@ def verify_l0_test_lists(llm_src): with open(test_list, "w") as f: f.writelines(f"{line}\n" for line in sorted(cleaned_lines)) - subprocess.run( - f"cd {llm_src}/tests/integration/defs && " - f"pytest --test-list={test_list} --output-dir={llm_src} -s --co -q", - shell=True, - check=True) + _run_collection_pytest(llm_src, test_list) def verify_qa_test_lists(llm_src): @@ -629,11 +650,7 @@ def verify_qa_test_lists(llm_src): test_def_files = subprocess.check_output( f"ls -d {test_qa_path}/*.txt", shell=True).decode().strip().split('\n') for test_def_file in test_def_files: - subprocess.run( - f"cd {llm_src}/tests/integration/defs && " - f"pytest --test-list={test_def_file} --output-dir={llm_src} -s --co -q", - shell=True, - check=True) + _run_collection_pytest(llm_src, test_def_file) # append all the test_def_file to qa_test.txt with open(f"{llm_src}/qa_test.txt", "a") as f: with open(test_def_file, "r") as test_file: @@ -744,11 +761,7 @@ def verify_waive_list(llm_src, args): with open(tmp_waives_file, "w") as f: f.writelines(f"{line}\n" for line in sorted(processed_lines)) - subprocess.run( - f"cd {llm_src}/tests/integration/defs && " - f"pytest --test-list={tmp_waives_file} --output-dir={llm_src} -s --co -q", - shell=True, - check=True) + _run_collection_pytest(llm_src, tmp_waives_file) def main(): diff --git a/tests/integration/defs/accuracy/test_cli_flow.py b/tests/integration/defs/accuracy/test_cli_flow.py deleted file mode 100644 index b61bdc21bc26..000000000000 --- a/tests/integration/defs/accuracy/test_cli_flow.py +++ /dev/null @@ -1,870 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 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. -import pytest - -from tensorrt_llm.llmapi import EagleDecodingConfig, MedusaDecodingConfig -from tensorrt_llm.quantization import QuantAlgo - -from ..conftest import (get_sm_version, llm_models_root, parametrize_with_ids, - skip_no_nvls, skip_post_blackwell, skip_pre_ada, - skip_pre_blackwell, skip_pre_hopper) -from .accuracy_core import (MMLU, CliFlowAccuracyTestHarness, CnnDailymail, - Humaneval, PassKeyRetrieval64k, ZeroScrolls) - -# skip trt flow cases on post-Blackwell-Ultra -if get_sm_version() >= 103: - pytest.skip( - "TRT workflow tests are not supported on post Blackwell-Ultra architecture", - allow_module_level=True) - - -class TestGpt2(CliFlowAccuracyTestHarness): - MODEL_NAME = "gpt2" - MODEL_PATH = f"{llm_models_root()}/gpt2" - EXAMPLE_FOLDER = "models/core/gpt" - - def test_auto_dtype(self): - # float16 - self.run(dtype='auto') - - @skip_post_blackwell - @pytest.mark.parametrize("precision", ["int8", "int4"]) - def test_weight_only(self, precision: str): - quant_algo = QuantAlgo.W8A16 if precision == "int8" else QuantAlgo.W4A16 - self.run(quant_algo=quant_algo) - - def test_beam_search(self): - self.run(extra_acc_spec="beam_width=4", - extra_build_args=["--max_beam_width=4"], - extra_summarize_args=["--num_beams=4", "--length_penalty=2.0"]) - - -class TestGpt2Medium(CliFlowAccuracyTestHarness): - MODEL_NAME = "gpt2-medium" - MODEL_PATH = f"{llm_models_root()}/gpt2-medium" - EXAMPLE_FOLDER = "models/core/gpt" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_pre_ada - def test_fp8(self): - self.run(quant_algo=QuantAlgo.FP8) - - -class TestStarcoder2_3B(CliFlowAccuracyTestHarness): - MODEL_NAME = "bigcode/starcoder2-3b" - MODEL_PATH = f"{llm_models_root()}/starcoder2-3b" - EXAMPLE_FOLDER = "models/core/gpt" - - def test_auto_dtype(self): - self.run(tasks=[Humaneval(self.MODEL_NAME)], dtype='auto') - - -class TestStarcoder2_15B(CliFlowAccuracyTestHarness): - MODEL_NAME = "bigcode/starcoder2-15b" - MODEL_PATH = f"{llm_models_root()}/starcoder2-model" - EXAMPLE_FOLDER = "models/core/gpt" - - -class TestGptNext(CliFlowAccuracyTestHarness): - MODEL_NAME = "gpt-next" - MODEL_PATH = f"{llm_models_root()}/gpt-next/megatron_converted_843m_tp1_pp1.nemo" - MODEL_FORMAT = "NEMO" - EXAMPLE_FOLDER = "models/core/gpt" - - def test_auto_dtype(self): - # bfloat16 - self.run(dtype='auto') - - -class TestMinitron4BBase(CliFlowAccuracyTestHarness): - MODEL_NAME = "nvidia/Minitron-4B-Base" - MODEL_PATH = f"{llm_models_root()}/nemotron/Minitron-4B-Base" - EXAMPLE_FOLDER = "models/core/gpt" - - def test_auto_dtype(self): - self.run(tasks=[Humaneval(self.MODEL_NAME)], dtype='auto') - - @skip_pre_ada - def test_fp8(self, mocker): - # Accuracy regression when using large batch size - mocker.patch.object(Humaneval, "MAX_BATCH_SIZE", 1) - self.run(tasks=[Humaneval(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8) - - -class TestNemotronMini4BInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "nvidia/Nemotron-Mini-4B-Instruct" - MODEL_PATH = f"{llm_models_root()}/nemotron/Nemotron-Mini-4B-Instruct" - EXAMPLE_FOLDER = "models/core/gpt" - - @skip_pre_ada - def test_fp8_prequantized(self, mocker): - mocker.patch.object( - self.__class__, "MODEL_PATH", - f"{llm_models_root()}/nemotron/nemotron-mini-4b-instruct_vfp8-fp8-bf16-export" - ) - self.run(quant_algo=QuantAlgo.FP8, kv_cache_quant_algo=QuantAlgo.FP8) - - -# TODO: Remove the CLI tests once NIMs use PyTorch backend -@pytest.mark.timeout(5400) -class TestLlama3_3NemotronSuper49Bv1(CliFlowAccuracyTestHarness): - MODEL_NAME = "nvidia/Llama-3_3-Nemotron-Super-49B-v1" - MODEL_PATH = f"{llm_models_root()}/nemotron-nas/Llama-3_3-Nemotron-Super-49B-v1" - EXAMPLE_FOLDER = "models/core/nemotron_nas" - - @pytest.mark.skip_less_device(2) - @pytest.mark.skip_less_device_memory(80000) - def test_auto_dtype_tp2(self): - self.run(tasks=[MMLU(self.MODEL_NAME)], tp_size=2, dtype='auto') - - -class TestLlama3_1NemotronNano8Bv1(CliFlowAccuracyTestHarness): - MODEL_NAME = "nvidia/Llama-3.1-Nemotron-Nano-8B-v1" - MODEL_PATH = f"{llm_models_root()}/Llama-3.1-Nemotron-Nano-8B-v1" - EXAMPLE_FOLDER = "models/core/llama" - - def test_auto_dtype(self): - self.run(tasks=[MMLU(self.MODEL_NAME)], dtype='auto') - - @skip_pre_hopper - @pytest.mark.skip_device_not_contain(["H100", "H200", "B200"]) - def test_fp8_prequantized(self, mocker): - mocker.patch.object( - self.__class__, "MODEL_PATH", - f"{llm_models_root()}/Llama-3.1-Nemotron-Nano-8B-v1-FP8") - - self.run(tasks=[MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8) - - -@pytest.mark.timeout(10800) -class TestNemotronUltra(CliFlowAccuracyTestHarness): - MODEL_NAME = "nvidia/Llama-3_1-Nemotron-Ultra-253B-v1" - MODEL_PATH = f"{llm_models_root()}/nemotron-nas/Llama-3_1-Nemotron-Ultra-253B-v1" - EXAMPLE_FOLDER = "models/core/nemotron_nas" - - @skip_pre_hopper - @pytest.mark.skip_less_device(8) - @pytest.mark.skip_less_device_memory(140000) - @parametrize_with_ids("cuda_graph", [False, True]) - @pytest.mark.parametrize("tp_size,pp_size", [(8, 1)], ids=["tp8"]) - def test_auto_dtype(self, cuda_graph, tp_size, pp_size): - extra_summarize_args = [] - if cuda_graph: - extra_summarize_args.append("--cuda_graph_mode") - - self.run(tasks=[MMLU(self.MODEL_NAME)], - tp_size=tp_size, - pp_size=pp_size, - extra_build_args=["--gemm_plugin=auto"], - extra_summarize_args=extra_summarize_args) - - @pytest.mark.skip( - reason="nemotron-nas scripts have to accommodate fp8 flags") - @skip_pre_hopper - @pytest.mark.skip_less_device(8) - @pytest.mark.skip_device_not_contain(["H100", "H200", "B200"]) - @parametrize_with_ids("cuda_graph", [False, True]) - @pytest.mark.parametrize("tp_size,pp_size", [(8, 1)], ids=["tp8"]) - def test_fp8_prequantized(self, cuda_graph, tp_size, pp_size, mocker): - mocker.patch.object( - self.__class__, "MODEL_PATH", - f"{llm_models_root()}/nemotron-nas/Llama-3_1-Nemotron-Ultra-253B-v1-FP8" - ) - - extra_summarize_args = [] - if cuda_graph: - extra_summarize_args.append("--cuda_graph_mode") - - self.run(tasks=[MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8, - tp_size=tp_size, - pp_size=pp_size, - extra_build_args=["--gemm_plugin=auto"], - extra_summarize_args=extra_summarize_args) - - -@skip_post_blackwell -class TestPhi2(CliFlowAccuracyTestHarness): - MODEL_NAME = "microsoft/phi-2" - MODEL_PATH = f"{llm_models_root()}/phi-2" - EXAMPLE_FOLDER = "models/core/phi" - - @skip_post_blackwell - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_post_blackwell - @pytest.mark.skip_less_device(2) - def test_tp2(self): - self.run(tp_size=2) - - -@skip_post_blackwell -class TestPhi3Mini4kInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct" - MODEL_PATH = f"{llm_models_root()}/Phi-3/Phi-3-mini-4k-instruct" - EXAMPLE_FOLDER = "models/core/phi" - - def test_auto_dtype(self): - self.run(dtype='auto') - - -@skip_post_blackwell -class TestPhi3Mini128kInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "microsoft/Phi-3-mini-128k-instruct" - MODEL_PATH = f"{llm_models_root()}/Phi-3/Phi-3-mini-128k-instruct" - EXAMPLE_FOLDER = "models/core/phi" - - def test_auto_dtype(self): - self.run(dtype='auto') - - -@skip_post_blackwell -class TestPhi3Small8kInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "microsoft/Phi-3-small-8k-instruct" - MODEL_PATH = f"{llm_models_root()}/Phi-3/Phi-3-small-8k-instruct" - EXAMPLE_FOLDER = "models/core/phi" - - def test_auto_dtype(self): - self.run(dtype='auto') - - -@skip_post_blackwell -class TestPhi3Small128kInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "microsoft/Phi-3-small-128k-instruct" - MODEL_PATH = f"{llm_models_root()}/Phi-3/Phi-3-small-128k-instruct" - EXAMPLE_FOLDER = "models/core/phi" - - def test_auto_dtype(self): - self.run(dtype='auto') - - -@skip_post_blackwell -class TestPhi3_5MiniInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "microsoft/Phi-3.5-mini-instruct" - MODEL_PATH = f"{llm_models_root()}/Phi-3.5/Phi-3.5-mini-instruct" - EXAMPLE_FOLDER = "models/core/phi" - - def test_auto_dtype(self): - self.run(dtype='auto') - - -class TestPhi4MiniInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "microsoft/Phi-4-mini-instruct" - MODEL_PATH = f"{llm_models_root()}/Phi-4-mini-instruct" - EXAMPLE_FOLDER = "models/core/phi" - - def test_auto_dtype(self): - self.run(tasks=[MMLU(self.MODEL_NAME)], dtype='auto') - - @pytest.mark.skip_less_device(2) - def test_tp2(self): - # Created a dummy accuracy to track tp_size=2 for phi4-mini model. - # TODO: update once https://nvbugs/5393849 is fixed. - MODEL_NAME = "microsoft/Phi-4-mini-instruct-tp2" - self.run(tasks=[MMLU(MODEL_NAME)], tp_size=2) - - -# Long sequence length test: -# Model FP16 7B + 32K tokens in KV cache = 14 * 1024 MB + 32K * 0.5 MB = 30720 MB + scratch memory -@pytest.mark.skip_less_device_memory(40000) -class TestLongAlpaca7B(CliFlowAccuracyTestHarness): - MODEL_NAME = "Yukang/LongAlpaca-7B" - MODEL_PATH = f"{llm_models_root()}/LongAlpaca-7B" - EXAMPLE_FOLDER = "models/core/llama" - - def test_auto_dtype(self): - self.run(tasks=[ZeroScrolls(self.MODEL_NAME)]) - - def test_multiblock_aggressive(self): - # MMHA + aggressive Multi_block_mode (export TRTLLM_ENABLE_MMHA_MULTI_BLOCK_DEBUG=1) - self.run(tasks=[ZeroScrolls(self.MODEL_NAME)], - extra_build_args=["--gemm_plugin=auto"], - env={ - "TRTLLM_ENABLE_MMHA_MULTI_BLOCK_DEBUG": "1", - "TRTLLM_MMHA_BLOCKS_PER_SEQUENCE": "32" - }) - - -class TestMamba130M(CliFlowAccuracyTestHarness): - MODEL_NAME = "state-spaces/mamba-130m-hf" - MODEL_PATH = f"{llm_models_root()}/mamba/mamba-130m-hf" - EXAMPLE_FOLDER = "models/core/mamba" - - def test_auto_dtype(self): - self.run(dtype='auto') - - -class TestVicuna7B(CliFlowAccuracyTestHarness): - MODEL_NAME = "lmsys/vicuna-7b-v1.3" - MODEL_PATH = f"{llm_models_root()}/vicuna-7b-v1.3" - EXAMPLE_FOLDER = "models/core/llama" - MEDUSA_MODEL_NAME = "FasterDecoding/medusa-vicuna-7b-v1.3" - MEDUSA_MODEL_PATH = f"{llm_models_root()}/medusa-vicuna-7b-v1.3" - EAGLE_MODEL_NAME = "yuhuili/EAGLE-Vicuna-7B-v1.3" - EAGLE_MODEL_PATH = f"{llm_models_root()}/EAGLE-Vicuna-7B-v1.3" - - @skip_post_blackwell - @parametrize_with_ids("cuda_graph", [False, True]) - def test_medusa(self, cuda_graph, mocker): - mocker.patch.object(self.__class__, "EXAMPLE_FOLDER", "medusa") - mocker.patch.object(CnnDailymail, "MAX_BATCH_SIZE", 8) - - extra_summarize_args = [ - "--medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" - ] - if cuda_graph: - extra_summarize_args.append("--cuda_graph_mode") - - self.run(dtype="float16", - spec_dec_algo=MedusaDecodingConfig. - model_fields["decoding_type"].default, - extra_convert_args=[ - f"--medusa_model_dir={self.MEDUSA_MODEL_PATH}", - "--num_medusa_heads=4" - ], - extra_build_args=["--speculative_decoding_mode=medusa"], - extra_summarize_args=extra_summarize_args) - - @skip_post_blackwell - @parametrize_with_ids("cuda_graph,chunked_context,typical_acceptance", - [(False, False, False), (True, False, False), - (True, True, False), (True, False, True)]) - def test_eagle(self, cuda_graph, chunked_context, typical_acceptance, - mocker): - mocker.patch.object(self.__class__, "EXAMPLE_FOLDER", "eagle") - mocker.patch.object(CnnDailymail, "MAX_BATCH_SIZE", 8) - - extra_summarize_args = [ - "--eagle_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" - ] - if cuda_graph: - extra_summarize_args.append("--cuda_graph_mode") - if chunked_context: - extra_summarize_args.append("--enable_chunked_context") - if typical_acceptance: - extra_summarize_args.extend( - ["--eagle_posterior_threshold=0.09", "--temperature=0.7"]) - - self.run(spec_dec_algo=EagleDecodingConfig. - model_fields["decoding_type"].default, - extra_convert_args=[ - f"--eagle_model_dir={self.EAGLE_MODEL_PATH}", - "--max_draft_len=63", "--num_eagle_layers=4", - "--max_non_leaves_per_layer=10" - ], - extra_build_args=[ - "--speculative_decoding_mode=eagle", "--max_draft_len=63" - ], - extra_summarize_args=extra_summarize_args) - - -class TestTinyLlama1_1BChat(CliFlowAccuracyTestHarness): - MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" - MODEL_PATH = f"{llm_models_root()}/llama-models-v2/TinyLlama-1.1B-Chat-v1.0" - EXAMPLE_FOLDER = "models/core/llama" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_post_blackwell - @pytest.mark.parametrize("precision", ["int8", "int4"]) - def test_weight_only(self, precision: str): - quant_algo = QuantAlgo.W8A16 if precision == "int8" else QuantAlgo.W4A16 - self.run(quant_algo=quant_algo) - - @skip_pre_ada - def test_fp8(self): - self.run(quant_algo=QuantAlgo.FP8, kv_cache_quant_algo=QuantAlgo.FP8) - - @pytest.mark.skip_less_device(4) - def test_pp4(self): - # Test num_hidden_layers (22) undivisible by pp_size (4) - self.run(extra_acc_spec="pp_size=4", pp_size=4) - - -class TestLlama3_8BInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct" - MODEL_PATH = f"{llm_models_root()}/llama-models-v3/llama-v3-8b-instruct-hf" - EXAMPLE_FOLDER = "models/core/llama" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_pre_ada - def test_fp8(self): - self.run(quant_algo=QuantAlgo.FP8, kv_cache_quant_algo=QuantAlgo.FP8) - - @skip_pre_blackwell - def test_nvfp4(self): - self.run(tasks=[MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.NVFP4, - kv_cache_quant_algo=QuantAlgo.FP8, - extra_build_args=["--gemm_plugin=disable"]) - - @pytest.mark.skip( - reason="Broken by modelopt. Will be fixed in next release") - @skip_pre_blackwell - @pytest.mark.parametrize("fuse_fp4_quant", [False, True], - ids=["disable_fused_quant", "enable_fused_quant"]) - @pytest.mark.parametrize( - "norm_quant_fusion", [False, True], - ids=["disable_norm_quant_fusion", "enable_norm_quant_fusion"]) - def test_nvfp4_gemm_plugin(self, fuse_fp4_quant: bool, - norm_quant_fusion: bool): - extra_build_args = ["--gemm_plugin=nvfp4"] - if fuse_fp4_quant: - extra_build_args.extend([ - "--use_paged_context_fmha=enable", - "--use_fp8_context_fmha=enable", "--fuse_fp4_quant=enable" - ]) - if norm_quant_fusion: - extra_build_args.append("--norm_quant_fusion=enable") - self.run(tasks=[MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.NVFP4, - kv_cache_quant_algo=QuantAlgo.FP8, - extra_build_args=extra_build_args) - - -class TestLlama3_8BInstructGradient1048k(CliFlowAccuracyTestHarness): - MODEL_NAME = "gradientai/Llama-3-8B-Instruct-Gradient-1048k" - MODEL_PATH = f"{llm_models_root()}/llama-models-v3/Llama-3-8B-Instruct-Gradient-1048k" - EXAMPLE_FOLDER = "models/core/llama" - - -class TestLlama3_1_8B(CliFlowAccuracyTestHarness): - MODEL_NAME = "meta-llama/Llama-3.1-8B" - MODEL_PATH = f"{llm_models_root()}/llama-3.1-model/Meta-Llama-3.1-8B" - EXAMPLE_FOLDER = "models/core/llama" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_pre_ada - def test_fp8(self): - self.run(quant_algo=QuantAlgo.FP8, kv_cache_quant_algo=QuantAlgo.FP8) - - @skip_pre_ada - @skip_post_blackwell - def test_fp8_rowwise(self): - self.run(tasks=[CnnDailymail(self.MODEL_NAME), - MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN) - - @skip_pre_ada - @skip_post_blackwell - def test_fp8_rowwise_meta_recipe(self): - self.run(quant_algo=QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, - extra_acc_spec="meta_recipe", - extra_convert_args=["--use_meta_fp8_rowwise_recipe"]) - - @pytest.mark.skip_less_device(4) - @pytest.mark.parametrize( - "gemm_allreduce", [False, pytest.param(True, marks=skip_no_nvls)], - ids=["disable_gemm_allreduce_plugin", "enable_gemm_allreduce_plugin"]) - def test_tp4(self, gemm_allreduce: bool): - extra_build_args = None - if gemm_allreduce: - extra_build_args = ["--gemm_allreduce_plugin=bfloat16"] - self.run( - tasks=[PassKeyRetrieval64k(self.MODEL_NAME), - MMLU(self.MODEL_NAME)], - tp_size=4, - extra_build_args=extra_build_args) - - @skip_pre_hopper - @skip_post_blackwell - @pytest.mark.skip_less_device(4) - @pytest.mark.parametrize( - "gemm_allreduce", [False, pytest.param(True, marks=skip_no_nvls)], - ids=["disable_gemm_allreduce_plugin", "enable_gemm_allreduce_plugin"]) - def test_fp8_rowwise_tp4(self, gemm_allreduce: bool): - extra_build_args = None - if gemm_allreduce: - extra_build_args = ["--gemm_allreduce_plugin=bfloat16"] - self.run( - tasks=[PassKeyRetrieval64k(self.MODEL_NAME), - MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, - tp_size=4, - extra_build_args=extra_build_args) - - -class TestLlama3_1_8BInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct" - MODEL_PATH = f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct" - EXAMPLE_FOLDER = "models/core/llama" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_pre_hopper - def test_fp8_prequantized(self, mocker): - mocker.patch.object( - self.__class__, "MODEL_PATH", - f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct-FP8") - self.run(quant_algo=QuantAlgo.FP8, kv_cache_quant_algo=QuantAlgo.FP8) - - @skip_pre_hopper - @skip_post_blackwell - def test_medusa_fp8_prequantized(self, mocker): - # nvidia/Llama-3.1-8B-Medusa-FP8 - mocker.patch.object(self.__class__, "MODEL_PATH", - f"{llm_models_root()}/llama3.1-medusa-8b-hf_v0.1") - mocker.patch.object(self.__class__, "EXAMPLE_FOLDER", "medusa") - mocker.patch.object(CnnDailymail, "MAX_BATCH_SIZE", 8) - - extra_summarize_args = [ - "--medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [1, 6], [0, 7, 0]]" - ] - self.run(dtype="float16", - spec_dec_algo=MedusaDecodingConfig. - model_fields["decoding_type"].default, - extra_build_args=["--speculative_decoding_mode=medusa"], - extra_summarize_args=extra_summarize_args) - - -class TestLlama3_2_1B(CliFlowAccuracyTestHarness): - MODEL_NAME = "meta-llama/Llama-3.2-1B" - MODEL_PATH = f"{llm_models_root()}/llama-3.2-models/Llama-3.2-1B" - EXAMPLE_FOLDER = "models/core/llama" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_pre_ada - def test_fp8(self): - self.run(quant_algo=QuantAlgo.FP8, kv_cache_quant_algo=QuantAlgo.FP8) - - @skip_pre_ada - @pytest.mark.skip_less_device(2) - @pytest.mark.parametrize( - "fp8_context_fmha", [False, True], - ids=["disable_fp8_context_fmha", "enable_fp8_context_fmha"]) - @pytest.mark.parametrize( - "reduce_fusion", [False, True], - ids=["disable_reduce_fusion", "enable_reduce_fusion"]) - def test_fp8_tp2(self, fp8_context_fmha: bool, reduce_fusion: bool): - if fp8_context_fmha: - extra_build_args = [ - "--use_fp8_context_fmha=enable", - "--use_paged_context_fmha=enable" - ] - else: - extra_build_args = [ - "--use_fp8_context_fmha=disable", - "--use_paged_context_fmha=disable" - ] - - if reduce_fusion: - extra_build_args.append("--reduce_fusion=enable") - else: - extra_build_args.append("--reduce_fusion=disable") - - self.run(quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8, - tp_size=2, - extra_build_args=extra_build_args) - - @skip_pre_ada - @skip_post_blackwell - def test_fp8_rowwise(self): - self.run(quant_algo=QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN) - - @skip_pre_ada - @skip_post_blackwell - def test_fp8_rowwise_meta_recipe(self): - self.run(quant_algo=QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, - extra_acc_spec="meta_recipe", - extra_convert_args=["--use_meta_fp8_rowwise_recipe"]) - - @pytest.mark.parametrize("max_gpu_percent", [0.1, 1.0]) - def test_weight_streaming(self, max_gpu_percent: float): - self.run(extra_build_args=["--weight_streaming"], - extra_summarize_args=["--gpu_weights_percent=0"]) - - for gpu_percent in [0.1, 0.5, 0.9, 1]: - if gpu_percent > max_gpu_percent: - break - self.extra_summarize_args = [f"--gpu_weights_percent={gpu_percent}"] - self.evaluate() - - -# TODO: Remove the CLI tests once NIMs use PyTorch backend -@pytest.mark.skip_less_device_memory(80000) -class TestLlama3_3_70BInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct" - MODEL_PATH = f"{llm_models_root()}/llama-3.3-models/Llama-3.3-70B-Instruct" - EXAMPLE_FOLDER = "models/core/llama" - - @pytest.mark.skip_less_device(8) - def test_auto_dtype_tp8(self): - self.run(tasks=[MMLU(self.MODEL_NAME)], tp_size=8, dtype='auto') - - -class TestMistral7B(CliFlowAccuracyTestHarness): - MODEL_NAME = "mistralai/Mistral-7B-v0.1" - MODEL_PATH = f"{llm_models_root()}/mistral-7b-v0.1" - EXAMPLE_FOLDER = "models/core/llama" - - @skip_pre_blackwell - def test_beam_search(self): - self.run(extra_acc_spec="beam_width=4", - extra_build_args=["--gemm_plugin=auto", "--max_beam_width=4"], - extra_summarize_args=["--num_beams=4"]) - import gc - - import torch - for num_beams in [1, 2]: - gc.collect() - torch.cuda.empty_cache() - self.extra_acc_spec = f"beam_width={num_beams}" - self.extra_summarize_args = [f"--num_beams={num_beams}"] - self.evaluate() - - -class TestMixtral8x7B(CliFlowAccuracyTestHarness): - MODEL_NAME = "mistralai/Mixtral-8x7B-v0.1" - MODEL_PATH = f"{llm_models_root()}/Mixtral-8x7B-v0.1" - EXAMPLE_FOLDER = "models/core/llama" - - @pytest.mark.skip_less_device(2) - @pytest.mark.skip_less_device_memory(80000) - def test_tp2(self): - self.run(dtype='auto', tp_size=2) - - @skip_pre_ada - @pytest.mark.skip_less_device(2) - @pytest.mark.skip_less_device_memory(80000) - def test_fp8_tp2(self): - self.run(quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8, - tp_size=2) - - @skip_pre_ada - @pytest.mark.skip_less_device(4) - @pytest.mark.skip_less_device_memory(40000) - def test_fp8_tp2pp2(self): - self.run(tasks=[CnnDailymail(self.MODEL_NAME), - MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8, - tp_size=2, - pp_size=2) - - @skip_pre_ada - @pytest.mark.skip_less_device(4) - @pytest.mark.skip_less_device_memory(40000) - def test_fp8_tp2pp2_manage_weights(self): - self.run(tasks=[CnnDailymail(self.MODEL_NAME), - MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8, - tp_size=2, - pp_size=2, - extra_build_args=["--fast_build"]) - - @skip_pre_blackwell - def test_nvfp4_prequantized(self, mocker): - mocker.patch.object( - self.__class__, "MODEL_PATH", - f"{llm_models_root()}/nvfp4-quantized/Mixtral-8x7B-Instruct-v0.1") - self.run(tasks=[MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.NVFP4, - kv_cache_quant_algo=QuantAlgo.FP8) - - -class TestMixtral8x22B(CliFlowAccuracyTestHarness): - MODEL_NAME = "mistralai/Mixtral-8x22B-v0.1" - MODEL_PATH = f"{llm_models_root()}/Mixtral-8x22B-v0.1" - EXAMPLE_FOLDER = "models/core/llama" - - @skip_pre_ada - @pytest.mark.skip_less_device(4) - @pytest.mark.skip_less_device_memory(80000) - def test_fp8_tp2pp2(self, timeout_manager): - self.run(tasks=[CnnDailymail(self.MODEL_NAME), - MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8, - tp_size=2, - pp_size=2, - extra_convert_args=["--calib_size=32"], - extra_build_args=["--gemm_plugin=auto"], - timeout_manager=timeout_manager) - - -class TestGemma2B(CliFlowAccuracyTestHarness): - MODEL_NAME = "google/gemma-2b" - MODEL_PATH = f"{llm_models_root()}/gemma/gemma-2b" - EXAMPLE_FOLDER = "models/core/gemma" - - def test_auto_dtype(self): - self.run(dtype='auto', extra_convert_args=["--ckpt-type=hf"]) - - @pytest.mark.parametrize("precision", ["int8"]) - def test_weight_only(self, precision: str): - quant_algo = QuantAlgo.W8A16 if precision == "int8" else QuantAlgo.W4A16 - self.run(quant_algo=quant_algo, extra_convert_args=["--ckpt-type=hf"]) - - @skip_pre_ada - def test_fp8(self): - self.run(quant_algo=QuantAlgo.FP8, kv_cache_quant_algo=QuantAlgo.FP8) - - -@pytest.mark.skip_less_device_memory(40000) -class TestGemma7B(CliFlowAccuracyTestHarness): - MODEL_NAME = "google/gemma-7b" - MODEL_PATH = f"{llm_models_root()}/gemma/gemma-7b" - EXAMPLE_FOLDER = "models/core/gemma" - - def test_auto_dtype(self): - self.run(dtype='auto', extra_convert_args=["--ckpt-type=hf"]) - - @pytest.mark.parametrize("precision", ["int8"]) - def test_weight_only(self, precision: str): - quant_algo = QuantAlgo.W8A16 if precision == "int8" else QuantAlgo.W4A16 - self.run(quant_algo=quant_algo, extra_convert_args=["--ckpt-type=hf"]) - - @skip_pre_ada - def test_fp8(self): - self.run(quant_algo=QuantAlgo.FP8, kv_cache_quant_algo=QuantAlgo.FP8) - - -@pytest.mark.skip_less_device_memory(40000) -class TestGemma2_9BIt(CliFlowAccuracyTestHarness): - MODEL_NAME = "google/gemma-2-9b-it" - MODEL_PATH = f"{llm_models_root()}/gemma/gemma-2-9b-it" - EXAMPLE_FOLDER = "models/core/gemma" - - @skip_post_blackwell - def test_auto_dtype(self): - self.run(tasks=[CnnDailymail(self.MODEL_NAME), - MMLU(self.MODEL_NAME)], - dtype='auto', - extra_convert_args=["--ckpt-type=hf"]) - - @skip_post_blackwell - @pytest.mark.parametrize("precision", ["int8", "int4"]) - def test_weight_only(self, precision: str): - quant_algo = QuantAlgo.W8A16 if precision == "int8" else QuantAlgo.W4A16 - self.run(quant_algo=quant_algo, extra_convert_args=["--ckpt-type=hf"]) - - @skip_pre_hopper - def test_fp8(self): - self.run(quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8, - extra_convert_args=["--device_map=sequential"]) - - -class TestQwen7BChat(CliFlowAccuracyTestHarness): - MODEL_NAME = "Qwen/Qwen-7B-Chat" - MODEL_PATH = f"{llm_models_root()}/Qwen-7B-Chat" - EXAMPLE_FOLDER = "models/core/qwen" - - def test_auto_dtype(self): - self.run(dtype='auto') - - def test_weight_only(self): - self.run(quant_algo=QuantAlgo.W8A16) - - -@pytest.mark.skip_less_device_memory(40000) -class TestQwen1_5MoeA2_7BChat(CliFlowAccuracyTestHarness): - MODEL_NAME = "Qwen/Qwen1.5-MoE-A2.7B-Chat" - MODEL_PATH = f"{llm_models_root()}/Qwen1.5-MoE-A2.7B-Chat" - EXAMPLE_FOLDER = "models/core/qwen" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @pytest.mark.skip(reason="https://nvbugs/5100102") - def test_weight_only(self): - self.run(quant_algo=QuantAlgo.W8A16) - - -class TestQwen2_0_5BInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "Qwen/Qwen2-0.5B-Instruct" - MODEL_PATH = f"{llm_models_root()}/Qwen2-0.5B-Instruct" - EXAMPLE_FOLDER = "models/core/qwen" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_post_blackwell - def test_weight_only(self): - self.run(quant_algo=QuantAlgo.W8A16) - - @skip_pre_ada - def test_fp8(self): - self.run(tasks=[CnnDailymail(self.MODEL_NAME), - MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8) - - -class TestQwen2_1_5B(CliFlowAccuracyTestHarness): - MODEL_NAME = "Qwen/Qwen2-1.5B" - MODEL_PATH = f"{llm_models_root()}/Qwen2-1.5B" - EXAMPLE_FOLDER = "models/core/qwen" - - -class TestQwen2_7BInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "Qwen/Qwen2-7B-Instruct" - MODEL_PATH = f"{llm_models_root()}/Qwen2-7B-Instruct" - EXAMPLE_FOLDER = "models/core/qwen" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_post_blackwell - def test_weight_only(self): - self.run(quant_algo=QuantAlgo.W8A16) - - -@pytest.mark.skip_less_device_memory(40000) -class TestQwen2_57B_A14B(CliFlowAccuracyTestHarness): - MODEL_NAME = "Qwen/Qwen2-57B-A14B" - MODEL_PATH = f"{llm_models_root()}/Qwen2-57B-A14B" - EXAMPLE_FOLDER = "models/core/qwen" - - @pytest.mark.skip(reason="https://nvbugs/5063469") - @pytest.mark.skip_less_device(4) - def test_tp4(self): - self.run(tp_size=4) - - -class TestQwen2_5_1_5BInstruct(CliFlowAccuracyTestHarness): - MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" - MODEL_PATH = f"{llm_models_root()}/Qwen2.5-1.5B-Instruct" - EXAMPLE_FOLDER = "models/core/qwen" - - def test_auto_dtype(self): - self.run(dtype='auto') - - @skip_post_blackwell - def test_weight_only(self): - self.run(quant_algo=QuantAlgo.W8A16) - - @skip_pre_ada - def test_fp8(self): - self.run(tasks=[CnnDailymail(self.MODEL_NAME), - MMLU(self.MODEL_NAME)], - quant_algo=QuantAlgo.FP8) diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 4e9742a89905..ec7047028aa4 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -48,9 +48,6 @@ # is harmless. from test_common import session_prefetcher_hooks as _prefetch_hooks -from tensorrt_llm.bindings import ipc_nvls_supported -from tensorrt_llm.llmapi.mpi_session import get_mpi_world_size - from .perf.gpu_clock_lock import GPUClockLock from .perf.session_data_writer import SessionDataWriter from .test_list_parser import (TestCorrectionMode, apply_waives, @@ -106,6 +103,12 @@ def llm_models_root() -> str: if not root.exists(): root = Path("/scratch.trt_llm_data/llm-models/") + if not root.exists() and os.environ.get("TRTLLM_BINDINGS_STUB") == "1": + # Collection-only (Check Test List): class attributes such as + # MODEL_PATH = f"{llm_models_root()}/..." just need a string, and no + # weights are ever read. Real runs keep the assert below. + return str(root) + assert root.exists(), ( "You shall set LLM_MODELS_ROOT env or be able to access scratch.trt_llm_data to run this test" ) @@ -1538,6 +1541,13 @@ def skip_by_device_count(request): f"Device count {device_count} is less than {expected_count}") +def get_mpi_world_size() -> int: + """Lazy wrapper: keeps bindings out of conftest import (also re-exported).""" + from tensorrt_llm.llmapi.mpi_session import \ + get_mpi_world_size as _get_mpi_world_size + return _get_mpi_world_size() + + @pytest.fixture(autouse=True) def skip_by_mpi_world_size(request): "fixture for skip less mpi world size" @@ -1612,6 +1622,7 @@ def is_ipc_nvls_supported(): if not torch.cuda.is_available(): return False try: + from tensorrt_llm.bindings import ipc_nvls_supported return ipc_nvls_supported() except RuntimeError: return False @@ -1661,8 +1672,19 @@ def is_ipc_nvls_supported(): reason="This test is not supported on GB200 or GB100", ) -skip_no_nvls = pytest.mark.skipif(not is_ipc_nvls_supported(), - reason="NVLS is not supported") + +def skip_no_nvls(func=None): + """Skip when NVLS is unsupported. + + Deferred so importing conftest does not call into bindings. + """ + mark = pytest.mark.skipif(not is_ipc_nvls_supported(), + reason="NVLS is not supported") + if func is not None: + return mark(func) + return mark + + skip_no_hopper = pytest.mark.skipif( get_sm_version() != 90, reason="This test is only supported in Hopper architecture") diff --git a/tests/integration/defs/stubify_bindings.py b/tests/integration/defs/stubify_bindings.py new file mode 100644 index 000000000000..112b292cd900 --- /dev/null +++ b/tests/integration/defs/stubify_bindings.py @@ -0,0 +1,344 @@ +# 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. +"""Pytest plugin: pure-Python stubs of TensorRT-LLM's compiled modules. + +Used by ``scripts/check_test_list.py`` so ``pytest --co`` can import product +code without a compiled ``bindings*.so`` or a prebuilt wheel. Covers the +build-generated Python modules listed in ``_STUB_ROOTS`` plus the +``libth_common.so`` surfaces torch exposes (see ``stub_torch_extensions``). + +Factory-first. The factory fabricates modules, classes and values on demand, +and deliberately exposes *empty* introspection surfaces (``dir()`` has no +public non-callable names, ``__members__`` is empty) so that PybindMirror's +field/enum mirroring in ``llm_args`` passes without per-symbol entries. + +Add ``_EXPLICIT`` entries only when Check Test List fails because a real +*value* is required at import time. Prefer removing +eager imports in tests/conftest over growing this table. + +Collection-only — never a substitute for running tests. +""" + +from __future__ import annotations + +import sys +import types +from abc import ABCMeta +from importlib.machinery import ModuleSpec +from pathlib import Path + +# --------------------------------------------------------------------------- +# Escape hatch +# --------------------------------------------------------------------------- +# Keyed by fully qualified name. The factory can fake any *shape*, but not a +# concrete value that product code consumes at import time. ``llm_args`` +# evaluates these lookahead getters in a class body to seed pydantic field +# defaults, so they must return the real C++ constants +# (``kDefaultLookaheadDecoding*`` in cpp/include/tensorrt_llm/executor/executor.h). + + +class LookaheadDecodingConfig: + """Explicit stub: real defaults consumed by ``llm_args`` at class-body time.""" + + @staticmethod + def get_default_lookahead_decoding_window(): + return 4 + + @staticmethod + def get_default_lookahead_decoding_ngram(): + return 3 + + @staticmethod + def get_default_lookahead_decoding_verification_set(): + return 4 + + +_EXPLICIT: dict[str, object] = { + "tensorrt_llm.bindings.executor.LookaheadDecodingConfig": LookaheadDecodingConfig, +} + +# --------------------------------------------------------------------------- +# Factory +# --------------------------------------------------------------------------- + +_BINDINGS = "tensorrt_llm.bindings" + +# Everything the C++ build generates under tensorrt_llm/ and that product code +# imports at module scope (see .gitignore / setup.py package_data). +_STUB_ROOTS = ( + _BINDINGS, + "tensorrt_llm.deep_ep", + "tensorrt_llm.deep_ep_cpp_tllm", + "tensorrt_llm.deep_gemm", + "tensorrt_llm.deep_gemm_cpp_tllm", + "tensorrt_llm.flash_mla", + "tensorrt_llm.flash_mla_cpp_tllm", + "tensorrt_llm.pg_utils_bindings", + "tensorrt_llm.tensorrt_llm_transfer_agent_binding", +) + +# Submodules that must resolve to modules even though they are not lower_case, +# so attribute access does not hit the CapWords "this is a class" rule. +_FORCED_SUBMODULES = (f"{_BINDINGS}.BuildInfo",) + +# torch::class_ namespaces registered by libs/libth_common.so, which +# TRT_LLM_NO_LIB_INIT=1 skips loading. Probing an unregistered class raises +# RuntimeError (not AttributeError), so hasattr() checks in product code blow +# up; an empty namespace makes them report "unavailable", which is the truth +# for a no-compile checkout. +_TORCH_CLASS_NAMESPACES = ("trtllm",) + + +class _StubMeta(ABCMeta): + """Metaclass for fabricated binding classes. + + Derives from ``ABCMeta`` so product classes can inherit from both a stubbed + binding type and an ABC without a metaclass conflict. + """ + + def __getattr__(cls, name: str): + # Never fabricate dunders: Python and pydantic probe them to decide + # which protocols a type supports. + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + + children = cls.__dict__.get("_stub_children") + if children is None: + children = {} + type.__setattr__(cls, "_stub_children", children) + if name not in children: + # Cached and distinct per name so enum-style members stay usable as + # dict keys (e.g. the DataType maps built in tensorrt_llm/_utils.py). + children[name] = _make_stub_class(f"{cls.__name__}.{name}", cls.__module__) + return children[name] + + @property + def __members__(cls): + # PybindMirror.mirror_pybind_enum iterates the C++ members and requires + # each to exist on the Python enum; an empty mapping trivially passes. + return {} + + def __iter__(cls): + # Product code materializes some binding sequences at import time + # (e.g. tuple(KVCacheIterationStatsDelta._field_names)). + return iter(()) + + def __int__(cls): + # Stubbed enum members are coerced at import time + # (e.g. int(BufferKind.DEFAULT) in cute_dsl_custom_ops.py). + return 0 + + def __index__(cls): + return 0 + + +class _StubBase(metaclass=_StubMeta): + """Base for fabricated binding classes; only dunders, so ``dir()`` is clean.""" + + def __init__(self, *args, **kwargs): + pass + + def __getattr__(self, name: str): + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + return _make_stub_class(f"{type(self).__name__}.{name}", type(self).__module__) + + def __call__(self, *args, **kwargs): + return False + + def __bool__(self): + return False + + def __iter__(self): + return iter(()) + + def __int__(self): + return 0 + + def __index__(self): + return 0 + + +def _make_stub_class(name: str, module: str) -> type: + return _StubMeta(name, (_StubBase,), {"__module__": module}) + + +class StubModule(types.ModuleType): + """Fake (sub)module of a stubbed root, resolving attributes on demand.""" + + def __init__(self, fullname: str): + super().__init__(fullname) + self.__file__ = f"" + self.__package__ = fullname + self.__path__ = [] # marks this as a package for nested imports + object.__setattr__(self, "_stub_cache", {}) + + def __getattr__(self, name: str): + cache = object.__getattribute__(self, "_stub_cache") + if name in cache: + return cache[name] + + fq = f"{self.__name__}.{name}" + if fq in _EXPLICIT: + value = _EXPLICIT[fq] + elif name.isupper(): + # Module-level constants such as BuildInfo.ENABLE_MULTI_DEVICE: + # falsy keeps mpi_session and communicator off their mpi4py paths. + value = 0 + elif name[:1].isupper(): + value = _make_stub_class(name, self.__name__) + else: + # Lower case: either a submodule or a free function. A module is + # both importable and callable, so it covers each case. + value = StubModule(fq) + sys.modules[fq] = value + + cache[name] = value + return value + + def __call__(self, *args, **kwargs): + return False + + def __bool__(self): + return False + + def __repr__(self): + return f"" + + +class _StubFinder: + """Meta path finder fabricating any module under a stubbed root. + + Attribute access alone is not enough: ``import a.b.c`` and + ``from a.b import c`` go through the import system, which never consults a + parent module's ``__getattr__``. + """ + + @staticmethod + def find_spec(fullname, path=None, target=None): + if not any(fullname == root or fullname.startswith(f"{root}.") for root in _STUB_ROOTS): + return None + return ModuleSpec(fullname, _StubLoader(), is_package=True) + + +class _StubLoader: + @staticmethod + def create_module(spec): + return StubModule(spec.name) + + @staticmethod + def exec_module(module): + pass + + +def _real_bindings_present() -> bool: + """True if a compiled bindings extension is loaded or present on disk. + + Avoids ``importlib.util.find_spec``, which would execute + ``tensorrt_llm/__init__.py`` before the stub is installed. + """ + existing = sys.modules.get(_BINDINGS) + if existing is not None and not isinstance(existing, StubModule): + return True + + for entry in sys.path: + pkg = Path(entry) / "tensorrt_llm" + if not pkg.is_dir(): + continue + for pattern in ("bindings*.so", "bindings*.pyd", "_bindings*.so"): + if list(pkg.glob(pattern)): + return True + return False + + +def install_bindings_stub() -> StubModule | None: + """Install stub modules for every ``_STUB_ROOTS`` entry into ``sys.modules``. + + Returns the ``tensorrt_llm.bindings`` stub, or None when real bindings are + present (the stub refuses to mask a real install). + """ + installed = sys.modules.get(_BINDINGS) + if isinstance(installed, StubModule): + # Idempotent: reinstalling would orphan the stub classes already + # captured by imported product modules. + return installed + + if _real_bindings_present(): + return None + + # Appended, so a real installation's finders always take precedence. + if not any(isinstance(f, _StubFinder) for f in sys.meta_path): + sys.meta_path.append(_StubFinder()) + + for root in _STUB_ROOTS: + sys.modules.setdefault(root, StubModule(root)) + + for sub in _FORCED_SUBMODULES: + if sub in sys.modules: + continue + child = StubModule(sub) + sys.modules[sub] = child + parent_name, _, attr = sub.rpartition(".") + parent = sys.modules.get(parent_name) + if isinstance(parent, StubModule): + object.__getattribute__(parent, "_stub_cache")[attr] = child + + return sys.modules[_BINDINGS] + + +def stub_torch_extensions() -> None: + """Neutralize the parts of torch that expect ``libth_common.so`` to be loaded.""" + if _real_bindings_present(): + return + + try: + import torch + except ImportError: + return + + for namespace in _TORCH_CLASS_NAMESPACES: + if namespace not in torch.classes.__dict__: + setattr(torch.classes, namespace, types.ModuleType(f"torch.classes.{namespace}")) + + # Product modules register fake kernels for C++ ops at import time; without + # the library the schemas are missing, so make those registrations no-ops. + register_fake = torch.library.register_fake + if getattr(register_fake, "_trtllm_collection_stub", False): + return + + def tolerant_register_fake(op, func=None, /, **kwargs): + def apply(fn): + try: + return register_fake(op, fn, **kwargs) + except RuntimeError as exc: + if "does not exist" not in str(exc): + raise + return fn + + return apply(func) if func is not None else apply + + tolerant_register_fake._trtllm_collection_stub = True + torch.library.register_fake = tolerant_register_fake + + +# Install on import so `pytest -p stubify_bindings` wins the race with +# conftest collection. +install_bindings_stub() + + +def pytest_configure(config): + """Re-assert the stubs early in the pytest session.""" + install_bindings_stub() + stub_torch_extensions() diff --git a/tests/integration/test_lists/test-db/README.md b/tests/integration/test_lists/test-db/README.md index 658b10f01577..bd92fc229c88 100644 --- a/tests/integration/test_lists/test-db/README.md +++ b/tests/integration/test_lists/test-db/README.md @@ -58,6 +58,15 @@ pytest -v --test-list=/TensorRT-LLM/src/l0_e2e.txt --output-dir=/tmp/logs This command runs the tests specified in the test list and outputs the results to the specified directory. +## Check Test List (CI) + +Jenkins **Check Test List** runs `scripts/check_test_list.py --l0 --qa --waive --validate`. +L0/QA/waive collection uses pure-Python stubs of the compiled modules +(`tests/integration/defs/stubify_bindings.py`, loaded only via +`-p stubify_bindings`), so no TensorRT-LLM wheel or C++ +compile is required. Pre-commit runs check_test_list.py only with `--validate` and +`--check-duplicate-waives`, which does not require stubs. + ## Additional Information - The `--context` parameter in the `trt-test-db` command specifies which context to search in the YAML files. - The `--match-exact` parameter provides system information used to filter tests based on the conditions defined in the YAML files. diff --git a/tests/unittest/utils/util.py b/tests/unittest/utils/util.py index 606a51cd11f5..63e4c036a1e6 100644 --- a/tests/unittest/utils/util.py +++ b/tests/unittest/utils/util.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -56,8 +56,16 @@ def ASSERT_DRV(err): # ref: https://github.com/NVIDIA/cuda-python/blob/main/examples/extra/jit_program_test.py def getSMVersion(): # Init - err, = cuda.cuInit(0) + try: + err, = cuda.cuInit(0) + except RuntimeError as exc: + # Broken or missing driver. Proceed without skipping to allow test collection on CPU hosts. + print(f"WARNING: CUDA driver init failed in getSMVersion(): {exc}") + return math.inf if err == cuda.CUresult.CUDA_ERROR_NO_DEVICE: + print( + "WARNING: CUDA reports no device in getSMVersion(). Tests that require a GPU will be skipped." + ) return -1 ASSERT_DRV(err) From 73b67011609379e3a1dfdd8ef06f2e8777a8ab6a Mon Sep 17 00:00:00 2001 From: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:33:38 -0700 Subject: [PATCH 3/4] lazily evaluate model paths, and remove TRTLLM_BINDINGS_STUB Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> --- scripts/check_test_list.py | 25 +++--- .../defs/accuracy/test_llm_api_autodeploy.py | 79 ++++++++----------- tests/integration/defs/conftest.py | 6 -- 3 files changed, 46 insertions(+), 64 deletions(-) diff --git a/scripts/check_test_list.py b/scripts/check_test_list.py index 703bce83a803..387f1a058f08 100755 --- a/scripts/check_test_list.py +++ b/scripts/check_test_list.py @@ -18,13 +18,10 @@ --validate: Run AST-based validation of test list entries against source files. Collection stub (``--l0`` / ``--qa`` / ``--waive``): - These modes run ``pytest --co`` with ``tests/integration/defs/stubify_bindings.py`` - (loaded only via ``-p stubify_bindings``, not by default) so TensorRT-LLM need - not be compiled and no ``tensorrt_llm`` wheel is downloaded. The stub fabricates - the compiled modules on demand; its ``_EXPLICIT`` table is only for symbols - whose real *value* is consumed at import time. Full local builds will not catch - stub gaps — watch Jenkins Check Test List. Pre-commit still runs only - ``--validate`` / waive duplicate checks (no stubbed ``--co``). + These modes run ``pytest --co`` with the plugin + ``tests/integration/defs/stubify_bindings.py`` to avoid compiling or + downloading TRT-LLM C++ binaries. Instead, the plugin creates a pure-Python + stub of the compiled modules. Note: All the perf tests will be excluded since they are generated dynamically. @@ -563,18 +560,18 @@ def install_python_dependencies(llm_src): def _collection_pytest_env(llm_src): - """Env for stubbed ``pytest --co``: PYTHONPATH + bindings stub flags. - - LLM_MODELS_ROOT is deliberately left alone: helpers degrade gracefully when - no models root exists, but an empty one makes model lookups fail. - """ + """Env for stubbed ``pytest --co``: PYTHONPATH + placeholder models root.""" + # The stubify_bindings plugin needs to be in the PYTHONPATH so it can be imported by pytest. + defs_dir = os.path.join(llm_src, "tests", "integration", "defs") existing = os.environ.get("PYTHONPATH", "") - pythonpath = os.pathsep.join(p for p in (llm_src, existing) if p) + pythonpath = os.pathsep.join(p for p in (llm_src, defs_dir, existing) if p) return { **os.environ, "PYTHONPATH": pythonpath, - "TRTLLM_BINDINGS_STUB": "1", "TRT_LLM_NO_LIB_INIT": "1", + # Collection only needs llm_models_root() to be a directory. + # Fixtures and weight loads do not run under pytest --co. + "LLM_MODELS_ROOT": "/tmp", } diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 1a77e5d48604..c76f740c000a 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -195,7 +195,6 @@ def reduced_model_kwargs(num_hidden_layers: int, class TestLlama3_1_8B(LlmapiAccuracyTestHarness): MODEL_NAME = "meta-llama/Llama-3.1-8B" - MODEL_PATH = hf_id_to_local_model_dir(MODEL_NAME) # Configuration presets for different attention backends ATTN_BACKEND_CONFIGS = { @@ -298,10 +297,11 @@ def get_default_sampling_params(self): ], ) def test_auto_dtype(self, world_size, enable_chunked_prefill, attn_backend): + model_path = hf_id_to_local_model_dir(self.MODEL_NAME) kwargs = self.get_default_kwargs(enable_chunked_prefill, attn_backend) sampling_params = self.get_default_sampling_params() - with AutoDeployLLM(model=self.MODEL_PATH, - tokenizer=self.MODEL_PATH, + with AutoDeployLLM(model=model_path, + tokenizer=model_path, world_size=world_size, **kwargs) as llm: task = CnnDailymail(self.MODEL_NAME) @@ -317,12 +317,13 @@ def test_auto_dtype(self, world_size, enable_chunked_prefill, attn_backend): ]) def test_attention_dp(self, world_size): """Test attention data parallelism mode where TP sharding is disabled.""" + model_path = hf_id_to_local_model_dir(self.MODEL_NAME) kwargs = self.get_default_kwargs(enable_chunked_prefill=True) # Enable attention DP - this disables TP sharding kwargs["transforms"]["detect_sharding"] = {"enable_attention_dp": True} sampling_params = self.get_default_sampling_params() - with AutoDeployLLM(model=self.MODEL_PATH, - tokenizer=self.MODEL_PATH, + with AutoDeployLLM(model=model_path, + tokenizer=model_path, world_size=world_size, **kwargs) as llm: task = CnnDailymail(self.MODEL_NAME) @@ -335,15 +336,13 @@ class TestLlama3_1_8B_Instruct_Eagle3(LlmapiAccuracyTestHarness): """Accuracy test for Eagle3 one-model speculative decoding with AutoDeploy.""" MODEL_NAME = "meta-llama/Llama-3.1-8B-Instruct" - MODEL_PATH = hf_id_to_local_model_dir(MODEL_NAME) - EAGLE_MODEL_PATH = hf_id_to_local_model_dir( - "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B") + EAGLE_MODEL_NAME = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" def get_default_kwargs(self, attn_backend="flashinfer"): yaml_paths, _ = _get_registry_yaml_extra(self.MODEL_NAME) speculative_config = Eagle3DecodingConfig( max_draft_len=3, - speculative_model=self.EAGLE_MODEL_PATH, + speculative_model=hf_id_to_local_model_dir(self.EAGLE_MODEL_NAME), eagle3_one_model=True, eagle3_layers_to_capture={1, 15, 28}, ) @@ -391,11 +390,12 @@ def check_acceptance_rate(self, llm, min_acceptance_rate: float): @pytest.mark.parametrize("attn_backend", ["flashinfer", "trtllm"]) def test_eagle3_one_model(self, attn_backend): """Test Eagle3 one-model speculative decoding accuracy on GSM8K.""" + model_path = hf_id_to_local_model_dir(self.MODEL_NAME) kwargs = self.get_default_kwargs(attn_backend=attn_backend) with AutoDeployLLM( - model=self.MODEL_PATH, - tokenizer=self.MODEL_PATH, + model=model_path, + tokenizer=model_path, **kwargs, ) as llm: task = GSM8K(self.MODEL_NAME) @@ -548,13 +548,10 @@ class TestNemotronNanoV3(LlmapiAccuracyTestHarness): CONFIG_YAML = str( Path(get_llm_root()) / "examples" / "auto_deploy" / "nano_v3.yaml") - MODEL_PATHS = { - "bf16": - hf_id_to_local_model_dir("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"), - "fp8": - hf_id_to_local_model_dir("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"), - "nvfp4": - hf_id_to_local_model_dir("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4"), + MODEL_NAMES = { + "bf16": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", + "fp8": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8", + "nvfp4": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4", } def get_default_sampling_params(self): @@ -587,7 +584,7 @@ def test_accuracy(self, model_id, world_size, enable_attention_dp, # max_dp_num_tokens path; on world_size=1 it's a no-op. if enable_attention_dp and world_size < 2: pytest.skip("attention_dp requires world_size >= 2") - model_path = self.MODEL_PATHS[model_id] + model_path = hf_id_to_local_model_dir(self.MODEL_NAMES[model_id]) kwargs = {} device_memory_mib = get_device_memory() # bf16 always needs low-memory overrides; below H100-class total @@ -622,16 +619,10 @@ class TestNemotronSuperV3(LlmapiAccuracyTestHarness): MODEL_NAME = "nvidia/Nemotron-Super-V3" CONFIG_YAML = str( Path(get_llm_root()) / "examples" / "auto_deploy" / "super_v3.yaml") - MODEL_PATHS = { - "bf16": - hf_id_to_local_model_dir( - "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16"), - "fp8": - hf_id_to_local_model_dir( - "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8"), - "nvfp4": - hf_id_to_local_model_dir( - "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4"), + MODEL_NAMES = { + "bf16": "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16", + "fp8": "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8", + "nvfp4": "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4", } def get_default_sampling_params(self): @@ -667,7 +658,7 @@ def test_accuracy(self, model_id, world_size, enable_attention_dp, if model_id == "bf16" and world_size < 4: pytest.skip("bf16 Super V3 requires at least 4 GPUs") - model_path = self.MODEL_PATHS[model_id] + model_path = hf_id_to_local_model_dir(self.MODEL_NAMES[model_id]) kwargs = {} if model_id == "bf16": low_memory_overrides(kwargs) @@ -712,7 +703,7 @@ def test_functional_small(self, dtype): No accuracy threshold is checked — the truncated model is not expected to produce meaningful text. """ - model_path = self.MODEL_PATHS[dtype] + model_path = hf_id_to_local_model_dir(self.MODEL_NAMES[dtype]) kwargs = {} kwargs.update( reduced_model_kwargs(num_hidden_layers=16, model_path=model_path)) @@ -783,7 +774,7 @@ def test_functional_small(self, dtype): ) def test_mtp(self, world_size, attn_backend, model_id): - model_path = self.MODEL_PATHS[model_id] + model_path = hf_id_to_local_model_dir(self.MODEL_NAMES[model_id]) kwargs = {} # TODO: gate for bf16 only after replay lands low_memory_overrides( @@ -837,8 +828,8 @@ class TestNemotronUltraV3(LlmapiAccuracyTestHarness): CONFIG_YAML = str( Path(get_llm_root()) / "examples" / "auto_deploy" / "model_registry" / "configs" / "ultra_v3.yaml") - MODEL_PATHS = { - "nvfp4": hf_id_to_local_model_dir("nvidia/Nemotron-Ultra-V3-NVFP4"), + MODEL_NAMES = { + "nvfp4": "nvidia/Nemotron-Ultra-V3-NVFP4", } def get_default_sampling_params(self): @@ -856,7 +847,7 @@ def test_accuracy(self, model_id, world_size): if get_device_count() < world_size: pytest.skip(f"Not enough devices for world_size={world_size}") - model_path = self.MODEL_PATHS[model_id] + model_path = hf_id_to_local_model_dir(self.MODEL_NAMES[model_id]) print_memory_usage("test start") with AutoDeployLLM( model=model_path, @@ -1018,7 +1009,6 @@ class TestQwen3_5_397B_MoE(LlmapiAccuracyTestHarness): MODEL_NAME = "Qwen/Qwen3.5-397B-A17B" MODEL_NAME_NVFP4 = "nvidia/Qwen3.5-397B-A17B-NVFP4" MODEL_NAME_SMALL = "Qwen/Qwen3.5-35B-A3B" - MODEL_PATH_SMALL = hf_id_to_local_model_dir(MODEL_NAME_SMALL) GSM8K_MAX_OUTPUT_LEN = 512 EXTRA_EVALUATOR_KWARGS = dict( apply_chat_template=True, @@ -1104,8 +1094,9 @@ def test_bf16_small(self, world_size): if get_device_count() < world_size: pytest.skip("Not enough devices for world size, skipping test") sampling_params = self.get_default_sampling_params() - with AutoDeployLLM(model=self.MODEL_PATH_SMALL, - tokenizer=self.MODEL_PATH_SMALL, + model_path = hf_id_to_local_model_dir(self.MODEL_NAME_SMALL) + with AutoDeployLLM(model=model_path, + tokenizer=model_path, dtype="bfloat16", world_size=world_size, **config) as llm: @@ -1129,7 +1120,6 @@ class TestMiniMaxM2(LlmapiAccuracyTestHarness): """ MODEL_NAME = "MiniMaxAI/MiniMax-M2" - MODEL_PATH = hf_id_to_local_model_dir(MODEL_NAME) # Set minimum possible seq len + small buffer, for test speed & memory usage MAX_SEQ_LEN = max(MMLU.MAX_INPUT_LEN + MMLU.MAX_OUTPUT_LEN, GSM8K.MAX_INPUT_LEN + GSM8K.MAX_OUTPUT_LEN) @@ -1158,9 +1148,10 @@ def get_default_kwargs(self): @skip_pre_hopper @pytest.mark.skip_less_device(4) def test_finegrained_fp8(self): + model_path = hf_id_to_local_model_dir(self.MODEL_NAME) kwargs = self.get_default_kwargs() - with AutoDeployLLM(model=self.MODEL_PATH, - tokenizer=self.MODEL_PATH, + with AutoDeployLLM(model=model_path, + tokenizer=model_path, world_size=4, **kwargs) as llm: task = MMLU(self.MODEL_NAME) @@ -1333,7 +1324,6 @@ class TestGemma4MoE(LlmapiAccuracyTestHarness): """Bench-run coverage for Gemma4 MoE via AutoDeploy.""" MODEL_NAME = "google/gemma-4-26B-A4B-it" - MODEL_PATH = hf_id_to_local_model_dir(MODEL_NAME) EXTRA_EVALUATOR_KWARGS = { "apply_chat_template": True, } @@ -1357,8 +1347,9 @@ def test_bf16(self): pytest.skip("Not enough devices for world size, skipping test") sampling_params = self.get_default_sampling_params() - with AutoDeployLLM(model=self.MODEL_PATH, - tokenizer=self.MODEL_PATH, + model_path = hf_id_to_local_model_dir(self.MODEL_NAME) + with AutoDeployLLM(model=model_path, + tokenizer=model_path, world_size=registry_world_size, yaml_extra=yaml_paths) as llm: task = MMMU(self.MODEL_NAME) # noqa: F821 diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index ec7047028aa4..aa80b720f393 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -103,12 +103,6 @@ def llm_models_root() -> str: if not root.exists(): root = Path("/scratch.trt_llm_data/llm-models/") - if not root.exists() and os.environ.get("TRTLLM_BINDINGS_STUB") == "1": - # Collection-only (Check Test List): class attributes such as - # MODEL_PATH = f"{llm_models_root()}/..." just need a string, and no - # weights are ever read. Real runs keep the assert below. - return str(root) - assert root.exists(), ( "You shall set LLM_MODELS_ROOT env or be able to access scratch.trt_llm_data to run this test" ) From 07cde59cb29d5f80a1ab27bf228f83e7d78f7e70 Mon Sep 17 00:00:00 2001 From: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:43:54 -0700 Subject: [PATCH 4/4] apply coding guidelines Signed-off-by: Tyler Burt <195370667+tburt-nv@users.noreply.github.com> --- scripts/check_test_list.py | 8 +- tests/integration/defs/conftest.py | 7 +- tests/integration/defs/stubify_bindings.py | 112 ++++++++++++++------- 3 files changed, 84 insertions(+), 43 deletions(-) diff --git a/scripts/check_test_list.py b/scripts/check_test_list.py index 387f1a058f08..b9f8c577a624 100755 --- a/scripts/check_test_list.py +++ b/scripts/check_test_list.py @@ -533,7 +533,7 @@ def validate_test_lists(test_lists_dir: str, test_base_dir: str): # ============================================================================= -def _get_trt_test_db_version(): +def _get_trt_test_db_version() -> str: """Read TRT_TEST_DB_VERSION from jenkins/ci_versions.properties.""" props_file = Path( __file__).resolve().parent.parent / "jenkins" / "ci_versions.properties" @@ -545,7 +545,7 @@ def _get_trt_test_db_version(): raise RuntimeError(f"TRT_TEST_DB_VERSION not found in {props_file}") -def install_python_dependencies(llm_src): +def install_python_dependencies(llm_src: str) -> None: """Install Python deps for collection — no TRT-LLM wheel or compile.""" subprocess.run(f"cd {llm_src} && pip3 install -r requirements-dev.txt", shell=True, @@ -559,7 +559,7 @@ def install_python_dependencies(llm_src): check=True) -def _collection_pytest_env(llm_src): +def _collection_pytest_env(llm_src: str) -> dict[str, str]: """Env for stubbed ``pytest --co``: PYTHONPATH + placeholder models root.""" # The stubify_bindings plugin needs to be in the PYTHONPATH so it can be imported by pytest. defs_dir = os.path.join(llm_src, "tests", "integration", "defs") @@ -575,7 +575,7 @@ def _collection_pytest_env(llm_src): } -def _run_collection_pytest(llm_src, test_list): +def _run_collection_pytest(llm_src: str, test_list: str) -> None: """Run pytest --co with the collection bindings stub plugin.""" env = _collection_pytest_env(llm_src) subprocess.run( diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index aa80b720f393..b94ee5653997 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -30,7 +30,7 @@ import warnings from functools import wraps from pathlib import Path -from typing import Iterable, Sequence +from typing import Callable, Iterable, Sequence import defs.ci_profiler import psutil @@ -39,6 +39,7 @@ import tqdm import yaml from _pytest.mark import ParameterSet +from _pytest.mark.structures import MarkDecorator # Dispatched explicitly (not via pytest_plugins, which pytest forbids in a # non-top-level conftest: a repo-root invocation like `pytest tests` loads # this file as a NESTED conftest and would fail collection; and not via "-p" @@ -1667,7 +1668,9 @@ def is_ipc_nvls_supported(): ) -def skip_no_nvls(func=None): +def skip_no_nvls( + func: Callable[..., object] | None = None +) -> Callable[..., object] | MarkDecorator: """Skip when NVLS is unsupported. Deferred so importing conftest does not call into bindings. diff --git a/tests/integration/defs/stubify_bindings.py b/tests/integration/defs/stubify_bindings.py index 112b292cd900..ac8d88f28931 100644 --- a/tests/integration/defs/stubify_bindings.py +++ b/tests/integration/defs/stubify_bindings.py @@ -36,8 +36,18 @@ import sys import types from abc import ABCMeta +from collections.abc import Callable, Iterator, Sequence +from importlib import abc from importlib.machinery import ModuleSpec from pathlib import Path +from types import ModuleType +from typing import TYPE_CHECKING, ParamSpec, TypeVar, cast + +if TYPE_CHECKING: + import pytest + +_P = ParamSpec("_P") +_R = TypeVar("_R") # --------------------------------------------------------------------------- # Escape hatch @@ -53,19 +63,19 @@ class LookaheadDecodingConfig: """Explicit stub: real defaults consumed by ``llm_args`` at class-body time.""" @staticmethod - def get_default_lookahead_decoding_window(): + def get_default_lookahead_decoding_window() -> int: return 4 @staticmethod - def get_default_lookahead_decoding_ngram(): + def get_default_lookahead_decoding_ngram() -> int: return 3 @staticmethod - def get_default_lookahead_decoding_verification_set(): + def get_default_lookahead_decoding_verification_set() -> int: return 4 -_EXPLICIT: dict[str, object] = { +_EXPLICIT: dict[str, type[LookaheadDecodingConfig]] = { "tensorrt_llm.bindings.executor.LookaheadDecodingConfig": LookaheadDecodingConfig, } @@ -108,13 +118,16 @@ class _StubMeta(ABCMeta): binding type and an ABC without a metaclass conflict. """ - def __getattr__(cls, name: str): + def __getattr__(cls, name: str) -> type[_StubBase]: # Never fabricate dunders: Python and pydantic probe them to decide # which protocols a type supports. if name.startswith("__") and name.endswith("__"): raise AttributeError(name) - children = cls.__dict__.get("_stub_children") + children = cast( + dict[str, type[_StubBase]] | None, + cls.__dict__.get("_stub_children"), + ) if children is None: children = {} type.__setattr__(cls, "_stub_children", children) @@ -125,72 +138,84 @@ def __getattr__(cls, name: str): return children[name] @property - def __members__(cls): + def __members__(cls) -> dict[str, type[_StubBase]]: # PybindMirror.mirror_pybind_enum iterates the C++ members and requires # each to exist on the Python enum; an empty mapping trivially passes. return {} - def __iter__(cls): + def __iter__(cls) -> Iterator[type[_StubBase]]: # Product code materializes some binding sequences at import time # (e.g. tuple(KVCacheIterationStatsDelta._field_names)). return iter(()) - def __int__(cls): + def __int__(cls) -> int: # Stubbed enum members are coerced at import time # (e.g. int(BufferKind.DEFAULT) in cute_dsl_custom_ops.py). return 0 - def __index__(cls): + def __index__(cls) -> int: return 0 class _StubBase(metaclass=_StubMeta): """Base for fabricated binding classes; only dunders, so ``dir()`` is clean.""" - def __init__(self, *args, **kwargs): + def __init__(self, *args: object, **kwargs: object) -> None: pass - def __getattr__(self, name: str): + def __getattr__(self, name: str) -> type[_StubBase]: if name.startswith("__") and name.endswith("__"): raise AttributeError(name) return _make_stub_class(f"{type(self).__name__}.{name}", type(self).__module__) - def __call__(self, *args, **kwargs): + def __call__(self, *args: object, **kwargs: object) -> bool: return False - def __bool__(self): + def __bool__(self) -> bool: return False - def __iter__(self): + def __iter__(self) -> Iterator[type[_StubBase]]: return iter(()) - def __int__(self): + def __int__(self) -> int: return 0 - def __index__(self): + def __index__(self) -> int: return 0 -def _make_stub_class(name: str, module: str) -> type: - return _StubMeta(name, (_StubBase,), {"__module__": module}) +def _make_stub_class(name: str, module: str) -> type[_StubBase]: + return cast( + type[_StubBase], + _StubMeta(name, (_StubBase,), {"__module__": module}), + ) class StubModule(types.ModuleType): """Fake (sub)module of a stubbed root, resolving attributes on demand.""" - def __init__(self, fullname: str): + def __init__(self, fullname: str) -> None: super().__init__(fullname) self.__file__ = f"" self.__package__ = fullname self.__path__ = [] # marks this as a package for nested imports object.__setattr__(self, "_stub_cache", {}) - def __getattr__(self, name: str): - cache = object.__getattribute__(self, "_stub_cache") + def __getattr__( + self, name: str + ) -> int | type[LookaheadDecodingConfig] | type[_StubBase] | StubModule: + cache = cast( + dict[ + str, + int | type[LookaheadDecodingConfig] | type[_StubBase] | StubModule, + ], + object.__getattribute__(self, "_stub_cache"), + ) if name in cache: return cache[name] fq = f"{self.__name__}.{name}" + value: int | type[LookaheadDecodingConfig] | type[_StubBase] | StubModule if fq in _EXPLICIT: value = _EXPLICIT[fq] elif name.isupper(): @@ -208,17 +233,17 @@ def __getattr__(self, name: str): cache[name] = value return value - def __call__(self, *args, **kwargs): + def __call__(self, *args: object, **kwargs: object) -> bool: return False - def __bool__(self): + def __bool__(self) -> bool: return False - def __repr__(self): + def __repr__(self) -> str: return f"" -class _StubFinder: +class _StubFinder(abc.MetaPathFinder): """Meta path finder fabricating any module under a stubbed root. Attribute access alone is not enough: ``import a.b.c`` and @@ -227,20 +252,25 @@ class _StubFinder: """ @staticmethod - def find_spec(fullname, path=None, target=None): + def find_spec( + fullname: str, + path: Sequence[str] | None = None, + target: ModuleType | None = None, + ) -> ModuleSpec | None: + del path, target if not any(fullname == root or fullname.startswith(f"{root}.") for root in _STUB_ROOTS): return None return ModuleSpec(fullname, _StubLoader(), is_package=True) -class _StubLoader: +class _StubLoader(abc.Loader): @staticmethod - def create_module(spec): + def create_module(spec: ModuleSpec) -> StubModule: return StubModule(spec.name) @staticmethod - def exec_module(module): - pass + def exec_module(module: ModuleType) -> None: + del module def _real_bindings_present() -> bool: @@ -295,7 +325,9 @@ def install_bindings_stub() -> StubModule | None: if isinstance(parent, StubModule): object.__getattribute__(parent, "_stub_cache")[attr] = child - return sys.modules[_BINDINGS] + bindings = sys.modules[_BINDINGS] + assert isinstance(bindings, StubModule) + return bindings def stub_torch_extensions() -> None: @@ -318,10 +350,15 @@ def stub_torch_extensions() -> None: if getattr(register_fake, "_trtllm_collection_stub", False): return - def tolerant_register_fake(op, func=None, /, **kwargs): - def apply(fn): + def tolerant_register_fake( + op: str, + func: Callable[_P, _R] | None = None, + /, + **kwargs: object, + ) -> Callable[_P, _R] | Callable[[Callable[_P, _R]], Callable[_P, _R]]: + def apply(fn: Callable[_P, _R]) -> Callable[_P, _R]: try: - return register_fake(op, fn, **kwargs) + return cast(Callable[_P, _R], register_fake(op, fn, **kwargs)) except RuntimeError as exc: if "does not exist" not in str(exc): raise @@ -329,7 +366,7 @@ def apply(fn): return apply(func) if func is not None else apply - tolerant_register_fake._trtllm_collection_stub = True + setattr(tolerant_register_fake, "_trtllm_collection_stub", True) torch.library.register_fake = tolerant_register_fake @@ -338,7 +375,8 @@ def apply(fn): install_bindings_stub() -def pytest_configure(config): +def pytest_configure(config: pytest.Config) -> None: """Re-assert the stubs early in the pytest session.""" + del config install_bindings_stub() stub_torch_extensions()