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 <ihar.hrachyshka@gmail.com>
This commit is contained in:
Ihar Hrachyshka 2025-03-27 17:13:22 -04:00 committed by GitHub
parent b5c27f77ad
commit 18bac27d4e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 102 additions and 36 deletions

View file

@ -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: 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) conda)
--image-name IMAGE_NAME --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) found. (default: None)
--print-deps-only Print the dependencies for the stack only, without building the stack (default: False) --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) --run Run the stack after building using the same image type, name, and other applicable arguments (default: False)

View file

@ -21,6 +21,7 @@ from prompt_toolkit.completion import WordCompleter
from prompt_toolkit.validation import Validator from prompt_toolkit.validation import Validator
from termcolor import cprint from termcolor import cprint
from llama_stack.cli.stack.utils import ImageType
from llama_stack.cli.table import print_table from llama_stack.cli.table import print_table
from llama_stack.distribution.build import ( from llama_stack.distribution.build import (
SERVER_DEPENDENCIES, SERVER_DEPENDENCIES,
@ -62,10 +63,10 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
if args.list_templates: if args.list_templates:
return _run_template_list_cmd() return _run_template_list_cmd()
if args.image_type == "venv": if args.image_type == ImageType.VENV.value:
current_venv = os.environ.get("VIRTUAL_ENV") current_venv = os.environ.get("VIRTUAL_ENV")
image_name = args.image_name or current_venv 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") current_conda_env = os.environ.get("CONDA_DEFAULT_ENV")
image_name = args.image_name or current_conda_env image_name = args.image_name or current_conda_env
else: else:
@ -84,7 +85,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
build_config.image_type = args.image_type build_config.image_type = args.image_type
else: else:
cprint( 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", color="red",
) )
sys.exit(1) sys.exit(1)
@ -98,15 +99,15 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
) )
image_type = prompt( 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( validator=Validator.from_callable(
lambda x: x in ["container", "conda", "venv"], lambda x: x in [e.value for e in ImageType],
error_message="Invalid image type, please enter conda or container or venv", 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: if not image_name:
cprint( cprint(
f"No current conda environment detected or specified, will create a new conda environment with the name `llamastack-{name}`", f"No current conda environment detected or specified, will create a new conda environment with the name `llamastack-{name}`",

View file

@ -6,6 +6,7 @@
import argparse import argparse
import textwrap import textwrap
from llama_stack.cli.stack.utils import ImageType
from llama_stack.cli.subcommand import Subcommand from llama_stack.cli.subcommand import Subcommand
@ -46,16 +47,16 @@ class StackBuild(Subcommand):
self.parser.add_argument( self.parser.add_argument(
"--image-type", "--image-type",
type=str, 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.", help="Image Type to use for the build. If not specified, will use the image type from the template config.",
choices=["conda", "container", "venv"], choices=[e.value for e in ImageType],
default="conda", default=ImageType.CONDA.value,
) )
self.parser.add_argument( self.parser.add_argument(
"--image-name", "--image-name",
type=str, type=str,
help=textwrap.dedent( 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. the build. If not specified, currently active Conda environment will be used if found.
""" """
), ),

View file

@ -8,6 +8,7 @@ import argparse
import os import os
from pathlib import Path from pathlib import Path
from llama_stack.cli.stack.utils import ImageType
from llama_stack.cli.subcommand import Subcommand from llama_stack.cli.subcommand import Subcommand
from llama_stack.log import get_logger from llama_stack.log import get_logger
@ -43,7 +44,7 @@ class StackRun(Subcommand):
self.parser.add_argument( self.parser.add_argument(
"--image-name", "--image-name",
type=str, type=str,
default=None, default=os.environ.get("CONDA_DEFAULT_ENV"),
help="Name of the image to run. Defaults to the current conda environment", help="Name of the image to run. Defaults to the current conda environment",
) )
self.parser.add_argument( self.parser.add_argument(
@ -56,7 +57,6 @@ class StackRun(Subcommand):
"--env", "--env",
action="append", action="append",
help="Environment variables to pass to the server in KEY=VALUE format. Can be specified multiple times.", help="Environment variables to pass to the server in KEY=VALUE format. Can be specified multiple times.",
default=[],
metavar="KEY=VALUE", metavar="KEY=VALUE",
) )
self.parser.add_argument( self.parser.add_argument(
@ -73,10 +73,24 @@ class StackRun(Subcommand):
"--image-type", "--image-type",
type=str, type=str,
help="Image Type used during the build. This can be either conda or container or venv.", help="Image Type used during the build. This can be either conda or container or venv.",
choices=["conda", "container", "venv"], choices=[e.value for e in ImageType],
default="conda",
) )
# 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: def _run_stack_run_cmd(self, args: argparse.Namespace) -> None:
import yaml import yaml
@ -120,20 +134,44 @@ class StackRun(Subcommand):
except AttributeError as e: except AttributeError as e:
self.parser.error(f"failed to parse config file '{config_file}':\n {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 neither image type nor image name is provided, assume the server should be run directly
if args.disable_ipv6: # using the current environment packages.
run_args.append("--disable-ipv6") 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: # Build the server args from the current args passed to the CLI
if "=" not in env_var: server_args = argparse.Namespace()
self.parser.error(f"Environment variable '{env_var}' must be in KEY=VALUE format") for arg in vars(args):
key, value = env_var.split("=", 1) # split on first = only # If this is a function, avoid passing it
if not key: # "args" contains:
self.parser.error(f"Environment variable '{env_var}' has empty key") # func=<bound method StackRun._run_stack_run_cmd of <llama_stack.cli.stack.run.StackRun object at 0x10484b010>>
run_args.extend(["--env", f"{key}={value}"]) if callable(getattr(args, arg)):
continue
setattr(server_args, arg, getattr(args, arg))
if args.tls_keyfile and args.tls_certfile: # Run the server
run_args.extend(["--tls-keyfile", args.tls_keyfile, "--tls-certfile", args.tls_certfile]) server_main(server_args)
run_command(run_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)

View file

@ -4,6 +4,14 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # 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): def print_subcommand_description(parser, subparsers):
"""Print descriptions of subcommands.""" """Print descriptions of subcommands."""

View file

@ -15,7 +15,7 @@ import warnings
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from importlib.metadata import version as parse_version from importlib.metadata import version as parse_version
from pathlib import Path from pathlib import Path
from typing import Any, List, Union from typing import Any, List, Optional, Union
import yaml import yaml
from fastapi import Body, FastAPI, HTTPException, Request from fastapi import Body, FastAPI, HTTPException, Request
@ -294,11 +294,17 @@ class ClientVersionMiddleware:
return await self.app(scope, receive, send) return await self.app(scope, receive, send)
def main(): def main(args: Optional[argparse.Namespace] = None):
"""Start the LlamaStack server.""" """Start the LlamaStack server."""
parser = argparse.ArgumentParser(description="Start the LlamaStack server.") parser = argparse.ArgumentParser(description="Start the LlamaStack server.")
parser.add_argument( parser.add_argument(
"--yaml-config", "--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", help="Path to YAML configuration file",
) )
parser.add_argument( parser.add_argument(
@ -328,12 +334,24 @@ def main():
required="--tls-keyfile" in sys.argv, 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 = "" log_line = ""
if args.yaml_config: if args.config:
# if the user provided a config file, use it, even if template was specified # 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(): if not config_file.exists():
raise ValueError(f"Config file {config_file} does not exist") raise ValueError(f"Config file {config_file} does not exist")
log_line = f"Using config file: {config_file}" log_line = f"Using config file: {config_file}"