Add retry to test

Author: mollyheamazon

sagemaker-serve/tests/integ/test_huggingface_integration.py

@@ -16,6 +16,8 @@
 import uuid
 import pytest
 import logging
+import time
+from botocore.exceptions import ClientError
 
 from sagemaker.serve.model_builder import ModelBuilder
 from sagemaker.serve.utils.types import ModelServer
@@ -116,11 +118,26 @@ def make_prediction(core_endpoint):
         "inputs": "What are falcons?", 
         "parameters": {"max_new_tokens": 32},
      }
-        
-    result = core_endpoint.invoke(
-        body=json.dumps(test_data),
-        content_type="application/json"
-     )
+    
+    # Retry logic to handle endpoint propagation delay in CodeBuild
+    max_retries = 5
+    for attempt in range(max_retries):
+        try:
+            result = core_endpoint.invoke(
+                body=json.dumps(test_data),
+                content_type="application/json"
+            )
+            break
+        except ClientError as e:
+            if e.response['Error']['Code'] == 'ValidationException' and 'not found' in str(e):
+                if attempt < max_retries - 1:
+                    wait_time = 2 ** attempt  # Exponential backoff: 1, 2, 4, 8 seconds
+                    logger.warning(f"Endpoint not found, retrying in {wait_time}s (attempt {attempt + 1}/{max_retries})")
+                    time.sleep(wait_time)
+                else:
+                    raise
+            else:
+                raise
 
     # Decode the output of the invocation and print the result
     prediction = json.loads(result.body.read().decode('utf-8'))