From 689e2ad3266ca76dd437e6eeee772aedf99e4918 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Mon, 2 Feb 2026 10:42:34 -0800 Subject: [PATCH 1/4] Fix mocking Signed-off-by: Anish Shanbhag --- tests/test_common/llm_data.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_common/llm_data.py b/tests/test_common/llm_data.py index a58d7dc6c3a0..22452cab29d2 100644 --- a/tests/test_common/llm_data.py +++ b/tests/test_common/llm_data.py @@ -104,11 +104,17 @@ def with_mocked_hf_download(func): When applied, any calls to snapshot_download will be redirected to use local model paths from LLM_MODELS_ROOT instead of downloading from HuggingFace. + + NOTE: We must patch snapshot_download at the location where it's actually imported + with 'from huggingface_hub import snapshot_download', since that creates a + local binding that won't be affected by patching huggingface_hub.snapshot_download. """ @wraps(func) def wrapper(*args, **kwargs): - with patch("huggingface_hub.snapshot_download", side_effect=mock_snapshot_download): + with patch( + "tensorrt_llm.llmapi.utils.snapshot_download", side_effect=mock_snapshot_download + ): return func(*args, **kwargs) return wrapper From c388b0461774f6c7787ccfb2fc227003c5610b3f Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Mon, 2 Feb 2026 11:11:37 -0800 Subject: [PATCH 2/4] Set HF_HUB_OFFLINE Signed-off-by: Anish Shanbhag --- tests/test_common/llm_data.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/test_common/llm_data.py b/tests/test_common/llm_data.py index 22452cab29d2..20a903294690 100644 --- a/tests/test_common/llm_data.py +++ b/tests/test_common/llm_data.py @@ -108,13 +108,25 @@ def with_mocked_hf_download(func): NOTE: We must patch snapshot_download at the location where it's actually imported with 'from huggingface_hub import snapshot_download', since that creates a local binding that won't be affected by patching huggingface_hub.snapshot_download. + + Additionally sets HF_HUB_OFFLINE=1 to ensure no network requests are made to + HuggingFace. """ @wraps(func) def wrapper(*args, **kwargs): - with patch( - "tensorrt_llm.llmapi.utils.snapshot_download", side_effect=mock_snapshot_download - ): - return func(*args, **kwargs) + original_offline = os.environ.get("HF_HUB_OFFLINE") + os.environ["HF_HUB_OFFLINE"] = "1" + + try: + with patch( + "tensorrt_llm.llmapi.utils.snapshot_download", side_effect=mock_snapshot_download + ): + return func(*args, **kwargs) + finally: + if original_offline is None: + os.environ.pop("HF_HUB_OFFLINE", None) + else: + os.environ["HF_HUB_OFFLINE"] = original_offline return wrapper From 77f7a0651a22a29d2e16902efe999bd2929a4be4 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Mon, 2 Feb 2026 11:23:33 -0800 Subject: [PATCH 3/4] Use patch.dict Signed-off-by: Anish Shanbhag --- tests/test_common/llm_data.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/tests/test_common/llm_data.py b/tests/test_common/llm_data.py index 20a903294690..59769e93b0fe 100644 --- a/tests/test_common/llm_data.py +++ b/tests/test_common/llm_data.py @@ -115,18 +115,12 @@ def with_mocked_hf_download(func): @wraps(func) def wrapper(*args, **kwargs): - original_offline = os.environ.get("HF_HUB_OFFLINE") - os.environ["HF_HUB_OFFLINE"] = "1" - - try: - with patch( + with ( + patch.dict(os.environ, {"HF_HUB_OFFLINE": "1"}), + patch( "tensorrt_llm.llmapi.utils.snapshot_download", side_effect=mock_snapshot_download - ): - return func(*args, **kwargs) - finally: - if original_offline is None: - os.environ.pop("HF_HUB_OFFLINE", None) - else: - os.environ["HF_HUB_OFFLINE"] = original_offline + ), + ): + return func(*args, **kwargs) return wrapper From fd7333adadc0f65ffd922e62d9b77256c3e1a041 Mon Sep 17 00:00:00 2001 From: Anish Shanbhag Date: Mon, 2 Feb 2026 12:00:24 -0800 Subject: [PATCH 4/4] Add warning for multi GPU Signed-off-by: Anish Shanbhag --- tests/test_common/llm_data.py | 5 ++++- .../unit/singlegpu/test_ad_speculative_decoding.py | 4 ++-- tests/unittest/_torch/speculative/test_eagle3.py | 4 ++-- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/test_common/llm_data.py b/tests/test_common/llm_data.py index 59769e93b0fe..a3855b8e8c46 100644 --- a/tests/test_common/llm_data.py +++ b/tests/test_common/llm_data.py @@ -99,7 +99,7 @@ def mock_snapshot_download(repo_id: str, **kwargs) -> str: return local_path -def with_mocked_hf_download(func): +def with_mocked_hf_download_for_single_gpu(func): """Decorator to mock huggingface_hub.snapshot_download for tests. When applied, any calls to snapshot_download will be redirected to use @@ -111,6 +111,9 @@ def with_mocked_hf_download(func): Additionally sets HF_HUB_OFFLINE=1 to ensure no network requests are made to HuggingFace. + + WARNING: This decorator only works for single-GPU tests. For multi-GPU tests, the + mock won't be applied in MPI worker processes. """ @wraps(func) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_speculative_decoding.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_speculative_decoding.py index 81481e8f51dd..cb3655411abb 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_speculative_decoding.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_speculative_decoding.py @@ -16,13 +16,13 @@ import pytest from _model_test_utils import get_small_model_config from build_and_run_ad import ExperimentConfig, main -from test_common.llm_data import with_mocked_hf_download +from test_common.llm_data import with_mocked_hf_download_for_single_gpu from tensorrt_llm.llmapi import DraftTargetDecodingConfig, KvCacheConfig @pytest.mark.parametrize("use_hf_speculative_model", [False, True]) -@with_mocked_hf_download +@with_mocked_hf_download_for_single_gpu def test_ad_speculative_decoding_smoke(use_hf_speculative_model: bool): """Test speculative decoding with AutoDeploy using the build_and_run_ad main().""" diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index c2d4cf50f4c1..7ef69a00be48 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -8,7 +8,7 @@ import pytest import torch -from test_common.llm_data import with_mocked_hf_download +from test_common.llm_data import with_mocked_hf_download_for_single_gpu from utils.llm_data import llm_models_root from tensorrt_llm import LLM, SamplingParams @@ -150,7 +150,7 @@ def test_kv_lens_runtime_with_eagle3_one_model(): [False, "TRTLLM", True, False, False, False, True, False, False, True], ]) @pytest.mark.high_cuda_memory -@with_mocked_hf_download +@with_mocked_hf_download_for_single_gpu def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, disable_overlap_scheduler: bool, enable_block_reuse: bool, use_one_model: bool, enable_chunked_prefill: bool,