fix:ResourceConflictException from AWS lambda on pipeline upsert

Author: Namrata Madan

src/sagemaker/lambda_helper.py

@@ -15,6 +15,7 @@
 
 from io import BytesIO
 import zipfile
+import time
 from botocore.exceptions import ClientError
 from sagemaker.session import Session
 
@@ -134,32 +135,35 @@ def update(self):
         Returns: boto3 response from Lambda's update_function method.
         """
         lambda_client = _get_lambda_client(self.session)
-
-        if self.script is not None:
-            try:
-                response = lambda_client.update_function_code(
-                    FunctionName=self.function_name, ZipFile=_zip_lambda_code(self.script)
-                )
-                return response
-            except ClientError as e:
-                error = e.response["Error"]
-                raise ValueError(error)
-        else:
+        retry_attempts = 7
+        for i in range(retry_attempts):
             try:
-                response = lambda_client.update_function_code(
-                    FunctionName=(self.function_name or self.function_arn),
-                    S3Bucket=self.s3_bucket,
-                    S3Key=_upload_to_s3(
-                        s3_client=_get_s3_client(self.session),
-                        function_name=self.function_name,
-                        zipped_code_dir=self.zipped_code_dir,
-                        s3_bucket=self.s3_bucket,
-                    ),
-                )
+                if self.script is not None:
+                    response = lambda_client.update_function_code(
+                        FunctionName=self.function_name, ZipFile=_zip_lambda_code(self.script)
+                    )
+                else:
+                    response = lambda_client.update_function_code(
+                        FunctionName=(self.function_name or self.function_arn),
+                        S3Bucket=self.s3_bucket,
+                        S3Key=_upload_to_s3(
+                            s3_client=_get_s3_client(self.session),
+                            function_name=self.function_name,
+                            zipped_code_dir=self.zipped_code_dir,
+                            s3_bucket=self.s3_bucket,
+                        ),
+                    )
                 return response
             except ClientError as e:
                 error = e.response["Error"]
-                raise ValueError(error)
+                code = error["Code"]
+                if code == "ResourceConflictException":
+                    if i == retry_attempts - 1:
+                        raise ValueError(error)
+                    # max wait time = 2**0 + 2**1 + .. + 2**6 = 127 seconds
+                    time.sleep(2**i)
+                else:
+                    raise ValueError(error)
 
     def upsert(self):
         """Method to create a lambda function or update it if it already exists

src/sagemaker/workflow/pipeline.py

@@ -215,30 +215,29 @@ def upsert(
         Returns:
             response dict from service
         """
+        exists = True
         try:
+            self.describe()
+        except ClientError:
+            exists = False
+
+        if not exists:
             response = self.create(role_arn, description, tags, parallelism_config)
-        except ClientError as e:
-            error = e.response["Error"]
-            if (
-                error["Code"] == "ValidationException"
-                and "Pipeline names must be unique within" in error["Message"]
-            ):
-                response = self.update(role_arn, description)
-                if tags is not None:
-                    old_tags = self.sagemaker_session.sagemaker_client.list_tags(
-                        ResourceArn=response["PipelineArn"]
-                    )["Tags"]
-
-                    tag_keys = [tag["Key"] for tag in tags]
-                    for old_tag in old_tags:
-                        if old_tag["Key"] not in tag_keys:
-                            tags.append(old_tag)
-
-                    self.sagemaker_session.sagemaker_client.add_tags(
-                        ResourceArn=response["PipelineArn"], Tags=tags
-                    )
-            else:
-                raise
+        else:
+            response = self.update(role_arn, description)
+            if tags is not None:
+                old_tags = self.sagemaker_session.sagemaker_client.list_tags(
+                    ResourceArn=response["PipelineArn"]
+                )["Tags"]
+
+                tag_keys = [tag["Key"] for tag in tags]
+                for old_tag in old_tags:
+                    if old_tag["Key"] not in tag_keys:
+                        tags.append(old_tag)
+
+                self.sagemaker_session.sagemaker_client.add_tags(
+                    ResourceArn=response["PipelineArn"], Tags=tags
+                )
         return response
 
     def delete(self) -> Dict[str, Any]:

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -178,20 +178,15 @@ def test_large_pipeline_update(sagemaker_session_mock, role_arn):
 
 
 def test_pipeline_upsert(sagemaker_session_mock, role_arn):
-    sagemaker_session_mock.side_effect = [
-        ClientError(
-            operation_name="CreatePipeline",
-            error_response={
-                "Error": {
-                    "Code": "ValidationException",
-                    "Message": "Pipeline names must be unique within ...",
-                }
-            },
-        ),
-        {"PipelineArn": "mock_pipeline_arn"},
-        [{"Key": "dummy", "Value": "dummy_tag"}],
-        {},
-    ]
+    sagemaker_session_mock.sagemaker_client.describe_pipeline.return_value = {
+        "PipelineArn": "pipeline-arn"
+    }
+    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"}]
+    }
 
     pipeline = Pipeline(
         name="MyPipeline",
@@ -205,9 +200,9 @@ def test_pipeline_upsert(sagemaker_session_mock, role_arn):
         {"Key": "bar", "Value": "xyz"},
     ]
     pipeline.upsert(role_arn=role_arn, tags=tags)
-    assert sagemaker_session_mock.sagemaker_client.create_pipeline.called_with(
-        PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn
-    )
+
+    sagemaker_session_mock.sagemaker_client.create_pipeline.assert_not_called()
+
     assert sagemaker_session_mock.sagemaker_client.update_pipeline.called_with(
         PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn
     )