retry if

Author: rafaelpierrehf

tests/integ/test_container.py

@@ -9,7 +9,10 @@
 from huggingface_inference_toolkit.utils import _is_gpu_available, _load_repository_from_hf
 from integ.config import task2input, task2model, task2output, task2validation
 from transformers.testing_utils import require_torch, slow, require_tf, _run_slow_tests
-from tenacity import retry
+import tenacity
+import logging
+
+logging.basicConfig(level = "DEBUG")
 
 IS_GPU = _run_slow_tests
 DEVICE = "gpu" if IS_GPU else "cpu"
@@ -40,7 +43,11 @@ def wait_for_container_to_be_ready(base_url):
             time.sleep(2)
     return True
 
-@retry
+@tenacity.retry(
+    wait = tenacity.wait_random(min=1, max=2),
+    retry = tenacity.retry_if_exception(requests.exception.ConnectionError),
+    stop = tenacity.retry.stop_after_attempt(5)
+)
 def verify_task(container: DockerClient, task: str, port: int = 5000, framework: str = "pytorch"):
     BASE_URL = f"http://localhost:{port}"
     input = task2input[task]