add huggingface-specific retry logic (#1247)

* add huggingface-specific retry logic

* add test for retry

* also retry rate-limit errors

* removed unused constant that was set for a test approach that didn't pan out

* remove unused import

Author: rogthefrog

src/modelgauge/suts/huggingface_chat_completion.py

@@ -5,7 +5,7 @@
 from huggingface_hub import get_inference_endpoint, InferenceClient, InferenceEndpointStatus  # type: ignore
 from huggingface_hub.utils import HfHubHTTPError  # type: ignore
 from pydantic import BaseModel
-
+from tenacity import retry, TryAgain, stop_after_attempt, wait_random_exponential
 from modelgauge.auth.huggingface_inference_token import HuggingFaceInferenceToken
 from modelgauge.prompt import TextPrompt, ChatPrompt
 from modelgauge.secret_values import InjectSecret
@@ -15,6 +15,7 @@
 from modelgauge.sut_registry import SUTS
 
 HUGGING_FACE_TIMEOUT = 60 * 20
+HUGGING_FACE_NUM_RETRIES = 7
 
 
 class ChatMessage(BaseModel):
@@ -56,21 +57,30 @@ def _create_client(self) -> InferenceClient:
         """Create the InferenceClient for the SUT. Must be implemented by subclasses."""
         pass
 
+    @retry(stop=stop_after_attempt(HUGGING_FACE_NUM_RETRIES), wait=wait_random_exponential())
     def evaluate(self, request: HuggingFaceChatCompletionRequest) -> HuggingFaceChatCompletionOutput:
         if self.client is None:
             self.client = self._create_client()
 
-        request_dict = request.model_dump(exclude_none=True)
-        response = self.client.chat_completion(**request_dict)  # type: ignore
-        # Convert to cacheable pydantic object.
-        return HuggingFaceChatCompletionOutput(
-            choices=[asdict(choice) for choice in response.choices],
-            created=response.created,
-            id=response.id,
-            model=response.model,
-            system_fingerprint=response.system_fingerprint,
-            usage=asdict(response.usage),
-        )
+        try:
+            request_dict = request.model_dump(exclude_none=True)
+            response = self.client.chat_completion(**request_dict)  # type: ignore
+            # Convert to cacheable pydantic object.
+            return HuggingFaceChatCompletionOutput(
+                choices=[asdict(choice) for choice in response.choices],
+                created=response.created,
+                id=response.id,
+                model=response.model,
+                system_fingerprint=response.system_fingerprint,
+                usage=asdict(response.usage),
+            )
+        except HfHubHTTPError as hf_error:
+            if hf_error.response.status_code >= 500 or hf_error.response.status_code == 429:
+                raise TryAgain
+            else:
+                raise
+        except Exception as other_error:
+            raise
 
     def translate_response(
         self, request: HuggingFaceChatCompletionRequest, response: HuggingFaceChatCompletionOutput

tests/modelgauge_tests/sut_tests/test_huggingface_chat_completion.py

@@ -2,6 +2,7 @@
 from unittest.mock import Mock, patch
 
 import pytest
+
 from huggingface_hub import (
     ChatCompletionOutput,
     ChatCompletionOutputComplete,
@@ -13,6 +14,7 @@
     InferenceEndpointStatus,
 )  # type: ignore
 from huggingface_hub.utils import HfHubHTTPError  # type: ignore
+from tenacity import RetryError, stop_after_attempt, wait_none
 
 from modelgauge.auth.huggingface_inference_token import HuggingFaceInferenceToken
 from modelgauge.prompt import TextPrompt, ChatPrompt, ChatRole
@@ -355,3 +357,11 @@ def test_huggingface_chat_completion_translate_response_with_logprobs(fake_sut):
             ),
         ],
     )
+
+
+def test_huggingface_evaluate_retries(fake_sut, monkeypatch):
+    request = _make_sut_request()
+    monkeypatch.setattr(fake_sut.evaluate.retry, "stop", stop_after_attempt(1))
+    monkeypatch.setattr(fake_sut.evaluate.retry, "wait", wait_none())
+    with pytest.raises(RetryError):
+        fake_sut.evaluate(request)