Skip to content

feat: export distribution container build artifacts #2186

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions .github/workflows/providers-build.yml
Original file line number Diff line number Diff line change
Expand Up @@ -198,3 +198,55 @@ jobs:
'source /etc/os-release && echo "$ID"' \
| grep -qE '^(rhel|ubi)$' \
|| { echo "Base image is not UBI 9!"; exit 1; }

export-build:
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2

- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: '3.10'

- name: Install uv
uses: astral-sh/setup-uv@0c5e2b8115b80b4c7c5ddf6ffdd634974642d182 # v5.4.1
with:
python-version: "3.10"

- name: Install LlamaStack
run: |
uv venv
source .venv/bin/activate
uv pip install -e .

- name: Pin template to UBI9 base
run: |
yq -i '
.image_type = "container" |
.image_name = "ubi9-test" |
.distribution_spec.container_image = "registry.access.redhat.com/ubi9:latest"
' llama_stack/templates/starter/build.yaml

- name: Test the export
run: |
# Support for USE_COPY_NOT_MOUNT=true LLAMA_STACK_DIR=. will be added in the future
uv run llama stack build --config llama_stack/templates/starter/build.yaml --export-dir export
for file in export/*; do
echo "File: $file"
if [[ "$file" == *.tar.gz ]]; then
echo "Tarball found"
tarball_found=1
tar -xzvf "$file" -C export
else
continue
fi
break
done
if [ -z "$tarball_found" ]; then
echo "Tarball not found"
exit 1
fi
cd export
docker build -t export-test -f ./Containerfile .
22 changes: 21 additions & 1 deletion docs/source/distributions/building_distro.md
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ The main points to consider are:

```
llama stack build -h
usage: llama stack build [-h] [--config CONFIG] [--template TEMPLATE] [--list-templates] [--image-type {conda,container,venv}] [--image-name IMAGE_NAME] [--print-deps-only] [--run]
usage: llama stack build [-h] [--config CONFIG] [--template TEMPLATE] [--list-templates] [--image-type {conda,container,venv}] [--image-name IMAGE_NAME] [--print-deps-only] [--run] [--export-dir EXPORT_DIR]

Build a Llama stack container

Expand All @@ -71,6 +71,8 @@ options:
found. (default: None)
--print-deps-only Print the dependencies for the stack only, without building the stack (default: False)
--run Run the stack after building using the same image type, name, and other applicable arguments (default: False)
--export-dir EXPORT_DIR
Export the build artifacts to a specified directory instead of building the container. This will create a tarball containing the Dockerfile and all necessary files to build the container. (default: None)

```

Expand Down Expand Up @@ -260,6 +262,24 @@ Containerfile created successfully in /tmp/tmp.viA3a3Rdsg/ContainerfileFROM pyth
You can now edit ~/meta-llama/llama-stack/tmp/configs/ollama-run.yaml and run `llama stack run ~/meta-llama/llama-stack/tmp/configs/ollama-run.yaml`
```

You can also export the build artifacts to a specified directory instead of building the container directly. This is useful when you want to:
- Build the container in a different environment
- Share the build configuration with others
- Customize the build process

To export the build artifacts, use the `--export-dir` flag:

```
llama stack build --config my-build.yaml --image-type container --export-dir ./my-build
```

This will create a tarball in the specified directory containing:
- The Dockerfile (named Containerfile)
- The run configuration file (if building from a config)
- Any external provider files (if specified in the config)

The tarball will be named with a timestamp to ensure uniqueness, for example: `<distro-name>_<timestamp>.tar.gz`

After this step is successful, you should be able to find the built container image and test it with `llama stack run <path/to/run.yaml>`.
:::

Expand Down
3 changes: 3 additions & 0 deletions llama_stack/cli/stack/_build.py
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
image_name=image_name,
config_path=args.config,
template_name=args.template,
export_dir=args.export_dir,
)

except (Exception, RuntimeError) as exc:
Expand Down Expand Up @@ -341,6 +342,7 @@ def _run_stack_build_command_from_build_config(
image_name: str | None = None,
template_name: str | None = None,
config_path: str | None = None,
export_dir: str | None = None,
) -> str:
image_name = image_name or build_config.image_name
if build_config.image_type == LlamaStackImageType.CONTAINER.value:
Expand Down Expand Up @@ -383,6 +385,7 @@ def _run_stack_build_command_from_build_config(
image_name,
template_or_config=template_name or config_path or str(build_file_path),
run_config=run_config_file,
export_dir=export_dir,
)
if return_code != 0:
raise RuntimeError(f"Failed to build image {image_name}")
Expand Down
8 changes: 8 additions & 0 deletions llama_stack/cli/stack/build.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
# the root directory of this source tree.
import argparse
import textwrap
from pathlib import Path

from llama_stack.cli.stack.utils import ImageType
from llama_stack.cli.subcommand import Subcommand
Expand Down Expand Up @@ -82,6 +83,13 @@ def _add_arguments(self):
help="Build a config for a list of providers and only those providers. This list is formatted like: api1=provider1,api2=provider2. Where there can be multiple providers per API.",
)

self.parser.add_argument(
"--export-dir",
type=Path,
default=None,
help="Export the build artifacts to a specified directory instead of building the container. This will create a directory containing the Dockerfile and all necessary files to build the container.",
)

def _run_stack_build_command(self, args: argparse.Namespace) -> None:
# always keep implementation completely silo-ed away from CLI so CLI
# can be fast to load and reduces dependencies
Expand Down
3 changes: 3 additions & 0 deletions llama_stack/distribution/build.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ def build_image(
image_name: str,
template_or_config: str,
run_config: str | None = None,
export_dir: str | None = None,
):
container_base = build_config.distribution_spec.container_image or "python:3.10-slim"

Expand All @@ -108,6 +109,8 @@ def build_image(
container_base,
" ".join(normal_deps),
]
if export_dir is not None:
args.append(export_dir)

# When building from a config file (not a template), include the run config path in the
# build arguments
Expand Down
89 changes: 68 additions & 21 deletions llama_stack/distribution/build_container.sh
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ BUILD_CONTEXT_DIR=$(pwd)

if [ "$#" -lt 4 ]; then
# This only works for templates
echo "Usage: $0 <template_or_config> <image_name> <container_base> <pip_dependencies> [<run_config>] [<special_pip_deps>]" >&2
echo "Usage: $0 <template_or_config> <image_name> <container_base> <pip_dependencies> [<run_config>] [<special_pip_deps>] [<export_dir>]" >&2
exit 1
fi
set -euo pipefail
Expand All @@ -43,23 +43,26 @@ shift
# Handle optional arguments
run_config=""
special_pip_deps=""

# Check if there are more arguments
# The logics is becoming cumbersom, we should refactor it if we can do better
if [ $# -gt 0 ]; then
# Check if the argument ends with .yaml
if [[ "$1" == *.yaml ]]; then
run_config="$1"
shift
# If there's another argument after .yaml, it must be special_pip_deps
if [ $# -gt 0 ]; then
special_pip_deps="$1"
fi
else
# If it's not .yaml, it must be special_pip_deps
special_pip_deps="$1"
fi
fi
export_dir=""

# Process remaining arguments
while [[ $# -gt 0 ]]; do
case "$1" in
*.yaml)
run_config="$1"
shift
;;
*)
# Check if argument is a valid directory path (export_dir) or special_pip_deps
if [[ "$1" == *" "* ]] || ! mkdir -p "$1" 2>/dev/null; then
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hm what is this? this looks very dangerous. why are we doing mkdir -p for checking?!

special_pip_deps="$1"
else
export_dir="$1"
fi
shift
;;
esac
done

# Define color codes
RED='\033[0;31m'
Expand All @@ -83,8 +86,8 @@ add_to_container() {
fi
}

# Check if container command is available
if ! is_command_available $CONTAINER_BINARY; then
# Check if container command is available only if not running in export mode
if ! is_command_available $CONTAINER_BINARY && [ -z "$export_dir" ]; then
printf "${RED}Error: ${CONTAINER_BINARY} command not found. Is ${CONTAINER_BINARY} installed and in your PATH?${NC}" >&2
exit 1
fi
Expand All @@ -96,7 +99,7 @@ FROM $container_base
WORKDIR /app

# We install the Python 3.11 dev headers and build tools so that any
# Cextension wheels (e.g. polyleven, faisscpu) can compile successfully.
# C-extension wheels (e.g. polyleven, faiss-cpu) can compile successfully.

RUN dnf -y update && dnf install -y iputils git net-tools wget \
vim-minimal python3.11 python3.11-pip python3.11-wheel \
Expand Down Expand Up @@ -270,6 +273,50 @@ printf "Containerfile created successfully in %s/Containerfile\n\n" "$TEMP_DIR"
cat "$TEMP_DIR"/Containerfile
printf "\n"

# If export_dir is specified, copy all necessary files and exit
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this script is getting very large. I think at the very least we can make this section into a function perhaps?

if [ -n "$export_dir" ]; then
mkdir -p "$export_dir"
timestamp=$(date '+%Y-%m-%d_%H-%M-%S')
tar_name="${image_name//[^a-zA-Z0-9]/_}_${timestamp}.tar.gz"

# If a run config is provided, copy it to the export directory otherwise it's a template build and
# we don't need to copy anything
if [ -n "$run_config" ]; then
mv "$run_config" "$TEMP_DIR"/run.yaml
fi

# Create the archive with all files
echo "Creating tarball with the following files:"
echo "- Containerfile"
[ -n "$run_config" ] && echo "- run.yaml"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nice one liners

edit: although the conditionals get repeated. can you add the echo as you are building the tar_cmd?

[ -n "$external_providers_dir" ] && [ -d "$external_providers_dir" ] && echo "- providers.d directory"

# Capture both stdout and stderr from tar command
tar_cmd="tar -czf \"$export_dir/$tar_name\" -C \"$TEMP_DIR\" Containerfile"

if [ -n "$run_config" ]; then
tar_cmd="$tar_cmd -C \"$BUILD_CONTEXT_DIR\" \"$(basename run.yaml)\""
fi

if [ -n "$external_providers_dir" ] && [ -d "$external_providers_dir" ]; then
tar_cmd="$tar_cmd -C \"$BUILD_CONTEXT_DIR\" providers.d"
fi

tar_output=$(eval "$tar_cmd" 2>&1)
tar_status=$?

if [ $tar_status -ne 0 ]; then
echo "ERROR: Failed to create tarball" >&2
echo "Tar command output:" >&2
echo "$tar_output" >&2
exit 1
fi
rm -rf providers.d run.yaml

echo "Build artifacts tarball created: $export_dir/$tar_name"
exit 0
fi

# Start building the CLI arguments
CLI_ARGS=()

Expand Down
3 changes: 2 additions & 1 deletion tests/unit/distribution/test_build_path.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,10 @@
def test_container_build_passes_path(monkeypatch, tmp_path):
called_with = {}

def spy_build_image(cfg, build_file_path, image_name, template_or_config, run_config=None):
def spy_build_image(cfg, build_file_path, image_name, template_or_config, run_config=None, export_dir=None):
called_with["path"] = template_or_config
called_with["run_config"] = run_config
called_with["export_dir"] = export_dir
return 0

monkeypatch.setattr(
Expand Down
Loading