fix: tag permission issue - remove describe before create

Author: mufaddal-rohawala

src/sagemaker/session.py

@@ -3223,14 +3223,11 @@ def create_model_package_from_containers(
 
         def submit(request):
             if model_package_group_name is not None:
-                try:
-                    self.sagemaker_client.describe_model_package_group(
-                        ModelPackageGroupName=request["ModelPackageGroupName"]
-                    )
-                except ClientError:
+                _create_resource(
                     self.sagemaker_client.create_model_package_group(
                         ModelPackageGroupName=request["ModelPackageGroupName"]
                     )
+                )
             return self.sagemaker_client.create_model_package(**request)
 
         return self._intercept_create_request(
@@ -3918,42 +3915,40 @@ def endpoint_from_model_data(
         name = name or name_from_image(image_uri)
         model_vpc_config = vpc_utils.sanitize(model_vpc_config)
 
-        if _deployment_entity_exists(
-            lambda: self.sagemaker_client.describe_endpoint(EndpointName=name)
-        ):
-            raise ValueError(
-                'Endpoint with name "{}" already exists; please pick a different name.'.format(name)
-            )
+        primary_container = container_def(
+            image_uri=image_uri,
+            model_data_url=model_s3_location,
+            env=model_environment_vars,
+        )
 
-        if not _deployment_entity_exists(
-            lambda: self.sagemaker_client.describe_model(ModelName=name)
-        ):
-            primary_container = container_def(
-                image_uri=image_uri,
-                model_data_url=model_s3_location,
-                env=model_environment_vars,
-            )
-            self.create_model(
-                name=name, role=role, container_defs=primary_container, vpc_config=model_vpc_config
-            )
+        self.create_model(
+            name=name, role=role, container_defs=primary_container, vpc_config=model_vpc_config
+        )
 
         data_capture_config_dict = None
         if data_capture_config is not None:
             data_capture_config_dict = data_capture_config._to_request_dict()
 
-        if not _deployment_entity_exists(
-            lambda: self.sagemaker_client.describe_endpoint_config(EndpointConfigName=name)
-        ):
-            self.create_endpoint_config(
+        _create_resource(
+            lambda: self.create_endpoint_config(
                 name=name,
                 model_name=name,
                 initial_instance_count=initial_instance_count,
                 instance_type=instance_type,
                 accelerator_type=accelerator_type,
                 data_capture_config_dict=data_capture_config_dict,
             )
+        )
+
+        # to make change backwards compatible
+        response = _create_resource(
+            lambda: self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+        )
+        if not response:
+            raise ValueError(
+                'Endpoint with name "{}" already exists; please pick a different name.'.format(name)
+            )
 
-        self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
         return name
 
     def endpoint_from_production_variants(
@@ -5452,6 +5447,34 @@ def _deployment_entity_exists(describe_fn):
         return False
 
 
+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.
+
+    Args:
+        create_fn: Create resource function.
+
+    Returns:
+        (bool): True if new resource was created, False if resource already exists.
+    """
+    try:
+        create_fn()
+        # create function succeeded, resource does not exist already
+        return True
+    except ClientError as ce:
+        error_code = ce.response["Error"]["Code"]
+        error_message = ce.response["Error"]["Message"]
+        already_exists_exceptions = ["ValidationException", "ResourceInUse"]
+        already_exists_msg_patterns = ["Cannot create already existing", "already exists"]
+        if not (
+            error_code in already_exists_exceptions
+            and any(p in error_message for p in already_exists_msg_patterns)
+        ):
+            raise ce
+        # no new resource created as resource already exists
+        return False
+
+
 def _train_done(sagemaker_client, job_name, last_desc):
     """Placeholder docstring"""
     in_progress_statuses = ["InProgress", "Created"]

src/sagemaker/workflow/pipeline.py

@@ -241,20 +241,19 @@ def upsert(
         Returns:
             response dict from service
         """
-        exists = True
         try:
-            self.describe()
-        except ClientError as e:
-            err = e.response.get("Error", {})
-            if err.get("Code", None) == "ResourceNotFound":
-                exists = False
-            else:
-                raise e
-
-        if not exists:
             response = self.create(role_arn, description, tags, parallelism_config)
-        else:
+        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
             response = self.update(role_arn, description)
+            # add new tags to existing resource
             if tags is not None:
                 old_tags = self.sagemaker_session.sagemaker_client.list_tags(
                     ResourceArn=response["PipelineArn"]

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -33,6 +33,8 @@
 from sagemaker.workflow.step_collections import StepCollection
 from tests.unit.sagemaker.workflow.helpers import ordered, CustomStep
 from sagemaker.local.local_session import LocalSession
+from botocore.exceptions import ClientError
+
 
 
 @pytest.fixture
@@ -173,10 +175,17 @@ def test_large_pipeline_update(sagemaker_session_mock, role_arn):
     )
 
 
-def test_pipeline_upsert(sagemaker_session_mock, role_arn):
-    sagemaker_session_mock.sagemaker_client.describe_pipeline.return_value = {
-        "PipelineArn": "pipeline-arn"
-    }
+def test_pipeline_upsert_resource_already_exists(sagemaker_session_mock, role_arn):
+
+    # case 1: resource already exists
+    def _raise_does_already_exists_client_error(**kwargs):
+        response = {"Error": {"Code": "ValidationException", "Message": "Resource already exists."}}
+        raise ClientError(error_response=response, operation_name="create_pipeline")
+
+    sagemaker_session_mock.sagemaker_client.create_pipeline = Mock(
+        name="create_pipeline", side_effect=_raise_does_already_exists_client_error
+    )
+
     sagemaker_session_mock.sagemaker_client.update_pipeline.return_value = {
         "PipelineArn": "pipeline-arn"
     }
@@ -197,9 +206,12 @@ def test_pipeline_upsert(sagemaker_session_mock, role_arn):
     ]
     pipeline.upsert(role_arn=role_arn, tags=tags)
 
-    sagemaker_session_mock.sagemaker_client.create_pipeline.assert_not_called()
+    sagemaker_session_mock.sagemaker_client.create_pipeline.assert_called_once_with(
+        PipelineName="MyPipeline", RoleArn=role_arn, PipelineDefinition=pipeline.definition(),
+        Tags=tags
+    )
 
-    assert sagemaker_session_mock.sagemaker_client.update_pipeline.called_with(
+    sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_once_with(
         PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn
     )
     assert sagemaker_session_mock.sagemaker_client.list_tags.called_with(
@@ -211,6 +223,89 @@ def test_pipeline_upsert(sagemaker_session_mock, role_arn):
         ResourceArn="mock_pipeline_arn", Tags=tags
     )
 
+def test_pipeline_upsert_create_unexpected_failure(sagemaker_session_mock, role_arn):
+
+    # case 2: unexpected failure on create
+    def _raise_unexpected_client_error(**kwargs):
+        response = {
+            "Error": {"Code": "ValidationException", "Message": "Name does not satisfy expression."}
+        }
+        raise ClientError(error_response=response, operation_name="foo")
+
+    sagemaker_session_mock.sagemaker_client.create_pipeline = Mock(
+        name="create_pipeline", side_effect=_raise_unexpected_client_error
+    )
+
+    sagemaker_session_mock.sagemaker_client.update_pipeline.return_value = {
+        "PipelineArn": "pipeline-arn"
+    }
+    sagemaker_session_mock.sagemaker_client.list_tags.return_value = {
+        "Tags": [{"Key": "dummy", "Value": "dummy_tag"}]
+    }
+
+    tags = [
+        {"Key": "foo", "Value": "abc"},
+        {"Key": "bar", "Value": "xyz"},
+    ]
+
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[],
+        steps=[],
+        sagemaker_session=sagemaker_session_mock,
+    )
+
+    with pytest.raises(ClientError):
+        pipeline.upsert(role_arn=role_arn, tags=tags)
+
+
+
+    sagemaker_session_mock.sagemaker_client.create_pipeline.assert_called_once_with(
+        PipelineName="MyPipeline", RoleArn=role_arn, PipelineDefinition=pipeline.definition(),
+        Tags=tags
+    )
+    sagemaker_session_mock.sagemaker_client.update_pipeline.assert_not_called()
+    sagemaker_session_mock.sagemaker_client.list_tags.assert_not_called()
+    sagemaker_session_mock.sagemaker_client.add_tags.assert_not_called()
+
+def test_pipeline_upsert_resourse_doesnt_exist(sagemaker_session_mock, role_arn):
+
+    # case 3: resource does not exist
+    sagemaker_session_mock.sagemaker_client.create_pipeline = Mock(name="create_pipeline")
+
+    sagemaker_session_mock.sagemaker_client.update_pipeline.return_value = {
+        "PipelineArn": "pipeline-arn"
+    }
+    sagemaker_session_mock.sagemaker_client.list_tags.return_value = {
+        "Tags": [{"Key": "dummy", "Value": "dummy_tag"}]
+    }
+
+    tags = [
+        {"Key": "foo", "Value": "abc"},
+        {"Key": "bar", "Value": "xyz"},
+    ]
+
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[],
+        steps=[],
+        sagemaker_session=sagemaker_session_mock,
+    )
+
+    try:
+        pipeline.upsert(role_arn=role_arn, tags=tags)
+    except ClientError:
+        assert False, f"Unexpected ClientError raised"
+
+    sagemaker_session_mock.sagemaker_client.create_pipeline.assert_called_once_with(
+        PipelineName="MyPipeline", RoleArn=role_arn, PipelineDefinition=pipeline.definition(),
+        Tags=tags
+    )
+
+    sagemaker_session_mock.sagemaker_client.update_pipeline.assert_not_called()
+    sagemaker_session_mock.sagemaker_client.list_tags.assert_not_called()
+    sagemaker_session_mock.sagemaker_client.add_tags.assert_not_called()
+
 
 def test_pipeline_delete(sagemaker_session_mock):
     pipeline = Pipeline(