From 14e60e3c02b4673f4b67bbfefaeb4be93a324f10 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Han?= Date: Thu, 24 Apr 2025 11:29:53 +0200 Subject: [PATCH] feat: include run.yaml in the container image (#2005) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit As part of the build process, we now include the generated run.yaml (based of the provided build configuration file) into the container. We updated the entrypoint to use this run configuration as well. Given this simple distribution configuration: ``` # build.yaml version: '2' distribution_spec: description: Use (an external) Ollama server for running LLM inference providers: inference: - remote::ollama vector_io: - inline::faiss safety: - inline::llama-guard agents: - inline::meta-reference telemetry: - inline::meta-reference eval: - inline::meta-reference datasetio: - remote::huggingface - inline::localfs scoring: - inline::basic - inline::llm-as-judge - inline::braintrust tool_runtime: - remote::brave-search - remote::tavily-search - inline::code-interpreter - inline::rag-runtime - remote::model-context-protocol - remote::wolfram-alpha container_image: "registry.access.redhat.com/ubi9" image_type: container image_name: test ``` Build it: ``` llama stack build --config build.yaml ``` Run it: ``` podman run --rm \ -p 8321:8321 \ -e OLLAMA_URL=http://host.containers.internal:11434 \ --name llama-stack-server \ localhost/leseb-test:0.2.2 ``` Signed-off-by: Sébastien Han --- .github/workflows/providers-build.yml | 38 +++++++++ llama_stack/cli/stack/_build.py | 22 ++++-- llama_stack/distribution/build.py | 6 ++ llama_stack/distribution/build_container.sh | 86 ++++++++++++++++++--- tests/unit/distribution/test_build_path.py | 4 +- 5 files changed, 139 insertions(+), 17 deletions(-) diff --git a/.github/workflows/providers-build.yml b/.github/workflows/providers-build.yml index 117c8b6d2..23257d7dc 100644 --- a/.github/workflows/providers-build.yml +++ b/.github/workflows/providers-build.yml @@ -107,3 +107,41 @@ jobs: - name: Build a single provider run: | USE_COPY_NOT_MOUNT=true LLAMA_STACK_DIR=. uv run llama stack build --image-type venv --image-name test --providers inference=remote::ollama + + build-custom-container-distribution: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + + - name: Set up Python + uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55 # v5.5.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: Build a single provider + run: | + yq -i '.image_type = "container"' llama_stack/templates/dev/build.yaml + yq -i '.image_name = "test"' llama_stack/templates/dev/build.yaml + USE_COPY_NOT_MOUNT=true LLAMA_STACK_DIR=. uv run llama stack build --config llama_stack/templates/dev/build.yaml + + - name: Inspect the container image entrypoint + run: | + IMAGE_ID=$(docker images --format "{{.Repository}}:{{.Tag}}" | head -n 1) + entrypoint=$(docker inspect --format '{{ .Config.Entrypoint }}' $IMAGE_ID) + echo "Entrypoint: $entrypoint" + if [ "$entrypoint" != "[python -m llama_stack.distribution.server.server --config /app/run.yaml]" ]; then + echo "Entrypoint is not correct" + exit 1 + fi diff --git a/llama_stack/cli/stack/_build.py b/llama_stack/cli/stack/_build.py index 26c09af4e..80ab0631b 100644 --- a/llama_stack/cli/stack/_build.py +++ b/llama_stack/cli/stack/_build.py @@ -317,11 +317,15 @@ def _generate_run_config( to_write = json.loads(run_config.model_dump_json()) f.write(yaml.dump(to_write, sort_keys=False)) - # this path is only invoked when no template is provided - cprint( - f"You can now run your stack with `llama stack run {run_config_file}`", - color="green", - ) + # Only print this message for non-container builds since it will be displayed before the + # container is built + # For non-container builds, the run.yaml is generated at the very end of the build process so it + # makes sense to display this message + if build_config.image_type != LlamaStackImageType.CONTAINER.value: + cprint( + f"You can now run your stack with `llama stack run {run_config_file}`", + color="green", + ) return run_config_file @@ -355,6 +359,13 @@ def _run_stack_build_command_from_build_config( build_file_path = build_dir / f"{image_name}-build.yaml" os.makedirs(build_dir, exist_ok=True) + run_config_file = None + # Generate the run.yaml so it can be included in the container image with the proper entrypoint + # Only do this if we're building a container image and we're not using a template + if build_config.image_type == LlamaStackImageType.CONTAINER.value and not template_name and config_path: + cprint("Generating run.yaml file", color="green") + run_config_file = _generate_run_config(build_config, build_dir, image_name) + with open(build_file_path, "w") as f: to_write = json.loads(build_config.model_dump_json()) f.write(yaml.dump(to_write, sort_keys=False)) @@ -364,6 +375,7 @@ def _run_stack_build_command_from_build_config( build_file_path, image_name, template_or_config=template_name or config_path or str(build_file_path), + run_config=run_config_file, ) if return_code != 0: raise RuntimeError(f"Failed to build image {image_name}") diff --git a/llama_stack/distribution/build.py b/llama_stack/distribution/build.py index 5b61ae081..9664449f3 100644 --- a/llama_stack/distribution/build.py +++ b/llama_stack/distribution/build.py @@ -93,6 +93,7 @@ def build_image( build_file_path: Path, image_name: str, template_or_config: str, + run_config: str | None = None, ): container_base = build_config.distribution_spec.container_image or "python:3.10-slim" @@ -108,6 +109,11 @@ def build_image( container_base, " ".join(normal_deps), ] + + # When building from a config file (not a template), include the run config path in the + # build arguments + if run_config is not None: + args.append(run_config) elif build_config.image_type == LlamaStackImageType.CONDA.value: script = str(importlib.resources.files("llama_stack") / "distribution/build_conda_env.sh") args = [ diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh index fb4780432..ad316d45e 100755 --- a/llama_stack/distribution/build_container.sh +++ b/llama_stack/distribution/build_container.sh @@ -19,12 +19,16 @@ UV_HTTP_TIMEOUT=${UV_HTTP_TIMEOUT:-500} # mounting is not supported by docker buildx, so we use COPY instead USE_COPY_NOT_MOUNT=${USE_COPY_NOT_MOUNT:-} +# Path to the run.yaml file in the container +RUN_CONFIG_PATH=/app/run.yaml + +BUILD_CONTEXT_DIR=$(pwd) + if [ "$#" -lt 4 ]; then # This only works for templates - echo "Usage: $0 []" >&2 + echo "Usage: $0 [] []" >&2 exit 1 fi - set -euo pipefail template_or_config="$1" @@ -35,8 +39,27 @@ container_base="$1" shift pip_dependencies="$1" shift -special_pip_deps="${1:-}" +# 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 # Define color codes RED='\033[0;31m' @@ -75,7 +98,7 @@ WORKDIR /app # We install the Python 3.11 dev headers and build tools so that any # C‑extension wheels (e.g. polyleven, faiss‑cpu) can compile successfully. -RUN dnf -y update && dnf install -y iputils net-tools wget \ +RUN dnf -y update && dnf install -y iputils git net-tools wget \ vim-minimal python3.11 python3.11-pip python3.11-wheel \ python3.11-setuptools python3.11-devel gcc make && \ ln -s /bin/pip3.11 /bin/pip && ln -s /bin/python3.11 /bin/python && dnf clean all @@ -119,6 +142,45 @@ EOF done fi +# Function to get Python command +get_python_cmd() { + if is_command_available python; then + echo "python" + elif is_command_available python3; then + echo "python3" + else + echo "Error: Neither python nor python3 is installed. Please install Python to continue." >&2 + exit 1 + fi +} + +if [ -n "$run_config" ]; then + # Copy the run config to the build context since it's an absolute path + cp "$run_config" "$BUILD_CONTEXT_DIR/run.yaml" + add_to_container << EOF +COPY run.yaml $RUN_CONFIG_PATH +EOF + + # Parse the run.yaml configuration to identify external provider directories + # If external providers are specified, copy their directory to the container + # and update the configuration to reference the new container path + python_cmd=$(get_python_cmd) + external_providers_dir=$($python_cmd -c "import yaml; config = yaml.safe_load(open('$run_config')); print(config.get('external_providers_dir') or '')") + if [ -n "$external_providers_dir" ]; then + echo "Copying external providers directory: $external_providers_dir" + add_to_container << EOF +COPY $external_providers_dir /app/providers.d +EOF + # Edit the run.yaml file to change the external_providers_dir to /app/providers.d + if [ "$(uname)" = "Darwin" ]; then + sed -i.bak -e 's|external_providers_dir:.*|external_providers_dir: /app/providers.d|' "$BUILD_CONTEXT_DIR/run.yaml" + rm -f "$BUILD_CONTEXT_DIR/run.yaml.bak" + else + sed -i 's|external_providers_dir:.*|external_providers_dir: /app/providers.d|' "$BUILD_CONTEXT_DIR/run.yaml" + fi + fi +fi + stack_mount="/app/llama-stack-source" client_mount="/app/llama-stack-client-source" @@ -178,15 +240,16 @@ fi RUN pip uninstall -y uv EOF -# if template_or_config ends with .yaml, it is not a template and we should not use the --template flag -if [[ "$template_or_config" != *.yaml ]]; then +# If a run config is provided, we use the --config flag +if [[ -n "$run_config" ]]; then + add_to_container << EOF +ENTRYPOINT ["python", "-m", "llama_stack.distribution.server.server", "--config", "$RUN_CONFIG_PATH"] +EOF +# If a template is provided (not a yaml file), we use the --template flag +elif [[ "$template_or_config" != *.yaml ]]; then add_to_container << EOF ENTRYPOINT ["python", "-m", "llama_stack.distribution.server.server", "--template", "$template_or_config"] EOF -else - add_to_container << EOF -ENTRYPOINT ["python", "-m", "llama_stack.distribution.server.server"] -EOF fi # Add other require item commands genearic to all containers @@ -258,9 +321,10 @@ $CONTAINER_BINARY build \ "${CLI_ARGS[@]}" \ -t "$image_tag" \ -f "$TEMP_DIR/Containerfile" \ - "." + "$BUILD_CONTEXT_DIR" # clean up tmp/configs +rm -f "$BUILD_CONTEXT_DIR/run.yaml" set +x echo "Success!" diff --git a/tests/unit/distribution/test_build_path.py b/tests/unit/distribution/test_build_path.py index a913bd88b..555cdda4a 100644 --- a/tests/unit/distribution/test_build_path.py +++ b/tests/unit/distribution/test_build_path.py @@ -16,8 +16,9 @@ from llama_stack.distribution.utils.image_types import LlamaStackImageType def test_container_build_passes_path(monkeypatch, tmp_path): called_with = {} - def spy_build_image(cfg, build_file_path, image_name, template_or_config): + def spy_build_image(cfg, build_file_path, image_name, template_or_config, run_config=None): called_with["path"] = template_or_config + called_with["run_config"] = run_config return 0 monkeypatch.setattr( @@ -36,3 +37,4 @@ def test_container_build_passes_path(monkeypatch, tmp_path): assert "path" in called_with assert isinstance(called_with["path"], str) assert Path(called_with["path"]).exists() + assert called_with["run_config"] is None