feature: update lambda code on pipeline create/update/upsert for Lamb… (aws#3052)

* feature: update lambda code on pipeline create/update/upsert for LambdaStep

* remove debugging statement

* fix pylint Unnecessary "else" after "return"

Co-authored-by: Namrata Madan <[email protected]>
Co-authored-by: Ahsan Khan <[email protected]>

Author: 3 people

src/sagemaker/lambda_helper.py

@@ -161,6 +161,18 @@ def update(self):
                 error = e.response["Error"]
                 raise ValueError(error)
 
+    def upsert(self):
+        """Method to create a lambda function or update it if it already exists
+
+        Returns: boto3 response from Lambda's methods.
+        """
+        try:
+            return self.create()
+        except ValueError as error:
+            if "ResourceConflictException" in str(error):
+                return self.update()
+            raise
+
     def invoke(self):
         """Method to invoke a lambda function.
 

src/sagemaker/workflow/lambda_step.py

@@ -149,28 +149,12 @@ def to_request(self) -> RequestType:
         return request_dict
 
     def _get_function_arn(self):
-        """Returns the lamba function arn
+        """Returns the lambda function arn
 
         Method creates a lambda function and returns it's arn.
         If the lambda is already present, it will build it's arn and return that.
         """
-        region = self.lambda_func.session.boto_region_name
-        if region.lower() == "cn-north-1" or region.lower() == "cn-northwest-1":
-            partition = "aws-cn"
-        else:
-            partition = "aws"
-
         if self.lambda_func.function_arn is None:
-            account_id = self.lambda_func.session.account_id()
-            try:
-                response = self.lambda_func.create()
-                return response["FunctionArn"]
-            except ValueError as error:
-                if "ResourceConflictException" not in str(error):
-                    raise
-                return (
-                    f"arn:{partition}:lambda:{region}:{account_id}:"
-                    f"function:{self.lambda_func.function_name}"
-                )
-        else:
-            return self.lambda_func.function_arn
+            response = self.lambda_func.upsert()
+            return response["FunctionArn"]
+        return self.lambda_func.function_arn

tests/unit/sagemaker/workflow/test_lambda_step.py

@@ -201,36 +201,40 @@ def test_lambda_step_no_inputs_outputs(sagemaker_session):
 
 
 def test_lambda_step_with_function_arn(sagemaker_session):
+    lambda_func = MagicMock(
+        function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
+        session=sagemaker_session,
+    )
     lambda_step = LambdaStep(
         name="MyLambdaStep",
         depends_on=["TestStep"],
-        lambda_func=Lambda(
-            function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
-            session=sagemaker_session,
-        ),
+        lambda_func=lambda_func,
         inputs={},
         outputs=[],
     )
-    lambda_step._get_function_arn()
-    sagemaker_session.account_id.assert_not_called()
+    function_arn = lambda_step._get_function_arn()
+    assert function_arn == "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda"
+    lambda_func.upsert.assert_not_called()
 
 
 def test_lambda_step_without_function_arn(sagemaker_session):
+    lambda_func = MagicMock(
+        function_arn=None,
+        function_name="name",
+        execution_role_arn="arn:aws:lambda:us-west-2:123456789012:execution_role",
+        zipped_code_dir="",
+        handler="",
+        session=sagemaker_session,
+    )
     lambda_step = LambdaStep(
         name="MyLambdaStep",
         depends_on=["TestStep"],
-        lambda_func=Lambda(
-            function_name="name",
-            execution_role_arn="arn:aws:lambda:us-west-2:123456789012:execution_role",
-            zipped_code_dir="",
-            handler="",
-            session=sagemaker_session,
-        ),
+        lambda_func=lambda_func,
         inputs={},
         outputs=[],
     )
     lambda_step._get_function_arn()
-    sagemaker_session.account_id.assert_called_once()
+    lambda_func.upsert.assert_called_once()
 
 
 def test_lambda_step_without_function_arn_and_with_error(sagemaker_session_cn):
@@ -242,26 +246,14 @@ def test_lambda_step_without_function_arn_and_with_error(sagemaker_session_cn):
         handler="",
         session=sagemaker_session_cn,
     )
-    # The raised ValueError contains ResourceConflictException
-    lambda_func.create.side_effect = ValueError("ResourceConflictException")
-    lambda_step1 = LambdaStep(
-        name="MyLambdaStep1",
-        depends_on=["TestStep"],
-        lambda_func=lambda_func,
-        inputs={},
-        outputs=[],
-    )
-    function_arn = lambda_step1._get_function_arn()
-    assert function_arn == "arn:aws-cn:lambda:cn-north-1:234567890123:function:name"
 
-    # The raised ValueError does not contain ResourceConflictException
-    lambda_func.create.side_effect = ValueError()
-    lambda_step2 = LambdaStep(
-        name="MyLambdaStep2",
+    lambda_func.upsert.side_effect = ValueError()
+    lambda_step = LambdaStep(
+        name="MyLambdaStep",
         depends_on=["TestStep"],
         lambda_func=lambda_func,
         inputs={},
         outputs=[],
     )
     with pytest.raises(ValueError):
-        lambda_step2._get_function_arn()
+        lambda_step._get_function_arn()

tests/unit/test_lambda_helper.py

@@ -294,6 +294,78 @@ def test_update_lambda_client_error(sagemaker_session):
     assert "Cannot update code" in str(error)
 
 
+@patch("sagemaker.lambda_helper._zip_lambda_code", return_value=ZIPPED_CODE)
+def test_upsert_lambda_happycase1(sagemaker_session):
+    lambda_obj = lambda_helper.Lambda(
+        function_name=FUNCTION_NAME,
+        execution_role_arn=EXECUTION_ROLE,
+        script=SCRIPT,
+        handler=HANDLER,
+        session=sagemaker_session,
+    )
+
+    code = {"ZipFile": ZIPPED_CODE}
+    lambda_obj.upsert()
+
+    sagemaker_session.lambda_client.create_function.assert_called_with(
+        FunctionName=FUNCTION_NAME,
+        Runtime="python3.8",
+        Handler=HANDLER,
+        Role=EXECUTION_ROLE,
+        Code=code,
+        Timeout=120,
+        MemorySize=128,
+    )
+
+
+@patch("sagemaker.lambda_helper._zip_lambda_code", return_value=ZIPPED_CODE)
+def test_upsert_lambda_happycase2(sagemaker_session):
+    lambda_obj = lambda_helper.Lambda(
+        function_name=FUNCTION_NAME,
+        execution_role_arn=EXECUTION_ROLE,
+        script=SCRIPT,
+        handler=HANDLER,
+        session=sagemaker_session,
+    )
+
+    sagemaker_session.lambda_client.create_function.side_effect = ClientError(
+        {"Error": {"Code": "ResourceConflictException", "Message": "Lambda already exists"}},
+        "CreateFunction",
+    )
+
+    lambda_obj.upsert()
+
+    sagemaker_session.lambda_client.update_function_code.assert_called_once_with(
+        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE
+    )
+
+
+@patch("sagemaker.lambda_helper._zip_lambda_code", return_value=ZIPPED_CODE)
+def test_upsert_lambda_client_error(sagemaker_session):
+    lambda_obj = lambda_helper.Lambda(
+        function_name=FUNCTION_NAME,
+        execution_role_arn=EXECUTION_ROLE,
+        script=SCRIPT,
+        handler=HANDLER,
+        session=sagemaker_session,
+    )
+
+    sagemaker_session.lambda_client.create_function.side_effect = ClientError(
+        {"Error": {"Code": "ResourceConflictException", "Message": "Lambda already exists"}},
+        "CreateFunction",
+    )
+
+    sagemaker_session.lambda_client.update_function_code.side_effect = ClientError(
+        {"Error": {"Code": "ResourceConflictException", "Message": "Cannot update code"}},
+        "UpdateFunctionCode",
+    )
+
+    with pytest.raises(ValueError) as error:
+        lambda_obj.upsert()
+
+    assert "Cannot update code" in str(error)
+
+
 def test_invoke_lambda_happycase(sagemaker_session):
     lambda_obj = lambda_helper.Lambda(function_arn=LAMBDA_ARN, session=sagemaker_session)
     lambda_obj.invoke()