From 18bac27d4ed15fa660cdf0e75c5b3712742c2289 Mon Sep 17 00:00:00 2001 From: Ihar Hrachyshka Date: Thu, 27 Mar 2025 17:13:22 -0400 Subject: [PATCH] fix: Use CONDA_DEFAULT_ENV presence as a flag to use conda mode (#1555) # What does this PR do? This is the second attempt to switch to system packages by default. Now with a hack to detect conda environment - in which case conda image-type is used. Note: Conda will only be used when --image-name is unset *and* CONDA_DEFAULT_ENV is set. This means that users without conda will correctly fall back to using system packages when no --image-* arguments are passed at all. [//]: # (If resolving an issue, uncomment and update the line below) [//]: # (Closes #[issue-number]) ## Test Plan Uses virtualenv: ``` $ llama stack build --template ollama --image-type venv $ llama stack run --image-type venv ~/.llama/distributions/ollama/ollama-run.yaml [...] Using virtual environment: /home/ec2-user/src/llama-stack/schedule/.local [...] ``` Uses system packages (virtualenv already initialized): ``` $ llama stack run ~/.llama/distributions/ollama/ollama-run.yaml [...] INFO 2025-03-27 20:46:22,882 llama_stack.cli.stack.run:142 server: No image type or image name provided. Assuming environment packages. [...] ``` Attempt to run from environment packages without necessary packages installed: ``` $ python -m venv barebones $ . ./barebones/bin/activate $ pip install -e . # to install llama command $ llama stack run ~/.llama/distributions/ollama/ollama-run.yaml [...] ModuleNotFoundError: No module named 'fastapi' ``` ^ failed as expected because the environment doesn't have necessary packages installed. Now install some packages in the new environment: ``` $ pip install fastapi opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp aiosqlite ollama openai datasets faiss-cpu mcp autoevals $ llama stack run ~/.llama/distributions/ollama/ollama-run.yaml [...] Uvicorn running on http://['::', '0.0.0.0']:8321 (Press CTRL+C to quit) ``` Now see if setting CONDA_DEFAULT_ENV will change what happens by default: ``` $ export CONDA_DEFAULT_ENV=base $ llama stack run ~/.llama/distributions/ollama/ollama-run.yaml [...] Using conda environment: base Conda environment base does not exist. [...] ``` --------- Signed-off-by: Ihar Hrachyshka --- docs/source/distributions/building_distro.md | 2 +- llama_stack/cli/stack/_build.py | 17 ++--- llama_stack/cli/stack/build.py | 9 +-- llama_stack/cli/stack/run.py | 74 +++++++++++++++----- llama_stack/cli/stack/utils.py | 8 +++ llama_stack/distribution/server/server.py | 28 ++++++-- 6 files changed, 102 insertions(+), 36 deletions(-) diff --git a/docs/source/distributions/building_distro.md b/docs/source/distributions/building_distro.md index 9b8c1b9ad..e1e38d7ce 100644 --- a/docs/source/distributions/building_distro.md +++ b/docs/source/distributions/building_distro.md @@ -67,7 +67,7 @@ options: Image Type to use for the build. This can be either conda or container or venv. If not specified, will use the image type from the template config. (default: conda) --image-name IMAGE_NAME - [for image-type=conda|venv] Name of the conda or virtual environment to use for the build. If not specified, currently active Conda environment will be used if + [for image-type=conda|container|venv] Name of the conda or virtual environment to use for the build. If not specified, currently active Conda environment will be used if found. (default: None) --print-deps-only Print the dependencies for the stack only, without building the stack (default: False) --run Run the stack after building using the same image type, name, and other applicable arguments (default: False) diff --git a/llama_stack/cli/stack/_build.py b/llama_stack/cli/stack/_build.py index d87e3bd0b..e440799d1 100644 --- a/llama_stack/cli/stack/_build.py +++ b/llama_stack/cli/stack/_build.py @@ -21,6 +21,7 @@ from prompt_toolkit.completion import WordCompleter from prompt_toolkit.validation import Validator from termcolor import cprint +from llama_stack.cli.stack.utils import ImageType from llama_stack.cli.table import print_table from llama_stack.distribution.build import ( SERVER_DEPENDENCIES, @@ -62,10 +63,10 @@ def run_stack_build_command(args: argparse.Namespace) -> None: if args.list_templates: return _run_template_list_cmd() - if args.image_type == "venv": + if args.image_type == ImageType.VENV.value: current_venv = os.environ.get("VIRTUAL_ENV") image_name = args.image_name or current_venv - elif args.image_type == "conda": + 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 else: @@ -84,7 +85,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None: build_config.image_type = args.image_type else: cprint( - f"Please specify a image-type (container | conda | venv) for {args.template}", + f"Please specify a image-type ({' | '.join(e.value for e in ImageType)}) for {args.template}", color="red", ) sys.exit(1) @@ -98,15 +99,15 @@ def run_stack_build_command(args: argparse.Namespace) -> None: ) image_type = prompt( - "> Enter the image type you want your Llama Stack to be built as (container or conda or venv): ", + f"> Enter the image type you want your Llama Stack to be built as ({' or '.join(e.value for e in ImageType)}): ", validator=Validator.from_callable( - lambda x: x in ["container", "conda", "venv"], - error_message="Invalid image type, please enter conda or container or venv", + lambda x: x in [e.value for e in ImageType], + error_message=f"Invalid image type, please enter {' or '.join(e.value for e in ImageType)}", ), - default="conda", + default=ImageType.CONDA.value, ) - if image_type == "conda": + if image_type == ImageType.CONDA.value: if not image_name: cprint( f"No current conda environment detected or specified, will create a new conda environment with the name `llamastack-{name}`", diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py index 70d74c620..0ada7c615 100644 --- a/llama_stack/cli/stack/build.py +++ b/llama_stack/cli/stack/build.py @@ -6,6 +6,7 @@ import argparse import textwrap +from llama_stack.cli.stack.utils import ImageType from llama_stack.cli.subcommand import Subcommand @@ -46,16 +47,16 @@ class StackBuild(Subcommand): self.parser.add_argument( "--image-type", type=str, - help="Image Type to use for the build. This can be either conda or container or venv. If not specified, will use the image type from the template config.", - choices=["conda", "container", "venv"], - default="conda", + help="Image Type to use for the build. If not specified, will use the image type from the template config.", + choices=[e.value for e in ImageType], + default=ImageType.CONDA.value, ) self.parser.add_argument( "--image-name", type=str, help=textwrap.dedent( - """[for image-type=conda|venv] Name of the conda or virtual environment to use for + f"""[for image-type={"|".join(e.value for e in ImageType)}] Name of the conda or virtual environment to use for the build. If not specified, currently active Conda environment will be used if found. """ ), diff --git a/llama_stack/cli/stack/run.py b/llama_stack/cli/stack/run.py index 88c3b061a..92015187b 100644 --- a/llama_stack/cli/stack/run.py +++ b/llama_stack/cli/stack/run.py @@ -8,6 +8,7 @@ import argparse import os from pathlib import Path +from llama_stack.cli.stack.utils import ImageType from llama_stack.cli.subcommand import Subcommand from llama_stack.log import get_logger @@ -43,7 +44,7 @@ class StackRun(Subcommand): self.parser.add_argument( "--image-name", type=str, - default=None, + default=os.environ.get("CONDA_DEFAULT_ENV"), help="Name of the image to run. Defaults to the current conda environment", ) self.parser.add_argument( @@ -56,7 +57,6 @@ class StackRun(Subcommand): "--env", action="append", help="Environment variables to pass to the server in KEY=VALUE format. Can be specified multiple times.", - default=[], metavar="KEY=VALUE", ) self.parser.add_argument( @@ -73,10 +73,24 @@ class StackRun(Subcommand): "--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", + choices=[e.value for e in ImageType], ) + # If neither image type nor image name is provided, but at the same time + # the current environment has conda breadcrumbs, then assume what the user + # wants to use conda mode and not the usual default mode (using + # pre-installed system packages). + # + # Note: yes, this is hacky. It's implemented this way to keep the existing + # conda users unaffected by the switch of the default behavior to using + # system packages. + def _get_image_type_and_name(self, args: argparse.Namespace) -> tuple[str, str]: + conda_env = os.environ.get("CONDA_DEFAULT_ENV") + if conda_env and args.image_name == conda_env: + logger.warning(f"Conda detected. Using conda environment {conda_env} for the run.") + return ImageType.CONDA.value, args.image_name + return args.image_type, args.image_name + def _run_stack_run_cmd(self, args: argparse.Namespace) -> None: import yaml @@ -120,20 +134,44 @@ class StackRun(Subcommand): except AttributeError as e: self.parser.error(f"failed to parse config file '{config_file}':\n {e}") - run_args = formulate_run_args(args.image_type, args.image_name, config, template_name) + image_type, image_name = self._get_image_type_and_name(args) - run_args.extend([str(config_file), str(args.port)]) - if args.disable_ipv6: - run_args.append("--disable-ipv6") + # If neither image type nor image name is provided, assume the server should be run directly + # using the current environment packages. + if not image_type and not image_name: + logger.info("No image type or image name provided. Assuming environment packages.") + from llama_stack.distribution.server.server import main as server_main - for env_var in args.env: - if "=" not in env_var: - self.parser.error(f"Environment variable '{env_var}' must be in KEY=VALUE format") - key, value = env_var.split("=", 1) # split on first = only - if not key: - self.parser.error(f"Environment variable '{env_var}' has empty key") - run_args.extend(["--env", f"{key}={value}"]) + # Build the server args from the current args passed to the CLI + server_args = argparse.Namespace() + for arg in vars(args): + # If this is a function, avoid passing it + # "args" contains: + # func=> + if callable(getattr(args, arg)): + continue + setattr(server_args, arg, getattr(args, arg)) - if args.tls_keyfile and args.tls_certfile: - run_args.extend(["--tls-keyfile", args.tls_keyfile, "--tls-certfile", args.tls_certfile]) - run_command(run_args) + # Run the server + server_main(server_args) + else: + run_args = formulate_run_args(image_type, image_name, config, template_name) + + run_args.extend([str(config_file), str(args.port)]) + if args.disable_ipv6: + run_args.append("--disable-ipv6") + + if args.env: + for env_var in args.env: + if "=" not in env_var: + self.parser.error(f"Environment variable '{env_var}' must be in KEY=VALUE format") + return + key, value = env_var.split("=", 1) # split on first = only + if not key: + self.parser.error(f"Environment variable '{env_var}' has empty key") + return + run_args.extend(["--env", f"{key}={value}"]) + + if args.tls_keyfile and args.tls_certfile: + run_args.extend(["--tls-keyfile", args.tls_keyfile, "--tls-certfile", args.tls_certfile]) + run_command(run_args) diff --git a/llama_stack/cli/stack/utils.py b/llama_stack/cli/stack/utils.py index 1e83a5cc8..74a606b2b 100644 --- a/llama_stack/cli/stack/utils.py +++ b/llama_stack/cli/stack/utils.py @@ -4,6 +4,14 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. +from enum import Enum + + +class ImageType(Enum): + CONDA = "conda" + CONTAINER = "container" + VENV = "venv" + def print_subcommand_description(parser, subparsers): """Print descriptions of subcommands.""" diff --git a/llama_stack/distribution/server/server.py b/llama_stack/distribution/server/server.py index b967b0269..7d4ec2a2f 100644 --- a/llama_stack/distribution/server/server.py +++ b/llama_stack/distribution/server/server.py @@ -15,7 +15,7 @@ import warnings from contextlib import asynccontextmanager from importlib.metadata import version as parse_version from pathlib import Path -from typing import Any, List, Union +from typing import Any, List, Optional, Union import yaml from fastapi import Body, FastAPI, HTTPException, Request @@ -294,11 +294,17 @@ class ClientVersionMiddleware: return await self.app(scope, receive, send) -def main(): +def main(args: Optional[argparse.Namespace] = None): """Start the LlamaStack server.""" parser = argparse.ArgumentParser(description="Start the LlamaStack server.") parser.add_argument( "--yaml-config", + dest="config", + help="(Deprecated) Path to YAML configuration file - use --config instead", + ) + parser.add_argument( + "--config", + dest="config", help="Path to YAML configuration file", ) parser.add_argument( @@ -328,12 +334,24 @@ def main(): required="--tls-keyfile" in sys.argv, ) - args = parser.parse_args() + # Determine whether the server args are being passed by the "run" command, if this is the case + # the args will be passed as a Namespace object to the main function, otherwise they will be + # parsed from the command line + if args is None: + args = parser.parse_args() + + # Check for deprecated argument usage + if "--yaml-config" in sys.argv: + warnings.warn( + "The '--yaml-config' argument is deprecated and will be removed in a future version. Use '--config' instead.", + DeprecationWarning, + stacklevel=2, + ) log_line = "" - if args.yaml_config: + if args.config: # if the user provided a config file, use it, even if template was specified - config_file = Path(args.yaml_config) + config_file = Path(args.config) if not config_file.exists(): raise ValueError(f"Config file {config_file} does not exist") log_line = f"Using config file: {config_file}"