Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions tests/test_common/llm_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,16 +99,31 @@ 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
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.

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)
def wrapper(*args, **kwargs):
with patch("huggingface_hub.snapshot_download", side_effect=mock_snapshot_download):
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)

return wrapper
Original file line number Diff line number Diff line change
Expand Up @@ -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()."""

Expand Down
4 changes: 2 additions & 2 deletions tests/unittest/_torch/speculative/test_eagle3.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down