From e466ec389bc00ac469bfdd0579c84bbf8c7aba1e Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Sun, 15 Sep 2024 23:16:18 -0700 Subject: [PATCH] address comments, configure within docker file --- llama_toolchain/cli/stack/build.py | 3 +- llama_toolchain/cli/stack/configure.py | 53 ++++++++++++------- .../generate_distribution_specs.py | 21 -------- llama_toolchain/core/build_container.sh | 9 ++-- llama_toolchain/core/configure_container.sh | 6 ++- llama_toolchain/core/datatypes.py | 4 +- llama_toolchain/memory/common/__init__.py | 5 ++ 7 files changed, 53 insertions(+), 48 deletions(-) delete mode 100644 llama_toolchain/configs/distributions/generate_distribution_specs.py create mode 100644 llama_toolchain/memory/common/__init__.py diff --git a/llama_toolchain/cli/stack/build.py b/llama_toolchain/cli/stack/build.py index 58d13b70a..3801637c0 100644 --- a/llama_toolchain/cli/stack/build.py +++ b/llama_toolchain/cli/stack/build.py @@ -84,7 +84,8 @@ class StackBuild(Subcommand): except Exception as e: self.parser.error(f"Could not parse config file {args.config}: {e}") return - build_config.name = args.name if args.name else build_config.name + if args.name: + build_config.name = args.name self._run_stack_build_command_from_build_config(build_config) return diff --git a/llama_toolchain/cli/stack/configure.py b/llama_toolchain/cli/stack/configure.py index afb2dd9fe..485200397 100644 --- a/llama_toolchain/cli/stack/configure.py +++ b/llama_toolchain/cli/stack/configure.py @@ -19,6 +19,8 @@ from termcolor import cprint from llama_toolchain.core.datatypes import * # noqa: F403 import os +from termcolor import cprint + class StackConfigure(Subcommand): """Llama cli for configuring llama toolchain configs""" @@ -41,53 +43,65 @@ class StackConfigure(Subcommand): help="Path to the build config file (e.g. ~/.llama/builds//-build.yaml). For docker, this could also be the name of the docker image. ", ) + self.parser.add_argument( + "--output-dir", + type=str, + help="Path to the output directory to store generated run.yaml config file. If not specified, will use ~/.llama/build//-run.yaml", + ) + def _run_stack_configure_cmd(self, args: argparse.Namespace) -> None: from llama_toolchain.core.package import ImageType + docker_image = None build_config_file = Path(args.config) if not build_config_file.exists(): cprint( f"Could not find {build_config_file}. Trying docker image name instead...", color="green", ) + docker_image = args.config - build_dir = ( - Path(os.path.expanduser("./.llama/distributions")) - / ImageType.docker.value - ) - build_config_file = build_dir / f"{args.config}-build.yaml" - - os.makedirs(build_dir, exist_ok=True) + builds_dir = BUILDS_BASE_DIR / ImageType.docker.value + if args.output_dir: + builds_dir = Path(output_dir) + os.makedirs(builds_dir, exist_ok=True) script = pkg_resources.resource_filename( "llama_toolchain", "core/configure_container.sh" ) - script_args = [ - script, - args.config, - str(build_config_file), - ] + script_args = [script, docker_image, str(builds_dir)] return_code = run_with_pty(script_args) # we have regenerated the build config file with script, now check if it exists - build_config_file = Path(str(build_config_file)) - if return_code != 0 or not build_config_file.exists(): + if return_code != 0: self.parser.error( f"Can not find {build_config_file}. Please run llama stack build first or check if docker image exists" ) - return + + build_name = docker_image.removeprefix("llamastack-") + cprint( + f"YAML configuration has been written to {builds_dir / f'{build_name}-run.yaml'}", + color="green", + ) + return with open(build_config_file, "r") as f: build_config = BuildConfig(**yaml.safe_load(f)) - self._configure_llama_distribution(build_config) + self._configure_llama_distribution(build_config, args.output_dir) - def _configure_llama_distribution(self, build_config: BuildConfig): + def _configure_llama_distribution( + self, + build_config: BuildConfig, + output_dir: Optional[str] = None, + ): from llama_toolchain.common.serialize import EnumEncoder from llama_toolchain.core.configure import configure_api_providers builds_dir = BUILDS_BASE_DIR / build_config.image_type + if output_dir: + builds_dir = Path(output_dir) os.makedirs(builds_dir, exist_ok=True) package_name = build_config.name.replace("::", "-") package_file = builds_dir / f"{package_name}-run.yaml" @@ -124,4 +138,7 @@ class StackConfigure(Subcommand): to_write = json.loads(json.dumps(config.dict(), cls=EnumEncoder)) f.write(yaml.dump(to_write, sort_keys=False)) - print(f"YAML configuration has been written to {package_file}") + cprint( + f"> YAML configuration has been written to {package_file}", + color="blue", + ) diff --git a/llama_toolchain/configs/distributions/generate_distribution_specs.py b/llama_toolchain/configs/distributions/generate_distribution_specs.py deleted file mode 100644 index 925ab8340..000000000 --- a/llama_toolchain/configs/distributions/generate_distribution_specs.py +++ /dev/null @@ -1,21 +0,0 @@ -from llama_toolchain.core.distribution_registry import * -import json - -import fire -import yaml -from llama_toolchain.common.serialize import EnumEncoder - - -def main(): - for d in available_distribution_specs(): - file_path = "./configs/distributions/distribution_registry/{}.yaml".format( - d.distribution_type - ) - - with open(file_path, "w") as f: - to_write = json.loads(json.dumps(d.dict(), cls=EnumEncoder)) - f.write(yaml.dump(to_write, sort_keys=False)) - - -if __name__ == "__main__": - fire.Fire(main) diff --git a/llama_toolchain/core/build_container.sh b/llama_toolchain/core/build_container.sh index ce7e3438f..d829e8399 100755 --- a/llama_toolchain/core/build_container.sh +++ b/llama_toolchain/core/build_container.sh @@ -25,7 +25,8 @@ set -euo pipefail SCRIPT_DIR=$(dirname "$(readlink -f "$0")") REPO_DIR=$(dirname $(dirname "$SCRIPT_DIR")) -BUILD_DIR=$(dirname $(dirname "$build_file_path")) +DOCKER_BINARY=${DOCKER_BINARY:-docker} +DOCKER_OPTS=${DOCKER_OPTS:-} TEMP_DIR=$(mktemp -d) @@ -93,7 +94,7 @@ add_to_docker <> ./$build_file_path +set -x +podman run -it -v $host_build_dir:$container_build_dir $docker_image llama stack configure ./llamastack-build.yaml --output-dir $container_build_dir diff --git a/llama_toolchain/core/datatypes.py b/llama_toolchain/core/datatypes.py index f44c4b70c..359ffe2ce 100644 --- a/llama_toolchain/core/datatypes.py +++ b/llama_toolchain/core/datatypes.py @@ -155,8 +155,8 @@ class DistributionSpec(BaseModel): default="local", description="Name of the distribution type. This can used to identify the distribution", ) - description: str = Field( - default="Use code from `llama_toolchain` itself to serve all llama stack APIs", + description: Optional[str] = Field( + default="", description="Description of the distribution", ) docker_image: Optional[str] = None diff --git a/llama_toolchain/memory/common/__init__.py b/llama_toolchain/memory/common/__init__.py new file mode 100644 index 000000000..756f351d8 --- /dev/null +++ b/llama_toolchain/memory/common/__init__.py @@ -0,0 +1,5 @@ +# 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.