fix: update lambda function when function arn is provided (#3668)

Co-authored-by: Rohan Gujarathi <gujrohan@amazon.com>

Author: rohangujarathi

src/sagemaker/lambda_helper.py

@@ -92,6 +92,10 @@ def __init__(
             if handler is None:
                 raise ValueError("Lambda handler must be provided.")
 
+        if function_arn is not None:
+            if zipped_code_dir and script:
+                raise ValueError("Provide either script or zipped_code_dir, not both.")
+
     def create(self):
         """Method to create a lambda function.
 
@@ -140,17 +144,29 @@ def update(self):
             try:
                 if self.script is not None:
                     response = lambda_client.update_function_code(
-                        FunctionName=self.function_name, ZipFile=_zip_lambda_code(self.script)
+                        FunctionName=self.function_name or self.function_arn,
+                        ZipFile=_zip_lambda_code(self.script),
                     )
                 else:
+                    bucket = self.s3_bucket or self.session.default_bucket()
+                    # get function name to be used in S3 upload path
+                    if self.function_arn:
+                        versioned_function_name = self.function_arn.split("funtion:")[-1]
+                        if ":" in versioned_function_name:
+                            function_name_for_s3 = versioned_function_name.split(":")[0]
+                        else:
+                            function_name_for_s3 = versioned_function_name
+                    else:
+                        function_name_for_s3 = self.function_name
+
                     response = lambda_client.update_function_code(
                         FunctionName=(self.function_name or self.function_arn),
-                        S3Bucket=self.s3_bucket,
+                        S3Bucket=bucket,
                         S3Key=_upload_to_s3(
                             s3_client=_get_s3_client(self.session),
-                            function_name=self.function_name,
+                            function_name=function_name_for_s3,
                             zipped_code_dir=self.zipped_code_dir,
-                            s3_bucket=self.s3_bucket,
+                            s3_bucket=bucket,
                         ),
                     )
                 return response

src/sagemaker/workflow/lambda_step.py

@@ -15,6 +15,7 @@
 
 from typing import List, Dict, Optional, Union
 from enum import Enum
+import warnings
 
 import attr
 
@@ -152,10 +153,20 @@ def to_request(self) -> RequestType:
     def _get_function_arn(self):
         """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.
+        It upserts a lambda function if function name is provided.
+        It updates a lambda function if lambda arn and code is provided.
+        It is a no-op if code is not provided but function arn is provided.
         """
         if self.lambda_func.function_arn is None:
             response = self.lambda_func.upsert()
             return response["FunctionArn"]
-        return self.lambda_func.function_arn
+
+        if self.lambda_func.zipped_code_dir is None and self.lambda_func.script is None:
+            warnings.warn(
+                "Lambda function won't be updated because zipped_code_dir \
+                or script is not provided."
+            )
+            return self.lambda_func.function_arn
+
+        response = self.lambda_func.update()
+        return response["FunctionArn"]

tests/unit/sagemaker/workflow/test_lambda_step.py

@@ -203,10 +203,12 @@ def test_lambda_step_no_inputs_outputs(sagemaker_session):
     }
 
 
-def test_lambda_step_with_function_arn(sagemaker_session):
+def test_lambda_step_with_function_arn_no_lambda_update(sagemaker_session):
     lambda_func = MagicMock(
         function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
         session=sagemaker_session,
+        zipped_code_dir=None,
+        script=None,
     )
     lambda_step = LambdaStep(
         name="MyLambdaStep",
@@ -220,6 +222,24 @@ def test_lambda_step_with_function_arn(sagemaker_session):
     lambda_func.upsert.assert_not_called()
 
 
+def test_lambda_step_with_function_arn_lambda_updated(sagemaker_session):
+    lambda_func = MagicMock(
+        function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
+        zipped_code_dir=None,
+        script="code",
+        session=sagemaker_session,
+    )
+    lambda_step = LambdaStep(
+        name="MyLambdaStep",
+        depends_on=["TestStep"],
+        lambda_func=lambda_func,
+        inputs={},
+        outputs=[],
+    )
+    lambda_step._get_function_arn()
+    lambda_func.update.assert_called_once()
+
+
 def test_lambda_step_without_function_arn(sagemaker_session):
     lambda_func = MagicMock(
         function_arn=None,

tests/unit/test_lambda_helper.py

@@ -41,6 +41,7 @@ def sagemaker_session():
         boto_region_name=REGION,
         config=None,
         local_mode=False,
+        # default_bucket=S3_BUCKET,
     )
     return session_mock
 
@@ -117,7 +118,7 @@ def test_lambda_object_no_code_error():
     assert "Either zipped_code_dir or script must be provided" in str(error)
 
 
-def test_lambda_object_both_script_and_code_dir_error():
+def test_lambda_object_both_script_and_code_dir_error_with_name():
     with pytest.raises(ValueError) as error:
         lambda_helper.Lambda(
             function_name=FUNCTION_NAME,
@@ -130,6 +131,17 @@ def test_lambda_object_both_script_and_code_dir_error():
     assert "Provide either script or zipped_code_dir, not both." in str(error)
 
 
+def test_lambda_object_both_script_and_code_dir_error_with_arn():
+    with pytest.raises(ValueError) as error:
+        lambda_helper.Lambda(
+            function_arn=LAMBDA_ARN,
+            script=SCRIPT,
+            zipped_code_dir=ZIPPED_CODE_DIR,
+            session=sagemaker_session,
+        )
+    assert "Provide either script or zipped_code_dir, not both." in str(error)
+
+
 def test_lambda_object_no_handler_error():
     with pytest.raises(ValueError) as error:
         lambda_helper.Lambda(
@@ -274,6 +286,25 @@ def test_update_lambda_happycase2(sagemaker_session):
     )
 
 
+@patch("sagemaker.lambda_helper._upload_to_s3", return_value=S3_KEY)
+def test_update_lambda_s3bucket_not_provided(s3_upload, sagemaker_session):
+    lambda_obj = lambda_helper.Lambda(
+        function_arn=LAMBDA_ARN,
+        execution_role_arn=EXECUTION_ROLE,
+        zipped_code_dir=ZIPPED_CODE_DIR,
+        handler=HANDLER,
+        session=sagemaker_session,
+    )
+
+    lambda_obj.update()
+
+    sagemaker_session.lambda_client.update_function_code.assert_called_with(
+        FunctionName=LAMBDA_ARN,
+        S3Bucket=sagemaker_session.default_bucket(),
+        S3Key=s3_upload.return_value,
+    )
+
+
 @patch("sagemaker.lambda_helper._zip_lambda_code", return_value=ZIPPED_CODE)
 def test_update_lambda_client_error(sagemaker_session):
     lambda_obj = lambda_helper.Lambda(