From 96c866180657abc564056acd4408b87490ddd86c Mon Sep 17 00:00:00 2001
From: Maxim Vafin <maxim.vafin@intel.com>
Date: Mon, 12 Aug 2024 19:19:18 +0200
Subject: [PATCH] [PT FE] Retry hf related tests (#26030)

### Details:
 - *item1*
 - *...*

### Tickets:
 - *CVS-149231*
 - *CVS-149229*
---
 tests/model_hub_tests/pytorch/test_hf_transformers.py          | 3 +--
 tests/model_hub_tests/pytorch/test_timm.py                     | 3 +--
 .../transformation_tests/test_pa_transformation.py             | 2 ++
 .../test_stateful_to_stateless_transformation.py               | 2 ++
 4 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tests/model_hub_tests/pytorch/test_hf_transformers.py b/tests/model_hub_tests/pytorch/test_hf_transformers.py
index 447077e371135c..d09df964b83ce0 100644
--- 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
index 69de16177cd0f7..8c11d32517c4d6 100644
--- 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
index f65213afd1bc59..dc65324d4f028b 100644
--- 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
index 8de04e8241a756..c3aaa082322fdf 100644
--- 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."