From 025f6158684bf647a94213cb76a5d5b3b23735f4 Mon Sep 17 00:00:00 2001 From: Charlie Doern Date: Wed, 12 Feb 2025 11:13:04 -0500 Subject: [PATCH] feat: add support for running in a venv (#1018) # What does this PR do? add --image-type to `llama stack run`. Which takes conda, container or venv also add start_venv.sh which start the stack using a venv resolves #1007 ## Test Plan running locally: `llama stack build --template ollama --image-type venv` `llama stack run --image-type venv ~/.llama/distributions/ollama/ollama-run.yaml` ... ``` llama stack run --image-type venv ~/.llama/distributions/ollama/ollama-run.yaml Using run configuration: /Users/charliedoern/.llama/distributions/ollama/ollama-run.yaml + python -m llama_stack.distribution.server.server --yaml-config /Users/charliedoern/.llama/distributions/ollama/ollama-run.yaml --port 8321 Using config file: /Users/charliedoern/.llama/distributions/ollama/ollama-run.yaml Run configuration: apis: - agents - datasetio ... ``` Signed-off-by: Charlie Doern --- docs/source/distributions/building_distro.md | 33 +++++++++ llama_stack/cli/stack/run.py | 20 +++++- llama_stack/distribution/start_venv.sh | 71 ++++++++++++++++++++ 3 files changed, 122 insertions(+), 2 deletions(-) create mode 100755 llama_stack/distribution/start_venv.sh diff --git a/docs/source/distributions/building_distro.md b/docs/source/distributions/building_distro.md index 5556d4aa1..90239cb4e 100644 --- a/docs/source/distributions/building_distro.md +++ b/docs/source/distributions/building_distro.md @@ -180,12 +180,45 @@ After this step is successful, you should be able to find the built container im ### Running your Stack server Now, let's start the Llama Stack Distribution Server. You will need the YAML configuration file which was written out at the end by the `llama stack build` step. +``` +llama stack run -h +usage: llama stack run [-h] [--port PORT] [--image-name IMAGE_NAME] [--disable-ipv6] [--env KEY=VALUE] [--tls-keyfile TLS_KEYFILE] + [--tls-certfile TLS_CERTFILE] [--image-type {conda,container,venv}] + config + +start the server for a Llama Stack Distribution. You should have already built (or downloaded) and configured the distribution. + +positional arguments: + config Path to config file to use for the run + +options: + -h, --help show this help message and exit + --port PORT Port to run the server on. Defaults to 8321 + --image-name IMAGE_NAME + Name of the image to run. Defaults to the current conda environment + --disable-ipv6 Disable IPv6 support + --env KEY=VALUE Environment variables to pass to the server in KEY=VALUE format. Can be specified multiple times. + --tls-keyfile TLS_KEYFILE + Path to TLS key file for HTTPS + --tls-certfile TLS_CERTFILE + Path to TLS certificate file for HTTPS + --image-type {conda,container,venv} + Image Type used during the build. This can be either conda or container or venv. + +``` + ``` # Start using template name llama stack run tgi # Start using config file llama stack run ~/.llama/distributions/llamastack-my-local-stack/my-local-stack-run.yaml + +# Start using a venv +llama stack run --image-type venv ~/.llama/distributions/llamastack-my-local-stack/my-local-stack-run.yaml + +# Start using a conda environment +llama stack run --image-type conda ~/.llama/distributions/llamastack-my-local-stack/my-local-stack-run.yaml ``` ``` diff --git a/llama_stack/cli/stack/run.py b/llama_stack/cli/stack/run.py index e7d6df292..c32e51fca 100644 --- a/llama_stack/cli/stack/run.py +++ b/llama_stack/cli/stack/run.py @@ -65,6 +65,13 @@ class StackRun(Subcommand): type=str, help="Path to TLS certificate file for HTTPS", ) + self.parser.add_argument( + "--image-type", + type=str, + help="Image Type used during the build. This can be either conda or container or venv.", + choices=["conda", "container", "venv"], + default="conda", + ) def _run_stack_run_cmd(self, args: argparse.Namespace) -> None: import importlib.resources @@ -118,11 +125,11 @@ class StackRun(Subcommand): config_dict = yaml.safe_load(config_file.read_text()) config = parse_and_maybe_upgrade_config(config_dict) - if config.container_image: + if args.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] - else: + elif args.image_type == ImageType.conda.value: current_conda_env = os.environ.get("CONDA_DEFAULT_ENV") image_name = args.image_name or current_conda_env if not image_name: @@ -167,6 +174,15 @@ class StackRun(Subcommand): script, image_name, ] + else: + # else must be venv since that is the only valid option left. + current_venv = os.environ.get("VIRTUAL_ENV") + venv = args.image_name or current_venv + script = importlib.resources.files("llama_stack") / "distribution/start_venv.sh" + run_args = [ + script, + venv, + ] run_args.extend([str(config_file), str(args.port)]) if args.disable_ipv6: diff --git a/llama_stack/distribution/start_venv.sh b/llama_stack/distribution/start_venv.sh new file mode 100755 index 000000000..1cfa7248f --- /dev/null +++ b/llama_stack/distribution/start_venv.sh @@ -0,0 +1,71 @@ +#!/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 + +venv_path="$1" +shift + +yaml_config="$1" +shift + +port="$1" +shift + +# Initialize env_vars as an empty array +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 + +# Activate virtual environment +if [ ! -d "$venv_path" ]; then + echo -e "${RED}Error: Virtual environment not found at $venv_path${NC}" >&2 + exit 1 +fi + +source "$venv_path/bin/activate" + +set -x +python -m llama_stack.distribution.server.server \ + --yaml-config "$yaml_config" \ + --port "$port" \ + $env_vars \ + $other_args