feature: Add VpcConfig, Architectures and Environment parameters for CreateLambda operation

Author: dipanjank

src/sagemaker/lambda_helper.py

@@ -36,6 +36,9 @@ def __init__(
         timeout: int = 120,
         memory_size: int = 128,
         runtime: str = "python3.8",
+        vpc_config: dict = None,
+        architectures: list = None,
+        environment: dict = None,
     ):
         """Constructs a Lambda instance.
 
@@ -66,6 +69,10 @@ def __init__(
             timeout (int): Timeout of the Lambda function in seconds. Default is 120 seconds.
             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.
         """
         self.function_arn = function_arn
         self.function_name = function_name
@@ -78,6 +85,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
 
         if function_arn is None and function_name is None:
             raise ValueError("Either function_arn or function_name must be provided.")
@@ -123,6 +133,9 @@ def create(self):
                 Code=code,
                 Timeout=self.timeout,
                 MemorySize=self.memory_size,
+                VpcConfig=self.vpc_config,
+                Environment=self.environment,
+                Architectures=self.architectures,
             )
             return response
         except ClientError as e:
@@ -152,6 +165,7 @@ def update(self):
                             zipped_code_dir=self.zipped_code_dir,
                             s3_bucket=self.s3_bucket,
                         ),
+                        Architectures=self.architectures
                     )
                 return response
             except ClientError as e:

tests/unit/test_lambda_helper.py

@@ -302,6 +302,12 @@ def test_upsert_lambda_happycase1(sagemaker_session):
         script=SCRIPT,
         handler=HANDLER,
         session=sagemaker_session,
+        architectures=["x86_64"],
+        environment={"Name": "my-test-lambda"},
+        vpc_config={
+            "SubnetIds": ["test-subnet-1"],
+            "SecurityGroupIds": ["sec-group-1"]
+        }
     )
 
     code = {"ZipFile": ZIPPED_CODE}
@@ -315,6 +321,12 @@ def test_upsert_lambda_happycase1(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
+        Architectures=["x86_64"],
+        Environment={"Name": "my-test-lambda"},
+        VpcConfig={
+            "SubnetIds": ["test-subnet-1"],
+            "SecurityGroupIds": ["sec-group-1"]
+        }
     )
 
 
@@ -326,6 +338,7 @@ def test_upsert_lambda_happycase2(sagemaker_session):
         script=SCRIPT,
         handler=HANDLER,
         session=sagemaker_session,
+        architectures=["x86_64"]
     )
 
     sagemaker_session.lambda_client.create_function.side_effect = ClientError(
@@ -336,7 +349,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
+        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE, Architectures=["x86_64"]
     )