Package Lambda functions and layers with CDK bundling, use cache to prevent unnecessary docker runs

package_lambda.py

from typing import List
import os
import shutil
import hashlib
import tempfile
from aws_cdk import core, aws_lambda
import jsii
from pathlib import Path


def create_lambda_function_deployment_package(
    dir_name: str, runtime: aws_lambda.Runtime
):
    commands = [
        "pip install -r requirements.txt --target /asset-output --no-dependencies --quiet",
        "cp -r /asset-input/* /asset-output",
        "cp -r /asset-output $ASSET_CACHE_DIR",
    ]
    return _create_deployment_package(dir_name, runtime, commands=commands)


def create_lambda_layer_deployment_package(dir_name: str, runtime: aws_lambda.Runtime):
    runtime_family = runtime.family.value.lower()
    commands = [
        f"pip install -r requirements.txt --target /asset-output/{runtime_family} --no-dependencies --quiet",
        f"rm -rf /asset-output/{runtime_family}/boto*",
        f"cp -r /asset-output/{runtime_family} $ASSET_CACHE_DIR",
    ]
    return _create_deployment_package(dir_name, runtime, commands=commands)


def _create_deployment_package(
    dir_name: str, runtime: aws_lambda.Runtime, commands: List[str] = []
):
    base_dir = Path(dir_name)
    hash_seed = "".join([runtime.family.value, *commands])
    bundler = LocalBundler(base_dir, hash_seed)
    code = aws_lambda.Code.from_asset(
        str(base_dir),
        bundling=core.BundlingOptions(
            image=runtime.bundling_docker_image,
            local=bundler,
            command=[
                "bash",
                "-c",
                " && ".join(commands),
            ],
            volumes=[
                core.DockerVolume(
                    host_path=str(bundler.cache_dir),
                    container_path=str(bundler.cache_dir),
                ),
            ],
            environment={
                "ASSET_CACHE_DIR": str(bundler.cache_dir),
            },
        ),
    )
    return code


@jsii.implements(core.ILocalBundling)
class LocalBundler:

    base_dir: Path
    cache_dir: str

    def __init__(self, base_dir: Path, hash_seed: str) -> None:
        self.base_dir = base_dir
        current_hash = self.get_dir_hash(self.base_dir, hash_seed)
        self.cache_dir = Path(os.getcwd()) / ".asset_cache" / current_hash
        print(f"Cache dir is {self.cache_dir}")

    def try_bundle(
        self,
        dir: str,
        options: core.BundlingOptions,
        *args,
        **kwargs,
    ):
        print("Running local bundler ... ")

        if os.path.isdir(self.cache_dir) and os.listdir(self.cache_dir):
            print("Found cache!")
            shutil.copytree(self.cache_dir, dir, dirs_exist_ok=True)
            return True
        print("No cache found!")
        return False

    @staticmethod
    def get_dir_hash(dir_name: str, hash_seed: str):
        with tempfile.TemporaryDirectory() as temp_dir:
            print(f"Hashing dir {dir_name} ...")
            shutil.make_archive(Path(temp_dir) / "tohash", "zip", dir_name)
            hash = hashlib.sha256()
            hash.update(hash_seed.encode())
            hash.update(open(Path(temp_dir) / "tohash.zip", "rb").read())
            return hash.hexdigest()