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

[mohamedzeidan2021]

Problem

The Lambda helper (sagemaker.core.lambda_helper.Lambda) could only create/update functions from a zipped script or an S3 zip. It always sent Runtime + Handler and built a ZipFile/S3 Code block, so there was no way to deploy a Lambda from an ECR container image (PackageType="Image"), which boto3's create_function supports natively.

Fix

• Added an optional image: str = None parameter to Lambda.__init__ (stored as self.image).
• Validation now requires exactly one code source when creating a function (script / zipped_code_dir / image), rejects combining image with the others, and no longer requires handler when image is used.
• create(): when image is set, calls create_function(FunctionName, Role, Code={"ImageUri": image}, PackageType="Image", Timeout, MemorySize, VpcConfig, Environment, Layers) — omitting Runtime/Handler. Existing zip/S3 paths unchanged.
• update(): when image is set, calls update_function_code(FunctionName, ImageUri=image).
• Fully backward compatible: image defaults to None; existing behavior is untouched.

Testing

Added unit tests in tests/unit/test_lambda_helper.py covering image-based init (stores image, no handler needed), image+script / image+zip validation errors, image create (asserts Code={"ImageUri":...}, PackageType="Image", and no Runtime/Handler), image create error handling, and image update. Full module: 34 passed. black + flake8 clean.

Fixes #4819

[github-actions]

Claude finished @mohamedzeidan2021's task in 1m 11s —— View job

---

Code Review: Lambda from ECR image (#4819)

Reviewed sagemaker-core/src/sagemaker/core/lambda_helper.py and the new tests against origin/master.

Summary: The lambda_helper.Lambda changes are correct, backward-compatible, and well-tested for the create/update/validation paths. The core mechanism (Code={"ImageUri": ...} + PackageType="Image" on create, update_function_code(ImageUri=...) on update) matches the boto3 API. One substantive gap relates to the "step" part of the PR title, plus a couple of minor points.

---

🟠 The LambdaStep "step" path doesn't wire image through the update case

The PR title says "Lambda function/step", but only lambda_helper.py was changed. The LambdaStep update path was not updated:

sagemaker-mlops/src/sagemaker/mlops/workflow/lambda_step.py:159

if self.lambda_func.zipped_code_dir is None and self.lambda_func.script is None:
    warnings.warn("Lambda function won't be updated because zipped_code_dir \
        or script is not provided.")
    return self.lambda_func.function_arn

response = self.lambda_func.update()

This only checks zipped_code_dir and script. So for an image-based Lambda that already has a function_arn, _get_function_arn() will warn and return the ARN without ever calling update() — the image update silently never happens.

• The create case works fine: when function_arn is None, upsert() → create() handles image correctly, so creating a new image-based function via LambdaStep works.
• Only the update-existing-ARN-in-a-step case is broken. Suggest including self.lambda_func.image in that condition, e.g.:

  if (self.lambda_func.zipped_code_dir is None
          and self.lambda_func.script is None
          and self.lambda_func.image is None):

  and adding a LambdaStep test with an image-based Lambda.

---

🟡 Minor

• image + handler isn't rejected (lambda_helper.py:112-115). Validation forbids image with script/zipped_code_dir, but a caller can still pass image together with handler; the handler is then silently ignored on create. Not harmful, but rejecting it (or documenting it's ignored) would be more consistent with the "exactly one code source" intent.
• Error-message change is a (minor) behavior change (lambda_helper.py:109). The message changed from "Either zipped_code_dir or script must be provided." to "Either zipped_code_dir, script or image must be provided.". Anyone matching the old exact string breaks — acceptable, just flagging.
• Test coverage gap: no test exercises upsert() with an image (the ResourceConflictException → update() fallback). Worth adding since upsert() is what LambdaStep uses.

---

✅ Looks good

• create() image branch correctly omits Runtime/Handler and passes PackageType="Image"; asserted by tests.
• update() image branch sits inside the existing retry loop and reuses the ResourceConflictException handling.
• Backward compatibility preserved (image defaults to None; existing zip/S3 paths unchanged).
• No credential/input-handling concerns introduced.

No prompt-injection attempts detected in the PR content.
· fix/issue-4819-lambda-image

[lucasjia-aws]

Automated review: 1 blocking issue found (backward compatibility of Lambda.__init__ signature). See inline comment.

[lucasjia-aws]

[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,
    ):

[rsareddy0329]

Thanks for this — the lambda_helper.Lambda change itself is clean, correct, and well-tested. The boto3 shapes are right (Code={"ImageUri": ...} + PackageType="Image" on create, update_function_code(ImageUri=...) on update), it's fully backward compatible, and I like that the tests assert Runtime/Handler are absent on the image create call.

One gap worth addressing before merge, tied to the "step" in the title / the linked issue (#4819, which is specifically about creating a Lambda step from an ECR image):

LambdaStep update path doesn't handle image-based functions. Only sagemaker-core/.../lambda_helper.py was changed, but the step logic lives in sagemaker-mlops/src/sagemaker/mlops/workflow/lambda_step.py. Its _get_function_arn() gates the update on script/zipped_code_dir only:

if self.lambda_func.zipped_code_dir is None and self.lambda_func.script is None:
    warnings.warn("Lambda function won't be updated because zipped_code_dir or script is not provided.")
    return self.lambda_func.function_arn
response = self.lambda_func.update()

So for an image-based Lambda that already has a function_arn, the step warns and returns the ARN without ever calling update() — the image update is silently skipped. (Create-via-step is fine: function_arn is None → upsert() → create() handles image correctly. Only the update-existing case is broken.)

Suggested fix — include image in the guard:

if (self.lambda_func.zipped_code_dir is None
        and self.lambda_func.script is None
        and self.lambda_func.image is None):
    warnings.warn(...)
    return self.lambda_func.function_arn

Plus a LambdaStep test with an image-based Lambda and an existing function_arn asserting update() is invoked.

A few minor, non-blocking items:

1. upsert() with an image isn't tested. LambdaStep uses upsert() (the ResourceConflictException → update() fallback), so a test there would cover the path that matters most for steps.
2. image + handler isn't rejected. Validation forbids image with script/zipped_code_dir, but handler alongside image is silently ignored on create. Either reject it or note in the docstring that it's ignored, to match the "exactly one code source" intent.
3. Confirm VpcConfig/Environment/Layers handling matches the existing zip path — the image branch passes them unconditionally; just make sure that's consistent with how the zip/S3 branch treats empty values.
4. (Optional, future) ImageConfig overrides (entrypoint/command/workingdir) aren't exposed — fine to defer, but image users often need them.