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 <cdoern@redhat.com>
Co-authored-by: Ashwin Bharambe <ashwin.bharambe@gmail.com>
Co-authored-by: Yuan Tang <terrytangyuan@gmail.com>

llama_stack/distribution/start_conda_env.sh

@@ -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 <build_name> <yaml_config> <port> <script_args...>"
-  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

llama_stack/distribution/start_container.sh

@@ -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 <build_name> <yaml_config> <port> <other_args...>"
-  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

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 <env_type> <env_path_or_name> <yaml_config> <port> <script_args...>"
+  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

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"
+        env_name = f"distribution-{template_name}" if template_name else config.container_image
-        image_name = f"distribution-{template_name}" if template_name else config.container_image
-        run_args = [script, image_name]
     elif image_type == ImageType.conda.value:
         current_conda_env = os.environ.get("CONDA_DEFAULT_ENV")
-        image_name = image_name or current_conda_env
+        env_name = image_name or current_conda_env
-        if not image_name:
+        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}")
+        print(f"Using conda environment: {env_name}")
-        conda_prefix = get_conda_prefix(image_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
+        env_name = image_name or current_venv
-        script = importlib.resources.files("llama_stack") / "distribution/start_venv.sh"
+        if not env_name:
-        run_args = [
+            cprint(
-            script,
+                "No current virtual environment detected, please specify a virtual environment name with --image-name",
-            venv,
+                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