implement throttle handling

Author: malavhs

tests/integ/sagemaker/jumpstart/conftest.py

@@ -34,6 +34,7 @@
     get_test_artifact_bucket,
     get_test_suite_id,
     get_sm_session,
+    with_exponential_backoff,
 )
 
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME, JUMPSTART_MODEL_HUB_NAME
@@ -62,6 +63,8 @@ def _teardown():
 
     test_suite_id = os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]
 
+    test_hub_name = os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]
+
     boto3_session = boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME)
 
     sagemaker_client = boto3_session.client(
@@ -133,36 +136,28 @@ def _teardown():
     bucket.objects.filter(Prefix=test_suite_id + "/").delete()
 
     # delete private hubs
-    _delete_hubs(sagemaker_session)
+    _delete_hubs(sagemaker_session, test_hub_name)
 
 
-def _delete_hubs(sagemaker_session):
-    # list Hubs created by PySDK integration tests
-    list_hub_response = sagemaker_session.list_hubs(
-        name_contains=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME]
+def _delete_hubs(sagemaker_session, hub_name):
+    # list and delete all hub contents first
+    list_hub_content_response = sagemaker_session.list_hub_contents(
+        hub_name=hub_name, hub_content_type=HubContentType.MODEL_REFERENCE.value
     )
+    for model in list_hub_content_response["HubContentSummaries"]:
+        _delete_hub_contents(sagemaker_session, hub_name, model)
 
-    for hub in list_hub_response["HubSummaries"]:
-        if hub["HubName"] != JUMPSTART_MODEL_HUB_NAME:
-            # delete all hub contents first
-            _delete_hub_contents(sagemaker_session, hub["HubName"])
-            sagemaker_session.delete_hub(hub["HubName"])
+    sagemaker_session.delete_hub(hub_name)
 
 
-def _delete_hub_contents(sagemaker_session, test_hub_name):
-    # list hub_contents for the given hub
-    list_hub_content_response = sagemaker_session.list_hub_contents(
-        hub_name=test_hub_name, hub_content_type=HubContentType.MODEL_REFERENCE.value
+@with_exponential_backoff()
+def _delete_hub_contents(sagemaker_session, hub_name, model):
+    sagemaker_session.delete_hub_content_reference(
+        hub_name=hub_name,
+        hub_content_type=HubContentType.MODEL_REFERENCE.value,
+        hub_content_name=model["HubContentName"],
     )
 
-    # delete hub_contents for the given hub
-    for models in list_hub_content_response["HubContentSummaries"]:
-        sagemaker_session.delete_hub_content_reference(
-            hub_name=test_hub_name,
-            hub_content_type=HubContentType.MODEL_REFERENCE.value,
-            hub_content_name=models["HubContentName"],
-        )
-
 
 @pytest.fixture(scope="session", autouse=True)
 def setup(request):

tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py

@@ -13,13 +13,15 @@
 from __future__ import absolute_import
 
 import os
+import random
 import time
 
 import pytest
 from sagemaker.enums import EndpointType
 from sagemaker.jumpstart.hub.hub import Hub
 from sagemaker.jumpstart.hub.utils import generate_hub_arn_for_init_kwargs
 from sagemaker.predictor import retrieve_default
+from botocore.exceptions import ClientError
 
 import tests.integ
 
@@ -32,6 +34,7 @@
 from tests.integ.sagemaker.jumpstart.utils import (
     get_public_hub_model_arn,
     get_sm_session,
+    with_exponential_backoff,
 )
 
 MAX_INIT_TIME_SECONDS = 5
@@ -45,14 +48,20 @@
 }
 
 
+@with_exponential_backoff()
+def create_model_reference(hub_instance, model_arn):
+    hub_instance.create_model_reference(model_arn=model_arn)
+
+
 @pytest.fixture(scope="session")
 def add_model_references():
     # Create Model References to test in Hub
     hub_instance = Hub(
         hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session()
     )
     for model in TEST_MODEL_IDS:
-        hub_instance.create_model_reference(model_arn=get_public_hub_model_arn(hub_instance, model))
+        model_arn = get_public_hub_model_arn(hub_instance, model)
+        create_model_reference(hub_instance, model_arn)
 
 
 def test_jumpstart_hub_model(setup, add_model_references):

tests/integ/sagemaker/jumpstart/utils.py

@@ -14,13 +14,16 @@
 import functools
 import json
 
+import random
+import time
 import uuid
 from typing import Any, Dict, List, Tuple
 import boto3
 import pandas as pd
 import os
 
 from botocore.config import Config
+from botocore.exceptions import ClientError
 import pytest
 
 
@@ -125,6 +128,32 @@ def get_public_hub_model_arn(hub: Hub, model_id: str) -> str:
     return models[0]["hub_content_arn"]
 
 
+def with_exponential_backoff(max_retries=5, initial_delay=1, max_delay=60):
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            retries = 0
+            while True:
+                try:
+                    return func(*args, **kwargs)
+                except ClientError as e:
+                    if retries >= max_retries or e.response["Error"]["Code"] not in [
+                        "ThrottlingException",
+                        "TooManyRequestsException",
+                    ]:
+                        raise
+                    delay = min(initial_delay * (2**retries) + random.random(), max_delay)
+                    print(
+                        f"Retrying {func.__name__} in {delay:.2f} seconds... (Attempt {retries + 1}/{max_retries})"
+                    )
+                    time.sleep(delay)
+                    retries += 1
+
+        return wrapper
+
+    return decorator
+
+
 class EndpointInvoker:
     def __init__(
         self,