From bda974e660e79f7e099f927bf28af5a04198884c Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Tue, 24 Sep 2024 14:18:57 -0700 Subject: [PATCH] Make the "all-remote" distribution lightweight in dependencies and size --- llama_stack/distribution/build.py | 17 +++++++++-- llama_stack/distribution/build_conda_env.sh | 31 +++++++++++++++------ llama_stack/distribution/build_container.sh | 19 +++++++++---- llama_stack/providers/registry/memory.py | 16 ++++++++++- 4 files changed, 65 insertions(+), 18 deletions(-) diff --git a/llama_stack/distribution/build.py b/llama_stack/distribution/build.py index e38f1af1a..828311ea8 100644 --- a/llama_stack/distribution/build.py +++ b/llama_stack/distribution/build.py @@ -66,6 +66,14 @@ def build_image(build_config: BuildConfig, build_file_path: Path): if provider_spec.docker_image: raise ValueError("A stack's dependencies cannot have a docker image") + special_deps = [] + deps = [] + for package in package_deps.pip_packages: + if "--no-deps" in package or "--index-url" in package: + special_deps.append(package) + else: + deps.append(package) + if build_config.image_type == ImageType.docker.value: script = pkg_resources.resource_filename( "llama_stack", "distribution/build_container.sh" @@ -75,7 +83,7 @@ def build_image(build_config: BuildConfig, build_file_path: Path): build_config.name, package_deps.docker_image, str(build_file_path), - " ".join(package_deps.pip_packages), + " ".join(deps), ] else: script = pkg_resources.resource_filename( @@ -84,14 +92,17 @@ def build_image(build_config: BuildConfig, build_file_path: Path): args = [ script, build_config.name, - " ".join(package_deps.pip_packages), + " ".join(deps), ] + if special_deps: + args.append("#".join(special_deps)) + return_code = run_with_pty(args) if return_code != 0: cprint( f"Failed to build target {build_config.name} with return code {return_code}", color="red", ) - + return return_code diff --git a/llama_stack/distribution/build_conda_env.sh b/llama_stack/distribution/build_conda_env.sh index abe59d978..65b2a8c0e 100755 --- a/llama_stack/distribution/build_conda_env.sh +++ b/llama_stack/distribution/build_conda_env.sh @@ -17,14 +17,16 @@ if [ -n "$LLAMA_MODELS_DIR" ]; then echo "Using llama-models-dir=$LLAMA_MODELS_DIR" fi -set -euo pipefail - -if [ "$#" -ne 2 ]; then - echo "Usage: $0 " >&2 +if [ "$#" -lt 2 ]; then + echo "Usage: $0 []" >&2 echo "Example: $0 mybuild 'numpy pandas scipy'" >&2 exit 1 fi +special_pip_deps="$3" + +set -euo pipefail + build_name="$1" env_name="llamastack-$build_name" pip_dependencies="$2" @@ -43,6 +45,7 @@ source "$SCRIPT_DIR/common.sh" ensure_conda_env_python310() { local env_name="$1" local pip_dependencies="$2" + local special_pip_deps="$3" local python_version="3.10" # Check if conda command is available @@ -78,7 +81,12 @@ ensure_conda_env_python310() { if [ -n "$TEST_PYPI_VERSION" ]; then # these packages are damaged in test-pypi, so install them first $CONDA_PREFIX/bin/pip install fastapi libcst - $CONDA_PREFIX/bin/pip install --extra-index-url https://test.pypi.org/simple/ llama-models==$TEST_PYPI_VERSION llama-stack==$TEST_PYPI_VERSION $pip_dependencies + $CONDA_PREFIX/bin/pip install --extra-index-url https://test.pypi.org/simple/ \ + llama-models==$TEST_PYPI_VERSION llama-stack==$TEST_PYPI_VERSION \ + $pip_dependencies + if [ -n "$special_pip_deps" ]; then + $CONDA_PREFIX/bin/pip install --no-deps "$special_pip_deps" + fi else # Re-installing llama-stack in the new conda environment if [ -n "$LLAMA_STACK_DIR" ]; then @@ -105,11 +113,16 @@ ensure_conda_env_python310() { fi # Install pip dependencies - if [ -n "$pip_dependencies" ]; then - printf "Installing pip dependencies: $pip_dependencies\n" - $CONDA_PREFIX/bin/pip install $pip_dependencies + printf "Installing pip dependencies\n" + $CONDA_PREFIX/bin/pip install $pip_dependencies + if [ -n "$special_pip_deps" ]; then + IFS='#' read -ra parts <<< "$special_pip_deps" + for part in "${parts[@]}"; do + echo "$part" + $CONDA_PREFIX/bin/pip install $part + done fi fi } -ensure_conda_env_python310 "$env_name" "$pip_dependencies" +ensure_conda_env_python310 "$env_name" "$pip_dependencies" "$special_pip_deps" diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh index 984e66afa..3efef6c97 100755 --- a/llama_stack/distribution/build_container.sh +++ b/llama_stack/distribution/build_container.sh @@ -4,12 +4,16 @@ LLAMA_MODELS_DIR=${LLAMA_MODELS_DIR:-} LLAMA_STACK_DIR=${LLAMA_STACK_DIR:-} TEST_PYPI_VERSION=${TEST_PYPI_VERSION:-} -if [ "$#" -ne 4 ]; then - echo "Usage: $0 - echo "Example: $0 my-fastapi-app python:3.9-slim 'fastapi uvicorn' +if [ "$#" -lt 4 ]; then + echo "Usage: $0 []" >&2 + echo "Example: $0 my-fastapi-app python:3.9-slim 'fastapi uvicorn' " >&2 exit 1 fi +special_pip_deps="$5" + +set -euo pipefail + build_name="$1" image_name="llamastack-$build_name" docker_base=$2 @@ -21,8 +25,6 @@ RED='\033[0;31m' GREEN='\033[0;32m' NC='\033[0m' # No Color -set -euo pipefail - SCRIPT_DIR=$(dirname "$(readlink -f "$0")") REPO_DIR=$(dirname $(dirname "$SCRIPT_DIR")) DOCKER_BINARY=${DOCKER_BINARY:-docker} @@ -85,6 +87,13 @@ if [ -n "$pip_dependencies" ]; then add_to_docker "RUN pip install $pip_dependencies" fi +if [ -n "$special_pip_deps" ]; then + IFS='#' read -ra parts <<< "$special_pip_deps" + for part in "${parts[@]}"; do + add_to_docker "RUN pip install $part" + done +fi + add_to_docker <