add try catch block

Author: Narrohag

tests/integ/sagemaker/jumpstart/private_hub/estimator/test_jumpstart_private_hub_estimator.py

@@ -34,10 +34,32 @@
     get_training_dataset_for_model_and_version,
 )
 
-from tests.integ.sagemaker.jumpstart.private_hub.setup import add_model_references
+MAX_INIT_TIME_SECONDS = 5
 
+TEST_MODEL_IDS = {
+    "huggingface-spc-bert-base-cased",
+    "meta-textgeneration-llama-2-7b",
+    "catboost-regression-model",
+}
 
-MAX_INIT_TIME_SECONDS = 5
+
+@with_exponential_backoff()
+def create_model_reference(hub_instance, model_arn):
+    try:
+        hub_instance.create_model_reference(model_arn=model_arn)
+    except:
+        pass
+
+
+@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:
+        model_arn = get_public_hub_model_arn(hub_instance, model)
+        create_model_reference(hub_instance, model_arn)
 
 
 def test_jumpstart_hub_estimator(setup, add_model_references):

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

@@ -35,10 +35,34 @@
     with_exponential_backoff,
 )
 
-from tests.integ.sagemaker.jumpstart.private_hub.setup import add_model_references
+MAX_INIT_TIME_SECONDS = 5
 
+TEST_MODEL_IDS = {
+    "catboost-classification-model",
+    "huggingface-txt2img-conflictx-complex-lineart",
+    "meta-textgeneration-llama-2-7b",
+    "meta-textgeneration-llama-3-2-1b",
+    "catboost-regression-model",
+}
 
-MAX_INIT_TIME_SECONDS = 5
+
+@with_exponential_backoff()
+def create_model_reference(hub_instance, model_arn):
+    try:
+        hub_instance.create_model_reference(model_arn=model_arn)
+    except:
+        pass
+
+
+@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:
+        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):