Skip to content

3826: Fix LambdaStep Creation #3828

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
May 9, 2023
13 changes: 3 additions & 10 deletions src/sagemaker/lambda_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
):
Expand Down Expand Up @@ -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.
"""
Expand All @@ -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.")
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand All @@ -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:
Expand Down
30 changes: 10 additions & 20 deletions tests/unit/test_lambda_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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=[],
)


Expand All @@ -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=[],
)


Expand All @@ -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"],
Expand All @@ -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"],
Expand Down Expand Up @@ -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,
)


Expand All @@ -335,7 +330,6 @@ def test_update_lambda_happycase2(sagemaker_session):
FunctionName=LAMBDA_ARN,
S3Bucket=S3_BUCKET,
S3Key=S3_KEY,
Architectures=None,
)


Expand All @@ -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"],
Expand All @@ -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"],
)


Expand All @@ -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,
)


Expand Down Expand Up @@ -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=[],
)


Expand All @@ -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
)


Expand Down