From fd3b65b7180a40a715c9701b37adb6a62a63de6a Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Wed, 28 Aug 2024 10:07:08 -0700 Subject: [PATCH] llama distribution -> llama stack + containers (WIP) --- llama_toolchain/cli/api/build.py | 27 ++-- llama_toolchain/cli/api/configure.py | 2 +- llama_toolchain/cli/api/start.py | 2 +- llama_toolchain/cli/llama.py | 4 +- .../cli/{distribution => stack}/__init__.py | 2 +- .../cli/{distribution => stack}/configure.py | 10 +- .../cli/{distribution => stack}/create.py | 4 +- .../cli/{distribution => stack}/install.py | 16 +-- .../cli/{distribution => stack}/list.py | 2 +- .../distribution.py => stack/stack.py} | 22 ++-- .../cli/{distribution => stack}/start.py | 2 +- .../distribution/build_conda_env.sh | 22 ++-- .../distribution/build_container.sh | 117 ++++++++++++++---- llama_toolchain/distribution/datatypes.py | 4 + .../distribution/start_conda_env.sh | 15 ++- .../distribution/start_container.sh | 33 ++++- 16 files changed, 204 insertions(+), 80 deletions(-) rename llama_toolchain/cli/{distribution => stack}/__init__.py (79%) rename llama_toolchain/cli/{distribution => stack}/configure.py (92%) rename llama_toolchain/cli/{distribution => stack}/create.py (92%) rename llama_toolchain/cli/{distribution => stack}/install.py (90%) rename llama_toolchain/cli/{distribution => stack}/list.py (97%) rename llama_toolchain/cli/{distribution/distribution.py => stack/stack.py} (58%) rename llama_toolchain/cli/{distribution => stack}/start.py (98%) diff --git a/llama_toolchain/cli/api/build.py b/llama_toolchain/cli/api/build.py index ae180fc20..248796ac6 100644 --- a/llama_toolchain/cli/api/build.py +++ b/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)) diff --git a/llama_toolchain/cli/api/configure.py b/llama_toolchain/cli/api/configure.py index 3fb00383d..fe68a2a12 100644 --- a/llama_toolchain/cli/api/configure.py +++ b/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( diff --git a/llama_toolchain/cli/api/start.py b/llama_toolchain/cli/api/start.py index 59b4d7dd4..e7dcfca18 100644 --- a/llama_toolchain/cli/api/start.py +++ b/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") diff --git a/llama_toolchain/cli/llama.py b/llama_toolchain/cli/llama.py index 7b6302f7e..737745ba5 100644 --- a/llama_toolchain/cli/llama.py +++ b/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: diff --git a/llama_toolchain/cli/distribution/__init__.py b/llama_toolchain/cli/stack/__init__.py similarity index 79% rename from llama_toolchain/cli/distribution/__init__.py rename to llama_toolchain/cli/stack/__init__.py index 81278f253..fd2cbcdff 100644 --- a/llama_toolchain/cli/distribution/__init__.py +++ b/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 diff --git a/llama_toolchain/cli/distribution/configure.py b/llama_toolchain/cli/stack/configure.py similarity index 92% rename from llama_toolchain/cli/distribution/configure.py rename to llama_toolchain/cli/stack/configure.py index 94764fce3..0e311feaa 100644 --- a/llama_toolchain/cli/distribution/configure.py +++ b/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 diff --git a/llama_toolchain/cli/distribution/create.py b/llama_toolchain/cli/stack/create.py similarity index 92% rename from llama_toolchain/cli/distribution/create.py rename to llama_toolchain/cli/stack/create.py index f4b6d3f20..037681355 100644 --- a/llama_toolchain/cli/distribution/create.py +++ b/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() diff --git a/llama_toolchain/cli/distribution/install.py b/llama_toolchain/cli/stack/install.py similarity index 90% rename from llama_toolchain/cli/distribution/install.py rename to llama_toolchain/cli/stack/install.py index cd75effc3..9e0a12de9 100644 --- a/llama_toolchain/cli/distribution/install.py +++ b/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", ) diff --git a/llama_toolchain/cli/distribution/list.py b/llama_toolchain/cli/stack/list.py similarity index 97% rename from llama_toolchain/cli/distribution/list.py rename to llama_toolchain/cli/stack/list.py index e214490ef..217d5ae03 100644 --- a/llama_toolchain/cli/distribution/list.py +++ b/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( diff --git a/llama_toolchain/cli/distribution/distribution.py b/llama_toolchain/cli/stack/stack.py similarity index 58% rename from llama_toolchain/cli/distribution/distribution.py rename to llama_toolchain/cli/stack/stack.py index 641f360e9..0584f62da 100644 --- a/llama_toolchain/cli/distribution/distribution.py +++ b/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) diff --git a/llama_toolchain/cli/distribution/start.py b/llama_toolchain/cli/stack/start.py similarity index 98% rename from llama_toolchain/cli/distribution/start.py rename to llama_toolchain/cli/stack/start.py index b854c79dc..9f35e5aa3 100644 --- a/llama_toolchain/cli/distribution/start.py +++ b/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( diff --git a/llama_toolchain/distribution/build_conda_env.sh b/llama_toolchain/distribution/build_conda_env.sh index e5f0e4a65..38d10a525 100755 --- a/llama_toolchain/distribution/build_conda_env.sh +++ b/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" diff --git a/llama_toolchain/distribution/build_container.sh b/llama_toolchain/distribution/build_container.sh index 5c4a7ccf3..95551abac 100755 --- a/llama_toolchain/distribution/build_container.sh +++ b/llama_toolchain/distribution/build_container.sh @@ -1,44 +1,113 @@ #!/bin/bash -if [ "$#" -ne 3 ]; then - echo "Usage: $0 - 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] + 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 <"$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 <&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 < Optional[str]: + return None + @property def module(self) -> str: if self.adapter: diff --git a/llama_toolchain/distribution/start_conda_env.sh b/llama_toolchain/distribution/start_conda_env.sh index 2d2bd57f7..2ab10cda4 100755 --- a/llama_toolchain/distribution/start_conda_env.sh +++ b/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 " +if [ $# -lt 3 ]; then + echo "Usage: $0 " 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" "$@" diff --git a/llama_toolchain/distribution/start_container.sh b/llama_toolchain/distribution/start_container.sh index d6e003d25..95ccf100d 100755 --- a/llama_toolchain/distribution/start_container.sh +++ b/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 " + 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 "$@"