fix: add tag fixes for pipelines and experiments

Author: mufaddal-rohawala

src/sagemaker/experiments/experiment.py

@@ -15,6 +15,8 @@
 
 import time
 
+from botocore.exceptions import ClientError
+
 from sagemaker.apiutils import _base_types
 from sagemaker.experiments.trial import _Trial
 from sagemaker.experiments.trial_component import _TrialComponent
@@ -154,17 +156,21 @@ def _load_or_create(
         Returns:
             experiments.experiment._Experiment: A SageMaker `_Experiment` object
         """
-        sagemaker_client = sagemaker_session.sagemaker_client
         try:
-            experiment = _Experiment.load(experiment_name, sagemaker_session)
-        except sagemaker_client.exceptions.ResourceNotFound:
             experiment = _Experiment.create(
                 experiment_name=experiment_name,
                 display_name=display_name,
                 description=description,
                 tags=tags,
                 sagemaker_session=sagemaker_session,
             )
+        except ClientError as ce:
+            error_code = ce.response["Error"]["Code"]
+            error_message = ce.response["Error"]["Message"]
+            if not (error_code == "ValidationException" and "already exists" in error_message):
+                raise ce
+            # already exists
+            experiment = _Experiment.load(experiment_name, sagemaker_session)
         return experiment
 
     def list_trials(self, created_before=None, created_after=None, sort_by=None, sort_order=None):

src/sagemaker/experiments/trial.py

@@ -13,6 +13,8 @@
 """Contains the Trial class."""
 from __future__ import absolute_import
 
+from botocore.exceptions import ClientError
+
 from sagemaker.apiutils import _base_types
 from sagemaker.experiments import _api_types
 from sagemaker.experiments.trial_component import _TrialComponent
@@ -268,8 +270,20 @@ def _load_or_create(
         Returns:
             experiments.trial._Trial: A SageMaker `_Trial` object
         """
-        sagemaker_client = sagemaker_session.sagemaker_client
         try:
+            trial = _Trial.create(
+                experiment_name=experiment_name,
+                trial_name=trial_name,
+                display_name=display_name,
+                tags=tags,
+                sagemaker_session=sagemaker_session,
+            )
+        except ClientError as ce:
+            error_code = ce.response["Error"]["Code"]
+            error_message = ce.response["Error"]["Message"]
+            if not (error_code == "ValidationException" and "already exists" in error_message):
+                raise ce
+            # already exists
             trial = _Trial.load(trial_name, sagemaker_session)
             if trial.experiment_name != experiment_name:  # pylint: disable=no-member
                 raise ValueError(
@@ -278,12 +292,4 @@ def _load_or_create(
                         trial.experiment_name  # pylint: disable=no-member
                     )
                 )
-        except sagemaker_client.exceptions.ResourceNotFound:
-            trial = _Trial.create(
-                experiment_name=experiment_name,
-                trial_name=trial_name,
-                display_name=display_name,
-                tags=tags,
-                sagemaker_session=sagemaker_session,
-            )
         return trial

src/sagemaker/experiments/trial_component.py

@@ -15,6 +15,8 @@
 
 import time
 
+from botocore.exceptions import ClientError
+
 from sagemaker.apiutils import _base_types
 from sagemaker.experiments import _api_types
 from sagemaker.experiments._api_types import TrialComponentSearchResult
@@ -326,16 +328,20 @@ def _load_or_create(
             experiments.trial_component._TrialComponent: A SageMaker `_TrialComponent` object.
             bool: A boolean variable indicating whether the trail component already exists
         """
-        sagemaker_client = sagemaker_session.sagemaker_client
         is_existed = False
         try:
-            run_tc = _TrialComponent.load(trial_component_name, sagemaker_session)
-            is_existed = True
-        except sagemaker_client.exceptions.ResourceNotFound:
             run_tc = _TrialComponent.create(
                 trial_component_name=trial_component_name,
                 display_name=display_name,
                 tags=tags,
                 sagemaker_session=sagemaker_session,
             )
+        except ClientError as ce:
+            error_code = ce.response["Error"]["Code"]
+            error_message = ce.response["Error"]["Message"]
+            if not (error_code == "ValidationException" and "already exists" in error_message):
+                raise ce
+            # already exists
+            run_tc = _TrialComponent.load(trial_component_name, sagemaker_session)
+            is_existed = True
         return run_tc, is_existed

src/sagemaker/session.py

@@ -3224,7 +3224,7 @@ def create_model_package_from_containers(
         def submit(request):
             if model_package_group_name is not None:
                 _create_resource(
-                    self.sagemaker_client.create_model_package_group(
+                    lambda: self.sagemaker_client.create_model_package_group(
                         ModelPackageGroupName=request["ModelPackageGroupName"]
                     )
                 )
@@ -5448,8 +5448,9 @@ def _deployment_entity_exists(describe_fn):
 
 
 def _create_resource(create_fn):
-    """Call create function and while doing so accepts/passes the resource already exists exception.
-    Throws an exception if any exception other than resource already exists.
+    """Call create function and accepts/pass when resource already exists.
+
+    This is a helper function to use an existing resource if found when creating.
 
     Args:
         create_fn: Create resource function.
@@ -5823,9 +5824,9 @@ def _wait_until_training_done(callable_fn, desc, poll=5):
             job_desc, finished = callable_fn(job_desc)
         except botocore.exceptions.ClientError as err:
             # For initial 5 mins we accept/pass AccessDeniedException.
-            # The reason is to await tag propagation to avoid false AccessDenied claims for an access
-            # policy based on resource tags, The caveat here is for true AccessDenied cases the routine
-            # will fail after 5 mins
+            # The reason is to await tag propagation to avoid false AccessDenied claims for an
+            # access policy based on resource tags, The caveat here is for true AccessDenied
+            # cases the routine will fail after 5 mins
             if err.response["Error"]["Code"] == "AccessDeniedException" and elapsed_time <= 300:
                 LOGGER.warning(
                     "Received AccessDeniedException. This could mean the IAM role does not "
@@ -5834,8 +5835,7 @@ def _wait_until_training_done(callable_fn, desc, poll=5):
                     "continuing to wait for tag propagation.."
                 )
                 continue
-            else:
-                raise err
+            raise err
     return job_desc
 
 
@@ -5861,8 +5861,7 @@ def _wait_until(callable_fn, poll=5):
                     "continuing to wait for tag propagation.."
                 )
                 continue
-            else:
-                raise err
+            raise err
     return result
 
 

src/sagemaker/utils.py

@@ -604,12 +604,17 @@ def retries(
     )
 
 
-def retry_with_backoff(callable_func, num_attempts=8):
+def retry_with_backoff(callable_func, num_attempts=8, botocore_client_error_code=None):
     """Retry with backoff until maximum attempts are reached
 
     Args:
         callable_func (callable): The callable function to retry.
-        num_attempts (int): The maximum number of attempts to retry.
+        num_attempts (int): The maximum number of attempts to retry.(Default: 8)
+        botocore_client_error_code (str): The specific Botocore ClientError exception error code
+            on which to retry on.
+            If provided other exceptions will be raised directly w/o retry.
+            If not provided, retry on any exception.
+            (Default: None)
     """
     if num_attempts < 1:
         raise ValueError(
@@ -619,7 +624,15 @@ def retry_with_backoff(callable_func, num_attempts=8):
         try:
             return callable_func()
         except Exception as ex:  # pylint: disable=broad-except
-            if i == num_attempts - 1:
+            if not botocore_client_error_code or (
+                botocore_client_error_code
+                and isinstance(ex, botocore.exceptions.ClientError)
+                and ex.response["Error"]["Code"]  # pylint: disable=no-member
+                == botocore_client_error_code
+            ):
+                if i == num_attempts - 1:
+                    raise ex
+            else:
                 raise ex
             logger.error("Retrying in attempt %s, due to %s", (i + 1), str(ex))
             time.sleep(2**i)

src/sagemaker/workflow/pipeline.py

@@ -26,6 +26,7 @@
 from sagemaker import s3
 from sagemaker._studio import _append_project_tags
 from sagemaker.session import Session
+from sagemaker.utils import retry_with_backoff
 from sagemaker.workflow.callback_step import CallbackOutput, CallbackStep
 from sagemaker.workflow.lambda_step import LambdaOutput, LambdaStep
 from sagemaker.workflow.entities import (
@@ -306,7 +307,12 @@ def start(
             update_args(kwargs, PipelineParameters=parameters)
             return self.sagemaker_session.sagemaker_client.start_pipeline_execution(**kwargs)
         update_args(kwargs, PipelineParameters=format_start_parameters(parameters))
-        response = self.sagemaker_session.sagemaker_client.start_pipeline_execution(**kwargs)
+
+        # retry on AccessDeniedException to cover case of tag propagation delay
+        response = retry_with_backoff(
+            lambda: self.sagemaker_session.sagemaker_client.start_pipeline_execution(**kwargs),
+            botocore_client_error_code="AccessDeniedException",
+        )
         return _PipelineExecution(
             arn=response["PipelineExecutionArn"],
             sagemaker_session=self.sagemaker_session,

tests/integ/test_inference_pipeline.py

@@ -28,7 +28,7 @@
 from sagemaker.predictor import Predictor
 from sagemaker.serializers import JSONSerializer
 from sagemaker.sparkml.model import SparkMLModel
-from sagemaker.utils import sagemaker_timestamp
+from sagemaker.utils import unique_name_from_base
 
 SPARKML_DATA_PATH = os.path.join(DATA_DIR, "sparkml_model")
 XGBOOST_DATA_PATH = os.path.join(DATA_DIR, "xgboost_model")
@@ -60,7 +60,7 @@ def test_inference_pipeline_batch_transform(sagemaker_session, cpu_instance_type
         path=os.path.join(XGBOOST_DATA_PATH, "xgb_model.tar.gz"),
         key_prefix="integ-test-data/xgboost/model",
     )
-    batch_job_name = "test-inference-pipeline-batch-{}".format(sagemaker_timestamp())
+    batch_job_name = unique_name_from_base("test-inference-pipeline-batch")
     sparkml_model = SparkMLModel(
         model_data=sparkml_model_data,
         env={"SAGEMAKER_SPARKML_SCHEMA": SCHEMA},
@@ -99,7 +99,7 @@ def test_inference_pipeline_batch_transform(sagemaker_session, cpu_instance_type
 def test_inference_pipeline_model_deploy(sagemaker_session, cpu_instance_type):
     sparkml_data_path = os.path.join(DATA_DIR, "sparkml_model")
     xgboost_data_path = os.path.join(DATA_DIR, "xgboost_model")
-    endpoint_name = "test-inference-pipeline-deploy-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-inference-pipeline-deploy")
     sparkml_model_data = sagemaker_session.upload_data(
         path=os.path.join(sparkml_data_path, "mleap_model.tar.gz"),
         key_prefix="integ-test-data/sparkml/model",
@@ -156,7 +156,7 @@ def test_inference_pipeline_model_deploy_and_update_endpoint(
 ):
     sparkml_data_path = os.path.join(DATA_DIR, "sparkml_model")
     xgboost_data_path = os.path.join(DATA_DIR, "xgboost_model")
-    endpoint_name = "test-inference-pipeline-deploy-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-inference-pipeline-deploy")
     sparkml_model_data = sagemaker_session.upload_data(
         path=os.path.join(sparkml_data_path, "mleap_model.tar.gz"),
         key_prefix="integ-test-data/sparkml/model",

tests/integ/test_mxnet.py

@@ -24,7 +24,7 @@
 from sagemaker.mxnet.model import MXNetModel
 from sagemaker.mxnet.processing import MXNetProcessor
 from sagemaker.serverless import ServerlessInferenceConfig
-from sagemaker.utils import sagemaker_timestamp
+from sagemaker.utils import unique_name_from_base
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.kms_utils import get_or_create_kms_key
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
@@ -98,7 +98,7 @@ def test_framework_processing_job_with_deps(
 
 @pytest.mark.release
 def test_attach_deploy(mxnet_training_job, sagemaker_session, cpu_instance_type):
-    endpoint_name = "test-mxnet-attach-deploy-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-mxnet-attach-deploy")
 
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         estimator = MXNet.attach(mxnet_training_job, sagemaker_session=sagemaker_session)
@@ -165,7 +165,7 @@ def test_deploy_model(
     mxnet_inference_latest_py_version,
     cpu_instance_type,
 ):
-    endpoint_name = "test-mxnet-deploy-model-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-mxnet-deploy-model")
 
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         desc = sagemaker_session.sagemaker_client.describe_training_job(
@@ -200,7 +200,7 @@ def test_register_model_package(
     mxnet_inference_latest_py_version,
     cpu_instance_type,
 ):
-    endpoint_name = "test-mxnet-deploy-model-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-mxnet-deploy-model")
 
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         desc = sagemaker_session.sagemaker_client.describe_training_job(
@@ -216,7 +216,7 @@ def test_register_model_package(
             sagemaker_session=sagemaker_session,
             framework_version=mxnet_inference_latest_version,
         )
-        model_package_name = "register-model-package-{}".format(sagemaker_timestamp())
+        model_package_name = unique_name_from_base("register-model-package")
         model_pkg = model.register(
             content_types=["application/json"],
             response_types=["application/json"],
@@ -239,13 +239,13 @@ def test_register_model_package_versioned(
     mxnet_inference_latest_py_version,
     cpu_instance_type,
 ):
-    endpoint_name = "test-mxnet-deploy-model-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-mxnet-deploy-model")
 
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         desc = sagemaker_session.sagemaker_client.describe_training_job(
             TrainingJobName=mxnet_training_job
         )
-        model_package_group_name = "register-model-package-{}".format(sagemaker_timestamp())
+        model_package_group_name = unique_name_from_base("register-model-package")
         sagemaker_session.sagemaker_client.create_model_package_group(
             ModelPackageGroupName=model_package_group_name
         )
@@ -287,7 +287,7 @@ def test_deploy_model_with_tags_and_kms(
     mxnet_inference_latest_py_version,
     cpu_instance_type,
 ):
-    endpoint_name = "test-mxnet-deploy-model-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-mxnet-deploy-model")
 
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         desc = sagemaker_session.sagemaker_client.describe_training_job(
@@ -347,7 +347,7 @@ def test_deploy_model_and_update_endpoint(
     cpu_instance_type,
     alternative_cpu_instance_type,
 ):
-    endpoint_name = "test-mxnet-deploy-model-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-mxnet-deploy-model")
 
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         desc = sagemaker_session.sagemaker_client.describe_training_job(
@@ -395,7 +395,7 @@ def test_deploy_model_with_accelerator(
     mxnet_eia_latest_py_version,
     cpu_instance_type,
 ):
-    endpoint_name = "test-mxnet-deploy-model-ei-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-mxnet-deploy-model-ei")
 
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         desc = sagemaker_session.sagemaker_client.describe_training_job(
@@ -426,7 +426,7 @@ def test_deploy_model_with_serverless_inference_config(
     mxnet_inference_latest_version,
     mxnet_inference_latest_py_version,
 ):
-    endpoint_name = "test-mxnet-deploy-model-serverless-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-mxnet-deploy-model-serverless")
 
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         desc = sagemaker_session.sagemaker_client.describe_training_job(
@@ -465,7 +465,7 @@ def test_async_fit(
     mxnet_inference_latest_py_version,
     cpu_instance_type,
 ):
-    endpoint_name = "test-mxnet-attach-deploy-{}".format(sagemaker_timestamp())
+    endpoint_name = unique_name_from_base("test-mxnet-attach-deploy")
 
     with timeout(minutes=5):
         script_path = os.path.join(DATA_DIR, "mxnet_mnist", "mnist.py")