llama distribution -> llama stack + containers (WIP)

llama_toolchain/cli/api/build.py

@@ -47,7 +47,7 @@ def get_dependencies(
 
     return Dependencies(
         docker_image=provider.docker_image,
-        pip_packages=pip_packages + SERVER_DEPENDENCIES
+        pip_packages=pip_packages + SERVER_DEPENDENCIES,
     )
 
 
@@ -161,18 +161,29 @@ class ApiBuild(Subcommand):
             },
             **provider_deps,
         }
-        with open(package_file, "w") as f:
+
+        # properly handle the case where package exists but has
+        # inconsistent configuration for the providers. if possible,
+        # we don't want to overwrite the existing configuration.
+        if package_file.exists():
+            cprint(
+                f"Build `{package_name}` exists; will reconfigure",
+                color="yellow",
+            )
+            c = PackageConfig(**yaml.safe_load(package_file.read_text()))
+        else:
             c = PackageConfig(
                 built_at=datetime.now(),
                 package_name=package_name,
-                docker_image=(
-                    package_name if args.type == BuildType.container.value else None
-                ),
-                conda_env=(
-                    package_name if args.type == BuildType.conda_env.value else None
-                ),
                 providers=stub_config,
             )
+
+        c.docker_image = (
+            package_name if args.type == BuildType.container.value else None
+        )
+        c.conda_env = package_name if args.type == BuildType.conda_env.value else None
+
+        with open(package_file, "w") as f:
             to_write = json.loads(json.dumps(c.dict(), cls=EnumEncoder))
             f.write(yaml.dump(to_write, sort_keys=False))
 

llama_toolchain/cli/api/configure.py

@@ -89,7 +89,7 @@ def configure_llama_provider(config_file: Path) -> None:
 
         try:
             existing_provider_config = config_type(**stub_config)
-        except KeyError:
+        except Exception:
             existing_provider_config = None
 
         provider_config = prompt_for_config(

llama_toolchain/cli/api/start.py

@@ -75,7 +75,7 @@ class ApiStart(Subcommand):
                 config.conda_env,
             ]
 
-        run_args.extend(["--yaml_config", str(config_file), "--port", str(args.port)])
+        run_args.extend([str(config_file), str(args.port)])
         if args.disable_ipv6:
             run_args.append("--disable-ipv6")
 

llama_toolchain/cli/llama.py

@@ -7,9 +7,9 @@
 import argparse
 
 from .api import ApiParser
-from .distribution import DistributionParser
 from .download import Download
 from .model import ModelParser
+from .stack import StackParser
 
 
 class LlamaCLIParser:
@@ -30,8 +30,8 @@ class LlamaCLIParser:
         # Add sub-commands
         Download.create(subparsers)
         ModelParser.create(subparsers)
-        DistributionParser.create(subparsers)
         ApiParser.create(subparsers)
+        StackParser.create(subparsers)
 
         # Import sub-commands from agentic_system if they exist
         try:

llama_toolchain/cli/stack/__init__.py

@@ -4,4 +4,4 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from .distribution import DistributionParser  # noqa
+from .stack import StackParser  # noqa

llama_toolchain/cli/stack/configure.py

@@ -9,13 +9,13 @@ import json
 import shlex
 
 import yaml
+from termcolor import cprint
 
 from llama_toolchain.cli.subcommand import Subcommand
 from llama_toolchain.common.config_dirs import DISTRIBS_BASE_DIR
-from termcolor import cprint
 
 
-class DistributionConfigure(Subcommand):
+class StackConfigure(Subcommand):
     """Llama cli for configuring llama toolchain configs"""
 
     def __init__(self, subparsers: argparse._SubParsersAction):
@@ -38,7 +38,7 @@ class DistributionConfigure(Subcommand):
         )
 
     def _run_distribution_configure_cmd(self, args: argparse.Namespace) -> None:
-        from llama_toolchain.distribution.datatypes import DistributionConfig
+        from llama_toolchain.distribution.datatypes import StackConfig
         from llama_toolchain.distribution.registry import resolve_distribution_spec
 
         config_file = DISTRIBS_BASE_DIR / args.name / "config.yaml"
@@ -50,7 +50,7 @@ class DistributionConfigure(Subcommand):
 
         # we need to find the spec from the name
         with open(config_file, "r") as f:
-            config = DistributionConfig(**yaml.safe_load(f))
+            config = StackConfig(**yaml.safe_load(f))
 
         dist = resolve_distribution_spec(config.spec)
         if dist is None:
@@ -59,7 +59,7 @@ class DistributionConfigure(Subcommand):
         configure_llama_distribution(dist, config)
 
 
-def configure_llama_distribution(dist: "Distribution", config: "DistributionConfig"):
+def configure_llama_distribution(dist: "Stack", config: "StackConfig"):
     from llama_toolchain.common.exec import run_command
     from llama_toolchain.common.prompt_for_config import prompt_for_config
     from llama_toolchain.common.serialize import EnumEncoder

llama_toolchain/cli/stack/create.py

@@ -9,7 +9,7 @@ import argparse
 from llama_toolchain.cli.subcommand import Subcommand
 
 
-class DistributionCreate(Subcommand):
+class StackCreate(Subcommand):
     def __init__(self, subparsers: argparse._SubParsersAction):
         super().__init__()
         self.parser = subparsers.add_parser(
@@ -37,7 +37,7 @@ class DistributionCreate(Subcommand):
 
         dist = resolve_distribution_spec(args.name)
         if dist is not None:
-            self.parser.error(f"Distribution with name {args.name} already exists")
+            self.parser.error(f"Stack with name {args.name} already exists")
             return
 
         raise NotImplementedError()

llama_toolchain/cli/stack/install.py

@@ -10,13 +10,13 @@ import os
 import pkg_resources
 import yaml
 
+from termcolor import cprint
+
 from llama_toolchain.cli.subcommand import Subcommand
 from llama_toolchain.common.config_dirs import DISTRIBS_BASE_DIR
 
-from termcolor import cprint
 
-
-class DistributionInstall(Subcommand):
+class StackInstall(Subcommand):
     """Llama cli for configuring llama toolchain configs"""
 
     def __init__(self, subparsers: argparse._SubParsersAction):
@@ -36,7 +36,7 @@ class DistributionInstall(Subcommand):
         self.parser.add_argument(
             "--spec",
             type=str,
-            help="Distribution spec to install (try local-ollama)",
+            help="Stack spec to install (try local-ollama)",
             required=True,
             choices=[d.spec_id for d in available_distribution_specs()],
         )
@@ -54,7 +54,7 @@ class DistributionInstall(Subcommand):
 
     def _run_distribution_install_cmd(self, args: argparse.Namespace) -> None:
         from llama_toolchain.common.exec import run_with_pty
-        from llama_toolchain.distribution.datatypes import DistributionConfig
+        from llama_toolchain.distribution.datatypes import StackConfig
         from llama_toolchain.distribution.distribution import distribution_dependencies
         from llama_toolchain.distribution.registry import resolve_distribution_spec
 
@@ -80,7 +80,7 @@ class DistributionInstall(Subcommand):
 
         config_file = distrib_dir / "config.yaml"
         if config_file.exists():
-            c = DistributionConfig(**yaml.safe_load(config_file.read_text()))
+            c = StackConfig(**yaml.safe_load(config_file.read_text()))
             if c.spec != dist.spec_id:
                 self.parser.error(
                     f"already installed distribution with `spec={c.spec}` does not match provided spec `{args.spec}`"
@@ -93,7 +93,7 @@ class DistributionInstall(Subcommand):
                 return
         else:
             with open(config_file, "w") as f:
-                c = DistributionConfig(
+                c = StackConfig(
                     spec=dist.spec_id,
                     name=args.name,
                     conda_env=conda_env,
@@ -106,6 +106,6 @@ class DistributionInstall(Subcommand):
             f"Failed to install distribution {dist.spec_id}", color="red"
         )
         cprint(
-            f"Distribution `{args.name}` (with spec {dist.spec_id}) has been installed successfully!",
+            f"Stack `{args.name}` (with spec {dist.spec_id}) has been installed successfully!",
             color="green",
         )

llama_toolchain/cli/stack/list.py

@@ -10,7 +10,7 @@ import json
 from llama_toolchain.cli.subcommand import Subcommand
 
 
-class DistributionList(Subcommand):
+class StackList(Subcommand):
     def __init__(self, subparsers: argparse._SubParsersAction):
         super().__init__()
         self.parser = subparsers.add_parser(

llama_toolchain/cli/stack/stack.py

@@ -8,14 +8,14 @@ import argparse
 
 from llama_toolchain.cli.subcommand import Subcommand
 
-from .configure import DistributionConfigure
-from .create import DistributionCreate
-from .install import DistributionInstall
-from .list import DistributionList
-from .start import DistributionStart
+from .configure import StackConfigure
+from .create import StackCreate
+from .install import StackInstall
+from .list import StackList
+from .start import StackStart
 
 
-class DistributionParser(Subcommand):
+class StackParser(Subcommand):
     def __init__(self, subparsers: argparse._SubParsersAction):
         super().__init__()
         self.parser = subparsers.add_parser(
@@ -27,8 +27,8 @@ class DistributionParser(Subcommand):
         subparsers = self.parser.add_subparsers(title="distribution_subcommands")
 
         # Add sub-commands
-        DistributionList.create(subparsers)
-        DistributionInstall.create(subparsers)
-        DistributionCreate.create(subparsers)
-        DistributionConfigure.create(subparsers)
-        DistributionStart.create(subparsers)
+        StackList.create(subparsers)
+        StackInstall.create(subparsers)
+        StackCreate.create(subparsers)
+        StackConfigure.create(subparsers)
+        StackStart.create(subparsers)

llama_toolchain/cli/stack/start.py

@@ -13,7 +13,7 @@ from llama_toolchain.cli.subcommand import Subcommand
 from llama_toolchain.common.config_dirs import DISTRIBS_BASE_DIR
 
 
-class DistributionStart(Subcommand):
+class StackStart(Subcommand):
     def __init__(self, subparsers: argparse._SubParsersAction):
         super().__init__()
         self.parser = subparsers.add_parser(

llama_toolchain/distribution/build_conda_env.sh

@@ -47,25 +47,25 @@ ensure_conda_env_python310() {
 
   # Check if conda command is available
   if ! command -v conda &>/dev/null; then
-    echo -e "${RED}Error: conda command not found. Is Conda installed and in your PATH?${NC}" >&2
+    printf "${RED}Error: conda command not found. Is Conda installed and in your PATH?${NC}" >&2
     exit 1
   fi
 
   # Check if the environment exists
   if conda env list | grep -q "^${env_name} "; then
-    echo "Conda environment '${env_name}' exists. Checking Python version..."
+    printf "Conda environment '${env_name}' exists. Checking Python version..."
 
     # Check Python version in the environment
     current_version=$(conda run -n "${env_name}" python --version 2>&1 | cut -d' ' -f2 | cut -d'.' -f1,2)
 
     if [ "$current_version" = "$python_version" ]; then
-      echo "Environment '${env_name}' already has Python ${python_version}. No action needed."
+      printf "Environment '${env_name}' already has Python ${python_version}. No action needed."
     else
-      echo "Updating environment '${env_name}' to Python ${python_version}..."
+      printf "Updating environment '${env_name}' to Python ${python_version}..."
       conda install -n "${env_name}" python="${python_version}" -y
     fi
   else
-    echo "Conda environment '${env_name}' does not exist. Creating with Python ${python_version}..."
+    printf "Conda environment '${env_name}' does not exist. Creating with Python ${python_version}..."
     conda create -n "${env_name}" python="${python_version}" -y
 
     ENVNAME="${env_name}"
@@ -83,11 +83,11 @@ ensure_conda_env_python310() {
     # Re-installing llama-toolchain in the new conda environment
     if [ -n "$LLAMA_TOOLCHAIN_DIR" ]; then
       if [ ! -d "$LLAMA_TOOLCHAIN_DIR" ]; then
-        echo -e "${RED}Warning: LLAMA_TOOLCHAIN_DIR is set but directory does not exist: $LLAMA_TOOLCHAIN_DIR${NC}" >&2
+        printf "${RED}Warning: LLAMA_TOOLCHAIN_DIR is set but directory does not exist: $LLAMA_TOOLCHAIN_DIR${NC}" >&2
         exit 1
       fi
 
-      echo "Installing from LLAMA_TOOLCHAIN_DIR: $LLAMA_TOOLCHAIN_DIR"
+      printf "Installing from LLAMA_TOOLCHAIN_DIR: $LLAMA_TOOLCHAIN_DIR"
       pip install --no-cache-dir -e "$LLAMA_TOOLCHAIN_DIR"
     else
       pip install --no-cache-dir llama-toolchain
@@ -95,18 +95,18 @@ ensure_conda_env_python310() {
 
     if [ -n "$LLAMA_MODELS_DIR" ]; then
       if [ ! -d "$LLAMA_MODELS_DIR" ]; then
-        echo -e "${RED}Warning: LLAMA_MODELS_DIR is set but directory does not exist: $LLAMA_MODELS_DIR${NC}" >&2
+        printf "${RED}Warning: LLAMA_MODELS_DIR is set but directory does not exist: $LLAMA_MODELS_DIR${NC}" >&2
         exit 1
       fi
 
-      echo "Installing from LLAMA_MODELS_DIR: $LLAMA_MODELS_DIR"
+      printf "Installing from LLAMA_MODELS_DIR: $LLAMA_MODELS_DIR"
       pip uninstall -y llama-models
       pip install --no-cache-dir -e "$LLAMA_MODELS_DIR"
     fi
 
     # Install pip dependencies
     if [ -n "$pip_dependencies" ]; then
-      echo "Installing pip dependencies: $pip_dependencies"
+      printf "Installing pip dependencies: $pip_dependencies"
       pip install $pip_dependencies
     fi
   fi
@@ -114,6 +114,6 @@ ensure_conda_env_python310() {
 
 ensure_conda_env_python310 "$env_name" "$pip_dependencies"
 
-echo -e "${GREEN}Successfully setup conda environment. Configuring build...${NC}"
+printf "${GREEN}Successfully setup conda environment. Configuring build...${NC}"
 
 $CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama api configure "$api_or_stack" --name "$env_name"

llama_toolchain/distribution/build_container.sh

@@ -1,44 +1,113 @@
 #!/bin/bash
 
-if [ "$#" -ne 3 ]; then
-  echo "Usage: $0 <image_name> <base_image> <pip_dependencies>
-  echo "Example: $0 my-fastapi-app python:3.9-slim 'fastapi uvicorn'
+LLAMA_MODELS_DIR=${LLAMA_MODELS_DIR:-}
+LLAMA_TOOLCHAIN_DIR=${LLAMA_TOOLCHAIN_DIR:-}
+TEST_PYPI_VERSION=${TEST_PYPI_VERSION:-}
+
+
+if [ "$#" -ne 4 ]; then
+  echo "Usage: $0 [api|stack] <image_name> <docker_base> <pip_dependencies>
+  echo "Example: $0 agentic_system my-fastapi-app python:3.9-slim 'fastapi uvicorn'
   exit 1
 fi
 
-IMAGE_NAME=$1
-BASE_IMAGE=$2
-PIP_DEPENDENCIES=$3
+api_or_stack=$1
+image_name=$2
+docker_base=$3
+pip_dependencies=$4
+
+# Define color codes
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+NC='\033[0m' # No Color
 
 set -euo pipefail
 
-PORT=8001
-
 SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
-SOURCE_DIR=$(dirname $(dirname $(dirname "$SCRIPT_DIR")))
-echo $SOURCE_DIR
+REPO_DIR=$(dirname $(dirname "$SCRIPT_DIR"))
 
 TEMP_DIR=$(mktemp -d)
-echo "Created temporary directory: $TEMP_DIR"
 
-cat <<EOF >"$TEMP_DIR/Dockerfile"
-FROM $BASE_IMAGE
+add_to_docker() {
+  local input
+  output_file="$TEMP_DIR/Dockerfile"
+  if [ -t 0 ]; then
+    printf '%s\n' "$1" >> "$output_file"
+  else
+    # If stdin is not a terminal, read from it (heredoc)
+    cat >> "$output_file"
+  fi
+}
+
+
+add_to_docker <<EOF
+FROM $docker_base
 WORKDIR /app
-COPY llama-stack/llama_toolchain /app/llama_toolchain
-COPY llama-models/models /app/llama_models
 
-RUN pip install $PIP_DEPENDENCIES
-
-EXPOSE $PORT
-ENTRYPOINT ["python3", "-m", "llama_toolchain.distribution.server", "--port", "$PORT"]
+RUN apt-get update && apt-get install -y \
+       iputils-ping net-tools iproute2 dnsutils telnet \
+       curl wget telnet \
+       procps psmisc lsof \
+       traceroute \
+       && rm -rf /var/lib/apt/lists/*
 
 EOF
 
-echo "Dockerfile created successfully in $TEMP_DIR/Dockerfile"
+toolchain_mount="/app/llama-toolchain-source"
+models_mount="/app/llama-models-source"
 
-podman build -t $IMAGE_NAME -f "$TEMP_DIR/Dockerfile" "$SOURCE_DIR"
+if [ -n "$LLAMA_TOOLCHAIN_DIR" ]; then
+  if [ ! -d "$LLAMA_TOOLCHAIN_DIR" ]; then
+    echo "${RED}Warning: LLAMA_TOOLCHAIN_DIR is set but directory does not exist: $LLAMA_TOOLCHAIN_DIR${NC}" >&2
+    exit 1
+  fi
+  add_to_docker "RUN pip install $toolchain_mount"
+else
+  add_to_docker "RUN pip install llama-toolchain"
+fi
 
-echo "Podman image '$IMAGE_NAME' built successfully."
-echo "You can run it with: podman run -p 8000:8000 $IMAGE_NAME"
+if [ -n "$LLAMA_MODELS_DIR" ]; then
+ if [ ! -d "$LLAMA_MODELS_DIR" ]; then
+    echo "${RED}Warning: LLAMA_MODELS_DIR is set but directory does not exist: $LLAMA_MODELS_DIR${NC}" >&2
+    exit 1
+  fi
 
-rm -rf "$TEMP_DIR"
+  add_to_docker <<EOF
+RUN pip uninstall -y llama-models
+RUN pip install $models_mount
+
+EOF
+fi
+
+if [ -n "$pip_dependencies" ]; then
+  add_to_docker "RUN pip install $pip_dependencies"
+fi
+
+add_to_docker <<EOF
+
+# This would be good in production but for debugging flexibility lets not add it right now
+# We need a more solid production ready entrypoint.sh anyway
+#
+# ENTRYPOINT ["python", "-m", "llama_toolchain.distribution.server"]
+
+EOF
+
+printf "Dockerfile created successfully in $TEMP_DIR/Dockerfile"
+cat $TEMP_DIR/Dockerfile
+printf "\n"
+
+mounts=""
+if [ -n "$LLAMA_TOOLCHAIN_DIR" ]; then
+  mounts="$mounts -v $(readlink -f $LLAMA_TOOLCHAIN_DIR):$toolchain_mount"
+fi
+if [ -n "$LLAMA_MODELS_DIR" ]; then
+  mounts="$mounts -v $(readlink -f $LLAMA_MODELS_DIR):$models_mount"
+fi
+set -x
+podman build -t $image_name -f "$TEMP_DIR/Dockerfile" "$REPO_DIR" $mounts
+set +x
+
+printf "${GREEN}Succesfully setup Podman image. Configuring build...${NC}"
+echo "You can run it with: podman run -p 8000:8000 $image_name"
+
+$CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama api configure "$api_or_stack" --name "$image_name"

llama_toolchain/distribution/datatypes.py

@@ -115,6 +115,10 @@ as being "Llama Stack compatible"
 """,
     )
 
+    @property
+    def docker_image(self) -> Optional[str]:
+        return None
+
     @property
     def module(self) -> str:
         if self.adapter:

llama_toolchain/distribution/start_conda_env.sh

@@ -18,15 +18,24 @@ error_handler() {
 
 trap 'error_handler ${LINENO}' ERR
 
-if [ $# -lt 2 ]; then
-  echo "Usage: $0 <environment_name> <script_args...>"
+if [ $# -lt 3 ]; then
+  echo "Usage: $0 <env_name> <yaml_config> <port> <script_args...>"
   exit 1
 fi
 
 env_name="$1"
 shift
 
+yaml_config="$1"
+shift
+
+port="$1"
+shift
+
 eval "$(conda shell.bash hook)"
 conda deactivate && conda activate "$env_name"
 
-$CONDA_PREFIX/bin/python -m llama_toolchain.distribution.server "$@"
+$CONDA_PREFIX/bin/python \
+  -m llama_toolchain.distribution.server \
+  --yaml_config "$yaml_config" \
+  --port "$port" "$@"

llama_toolchain/distribution/start_container.sh

@@ -8,4 +8,35 @@
 
 set -euo pipefail
 
-podman run -it -p 8001:8001 -v ~/.llama/test.yaml:/app/test.yaml test-image --yaml_config /app/test.yaml --disable-ipv6
+RED='\033[0;31m'
+NC='\033[0m' # No Color
+
+error_handler() {
+  echo "Error occurred in script at line: ${1}" >&2
+  exit 1
+}
+
+trap 'error_handler ${LINENO}' ERR
+
+if [ $# -lt 3 ]; then
+  echo "Usage: $0 <docker_image> <yaml_config> <port> <other_args...>"
+  exit 1
+fi
+
+docker_image="$1"
+shift
+
+yaml_config="$1"
+shift
+
+port="$1"
+shift
+
+set -x
+podman run -it \
+  -p $port:$port \
+  -v "$yaml_config:/app/config.yaml" \
+  $docker_image \
+  python -m llama_toolchain.distribution.server \
+  --yaml_config /app/config.yaml \
+  --port $port "$@"