tenacity

Author: rafaelpierrehf

.github/workflows/gpu-integ-new.yaml

@@ -20,10 +20,10 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v2
-      - name: Set up Python 3.9
+      - name: Set up Python 3.11
         uses: actions/setup-python@v2
         with:
-          python-version: 3.9
+          python-version: 3.11
       - name: Install Python dependencies
         run: pip install -e .[test,dev,torch]
       - name: Build Docker

.github/workflows/gpu-integ-test.yaml

@@ -50,6 +50,7 @@
     "mock==2.0.0",
     "docker",
     "requests",
+    "tenacity"
 ]
 extras["quality"] = [
     "black",

setup.py

@@ -9,6 +9,7 @@
 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
 
 IS_GPU = _run_slow_tests
 DEVICE = "gpu" if IS_GPU else "cpu"
@@ -39,7 +40,7 @@ def wait_for_container_to_be_ready(base_url):
             time.sleep(2)
     return True
 
-
+@retry
 def verify_task(container: DockerClient, task: str, port: int = 5000, framework: str = "pytorch"):
     BASE_URL = f"http://localhost:{port}"
     input = task2input[task]