[PT FE] Retry hf related tests (#26030)

### Details: - *item1* - *...* ### Tickets: - *CVS-149231* - *CVS-149229*
openvinotoolkit · Aug 12, 2024 · 96c8661 · 96c8661
1 parent 7dbceea
commit 96c8661
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 4 deletions.
diff --git a/tests/model_hub_tests/pytorch/test_hf_transformers.py b/tests/model_hub_tests/pytorch/test_hf_transformers.py
@@ -5,7 +5,6 @@
 
 from datasets import Audio, load_dataset
 from huggingface_hub import hf_hub_download, model_info
-from huggingface_hub.utils import HfHubHTTPError, LocalEntryNotFoundError
 from PIL import Image
 import pytest
 import torch
@@ -56,7 +55,7 @@ def setup_class(self):
         url = "http://images.cocodataset.org/val2017/000000039769.jpg"
         self.image = Image.open(requests.get(url, stream=True).raw)
 
-    @retry(3, exceptions=(HfHubHTTPError, LocalEntryNotFoundError), delay=1)
+    @retry(3, exceptions=(OSError,), delay=1)
     def load_model(self, name, type):
         name, _, name_suffix = name.partition(':')
 

diff --git a/tests/model_hub_tests/pytorch/test_timm.py b/tests/model_hub_tests/pytorch/test_timm.py
@@ -6,7 +6,6 @@
 import pytest
 import timm
 import torch
-from huggingface_hub.utils import HfHubHTTPError, LocalEntryNotFoundError
 from models_hub_common.utils import get_models_list, retry
 
 from torch_utils import TestTorchConvertModel
@@ -48,7 +47,7 @@ def filter_timm(timm_list: list) -> list:
 
 
 class TestTimmConvertModel(TestTorchConvertModel):
-    @retry(3, exceptions=(HfHubHTTPError, LocalEntryNotFoundError), delay=1)
+    @retry(3, exceptions=(OSError,), delay=1)
     def load_model(self, model_name, model_link):
         m = timm.create_model(model_name, pretrained=True)
         cfg = timm.get_pretrained_cfg(model_name)

diff --git a/tests/model_hub_tests/transformation_tests/test_pa_transformation.py b/tests/model_hub_tests/transformation_tests/test_pa_transformation.py
@@ -4,11 +4,13 @@
 from openvino._offline_transformations import paged_attention_transformation
 from openvino._pyopenvino.op import _PagedAttentionExtension
 from optimum.intel import OVModelForCausalLM
+from models_hub_common.utils import retry
 import models_hub_common.utils as utils
 import pytest
 import os
 import re
 
+@retry(3, exceptions=(OSError,), delay=1)
 def run_pa(tmp_path, model_id, model_link, use_block_indices_inputs, use_score_outputs):
     model = OVModelForCausalLM.from_pretrained(model_id, export=True, trust_remote_code=True)
 

diff --git a/tests/model_hub_tests/transformation_tests/test_stateful_to_stateless_transformation.py b/tests/model_hub_tests/transformation_tests/test_stateful_to_stateless_transformation.py
@@ -4,6 +4,7 @@
 import openvino as ov
 from openvino._offline_transformations import stateful_to_stateless_transformation
 from optimum.intel import OVModelForCausalLM
+from models_hub_common.utils import retry
 import models_hub_common.utils as utils
 import pytest
 import os
@@ -23,6 +24,7 @@ def check_desc_tensors(tensors1, tensors2):
         assert t1.get_partial_shape() == t2.get_partial_shape()
         assert t1.get_element_type() == t2.get_element_type()
 
+@retry(3, exceptions=(OSError,), delay=1)
 def run_stateful_to_stateless_in_runtime(tmp_path, model_id, model_link):
     model = OVModelForCausalLM.from_pretrained(model_id, export=True, stateful=True, compile=False)
     assert len(model.model.get_sinks()), f"Input model is not in the expected stateful form because it doesn't have any sinks."