Various docker caching, CLI and config validation improvements

README.md

@@ -18,12 +18,13 @@
 
 📦 **`agi-pack`** allows you to define your Docker images using a simple YAML format, and then generate them on-the-fly using [Jinja2](https://jinja.palletsprojects.com/en/3.1.x/) templates with [Pydantic](https://docs.pydantic.dev/latest/)-based validation. It's a simple tool that aims to simplify the process of building Docker images for ML.
 
-🚨 **Disclaimer:** More than **75%** of this initial implementation was generated by GPT-4 and [Github Co-Pilot](https://github.com/features/copilot). See [attribution](#inspiration-and-attribution-🌟) section below for more details.
+🚨 **Disclaimer:** More than **75%** of the original implementation was generated by GPT-4 and [Github Co-Pilot](https://github.com/features/copilot). I've spent enough time building on it that it's probably **<30%** auto-generated code now. See [attribution](#inspiration-and-attribution-🌟) section below for more details.
 
 ## Goals 🎯
 
 - 😇 **Simplicity**: Make it easy to define and build docker images for ML.
 - 📦 **Best-practices**: Bring best-practices to building docker images for ML -- good base images, multi-stage builds, minimal image sizes, etc.
+- ⚡️ **Fast**: Make it lightning-fast to build and re-build docker images with out-of-the-box caching for apt, conda and pip packages.
 - 🧩 **Modular, Re-usable, Composable**: Define `base`, `dev` and `prod` targets with multi-stage builds, and re-use them wherever possible.
 - 👩‍💻 **Extensible**: Make the YAML / DSL easily hackable and extensible to support the ML ecosystem, as more libraries, drivers, HW vendors, come into the market.
 - ☁️ **Vendor-agnostic**: `agi-pack` is not intended to be built for any specific vendor -- I need this tool for internal purposes, but I decided to build it in the open and keep it simple.
@@ -150,7 +151,7 @@ Here's the corresponding [`Dockerfile`](./examples/generated/Dockerfile-multista
 
 📦 **agi-pack** is simply a weekend project I hacked together, that started with a conversation with [ChatGPT / GPT-4](#chatgpt-prompt).
 
-🚨 **Disclaimer:** More than **75%** of this initial implementation was generated by GPT-4 and [Github Co-Pilot](https://github.com/features/copilot).
+🚨 **Disclaimer:** More than **75%** of this initial implementation was generated by GPT-4 and [Github Co-Pilot](https://github.com/features/copilot). I've spent enough time building on it that it's probably **<30%** auto-generated code now.
 
 ### ChatGPT Prompt
 ---

agipack/builder.py

@@ -28,6 +28,9 @@ class AGIPackRenderOptions:
     env: str = field(default=AGIPACK_ENV)
     """AGIPack environment to use for the build."""
 
+    skip_base_builds: bool = field(default=False)
+    """Skip building the base images."""
+
     def is_prod(self) -> bool:
         """Check if the build is for production."""
         return self.env == "prod"
@@ -96,7 +99,10 @@ def _render_one(self, target: str, image_config: ImageConfig, options: AGIPackRe
 
         # Render the Dockerfile template
         image_dict["target"] = target
-        image_dict["is_base_image"] = self.config.is_root(target)
+        if options.skip_base_builds:
+            image_dict["is_base_image"] = False
+        else:
+            image_dict["is_base_image"] = self.config.is_root(target)
         image_dict["is_prod"] = options.is_prod()
         image_dict["agipack_version"] = __version__
         content = template.render(image_dict)

agipack/cli.py

@@ -55,10 +55,14 @@ def generate(
     base_image: str = typer.Option(
         None, "--base", "-b", help="Base image to use for the root/base target.", show_default=False
     ),
-    tag: str = typer.Option("agipack:{target}", "--tag", "-t", help="Image tag f-string.", show_default=True),
-    prod: bool = typer.Option(False, "--prod", "-p", help="Generate a production Dockerfile.", show_default=False),
-    lint: bool = typer.Option(False, "--lint", "-l", help="Lint the generated Dockerfile.", show_default=False),
-    build: bool = typer.Option(False, "--build", "-b", help="Build the Docker image after generating the Dockerfile."),
+    tag: str = typer.Option("{name}:{target}", "--tag", "-t", help="Image tag f-string.", show_default=True),
+    _target: str = typer.Option(None, "--target", help="Build specific target.", show_default=False),
+    prod: bool = typer.Option(False, "--prod", help="Generate a production Dockerfile.", show_default=False),
+    lint: bool = typer.Option(False, "--lint", help="Lint the generated Dockerfile.", show_default=False),
+    build: bool = typer.Option(False, "--build", help="Build the Docker image after generating the Dockerfile."),
+    skip_base_builds: bool = typer.Option(
+        False, "--skip-base", help="Skip building the base image.", show_default=False
+    ),
 ):
     """Generate the Dockerfile with optional overrides.
 
@@ -81,17 +85,25 @@ def generate(
         config.images[root].base = base_image
 
     # Render the Dockerfiles with the new filename and configuration
+    trees = []
     builder = AGIPack(config)
-    dockerfiles = builder.render(filename=filename, env="prod" if prod else "dev")
+    dockerfiles = builder.render(filename=filename, env="prod" if prod else "dev", skip_base_builds=skip_base_builds)
     for target, filename in dockerfiles.items():
+        # Skip if the target is not the one we want to build
+        if _target is not None and target != _target:
+            continue
         image_config = config.images[target]
 
         # Build the Dockerfile using the generated filename and target
-        tag_name = f"{image_config.name}:{target}" if tag is None else tag.format(target=target)
+        tag_name = (
+            f"{image_config.name}:{target}" if tag is None else tag.format(name=image_config.name, target=target)
+        )
         cmd = f"docker build -f {filename} --target {target} -t {tag_name} ."
+
+        # Print the command to build the Dockerfile
         tree = Tree(f"📦 [bold white]{target}[/bold white]")
         tree.add(
-            f"🎉 Successfully generated Dockerfile (target=[bold white]{target}[/bold white], filename=[bold white]{filename}[/bold white])."
+            f"[bold green]✓[/bold green] Successfully generated Dockerfile (target=[bold white]{target}[/bold white], filename=[bold white]{filename}[/bold white])."
         ).add(f"[green]`{cmd}`[/green]")
         print(tree)
 
@@ -105,6 +117,15 @@ def generate(
             print(f"🚀 Building Docker image for target [{target}]")
             builder.build(filename=filename, target=target, tags=[tag_name])
 
+            tree.add(
+                f"[bold green]✓[/bold green] Successfully built image (target=[bold white]{target}[/bold white], image=[bold white]{tag_name}[/bold white])."
+            )
+            trees.append(tree)
+
+    # Re-render the tree
+    for tree in trees:
+        print(tree)
+
 
 @app.command()
 def build(
@@ -120,9 +141,13 @@ def build(
     base_image: str = typer.Option(
         None, "--base", "-b", help="Base image to use for the root/base target.", show_default=False
     ),
-    tag: str = typer.Option("agipack:{target}", "--tag", "-t", help="Image tag f-string.", show_default=True),
-    prod: bool = typer.Option(False, "--prod", "-p", help="Generate a production Dockerfile.", show_default=False),
-    lint: bool = typer.Option(False, "--lint", "-l", help="Lint the generated Dockerfile.", show_default=False),
+    tag: str = typer.Option("{name}:{target}", "--tag", "-t", help="Image tag f-string.", show_default=True),
+    target: str = typer.Option(None, "--target", help="Build specific target.", show_default=False),
+    prod: bool = typer.Option(False, "--prod", help="Generate a production Dockerfile.", show_default=False),
+    lint: bool = typer.Option(False, "--lint", help="Lint the generated Dockerfile.", show_default=False),
+    skip_base_builds: bool = typer.Option(
+        False, "--skip-base", help="Skip building the base image.", show_default=False
+    ),
 ):
     """Generate the Dockerfile with optional overrides.
 
@@ -132,9 +157,21 @@ def build(
         agi-pack build -c agibuild.yaml -p 3.8.10
         agi-pack build -c agibuild.yaml -b python:3.8.10-slim
         agi-pack build -c agibuild.yaml -t "my-image-name:{target}"
+        agi-pack build -c agibuild.yaml -t "my-image-name:my-target"
         agi-pack build -c agibuild.yaml --prod --lint
     """
-    generate(config_filename, filename, python, base_image, tag, prod, lint, build=True)
+    generate(
+        config_filename,
+        filename,
+        python,
+        base_image,
+        tag,
+        target,
+        prod,
+        lint,
+        build=True,
+        skip_base_builds=skip_base_builds,
+    )
 
 
 if __name__ == "__main__":

agipack/config.py

@@ -4,7 +4,7 @@
 from typing import Dict, List, Optional, Tuple, Union
 
 import yaml
-from pydantic import validator
+from pydantic import Extra, validator
 from pydantic.dataclasses import dataclass
 
 logger = logging.getLogger(__name__)
@@ -28,7 +28,13 @@ def is_leaf(self) -> bool:
         return not len(self.children)
 
 
-@dataclass
+class _ForbidExtrasConfig:
+    """Pydantic config to forbid extra fields."""
+
+    extra = Extra.forbid
+
+
+@dataclass(config=_ForbidExtrasConfig)
 class ImageConfig:
     """AGIPack configuration for a docker target specified in `agibuild.yaml`
 
@@ -50,16 +56,19 @@ class ImageConfig:
 
     """
 
-    image: str = field(default="agi:latest")
+    image: str = field(default=None)
     """Name of the image repository.
-    Defaults to the <target> name if not provided.
+    Defaults to <name>:<target> if not specified.
     """
 
-    name: str = field(default="agi")
+    name: str = field(default="agipack")
     """Pretty-name of the project in the image and
     reflected in the `AGIPACK_ENV` environment variable in the image.
     """
 
+    target: str = field(default=None)
+    """Name of the target."""
+
     base: str = field(default="debian:buster-slim")
     """Base docker image / target to use (FROM clause in the Dockerfile)."""
 
@@ -96,6 +105,12 @@ class ImageConfig:
     command: Optional[Union[str, List[str]]] = field(default_factory=list)
     """Command to run in the image."""
 
+    def __post_init__(self):
+        if self.target is None:
+            self.target = "latest"
+        if self.image is None:
+            self.image = f"{self.name}:{self.target}"
+
     def additional_kwargs(self):
         """Additional kwargs to pass to the Jinja2 Dockerfile template."""
         python_alias = f"py{''.join(self.python.split('.')[:2])}"
@@ -211,9 +226,15 @@ def load_yaml(cls, filename: Union[str, Path]) -> "AGIPackConfig":
         Returns:
             AGIPackConfig: AGIPack configuration.
         """
+        path = Path(filename)
+        logger.debug(f"Loading AGIPack configuration from {path}")
+        if not path.exists():
+            raise ValueError(f"YAML file {path.name} does not exist")
+        if not (path.name.endswith(".yaml") or path.name.endswith(".yml")):
+            raise ValueError(f"YAML file {path.name} must have a .yaml or .yml extension")
+
         # Load the YAML file
-        logger.debug(f"Loading AGIPack configuration from {filename}")
-        with open(filename, "r") as f:
+        with path.open("r") as f:
             data = yaml.safe_load(f)
         logger.debug(f"AGIPack configuration: {data}")
 

agipack/templates/Dockerfile.j2

@@ -2,7 +2,6 @@
 # >>>>>>>>>>>>>>>>>>>>>>>>>>>
 # Auto-generated by agi-pack (version={{ agipack_version }}).
 {%- else %}
-
 # >>>>>>>>>>>>>>>>>>>>>>>>>>>
 {%- endif %}
 FROM {{ base }} AS {{ target }}
@@ -15,6 +14,7 @@ ENV AGIPACK_PROJECT {{ name }}
 ENV AGIPACK_PYENV {{ name }}-{{ python_alias }}
 ENV AGIPACK_PATH /opt/agi-pack
 
+ENV DEBIAN_FRONTEND="noninteractive"
 ENV PYTHON_VERSION {{ python }}
 ENV PYTHONDONTWRITEBYTECODE 1
 ENV PYTHONUNBUFFERED 1
@@ -178,7 +178,8 @@ RUN apt-get -y autoclean \
 
 # Setup environment variables
 {% for key, value in env.items() -%}
-ENV {{ key }}={{ value }} {% endfor %}
+ENV {{ key }}={{ value }}
+{% endfor %}
 
 {%- endif %}
 

agipack/version.py

@@ -1 +1 @@
-__version__ = "0.1.15"
+__version__ = "0.1.16"

examples/generated/Dockerfile-base-cpu

@@ -1,12 +1,13 @@
 # >>>>>>>>>>>>>>>>>>>>>>>>>>>
-# Auto-generated by agi-pack (version=0.1.15).
+# Auto-generated by agi-pack (version=0.1.16).
 FROM debian:buster-slim AS base-cpu
 
 # Setup environment variables
-ENV AGIPACK_PROJECT agi
-ENV AGIPACK_PYENV agi-py38
+ENV AGIPACK_PROJECT agipack
+ENV AGIPACK_PYENV agipack-py38
 ENV AGIPACK_PATH /opt/agi-pack
 
+ENV DEBIAN_FRONTEND="noninteractive"
 ENV PYTHON_VERSION 3.8
 ENV PYTHONDONTWRITEBYTECODE 1
 ENV PYTHONUNBUFFERED 1
@@ -94,4 +95,4 @@ RUN apt-get -y autoclean \
     && echo "pip cleanup complete"
 
 # Setup environment variables
-ENV MY_ENV_VAR=value 
+ENV MY_ENV_VAR=value

examples/generated/Dockerfile-base-cu118

@@ -1,12 +1,13 @@
 # >>>>>>>>>>>>>>>>>>>>>>>>>>>
-# Auto-generated by agi-pack (version=0.1.15).
+# Auto-generated by agi-pack (version=0.1.16).
 FROM nvidia/cuda:11.8.0-base-ubuntu22.04 AS base-gpu
 
 # Setup environment variables
-ENV AGIPACK_PROJECT agi
-ENV AGIPACK_PYENV agi-py38
+ENV AGIPACK_PROJECT agipack
+ENV AGIPACK_PYENV agipack-py38
 ENV AGIPACK_PATH /opt/agi-pack
 
+ENV DEBIAN_FRONTEND="noninteractive"
 ENV PYTHON_VERSION 3.8
 ENV PYTHONDONTWRITEBYTECODE 1
 ENV PYTHONUNBUFFERED 1
@@ -104,4 +105,4 @@ RUN apt-get -y autoclean \
     && echo "pip cleanup complete"
 
 # Setup environment variables
-ENV MY_ENV_VAR=value 
+ENV MY_ENV_VAR=value

examples/generated/Dockerfile-builder

@@ -1,12 +1,13 @@
 # >>>>>>>>>>>>>>>>>>>>>>>>>>>
-# Auto-generated by agi-pack (version=0.1.15).
+# Auto-generated by agi-pack (version=0.1.16).
 FROM debian:buster-slim AS agipack-builder
 
 # Setup environment variables
-ENV AGIPACK_PROJECT agi
-ENV AGIPACK_PYENV agi-py38
+ENV AGIPACK_PROJECT agipack
+ENV AGIPACK_PYENV agipack-py38
 ENV AGIPACK_PATH /opt/agi-pack
 
+ENV DEBIAN_FRONTEND="noninteractive"
 ENV PYTHON_VERSION 3.8
 ENV PYTHONDONTWRITEBYTECODE 1
 ENV PYTHONUNBUFFERED 1