Simplify and generalize llama api build yay

This commit is contained in:
Ashwin Bharambe 2024-08-30 14:51:40 -07:00
parent 297d51b183
commit f8517e4688
9 changed files with 103 additions and 151 deletions

View file

@ -20,13 +20,14 @@ fi
set -euo pipefail
if [ "$#" -ne 3 ]; then
echo "Usage: $0 <api_or_stack> <environment_name> <pip_dependencies>" >&2
echo "Example: $0 [api|stack] conda-env 'numpy pandas scipy'" >&2
echo "Usage: $0 <distribution_id> <build_name> <pip_dependencies>" >&2
echo "Example: $0 <distribution_id> mybuild 'numpy pandas scipy'" >&2
exit 1
fi
api_or_stack="$1"
env_name="$2"
distribution_id="$1"
build_name="$2"
env_name="llamastack-$build_name"
pip_dependencies="$3"
# Define color codes
@ -53,19 +54,19 @@ ensure_conda_env_python310() {
# Check if the environment exists
if conda env list | grep -q "^${env_name} "; then
printf "Conda environment '${env_name}' exists. Checking Python version..."
printf "Conda environment '${env_name}' exists. Checking Python version...\n"
# 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
printf "Environment '${env_name}' already has Python ${python_version}. No action needed."
printf "Environment '${env_name}' already has Python ${python_version}. No action needed.\n"
else
printf "Updating environment '${env_name}' to Python ${python_version}..."
printf "Updating environment '${env_name}' to Python ${python_version}...\n"
conda install -n "${env_name}" python="${python_version}" -y
fi
else
printf "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}...\n"
conda create -n "${env_name}" python="${python_version}" -y
ENVNAME="${env_name}"
@ -83,11 +84,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
printf "${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}\n" >&2
exit 1
fi
printf "Installing from LLAMA_TOOLCHAIN_DIR: $LLAMA_TOOLCHAIN_DIR"
printf "Installing from LLAMA_TOOLCHAIN_DIR: $LLAMA_TOOLCHAIN_DIR\n"
pip install --no-cache-dir -e "$LLAMA_TOOLCHAIN_DIR"
else
pip install --no-cache-dir llama-toolchain
@ -95,18 +96,18 @@ ensure_conda_env_python310() {
if [ -n "$LLAMA_MODELS_DIR" ]; then
if [ ! -d "$LLAMA_MODELS_DIR" ]; then
printf "${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}\n" >&2
exit 1
fi
printf "Installing from LLAMA_MODELS_DIR: $LLAMA_MODELS_DIR"
printf "Installing from LLAMA_MODELS_DIR: $LLAMA_MODELS_DIR\n"
pip uninstall -y llama-models
pip install --no-cache-dir -e "$LLAMA_MODELS_DIR"
fi
# Install pip dependencies
if [ -n "$pip_dependencies" ]; then
printf "Installing pip dependencies: $pip_dependencies"
printf "Installing pip dependencies: $pip_dependencies\n"
pip install $pip_dependencies
fi
fi
@ -114,14 +115,14 @@ ensure_conda_env_python310() {
ensure_conda_env_python310 "$env_name" "$pip_dependencies"
printf "${GREEN}Successfully setup conda environment. Configuring build...${NC}"
printf "${GREEN}Successfully setup conda environment. Configuring build...${NC}\n"
if [ "$api_or_stack" = "stack" ]; then
subcommand="stack"
if [ "$distribution_id" = "adhoc" ]; then
subcommand="api"
target=""
else
subcommand="api"
target="$api_or_stack"
subcommand="stack"
target="$distribution_id"
fi
$CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama $subcommand configure $target --build-name "$env_name"
$CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama $subcommand configure $target --build-name "$build_name" --build-type conda_env

View file

@ -4,15 +4,15 @@ 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'
echo "Usage: $0 <distribution_id> <build_name> <docker_base> <pip_dependencies>
echo "Example: $0 distribution_id my-fastapi-app python:3.9-slim 'fastapi uvicorn'
exit 1
fi
api_or_stack=$1
image_name=$2
distribution_id=$1
build_name="$2"
image_name="llamastack-$build_name"
docker_base=$3
pip_dependencies=$4
@ -32,14 +32,13 @@ add_to_docker() {
local input
output_file="$TEMP_DIR/Dockerfile"
if [ -t 0 ]; then
printf '%s\n' "$1" >> "$output_file"
printf '%s\n' "$1" >>"$output_file"
else
# If stdin is not a terminal, read from it (heredoc)
cat >> "$output_file"
cat >>"$output_file"
fi
}
add_to_docker <<EOF
FROM $docker_base
WORKDIR /app
@ -67,7 +66,7 @@ else
fi
if [ -n "$LLAMA_MODELS_DIR" ]; then
if [ ! -d "$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
@ -110,12 +109,12 @@ set +x
printf "${GREEN}Succesfully setup Podman image. Configuring build...${NC}"
echo "You can run it with: podman run -p 8000:8000 $image_name"
if [ "$api_or_stack" = "stack" ]; then
subcommand="stack"
if [ "$distribution_id" = "adhoc" ]; then
subcommand="api"
target=""
else
subcommand="api"
target="$api_or_stack"
subcommand="stack"
target="$distribution_id"
fi
$CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama $subcommand configure $target --build-name "$image_name"
$CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama $subcommand configure $target --build-name "$build_name" --build-type container

View file

@ -8,7 +8,7 @@ import json
import os
from datetime import datetime
from enum import Enum
from typing import Dict, List, Optional
from typing import List, Optional
import pkg_resources
import yaml
@ -17,11 +17,11 @@ from pydantic import BaseModel
from termcolor import cprint
from llama_toolchain.common.config_dirs import BUILDS_BASE_DIR
from llama_toolchain.core.datatypes import * # noqa: F403
from llama_toolchain.common.exec import run_with_pty
from llama_toolchain.common.serialize import EnumEncoder
from llama_toolchain.core.distribution import api_providers
from llama_toolchain.core.datatypes import * # noqa: F403
from llama_toolchain.core.distribution import api_providers, SERVER_DEPENDENCIES
class BuildType(Enum):
@ -29,7 +29,7 @@ class BuildType(Enum):
conda_env = "conda_env"
def descriptor(self) -> str:
return "image" if self == self.container else "env"
return "docker" if self == self.container else "conda"
class Dependencies(BaseModel):
@ -37,29 +37,9 @@ class Dependencies(BaseModel):
docker_image: Optional[str] = None
def get_dependencies(
provider: ProviderSpec, dependencies: Dict[str, ProviderSpec]
) -> Dependencies:
from llama_toolchain.core.distribution import SERVER_DEPENDENCIES
pip_packages = provider.pip_packages
for dep in dependencies.values():
if dep.docker_image:
raise ValueError(
"You can only have the root provider specify a docker image"
)
pip_packages.extend(dep.pip_packages)
return Dependencies(
docker_image=provider.docker_image,
pip_packages=pip_packages + SERVER_DEPENDENCIES,
)
class ApiInput(BaseModel):
api: Api
provider: str
dependencies: Dict[str, ProviderSpec]
def build_package(
@ -69,31 +49,22 @@ def build_package(
distribution_id: Optional[str] = None,
docker_image: Optional[str] = None,
):
is_stack = len(api_inputs) > 1
if is_stack:
if not distribution_id:
raise ValueError(
"You must specify a distribution name when building the Llama Stack"
)
if not distribution_id:
distribution_id = "adhoc"
api1 = api_inputs[0]
provider = distribution_id if is_stack else api1.provider
api_or_stack = "stack" if is_stack else api1.api.value
build_dir = BUILDS_BASE_DIR / api_or_stack
build_dir = BUILDS_BASE_DIR / distribution_id / build_type.descriptor()
os.makedirs(build_dir, exist_ok=True)
package_name = f"{build_type.descriptor()}-{provider}-{name}"
package_name = package_name.replace("::", "-")
package_name = name.replace("::", "-")
package_file = build_dir / f"{package_name}.yaml"
all_providers = api_providers()
package_deps = Dependencies(
docker_image=docker_image or "python:3.10-slim",
pip_packages=[],
pip_packages=SERVER_DEPENDENCIES,
)
stub_config = {}
for api_input in api_inputs:
api = api_input.api
@ -103,18 +74,12 @@ def build_package(
f"Provider `{api_input.provider}` is not available for API `{api}`"
)
deps = get_dependencies(
providers_for_api[api_input.provider],
api_input.dependencies,
)
if deps.docker_image:
provider = providers_for_api[api_input.provider]
package_deps.pip_packages.extend(provider.pip_packages)
if provider.docker_image:
raise ValueError("A stack's dependencies cannot have a docker image")
package_deps.pip_packages.extend(deps.pip_packages)
stub_config[api.value] = {"provider_id": api_input.provider}
for dep_api, dep_spec in api_input.dependencies.items():
if dep_api not in stub_config:
stub_config[dep_api] = {"provider_id": dep_spec.provider_id}
if package_file.exists():
cprint(
@ -154,7 +119,7 @@ def build_package(
)
args = [
script,
api_or_stack,
distribution_id,
package_name,
package_deps.docker_image,
" ".join(package_deps.pip_packages),
@ -165,7 +130,7 @@ def build_package(
)
args = [
script,
api_or_stack,
distribution_id,
package_name,
" ".join(package_deps.pip_packages),
]

View file

@ -19,11 +19,12 @@ error_handler() {
trap 'error_handler ${LINENO}' ERR
if [ $# -lt 3 ]; then
echo "Usage: $0 <env_name> <yaml_config> <port> <script_args...>"
echo "Usage: $0 <build_name> <yaml_config> <port> <script_args...>"
exit 1
fi
env_name="$1"
build_name="$1"
env_name="llamastack-$build_name"
shift
yaml_config="$1"

View file

@ -19,11 +19,12 @@ error_handler() {
trap 'error_handler ${LINENO}' ERR
if [ $# -lt 3 ]; then
echo "Usage: $0 <docker_image> <yaml_config> <port> <other_args...>"
echo "Usage: $0 <build_name> <yaml_config> <port> <other_args...>"
exit 1
fi
docker_image="$1"
build_name="$1"
docker_image="llamastack-$build_name"
shift
yaml_config="$1"