fix: Fix LambdaStep Creation (#3828)

Co-authored-by: Michael Trinh <[email protected]>
Co-authored-by: Michael <[email protected]>

Author: 3 people

src/sagemaker/lambda_helper.py

@@ -37,7 +37,6 @@ def __init__(
         memory_size: int = 128,
         runtime: str = "python3.8",
         vpc_config: dict = None,
-        architectures: list = None,
         environment: dict = None,
         layers: list = None,
     ):
@@ -71,8 +70,6 @@ def __init__(
             memory_size (int): Memory of the Lambda function in megabytes. Default is 128 MB.
             runtime (str): Runtime of the Lambda function. Default is set to python3.8.
             vpc_config (dict): VPC to deploy the Lambda function to. Default is None.
-            architectures (list): Which architecture to deploy to. Valid Values are
-                'x86_64' and 'arm64', default is None.
             environment (dict): Environment Variables for the Lambda function. Default is None.
             layers (list): List of Lambda layers for the Lambda function. Default is None.
         """
@@ -87,10 +84,9 @@ def __init__(
         self.timeout = timeout
         self.memory_size = memory_size
         self.runtime = runtime
-        self.vpc_config = vpc_config
-        self.environment = environment
-        self.architectures = architectures
-        self.layers = layers
+        self.vpc_config = vpc_config or {}
+        self.environment = environment or {}
+        self.layers = layers or []
 
         if function_arn is None and function_name is None:
             raise ValueError("Either function_arn or function_name must be provided.")
@@ -142,7 +138,6 @@ def create(self):
                 MemorySize=self.memory_size,
                 VpcConfig=self.vpc_config,
                 Environment=self.environment,
-                Architectures=self.architectures,
                 Layers=self.layers,
             )
             return response
@@ -163,7 +158,6 @@ def update(self):
                     response = lambda_client.update_function_code(
                         FunctionName=self.function_name or self.function_arn,
                         ZipFile=_zip_lambda_code(self.script),
-                        Architectures=self.architectures,
                     )
                 else:
                     bucket = self.s3_bucket or self.session.default_bucket()
@@ -186,7 +180,6 @@ def update(self):
                             zipped_code_dir=self.zipped_code_dir,
                             s3_bucket=bucket,
                         ),
-                        Architectures=self.architectures,
                     )
                 return response
             except ClientError as e:

tests/unit/test_lambda_helper.py

@@ -187,10 +187,9 @@ def test_create_lambda_happycase1(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
-        Architectures=None,
-        VpcConfig=None,
-        Environment=None,
-        Layers=None,
+        VpcConfig={},
+        Environment={},
+        Layers=[],
     )
 
 
@@ -216,10 +215,9 @@ def test_create_lambda_happycase2(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
-        Architectures=None,
-        VpcConfig=None,
-        Environment=None,
-        Layers=None,
+        VpcConfig={},
+        Environment={},
+        Layers=[],
     )
 
 
@@ -231,7 +229,6 @@ def test_create_lambda_happycase3(sagemaker_session):
         script=SCRIPT,
         handler=HANDLER,
         session=sagemaker_session,
-        architectures=["x86_64"],
         environment={"Name": "my-test-lambda"},
         vpc_config={
             "SubnetIds": ["test-subnet-1"],
@@ -251,7 +248,6 @@ def test_create_lambda_happycase3(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
-        Architectures=["x86_64"],
         VpcConfig={"SubnetIds": ["test-subnet-1"], "SecurityGroupIds": ["sec-group-1"]},
         Environment={"Name": "my-test-lambda"},
         Layers=["my-test-layer-1", "my-test-layer-2"],
@@ -314,7 +310,6 @@ def test_update_lambda_happycase1(sagemaker_session):
     sagemaker_session.lambda_client.update_function_code.assert_called_with(
         FunctionName=FUNCTION_NAME,
         ZipFile=ZIPPED_CODE,
-        Architectures=None,
     )
 
 
@@ -335,7 +330,6 @@ def test_update_lambda_happycase2(sagemaker_session):
         FunctionName=LAMBDA_ARN,
         S3Bucket=S3_BUCKET,
         S3Key=S3_KEY,
-        Architectures=None,
     )
 
 
@@ -347,7 +341,6 @@ def test_update_lambda_happycase3(sagemaker_session):
         script=SCRIPT,
         handler=HANDLER,
         session=sagemaker_session,
-        architectures=["x86_64"],
         environment={"Name": "my-test-lambda"},
         vpc_config={
             "SubnetIds": ["test-subnet-1"],
@@ -360,7 +353,6 @@ def test_update_lambda_happycase3(sagemaker_session):
     sagemaker_session.lambda_client.update_function_code.assert_called_with(
         FunctionName=FUNCTION_NAME,
         ZipFile=ZIPPED_CODE,
-        Architectures=["x86_64"],
     )
 
 
@@ -380,7 +372,6 @@ def test_update_lambda_s3bucket_not_provided(s3_upload, sagemaker_session):
         FunctionName=LAMBDA_ARN,
         S3Bucket=sagemaker_session.default_bucket(),
         S3Key=s3_upload.return_value,
-        Architectures=None,
     )
 
 
@@ -425,10 +416,9 @@ def test_upsert_lambda_happycase1(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
-        Architectures=None,
-        VpcConfig=None,
-        Environment=None,
-        Layers=None,
+        VpcConfig={},
+        Environment={},
+        Layers=[],
     )
 
 
@@ -455,7 +445,7 @@ def test_upsert_lambda_happycase2(sagemaker_session):
     lambda_obj.upsert()
 
     sagemaker_session.lambda_client.update_function_code.assert_called_once_with(
-        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE, Architectures=None
+        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE
     )