From 4684fd3f8dea660ade15f59d27318b9183fb7e53 Mon Sep 17 00:00:00 2001 From: Charlie Doern Date: Mon, 24 Feb 2025 19:53:31 -0500 Subject: [PATCH] refactor: combine start scripts for each env (#1139) # What does this PR do? now that llama stack supports running in venv, conda, and container modes and the 3 scripts overlap alot, combine these three into ons `start_stack.sh` script ## Test Plan tested this locally on venv, conda, and container --------- Signed-off-by: Charlie Doern Co-authored-by: Ashwin Bharambe Co-authored-by: Yuan Tang --- llama_stack/distribution/start_conda_env.sh | 67 --------- llama_stack/distribution/start_container.sh | 105 -------------- llama_stack/distribution/start_stack.sh | 150 ++++++++++++++++++++ llama_stack/distribution/utils/exec.py | 43 +++--- 4 files changed, 173 insertions(+), 192 deletions(-) delete mode 100755 llama_stack/distribution/start_conda_env.sh delete mode 100755 llama_stack/distribution/start_container.sh create mode 100755 llama_stack/distribution/start_stack.sh diff --git a/llama_stack/distribution/start_conda_env.sh b/llama_stack/distribution/start_conda_env.sh deleted file mode 100755 index fe830059f..000000000 --- a/llama_stack/distribution/start_conda_env.sh +++ /dev/null @@ -1,67 +0,0 @@ -#!/bin/bash - -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the terms described in the LICENSE file in -# the root directory of this source tree. - -set -euo pipefail - -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 - -env_name="$1" -shift - -yaml_config="$1" -shift - -port="$1" -shift - -# Process environment variables from --env arguments -env_vars="" -other_args="" -while [[ $# -gt 0 ]]; do - case "$1" in - --env) - - if [[ -n "$2" ]]; then - # collect environment variables so we can set them after activating the conda env - env_vars="$env_vars --env $2" - shift 2 - else - echo -e "${RED}Error: --env requires a KEY=VALUE argument${NC}" >&2 - exit 1 - fi - ;; - *) - other_args="$other_args $1" - shift - ;; - esac -done - -eval "$(conda shell.bash hook)" -conda deactivate && conda activate "$env_name" - -set -x -$CONDA_PREFIX/bin/python \ - -m llama_stack.distribution.server.server \ - --yaml-config "$yaml_config" \ - --port "$port" \ - $env_vars \ - $other_args diff --git a/llama_stack/distribution/start_container.sh b/llama_stack/distribution/start_container.sh deleted file mode 100755 index a5f543fb4..000000000 --- a/llama_stack/distribution/start_container.sh +++ /dev/null @@ -1,105 +0,0 @@ -#!/bin/bash - -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the terms described in the LICENSE file in -# the root directory of this source tree. - -CONTAINER_BINARY=${CONTAINER_BINARY:-docker} -CONTAINER_OPTS=${CONTAINER_OPTS:-} -LLAMA_CHECKPOINT_DIR=${LLAMA_CHECKPOINT_DIR:-} -LLAMA_STACK_DIR=${LLAMA_STACK_DIR:-} -TEST_PYPI_VERSION=${TEST_PYPI_VERSION:-} -PYPI_VERSION=${PYPI_VERSION:-} - -set -euo pipefail - -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 - -image_name="$1" -container_image="localhost/$image_name" -shift - -yaml_config="$1" -shift - -port="$1" -shift - -# Initialize other_args -other_args="" - -# Process environment variables from --env arguments -env_vars="" - -while [[ $# -gt 0 ]]; do - case "$1" in - --env) - echo "env = $2" - if [[ -n "$2" ]]; then - env_vars="$env_vars -e $2" - shift 2 - else - echo -e "${RED}Error: --env requires a KEY=VALUE argument${NC}" >&2 - exit 1 - fi - ;; - *) - other_args="$other_args $1" - shift - ;; - esac -done - -set -x - -if command -v selinuxenabled &> /dev/null && selinuxenabled; then - # Disable SELinux labels - CONTAINER_OPTS="$CONTAINER_OPTS --security-opt label=disable" -fi - -mounts="" -if [ -n "$LLAMA_STACK_DIR" ]; then - mounts="$mounts -v $(readlink -f $LLAMA_STACK_DIR):/app/llama-stack-source" -fi -if [ -n "$LLAMA_CHECKPOINT_DIR" ]; then - mounts="$mounts -v $LLAMA_CHECKPOINT_DIR:/root/.llama" - CONTAINER_OPTS="$CONTAINER_OPTS --gpus=all" -fi - -if [ -n "$PYPI_VERSION" ]; then - version_tag="$PYPI_VERSION" -elif [ -n "$LLAMA_STACK_DIR" ]; then - version_tag="dev" -elif [ -n "$TEST_PYPI_VERSION" ]; then - version_tag="test-$TEST_PYPI_VERSION" -else - URL="https://pypi.org/pypi/llama-stack/json" - version_tag=$(curl -s $URL | jq -r '.info.version') -fi - -$CONTAINER_BINARY run $CONTAINER_OPTS -it \ - -p $port:$port \ - $env_vars \ - -v "$yaml_config:/app/config.yaml" \ - $mounts \ - --env LLAMA_STACK_PORT=$port \ - --entrypoint python \ - $container_image:$version_tag \ - -m llama_stack.distribution.server.server \ - --yaml-config /app/config.yaml \ - $other_args diff --git a/llama_stack/distribution/start_stack.sh b/llama_stack/distribution/start_stack.sh new file mode 100755 index 000000000..901af1ce0 --- /dev/null +++ b/llama_stack/distribution/start_stack.sh @@ -0,0 +1,150 @@ +#!/usr/bin/env bash + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the terms described in the LICENSE file in +# the root directory of this source tree. + + +CONTAINER_BINARY=${CONTAINER_BINARY:-docker} +CONTAINER_OPTS=${CONTAINER_OPTS:-} +LLAMA_CHECKPOINT_DIR=${LLAMA_CHECKPOINT_DIR:-} +LLAMA_STACK_DIR=${LLAMA_STACK_DIR:-} +TEST_PYPI_VERSION=${TEST_PYPI_VERSION:-} +PYPI_VERSION=${PYPI_VERSION:-} + +set -euo pipefail + +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 + +env_type="$1" +shift + +env_path_or_name="$1" +container_image="localhost/$env_path_or_name" +shift + +yaml_config="$1" +shift + +port="$1" +shift + +SCRIPT_DIR=$(dirname "$(readlink -f "$0")") +source "$SCRIPT_DIR/common.sh" + +# Initialize env_vars as an string +env_vars="" +other_args="" +# Process environment variables from --env arguments +while [[ $# -gt 0 ]]; do + case "$1" in + --env) + + if [[ -n "$2" ]]; then + env_vars="$env_vars --env $2" + shift 2 + else + echo -e "${RED}Error: --env requires a KEY=VALUE argument${NC}" >&2 + exit 1 + fi + ;; + *) + other_args="$other_args $1" + shift + ;; + esac +done + +PYTHON_BINARY="python" +case "$env_type" in + "venv") + # Activate virtual environment + if [ ! -d "$env_path_or_name" ]; then + echo -e "${RED}Error: Virtual environment not found at $env_path_or_name${NC}" >&2 + exit 1 + fi + + if [ ! -f "$env_path_or_name/bin/activate" ]; then + echo -e "${RED}Error: Virtual environment activate binary not found at $env_path_or_name/bin/activate" >&2 + exit 1 + fi + + source "$env_path_or_name/bin/activate" + ;; + "conda") + if ! is_command_available conda; then + echo -e "${RED}Error: conda not found" >&2 + exit 1 + fi + eval "$(conda shell.bash hook)" + conda deactivate && conda activate "$env_path_or_name" + PYTHON_BINARY="$CONDA_PREFIX/bin/python" + ;; + *) +esac + +set -x + +if [[ "$env_type" == "venv" || "$env_type" == "conda" ]]; then + $PYTHON_BINARY -m llama_stack.distribution.server.server \ + --yaml-config "$yaml_config" \ + --port "$port" \ + $env_vars \ + $other_args +elif [[ "$env_type" == "container" ]]; then + if is_command_available selinuxenabled &> /dev/null && selinuxenabled; then + # Disable SELinux labels + CONTAINER_OPTS="$CONTAINER_OPTS --security-opt label=disable" + fi + + mounts="" + if [ -n "$LLAMA_STACK_DIR" ]; then + mounts="$mounts -v $(readlink -f $LLAMA_STACK_DIR):/app/llama-stack-source" + fi + if [ -n "$LLAMA_CHECKPOINT_DIR" ]; then + mounts="$mounts -v $LLAMA_CHECKPOINT_DIR:/root/.llama" + CONTAINER_OPTS="$CONTAINER_OPTS --gpus=all" + fi + + if [ -n "$PYPI_VERSION" ]; then + version_tag="$PYPI_VERSION" + elif [ -n "$LLAMA_STACK_DIR" ]; then + version_tag="dev" + elif [ -n "$TEST_PYPI_VERSION" ]; then + version_tag="test-$TEST_PYPI_VERSION" + else + if ! is_command_available jq; then + echo -e "${RED}Error: jq not found" >&2 + exit 1 + fi + URL="https://pypi.org/pypi/llama-stack/json" + version_tag=$(curl -s $URL | jq -r '.info.version') + fi + + $CONTAINER_BINARY run $CONTAINER_OPTS -it \ + -p $port:$port \ + $env_vars \ + -v "$yaml_config:/app/config.yaml" \ + $mounts \ + --env LLAMA_STACK_PORT=$port \ + --entrypoint python \ + $container_image:$version_tag \ + -m llama_stack.distribution.server.server \ + --yaml-config /app/config.yaml \ + $other_args +fi diff --git a/llama_stack/distribution/utils/exec.py b/llama_stack/distribution/utils/exec.py index 00afdadbe..82bf00e3c 100644 --- a/llama_stack/distribution/utils/exec.py +++ b/llama_stack/distribution/utils/exec.py @@ -24,14 +24,13 @@ from llama_stack.distribution.utils.image_types import ImageType def formulate_run_args(image_type, image_name, config, template_name) -> list: + env_name = "" if image_type == ImageType.container.value or config.container_image: - script = importlib.resources.files("llama_stack") / "distribution/start_container.sh" - image_name = f"distribution-{template_name}" if template_name else config.container_image - run_args = [script, image_name] + env_name = f"distribution-{template_name}" if template_name else config.container_image elif image_type == ImageType.conda.value: current_conda_env = os.environ.get("CONDA_DEFAULT_ENV") - image_name = image_name or current_conda_env - if not image_name: + env_name = image_name or current_conda_env + if not env_name: cprint( "No current conda environment detected, please specify a conda environment name with --image-name", color="red", @@ -51,11 +50,11 @@ def formulate_run_args(image_type, image_name, config, template_name) -> list: return envpath return None - print(f"Using conda environment: {image_name}") - conda_prefix = get_conda_prefix(image_name) + print(f"Using conda environment: {env_name}") + conda_prefix = get_conda_prefix(env_name) if not conda_prefix: cprint( - f"Conda environment {image_name} does not exist.", + f"Conda environment {env_name} does not exist.", color="red", ) return @@ -67,21 +66,25 @@ def formulate_run_args(image_type, image_name, config, template_name) -> list: color="red", ) return - - script = importlib.resources.files("llama_stack") / "distribution/start_conda_env.sh" - run_args = [ - script, - image_name, - ] else: # else must be venv since that is the only valid option left. current_venv = os.environ.get("VIRTUAL_ENV") - venv = image_name or current_venv - script = importlib.resources.files("llama_stack") / "distribution/start_venv.sh" - run_args = [ - script, - venv, - ] + env_name = image_name or current_venv + if not env_name: + cprint( + "No current virtual environment detected, please specify a virtual environment name with --image-name", + color="red", + ) + return + print(f"Using virtual environment: {env_name}") + + script = importlib.resources.files("llama_stack") / "distribution/start_stack.sh" + run_args = [ + script, + image_type, + env_name, + ] + return run_args