[CI] Stop flaky test from downloading model every time (#49800)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jul 25, 2026 at 14:29 UTC
9a504646985848c1f5785fa57b5f0afc297a29e6
1 file changed
+16
-20
tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py
+16
-20
@@ -2,9 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import glob
5
-import tempfile
5
7
-import huggingface_hub.constants
6
import pytest
7
import torch
8
@@ -21,29 +19,27 @@ from vllm.platforms import current_platform
19
reason="InstantTensor requires NVIDIA GPUs",
20
)
21
def test_instanttensor_model_loader():
24
- with tempfile.TemporaryDirectory() as tmpdir:
25
- huggingface_hub.constants.HF_HUB_OFFLINE = False
26
- download_weights_from_hf(
27
- "openai-community/gpt2", allow_patterns=["*.safetensors"], cache_dir=tmpdir
28
- )
29
- safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
30
- assert len(safetensors) > 0
22
+ model_dir = download_weights_from_hf(
23
+ "openai-community/gpt2", cache_dir=None, allow_patterns=["*.safetensors"]
24
+ )
25
+ safetensors = glob.glob(f"{model_dir}/*.safetensors")
26
+ assert len(safetensors) > 0
27
32
- instanttensor_tensors = {}
33
- hf_safetensors_tensors = {}
28
+ instanttensor_tensors = {}
29
+ hf_safetensors_tensors = {}
30
35
- for name, tensor in instanttensor_weights_iterator(safetensors, True):
36
- instanttensor_tensors[name] = tensor.to("cpu")
31
+ for name, tensor in instanttensor_weights_iterator(safetensors, True):
32
+ instanttensor_tensors[name] = tensor.to("cpu")
33
38
- for name, tensor in safetensors_weights_iterator(safetensors, True):
39
- hf_safetensors_tensors[name] = tensor
34
+ for name, tensor in safetensors_weights_iterator(safetensors, True):
35
+ hf_safetensors_tensors[name] = tensor
36
41
- assert len(instanttensor_tensors) == len(hf_safetensors_tensors)
37
+ assert len(instanttensor_tensors) == len(hf_safetensors_tensors)
38
43
- for name, instanttensor_tensor in instanttensor_tensors.items():
44
- assert instanttensor_tensor.dtype == hf_safetensors_tensors[name].dtype
45
- assert instanttensor_tensor.shape == hf_safetensors_tensors[name].shape
46
- assert torch.all(instanttensor_tensor.eq(hf_safetensors_tensors[name]))
39
+ for name, instanttensor_tensor in instanttensor_tensors.items():
40
+ assert instanttensor_tensor.dtype == hf_safetensors_tensors[name].dtype
41
+ assert instanttensor_tensor.shape == hf_safetensors_tensors[name].shape
42
+ assert torch.all(instanttensor_tensor.eq(hf_safetensors_tensors[name]))
43
44
45
if __name__ == "__main__":