Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

CLI - add back build wizard, configure with name instead of build.yaml #74

Merged
merged 17 commits into from
Sep 18, 2024
74 changes: 67 additions & 7 deletions llama_stack/cli/stack/build.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
from pathlib import Path

import yaml
from prompt_toolkit.validation import Validator
yanxi0830 marked this conversation as resolved.
Show resolved Hide resolved
from termcolor import cprint


class StackBuild(Subcommand):
Expand All @@ -29,7 +31,9 @@ def _add_arguments(self):
self.parser.add_argument(
"config",
type=str,
help="Path to a config file to use for the build. You may find example configs in llama_stack/distribution/example_configs",
default=None,
nargs="*",
yanxi0830 marked this conversation as resolved.
Show resolved Hide resolved
help="Path to a config file to use for the build. You can find example configs in llama_stack/distribution/example_configs. If this argument is not provided, you will be prompted to enter information interactively",
)

self.parser.add_argument(
Expand All @@ -44,10 +48,10 @@ def _run_stack_build_command_from_build_config(
import json
import os

from llama_stack.distribution.build import ApiInput, build_image, ImageType

from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR
from llama_stack.distribution.utils.serialize import EnumEncoder
from llama_stack.distribution.build import ApiInput, build_image, ImageType
from termcolor import cprint

# save build.yaml spec for building same distribution again
if build_config.image_type == ImageType.docker.value:
Expand All @@ -57,7 +61,10 @@ def _run_stack_build_command_from_build_config(
llama_stack_path / "configs/distributions" / build_config.image_type
)
else:
build_dir = DISTRIBS_BASE_DIR / build_config.image_type
build_dir = (
Path(os.getenv("CONDA_PREFIX")).parent
/ f"llamastack-{build_config.name}"
)

os.makedirs(build_dir, exist_ok=True)
build_file_path = build_dir / f"{build_config.name}-build.yaml"
Expand All @@ -74,13 +81,66 @@ def _run_stack_build_command_from_build_config(
)

def _run_stack_build_command(self, args: argparse.Namespace) -> None:
from llama_stack.distribution.utils.prompt_for_config import prompt_for_config
from llama_stack.distribution.distribution import Api, api_providers
from llama_stack.distribution.utils.dynamic import instantiate_class_type
from prompt_toolkit import prompt

if not args.config:
self.parser.error(
"No config file specified. Please use `llama stack build /path/to/*-build.yaml`. Example config files can be found in llama_stack/distribution/example_configs"
name = prompt(
"> Enter an unique name for identifying your Llama Stack build distribution (e.g. my-local-stack): "
yanxi0830 marked this conversation as resolved.
Show resolved Hide resolved
)
image_type = prompt(
"> Enter the image type you want your distribution to be built with (docker or conda): ",
yanxi0830 marked this conversation as resolved.
Show resolved Hide resolved
validator=Validator.from_callable(
lambda x: x in ["docker", "conda"],
error_message="Invalid image type, please enter (conda|docker)",
yanxi0830 marked this conversation as resolved.
Show resolved Hide resolved
),
default="conda",
)

cprint(
yanxi0830 marked this conversation as resolved.
Show resolved Hide resolved
f"\n Llama Stack is composed of several APIs working together. Let's configure the providers (implementations) you want to use for these APIs.",
color="green",
)

providers = dict()
for api in Api:
all_providers = api_providers()
providers_for_api = all_providers[api]

api_provider = prompt(
"> Enter the API provider for the {} API: (default=meta-reference): ".format(
api.value
),
validator=Validator.from_callable(
lambda x: x in providers_for_api,
error_message="Invalid provider, please enter one of the following: {}".format(
providers_for_api.keys()
),
),
default=(
"meta-reference"
yanxi0830 marked this conversation as resolved.
Show resolved Hide resolved
if "meta-reference" in providers_for_api
else list(providers_for_api.keys())[0]
),
)

providers[api.value] = api_provider

description = prompt(
yanxi0830 marked this conversation as resolved.
Show resolved Hide resolved
"\n > (Optional) Enter a short description for your Llama Stack distribution: ",
default="",
)

distribution_spec = DistributionSpec(
providers=providers,
description=description,
)

build_config = BuildConfig(
name=name, image_type=image_type, distribution_spec=distribution_spec
)
self._run_stack_build_command_from_build_config(build_config)
return

with open(args.config, "r") as f:
Expand Down
16 changes: 12 additions & 4 deletions llama_stack/cli/stack/configure.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,12 +9,12 @@
from pathlib import Path

import yaml
from termcolor import cprint

from llama_stack.cli.subcommand import Subcommand
from llama_stack.distribution.utils.config_dirs import BUILDS_BASE_DIR

from llama_stack.distribution.utils.exec import run_with_pty
from termcolor import cprint
yanxi0830 marked this conversation as resolved.
Show resolved Hide resolved
from llama_stack.distribution.datatypes import * # noqa: F403
import os

Expand Down Expand Up @@ -52,7 +52,9 @@ def _run_stack_configure_cmd(self, args: argparse.Namespace) -> None:
from llama_stack.distribution.build import ImageType

docker_image = None
build_config_file = Path(args.config)
conda_dir = Path(os.getenv("CONDA_PREFIX")).parent / f"llamastack-{args.config}"
build_config_file = Path(conda_dir) / f"{args.config}-build.yaml"

if not build_config_file.exists():
cprint(
f"Could not find {build_config_file}. Trying docker image name instead...",
Expand All @@ -79,8 +81,9 @@ def _run_stack_configure_cmd(self, args: argparse.Namespace) -> None:
)

build_name = docker_image.removeprefix("llamastack-")
saved_file = str(builds_dir / f"{build_name}-run.yaml")
cprint(
f"YAML configuration has been written to {builds_dir / f'{build_name}-run.yaml'}",
f"YAML configuration has been written to {saved_file}. You can now run `llama stack run {saved_file}`",
color="green",
)
return
Expand Down Expand Up @@ -132,6 +135,11 @@ def _configure_llama_distribution(
f.write(yaml.dump(to_write, sort_keys=False))

cprint(
f"> YAML configuration has been written to {run_config_file}",
f"> YAML configuration has been written to {run_config_file}.",
color="blue",
)

cprint(
f"You can now run `llama stack run {image_name} --port PORT` or `llama stack run {run_config_file} --port PORT`",
color="green",
)
18 changes: 15 additions & 3 deletions llama_stack/cli/stack/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,19 +47,31 @@ def _add_arguments(self):

def _run_stack_run_cmd(self, args: argparse.Namespace) -> None:
import pkg_resources
from llama_stack.distribution.build import ImageType
from llama_stack.distribution.utils.config_dirs import BUILDS_BASE_DIR

from llama_stack.distribution.utils.exec import run_with_pty

if not args.config:
self.parser.error("Must specify a config file to run")
return

path = args.config
config_file = Path(path)
config_file = Path(args.config)
if not config_file.exists() and not args.config.endswith(".yaml"):
# check if it's a build config saved to conda dir
config_file = Path(
BUILDS_BASE_DIR / ImageType.conda.value / f"{args.config}-run.yaml"
)

if not config_file.exists() and not args.config.endswith(".yaml"):
# check if it's a build config saved to docker dir
config_file = Path(
BUILDS_BASE_DIR / ImageType.docker.value / f"{args.config}-run.yaml"
)

if not config_file.exists():
self.parser.error(
f"File {str(config_file)} does not exist. Did you run `llama stack build`?"
f"File {str(config_file)} does not exist. Please run `llama stack build` and `llama stack configure <name>` to generate a run.yaml file"
)
return

Expand Down