Skip to content

feat: support creating a Lambda function/step from an ECR image (#4819) #6349

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.

Sign up for GitHub

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

Open
mohamedzeidan2021 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: support creating a Lambda function/step from an ECR image (#4819) #6349

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 40 additions & 6 deletions sagemaker-core/src/sagemaker/core/lambda_helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ def __init__(
zipped_code_dir: str = None,
s3_bucket: str = None,
script: str = None,
image: str = None,

@lucasjia-aws lucasjia-aws Sep 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[blocker] backward-compat

The new image parameter is inserted between script and handler rather than appended to the end of the signature. This shifts the positional binding of every subsequent parameter (handler, session, timeout, memory_size, runtime, vpc_config, environment, layers), so any existing caller that passes handler (or a later arg) positionally now silently binds that value to image. Positional argument order is public API contract for pip install sagemaker users. Append image after the existing parameters instead:

    def __init__(
        self,
        ...
        layers: list = None,
        image: str = None,
    ):

handler: str = None,
session: Session = None,
timeout: int = 120,
Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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()
Expand All @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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),
Expand Down
107 changes: 106 additions & 1 deletion sagemaker-core/tests/unit/test_lambda_helper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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(
Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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):
Expand Down
Loading