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]: