From 1392222f53669dc8d3e6c149b0b1ee5a6e8fac44 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Mon, 28 Sep 2026 14:23:47 -0700 Subject: [PATCH] feat: support creating a Lambda function/step from an ECR image (#4819) --- .../src/sagemaker/core/lambda_helper.py | 46 +++++++- .../tests/unit/test_lambda_helper.py | 107 +++++++++++++++++- 2 files changed, 146 insertions(+), 7 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/lambda_helper.py b/sagemaker-core/src/sagemaker/core/lambda_helper.py index 41ee4d5919..9d2b5be3e6 100644 --- a/sagemaker-core/src/sagemaker/core/lambda_helper.py +++ b/sagemaker-core/src/sagemaker/core/lambda_helper.py @@ -34,6 +34,7 @@ def __init__( zipped_code_dir: str = None, s3_bucket: str = None, script: str = None, + image: str = None, handler: str = None, session: Session = None, timeout: int = 120, @@ -50,8 +51,10 @@ def __init__( This class can be used either for creating a new Lambda function or using an existing one. When using an existing Lambda function, only the function_arn argument is required. - When creating a new one the function_name, execution_role_arn and handler arguments - are required, as well as either script or zipped_code_dir. + When creating a new one the function_name and execution_role_arn arguments are + required, as well as exactly one code source: script, zipped_code_dir or image. + A handler is additionally required for zip-based functions (script or + zipped_code_dir); it is not used for image-based functions. Args: function_arn (str): The arn of the Lambda function. @@ -62,6 +65,10 @@ def __init__( s3_bucket (str): The bucket where zipped code is uploaded. If not provided, default session bucket is used to upload zipped_code_dir. script (str): The path of Lambda function script for direct zipped upload + image (str): The URI of an ECR container image to deploy the Lambda function + from. When provided, the function is created with PackageType "Image" and + no handler/runtime are required. Cannot be combined with script or + zipped_code_dir. handler (str): The Lambda handler. The format for handler should be file_name.function_name. For ex: if the name of the Lambda script is hello_world.py and Lambda function definition in that script is @@ -81,6 +88,7 @@ def __init__( self.zipped_code_dir = zipped_code_dir self.s3_bucket = s3_bucket self.script = script + self.image = image self.handler = handler self.execution_role_arn = execution_role_arn self.session = session if session is not None else Session() @@ -97,16 +105,20 @@ def __init__( if function_name is not None: if execution_role_arn is None: raise ValueError("execution_role_arn must be provided.") - if zipped_code_dir is None and script is None: - raise ValueError("Either zipped_code_dir or script must be provided.") + if zipped_code_dir is None and script is None and image is None: + raise ValueError("Either zipped_code_dir, script or image must be provided.") if zipped_code_dir and script: raise ValueError("Provide either script or zipped_code_dir, not both.") - if handler is None: + if image and (zipped_code_dir or script): + raise ValueError("Provide either image or zipped_code_dir/script, not both.") + if image is None and 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.") + if image and (zipped_code_dir or script): + raise ValueError("Provide either image or zipped_code_dir/script, not both.") def create(self): """Method to create a lambda function. @@ -118,6 +130,23 @@ def create(self): if self.function_name is None: raise ValueError("FunctionName must be provided to create a Lambda function.") + if self.image is not None: + try: + return lambda_client.create_function( + FunctionName=self.function_name, + Role=self.execution_role_arn, + Code={"ImageUri": self.image}, + PackageType="Image", + Timeout=self.timeout, + MemorySize=self.memory_size, + VpcConfig=self.vpc_config, + Environment=self.environment, + Layers=self.layers, + ) + except ClientError as e: + error = e.response["Error"] + raise ValueError(error) + if self.script is not None: code = {"ZipFile": _zip_lambda_code(self.script)} else: @@ -165,7 +194,12 @@ def update(self): retry_attempts = 7 for i in range(retry_attempts): try: - if self.script is not None: + if self.image is not None: + response = lambda_client.update_function_code( + FunctionName=self.function_name or self.function_arn, + ImageUri=self.image, + ) + elif self.script is not None: response = lambda_client.update_function_code( FunctionName=self.function_name or self.function_arn, ZipFile=_zip_lambda_code(self.script), diff --git a/sagemaker-core/tests/unit/test_lambda_helper.py b/sagemaker-core/tests/unit/test_lambda_helper.py index fb075e1ba9..e4ef632044 100644 --- a/sagemaker-core/tests/unit/test_lambda_helper.py +++ b/sagemaker-core/tests/unit/test_lambda_helper.py @@ -71,7 +71,9 @@ def test_lambda_init_missing_execution_role(self): def test_lambda_init_missing_code(self): """Test initialization fails without code when creating new function.""" - with pytest.raises(ValueError, match="Either zipped_code_dir or script must be provided"): + with pytest.raises( + ValueError, match="Either zipped_code_dir, script or image must be provided" + ): Lambda( function_name="my-function", execution_role_arn="arn:aws:iam::123456789012:role/my-role", @@ -98,6 +100,41 @@ def test_lambda_init_missing_handler(self): script="/path/to/script.py", ) + def test_lambda_init_with_image(self): + """Test initialization with an ECR image (no handler/runtime required).""" + lambda_obj = Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + image="123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest", + ) + assert lambda_obj.image == "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest" + assert lambda_obj.handler is None + + def test_lambda_init_image_and_script_fails(self): + """Test initialization fails when both image and script are provided.""" + with pytest.raises( + ValueError, match="Provide either image or zipped_code_dir/script, not both" + ): + Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + image="123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest", + script="/path/to/script.py", + handler="script.handler", + ) + + def test_lambda_init_image_and_zipped_code_fails(self): + """Test initialization fails when both image and zipped_code_dir are provided.""" + with pytest.raises( + ValueError, match="Provide either image or zipped_code_dir/script, not both" + ): + Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + image="123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest", + zipped_code_dir="/path/to/code.zip", + ) + def test_lambda_init_with_optional_params(self): """Test initialization with optional parameters.""" lambda_obj = Lambda( @@ -179,6 +216,51 @@ def test_create_with_zipped_code( "S3Key": "prefix/lambda/my-function/code", } + @patch("sagemaker.core.lambda_helper._get_lambda_client") + def test_create_with_image(self, mock_get_client): + """Test creating a Lambda function from an ECR image.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + mock_client.create_function.return_value = { + "FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:my-function" + } + + image_uri = "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest" + lambda_obj = Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + image=image_uri, + ) + result = lambda_obj.create() + + assert result["FunctionArn"] == "arn:aws:lambda:us-west-2:123456789012:function:my-function" + mock_client.create_function.assert_called_once() + call_args = mock_client.create_function.call_args[1] + assert call_args["Code"] == {"ImageUri": image_uri} + assert call_args["PackageType"] == "Image" + assert "Runtime" not in call_args + assert "Handler" not in call_args + + @patch("sagemaker.core.lambda_helper._get_lambda_client") + def test_create_with_image_client_error(self, mock_get_client): + """Test image-based create handles ClientError.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + error = ClientError( + {"Error": {"Code": "InvalidParameterValue", "Message": "Invalid parameter"}}, + "CreateFunction", + ) + mock_client.create_function.side_effect = error + + lambda_obj = Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + image="123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest", + ) + + with pytest.raises(ValueError): + lambda_obj.create() + @patch("sagemaker.core.lambda_helper._get_lambda_client") def test_create_without_function_name(self, mock_get_client): """Test create fails without function name.""" @@ -238,6 +320,29 @@ def test_update_with_script(self, mock_zip, mock_get_client): assert result["FunctionArn"] == "arn:aws:lambda:us-west-2:123456789012:function:my-function" mock_client.update_function_code.assert_called_once() + @patch("sagemaker.core.lambda_helper._get_lambda_client") + def test_update_with_image(self, mock_get_client): + """Test updating a Lambda function's code from an ECR image.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + mock_client.update_function_code.return_value = { + "FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:my-function" + } + + image_uri = "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-image:latest" + lambda_obj = Lambda( + function_name="my-function", + execution_role_arn="arn:aws:iam::123456789012:role/my-role", + image=image_uri, + ) + result = lambda_obj.update() + + assert result["FunctionArn"] == "arn:aws:lambda:us-west-2:123456789012:function:my-function" + mock_client.update_function_code.assert_called_once_with( + FunctionName="my-function", + ImageUri=image_uri, + ) + @patch("sagemaker.core.lambda_helper._get_lambda_client") @patch("sagemaker.core.lambda_helper._zip_lambda_code") def test_update_with_retry_on_resource_conflict(self, mock_zip, mock_get_client):