From 5bed6c276c679088a8d8617223d085f313cc8a98 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Fri, 25 Oct 2024 09:17:33 -0700 Subject: [PATCH] Move function around --- .pre-commit-config.yaml | 2 +- llama_stack/cli/stack/build.py | 192 ++++++++++++++++----------------- 2 files changed, 97 insertions(+), 97 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1c85436c4..3707d4671 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,4 +1,4 @@ -exclude: 'build' +exclude: 'build/' default_language_version: python: python3 diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py index 26aa35e16..a425a6c99 100644 --- a/llama_stack/cli/stack/build.py +++ b/llama_stack/cli/stack/build.py @@ -78,102 +78,6 @@ class StackBuild(Subcommand): choices=["conda", "docker"], ) - def _get_build_config_from_name(self, args: argparse.Namespace) -> Optional[Path]: - if os.getenv("CONDA_PREFIX", ""): - conda_dir = ( - Path(os.getenv("CONDA_PREFIX")).parent / f"llamastack-{args.name}" - ) - else: - cprint( - "Cannot find CONDA_PREFIX. Trying default conda path ~/.conda/envs...", - color="green", - ) - conda_dir = ( - Path(os.path.expanduser("~/.conda/envs")) / f"llamastack-{args.name}" - ) - build_config_file = Path(conda_dir) / f"{args.name}-build.yaml" - if build_config_file.exists(): - return build_config_file - - return None - - def _run_stack_build_command_from_build_config( - self, build_config: BuildConfig - ) -> None: - import json - import os - - import yaml - - from llama_stack.distribution.build import build_image, ImageType - from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR - from llama_stack.distribution.utils.serialize import EnumEncoder - from termcolor import cprint - - # save build.yaml spec for building same distribution again - if build_config.image_type == ImageType.docker.value: - # docker needs build file to be in the llama-stack repo dir to be able to copy over to the image - llama_stack_path = Path( - os.path.abspath(__file__) - ).parent.parent.parent.parent - build_dir = llama_stack_path / "tmp/configs/" - else: - build_dir = DISTRIBS_BASE_DIR / f"llamastack-{build_config.name}" - - os.makedirs(build_dir, exist_ok=True) - build_file_path = build_dir / f"{build_config.name}-build.yaml" - - with open(build_file_path, "w") as f: - to_write = json.loads(json.dumps(build_config.dict(), cls=EnumEncoder)) - f.write(yaml.dump(to_write, sort_keys=False)) - - return_code = build_image(build_config, build_file_path) - if return_code != 0: - return - - configure_name = ( - build_config.name - if build_config.image_type == "conda" - else (f"llamastack-{build_config.name}") - ) - if build_config.image_type == "conda": - cprint( - f"You can now run `llama stack configure {configure_name}`", - color="green", - ) - else: - cprint( - f"You can now run `llama stack run {build_config.name}`", - color="green", - ) - - def _run_template_list_cmd(self, args: argparse.Namespace) -> None: - import json - - from llama_stack.cli.table import print_table - - # eventually, this should query a registry at llama.meta.com/llamastack/distributions - headers = [ - "Template Name", - "Providers", - "Description", - ] - - rows = [] - for spec in available_templates_specs(): - rows.append( - [ - spec.name, - json.dumps(spec.distribution_spec.providers, indent=2), - spec.distribution_spec.description, - ] - ) - print_table( - rows, - headers, - separate_rows=True, - ) - def _run_stack_build_command(self, args: argparse.Namespace) -> None: import textwrap @@ -297,3 +201,99 @@ class StackBuild(Subcommand): self.parser.error(f"Could not parse config file {args.config}: {e}") return self._run_stack_build_command_from_build_config(build_config) + def _get_build_config_from_name(self, args: argparse.Namespace) -> Optional[Path]: + if os.getenv("CONDA_PREFIX", ""): + conda_dir = ( + Path(os.getenv("CONDA_PREFIX")).parent / f"llamastack-{args.name}" + ) + else: + cprint( + "Cannot find CONDA_PREFIX. Trying default conda path ~/.conda/envs...", + color="green", + ) + conda_dir = ( + Path(os.path.expanduser("~/.conda/envs")) / f"llamastack-{args.name}" + ) + build_config_file = Path(conda_dir) / f"{args.name}-build.yaml" + if build_config_file.exists(): + return build_config_file + + return None + + def _run_stack_build_command_from_build_config( + self, build_config: BuildConfig + ) -> None: + import json + import os + + import yaml + + from llama_stack.distribution.build import build_image, ImageType + from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR + from llama_stack.distribution.utils.serialize import EnumEncoder + from termcolor import cprint + + # save build.yaml spec for building same distribution again + if build_config.image_type == ImageType.docker.value: + # docker needs build file to be in the llama-stack repo dir to be able to copy over to the image + llama_stack_path = Path( + os.path.abspath(__file__) + ).parent.parent.parent.parent + build_dir = llama_stack_path / "tmp/configs/" + else: + build_dir = DISTRIBS_BASE_DIR / f"llamastack-{build_config.name}" + + os.makedirs(build_dir, exist_ok=True) + build_file_path = build_dir / f"{build_config.name}-build.yaml" + + with open(build_file_path, "w") as f: + to_write = json.loads(json.dumps(build_config.dict(), cls=EnumEncoder)) + f.write(yaml.dump(to_write, sort_keys=False)) + + return_code = build_image(build_config, build_file_path) + if return_code != 0: + return + + configure_name = ( + build_config.name + if build_config.image_type == "conda" + else (f"llamastack-{build_config.name}") + ) + if build_config.image_type == "conda": + cprint( + f"You can now run `llama stack configure {configure_name}`", + color="green", + ) + else: + cprint( + f"You can now run `llama stack run {build_config.name}`", + color="green", + ) + + def _run_template_list_cmd(self, args: argparse.Namespace) -> None: + import json + + from llama_stack.cli.table import print_table + + # eventually, this should query a registry at llama.meta.com/llamastack/distributions + headers = [ + "Template Name", + "Providers", + "Description", + ] + + rows = [] + for spec in available_templates_specs(): + rows.append( + [ + spec.name, + json.dumps(spec.distribution_spec.providers, indent=2), + spec.distribution_spec.description, + ] + ) + print_table( + rows, + headers, + separate_rows=True, + ) +