Merge branch 'topic/Add-support-for-arm64-lambda-functions-and-docker-buildkit' into 'master'

Improve support for Docker based Lambda functions

See merge request it/e3-aws!105

Author: adanaja

pyproject.toml

@@ -10,11 +10,9 @@ description = "E3 Cloud Formation Extensions"
 dependencies = [
   "boto3",
   "botocore",
-  # Only require docker for linux as there is a known dependency issue
-  # with pywin32 on windows
-  "docker; platform_system=='Linux'",
   "e3-core",
   "pyyaml",
+  "python-on-whales",
   "troposphere"
 ]
 
@@ -91,7 +89,6 @@ module = [
     "botocore.*",
     "boto3.*",
     "requests.*",
-    "docker.*",
     "troposphere.*",
 ]
 ignore_missing_imports = true

src/e3/aws/troposphere/awslambda/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from datetime import datetime
+from datetime import UTC, datetime
+from enum import Enum
 import logging
 import os
 import sys
@@ -34,10 +35,23 @@
     from troposphere import AWSObject
     import botocore.client
     from e3.aws.troposphere import Stack
+    from python_on_whales import DockerClient
 
 logger = logging.getLogger("e3.aws.troposphere.awslambda")
 
 
+class Architecture(Enum):
+    X86_64 = "x86_64"
+    ARM64 = "arm64"
+
+
+class UnknownPlatform(Exception):
+    """Unknown platform exception."""
+
+    def __init__(self, architecture: Architecture):
+        super().__init__(f"Unknown platform for architecture {architecture}.")
+
+
 def package_pyfunction_code(
     filename: str,
     /,
@@ -263,6 +277,7 @@ def __init__(
         code_version: int | None = None,
         timeout: int = 3,
         runtime: str | None = None,
+        architecture: Architecture | None = None,
         memory_size: int | None = None,
         ephemeral_storage_size: int | None = None,
         logs_retention_in_days: int | None = 731,
@@ -288,6 +303,7 @@ def __init__(
         :param code_version: code version
         :param timeout: maximum execution time (default: 3s)
         :param runtime: runtime to use
+        :param architecture: x86_64 or arm64. (default: x86_64)
         :param memory_size: the amount of memory available to the function at
             runtime. The value can be any multiple of 1 MB.
         :param ephemeral_storage_size: The size of the function’s /tmp directory
@@ -317,6 +333,7 @@ def __init__(
         self.runtime = runtime
         self.role = role
         self.handler = handler
+        self.architecture = architecture
         self.memory_size = memory_size
         self.ephemeral_storage_size = ephemeral_storage_size
         self.logs_retention_in_days = logs_retention_in_days
@@ -450,6 +467,9 @@ def lambda_resources(
         if self.handler is not None:
             params["Handler"] = self.handler
 
+        if self.architecture is not None:
+            params["Architectures"] = [self.architecture.value]
+
         if self.memory_size is not None:
             params["MemorySize"] = self.memory_size
 
@@ -592,9 +612,14 @@ def __init__(
         repository_name: str,
         image_tag: str,
         timeout: int = 3,
+        architecture: Architecture | None = None,
         memory_size: int | None = None,
+        logs_retention_in_days: int | None = 731,
+        environment: dict[str, str] | None = None,
         logging_config: awslambda.LoggingConfig | None = None,
         dl_config: awslambda.DeadLetterConfig | None = None,
+        docker_client: DockerClient | None = None,
+        **build_args: Any,
     ):
         """Initialize an AWS lambda function using a Docker image.
 
@@ -605,26 +630,48 @@ def __init__(
         :param repository_name: ECR repository name
         :param image_tag: docker image version
         :param timeout: maximum execution time (default: 3s)
+        :param architecture: x86_64 or arm64. (default: x86_64)
         :param memory_size: the amount of memory available to the function at
             runtime. The value can be any multiple of 1 MB.
+        :param logs_retention_in_days: The number of days to retain the log
+            events in the lambda log group
+        :param environment: Environment variables that are accessible from
+            function code during execution
         :param logging_config: The function's Amazon CloudWatch Logs settings
-        :param dl_config: The dead letter config that specifies the topic or queue where
-            lambda sends asynchronous events when they fail processing
+        :param dl_config: The dead letter config that specifies the topic or
+            queue where lambda sends asynchronous events when they fail processing
+        :param docker_client: Docker client to use for building and pushing.
+            This is here in case the user wants to customize the Docker client,
+            for example to use podman.
+        :param build_args: args to pass to docker build
         """
         super().__init__(
             name=name,
             description=description,
             role=role,
             timeout=timeout,
+            architecture=architecture,
             memory_size=memory_size,
+            logs_retention_in_days=logs_retention_in_days,
+            environment=environment,
             logging_config=logging_config,
             dl_config=dl_config,
         )
         self.source_dir: str = source_dir
         self.repository_name: str = repository_name
-        timestamp = datetime.utcnow().strftime("%Y-%m-%d-%H-%M-%S-%f")
+        timestamp = datetime.now(UTC).strftime("%Y-%m-%d-%H-%M-%S-%f")
         self.image_tag: str = f"{image_tag}-{timestamp}"
         self.image_uri: str | None = None
+        self.docker_client = docker_client
+        self.build_args = build_args
+        if "platforms" not in self.build_args:
+            match self.architecture:
+                case Architecture.ARM64:
+                    self.build_args["platforms"] = ["linux/arm64"]
+                case Architecture.X86_64 | None:
+                    self.build_args["platforms"] = ["linux/amd64"]
+                case _:
+                    raise UnknownPlatform(self.architecture)
 
     def resources(self, stack: Stack) -> list[AWSObject]:
         """Compute AWS resources for the construct.
@@ -641,6 +688,9 @@ def resources(self, stack: Stack) -> list[AWSObject]:
                 self.repository_name,
                 self.image_tag,
                 stack.deploy_session,
+                push=True,
+                docker_client=self.docker_client,
+                **self.build_args,
             )
 
         return self.lambda_resources(image_uri=self.image_uri)

src/e3/aws/util/ecr.py

@@ -5,64 +5,89 @@
 import base64
 import logging
 import tempfile
-from typing import TYPE_CHECKING
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any
 
-import docker
-from e3.fs import sync_tree, rm
+from e3.fs import sync_tree
+from python_on_whales import DockerClient
 
 if TYPE_CHECKING:
     from e3.aws import Session
 
 logger = logging.getLogger("e3.aws.utils.ecr")
 
 
+def get_ecr_credentials(session: Session) -> tuple[str, str, str]:
+    """Get ECR credentials (username, password, registry URL).
+
+    :param session: AWS session to get ECR credentials
+    :return: tuple of (username, password, registry URL)
+    """
+    ecr_client = session.client("ecr")
+    ecr_creds = ecr_client.get_authorization_token()["authorizationData"][0]
+    ecr_username, ecr_password = (
+        base64.b64decode(ecr_creds["authorizationToken"]).decode().split(":", 1)
+    )
+    ecr_url = ecr_creds["proxyEndpoint"]
+    return ecr_username, ecr_password, ecr_url
+
+
 def build_and_push_image(
-    source_dir: str, repository_name: str, image_tag: str, session: Session
+    source_dir: str,
+    repository_name: str,
+    image_tag: str,
+    session: Session,
+    docker_client: DockerClient | None = None,
+    **build_args: Any,
 ) -> str:
     """Build and push image to an ECR repository and return image URI.
 
     :param source_dir: directory where to find Dockerfile
     :param repository_name: ECR repository name
     :param image_tag: Docker image tag
     :param session: AWS session to push docker image to ECR
+    :param build_args: Keyword arguments to pass to ``docker_client.build``.
+        See the ``python-on-whales`` documentation for a complete list of
+        possible arguments. Note that this function sets default values for
+        some arguments if they are not provided: ``push=True``,
+        ``progress="plain"``, ``stream_logs=True``, and ``provenance=False``.
+    :param docker_client: Docker client to use for building and pushing.
+        This is here in case the user wants to customize the Docker client,
+        for example to use podman.
+    :raises DockerException: if there is an error building or pushing the image
+    :return: image URI
     """
-    temp_dir = tempfile.mkdtemp()
-    try:
-        sync_tree(source_dir, temp_dir)
-
-        docker_client = docker.from_env()
-        # Build image
-        image, build_log = docker_client.images.build(
-            path=temp_dir, tag=f"{repository_name}:{image_tag}", rm=True
-        )
-        logging.info(f"Building Docker image from {source_dir}")
-        for chunk in build_log:
-            if "stream" in chunk:
-                for line in chunk["stream"].splitlines():
-                    logging.debug(line)
+    # Create a Docker client and login to ECR
+    if docker_client is None:
+        docker_client = DockerClient()
+    ecr_username, ecr_password, ecr_url = get_ecr_credentials(session)
+    docker_client.login(username=ecr_username, password=ecr_password, server=ecr_url)
+    ecr_repo_name = f"{ecr_url.replace('https://', '')}/{repository_name}"
+    image_uri = f"{ecr_repo_name}:{image_tag}"
 
-        # Push image to registry
-        ecr_client = session.client("ecr")
-        ecr_creds = ecr_client.get_authorization_token()["authorizationData"][0]
-        ecr_username, ecr_password = (
-            base64.b64decode(ecr_creds["authorizationToken"])
-            .decode("utf-8")
-            .split(":", 1)
-        )
-        ecr_url = ecr_creds["proxyEndpoint"]
-        docker_client.login(
-            username=ecr_username, password=ecr_password, registry=ecr_url
-        )
-        ecr_repo_name = f"{ecr_url.replace('https://', '')}/{repository_name}"
-        image.tag(ecr_repo_name, tag=image_tag)
-
-        push_log = docker_client.images.push(ecr_repo_name, tag=image_tag)
-        logging.info(f"Pushing Docker image to {ecr_repo_name} with tag {image_tag}")
-        logging.debug(push_log)
+    with tempfile.TemporaryDirectory() as temp_dir:
+        sync_tree(source_dir, temp_dir)
 
-        image_uri = f"{ecr_repo_name}:{image_tag}"
+        # Set various build args, unless the user has already specified a value
+        # Disable provenance to avoid manifest list issues with AWS Lambda
+        # See: https://github.com/docker/buildx/issues/1533
+        defaults: dict[str, Any] = {
+            "push": True,
+            "context_path": temp_dir,
+            "tags": [image_uri],
+            "progress": "plain",
+            "stream_logs": True,
+            "provenance": False,
+        }
+        build_args = defaults | build_args
 
-    finally:
-        rm(temp_dir, recursive=True)
+        # Build and push image
+        logger.info(f"Building Docker image from {source_dir}")
+        logger.debug(f"Build args {build_args}")
+        build_result = docker_client.build(**build_args)
+        if isinstance(build_result, Iterator):
+            # stream_logs=True, causes docker_client.build to return an iterator
+            for log_text in build_result:
+                logger.debug(log_text)
 
     return image_uri