build_pipeline_generator.py

import os
import re
from pathlib import Path
from argparse import ArgumentParser


class ModuleBuildConfig:
    
    def __init__(self, context: "BuildContext", module_dir: Path):
        super().__init__()
        
        self.context = context
        self.module_dir = module_dir.resolve()
        self.config_file = self.module_dir.joinpath("build_config.py")
        self.name = str(self.module_dir.relative_to(context.build_dir))
        
        globals = {}
        exec(self.config_file.read_text(), globals)
        
        self.target_image_name = globals.pop("TARGET_IMAGE_NAME", None)
        if self.target_image_name is not None and  not isinstance(self.target_image_name, str):
            raise TypeError("TARGET_IMAGE_NAME must be a str (or absent)")
        
        context.add_module(self)
        
        self.dependencies: list[ModuleBuildConfig] = []
        for dependency in globals.pop("BUILD_DEPENDENCIES", []):
            if not isinstance(dependency, str):
                raise ValueError("BUILD_DEPENDENCIES must be list of str")
            dependency_name = str(module_dir.joinpath(Path(dependency)).resolve().relative_to(context.build_dir))
            self.dependencies.append(context.get_or_load_module(dependency_name))
        
        self.pip_req_file = None
        pip_req_file_name = globals.pop("PIP_REQUIREMENTS_FILE", None)
        if pip_req_file_name is not None:
            assert isinstance(pip_req_file_name, str)
            self.pip_req_file = self.module_dir.joinpath(Path(pip_req_file_name))
            if not self.pip_req_file.is_file():
                raise ValueError(f"Cannot find pip requirements file {self.pip_req_file}")
        
        self.apt_runtime_dependencies = globals.pop("APT_RUNTIME_DEPENDENCIES", [])
        if any(not isinstance(r, str) for r in self.apt_runtime_dependencies):
            raise TypeError(f"APT_RUNTIME_DEPENDENCIES must be list of str")
        
        self.apt_build_dependencies = globals.pop("APT_BUILD_DEPENDENCIES", [])
        if any(not isinstance(r, str) for r in self.apt_build_dependencies):
            raise TypeError(f"APT_BUILD_DEPENDENCIES must be list of str")
        
        self.dockerfile_extra = globals.pop("DOCKERFILE_EXTRA", "")
        if not isinstance(self.dockerfile_extra, str):
            raise TypeError("DOCKERFILE_EXTRA must be str (or absent)")
        self.dockerfile_extra = self.dockerfile_extra.replace(
            "$MODULE_DIR",
            str(self.module_dir.relative_to(context.build_dir))
        )
        
        self.ci_test_job_template: str or None = globals.pop("CI_TEST_JOB_TEMPLATE", None)
        if self.ci_test_job_template is not None and not isinstance(self.dockerfile_extra, str):
            raise TypeError("CI_TEST_JOB_TEMPLATE must be str (or absent)")
        
        for g in globals.keys():
            assert isinstance(g, str)
            if g.isupper():
                raise ValueError(f"Unknown key {g} in config file")
    
    def check_cyclic_dependency(self, dependents_stack: list[str]):
        if self.name in dependents_stack:
            raise ValueError(f"Dependency cycle involving {self.name} detected")
        dependents_stack = dependents_stack + [self.name]
        for dependency in self.dependencies:
            dependency.check_cyclic_dependency(dependents_stack)
    
    def collect_docker_dependencies(
            self,
            pip_requirement_files: list[Path],
            apt_runtime_dependencies: list[str],
            apt_build_dependencies: list[str],
            dockerfile_extras: list[str]
    ):
        for dep in self.dependencies:
            dep.collect_docker_dependencies(
                pip_requirement_files,
                apt_runtime_dependencies,
                apt_build_dependencies,
                dockerfile_extras
            )
        
        if self.pip_req_file is not None:
            pip_requirement_files.append(self.pip_req_file)
        
        apt_runtime_dependencies.extend(self.apt_runtime_dependencies)
        apt_build_dependencies.extend(self.apt_build_dependencies)
        dockerfile_extras.append(self.dockerfile_extra)
    
    def generate_dockerfile(self):
        pip_requirement_files: list[Path] = []
        apt_runtime_dependencies: list[str] = []
        apt_build_dependencies: list[str] = []
        dockerfile_extras: list[str] = []
        self.collect_docker_dependencies(
            pip_requirement_files,
            apt_runtime_dependencies,
            apt_build_dependencies,
            dockerfile_extras
        )
        # Remove duplicates and ensure same order every time
        apt_runtime_dependencies = list(set(apt_runtime_dependencies))
        apt_runtime_dependencies.sort()
        apt_build_dependencies = list(set(apt_build_dependencies))
        apt_build_dependencies.sort()
        
        res = f"""\
#####################################################################
### WARNING: THIS FILE WAS AUTOMATICALLY GENERATED. DO NOT EDIT ! ###
#####################################################################
FROM python:3.13-slim AS base

RUN mkdir -p /code
WORKDIR /code

ENV PIP_CACHE_DIR=/tmp/pip-cache
RUN apt-get update && apt-get --no-install-recommends install -y {' '.join(apt_runtime_dependencies)}
"""
        for path, i in zip(pip_requirement_files, range(len(pip_requirement_files))):
            res += f"COPY {str(path.relative_to(self.context.build_dir))} /tmp/req-{i}.txt\n"
        
        pip_requirement_list_arg = " ".join(f"-r /tmp/req-{i}.txt" for i in range(len(pip_requirement_files)))
        
        res += f"""
FROM base AS builder

# This step builds (and installs) the requirements and also puts the build result into the pip cache
RUN apt-get update && apt-get --no-install-recommends install -y {' '.join(apt_build_dependencies)}
RUN pip install {pip_requirement_list_arg}

FROM base AS final

# Here we copy the pip cache with the built packages and install them again. Pip will use the cache and won't need the
# apt build dependencies. This saves a lot of space, compared to leaving the build dependencies in the final image
# (reduces the final image size by about half)
COPY --from=builder /tmp/pip-cache /tmp/pip-cache
RUN pip install {pip_requirement_list_arg}
"""
        for docker in dockerfile_extras:
            res += "\n" + docker + "\n"
        
        return res
    
    def generate_ci_jobs(self):
        self.context.ensure_in_ci()
        if self.target_image_name is None:
            raise ValueError("This module has no target image name and therefore cannot be built")
        return (self._generate_ci_build_job()
                + "\n" + self._generate_ci_test_job()
                + "\n" + self._generate_ci_deploy_job())
    
    def output_dockerfile_path(self) -> Path:
        return self.context.build_dir.joinpath(".dockerfiles").joinpath(self.target_image_name)
    
    def image_full_name(self):
        return f"{self.context.env_type()}/{self.target_image_name}"
    
    def ci_build_job_name(self):
        return f"build-{self.target_image_name}"
    
    @staticmethod
    def _get_auth_echo():
        return """\
echo "{\\"auths\\":{\\"$CI_REGISTRY\\":{\\"username\\":\\"$CI_REGISTRY_USER\\",\\"password\\":\\"$CI_REGISTRY_PASSWORD\\"}}}" > /kaniko/.docker/config.json\
"""

    def _generate_ci_build_job(self):
        kaniko_args = [
            f"--context=git://git.fsmpi.rwth-aachen.de/videoag/backend.git#{self.context.commit_sha}",
            f"--dockerfile={str(self.output_dockerfile_path().relative_to(self.context.build_dir))}",
            f"--git recurse-submodules=true",
            f"--destination=$CI_REGISTRY_IMAGE/{self.image_full_name()}#{self.context.commit_sha}",
            f"--build-arg=GIT_COMMIT_SHA={self.context.commit_sha}",
            f"--cache=true",
        ]
        
        if self.context.commit_tag is not None:
            kaniko_args.append(f"--build-arg=GIT_COMMIT_TAG={self.context.commit_tag}")
        
        return f"""
{self.ci_build_job_name()}:
    stage: build-and-test
    timeout: 1h
    needs:
        - pipeline: $PARENT_PIPELINE_ID
          job: generate-pipeline
    image:
        name: gcr.io/kaniko-project/executor:v1.23.2-debug
        entrypoint: [""]
    script:
        - {self._get_auth_echo()}
        - echo {self.context.commit_sha}
        - >-
            /kaniko/executor
            {"\n            ".join(kaniko_args)}
"""
    
    def _generate_ci_test_job(self):
        if self.ci_test_job_template is None:
            return ""
        assert isinstance(self.ci_test_job_template, str)
        res = f"""
test-{self.target_image_name}:
    stage: build-and-test
    needs: [{self.ci_build_job_name()}]
    image:
        name: $CI_REGISTRY_IMAGE/{self.image_full_name()}#{self.context.commit_sha}
"""
        res += "    " + "\n    ".join(self.ci_test_job_template.splitlines()) + "\n"
        return res
    
    def _generate_ci_deploy_job(self):
        destination_args = [
            f"--destination=$CI_REGISTRY_IMAGE/{self.image_full_name()}:latest"
        ]
        if self.context.is_production:
            destination_args.append(f"--destination=$CI_REGISTRY_IMAGE/{self.image_full_name()}:{self.context.commit_tag}")
        
        return f"""
deploy-{self.target_image_name}:
    stage: deploy
    timeout: 1h
    image:
        name: gcr.io/kaniko-project/executor:v1.23.2-debug
        entrypoint: [""]
    script:
        - {self._get_auth_echo()}
        - mkdir /workdir
        - echo "FROM $CI_REGISTRY_IMAGE/{self.image_full_name()} > /workdir/Dockerfile"
        - echo {self.context.commit_sha}
        - >-
            /kaniko/executor
            --context=dir:///workdir
            {"\n            ".join(destination_args)}
"""


class BuildContext:
    
    def __init__(self, build_dir: Path, commit_sha: str or None, commit_tag: str or None):
        super().__init__()
        self.build_dir = build_dir
        self.modules: dict[str, ModuleBuildConfig] = {}
        self._target_image_names: set[str] = set()
        self.commit_sha = commit_sha or None  # Make empty string to None
        self.commit_tag = commit_tag or None
        self.is_production = commit_tag is not None and re.fullmatch("v.*", commit_tag) is not None
    
    def env_type(self) -> str:
        return "production" if self.is_production else "development"
    
    def ensure_in_ci(self):
        if self.commit_sha is None:
            raise Exception("Not in GitLab CI. No commit sha given")
    
    def add_module(self, module: ModuleBuildConfig):
        if module.target_image_name in self._target_image_names:
            raise ValueError(f"Duplicate target image name {module.target_image_name}")
        self.modules[module.name] = module
    
    def get_or_load_module(self, name: str):
        if name in self.modules:
            return self.modules[name]
        module_dir = self.build_dir.joinpath(name)
        try:
            return ModuleBuildConfig(self, module_dir)
        except Exception as e:
            raise Exception(f"Exception while loading module {module_dir}", e)
    
    def generate_ci_pipeline(self):
        self.ensure_in_ci()
        
        pipeline = """
####################################################################
##### AUTOMATICALLY GENERATED PIPELINE. DO NOT CHANGE MANUALLY! ####
####################################################################

stages:
    - build-and-test
    - deploy
"""
        for module in self.modules.values():
            if module.target_image_name is not None:
                pipeline += module.generate_ci_jobs()
        
        return pipeline


def main():
    parser = ArgumentParser()
    parser.add_argument("--build-dir", type=Path, default=Path("."))
    parser.add_argument("--commit-sha", type=str, required=False)
    parser.add_argument("--commit-tag", type=str, required=False)
    parser.add_argument("--ci-pipeline-dest", type=Path, required=False)
    parser.add_argument("modules", nargs="+", type=Path)
    
    args = parser.parse_args()
    
    context = BuildContext(args.build_dir.resolve(), args.commit_sha, args.commit_tag)
    for module in args.modules:
        context.get_or_load_module(str(module.resolve().relative_to(context.build_dir)))
    
    for module in context.modules.values():
        if module.target_image_name is None:
            continue
        module.output_dockerfile_path().parent.mkdir(parents=True, exist_ok=True)
        module.output_dockerfile_path().write_text(module.generate_dockerfile())
    
    if args.ci_pipeline_dest is not None:
        args.ci_pipeline_dest.parent.mkdir(parents=True, exist_ok=True)
        args.ci_pipeline_dest.write_text(context.generate_ci_pipeline())


if __name__ == "__main__":
    main()