address comments, configure within docker file

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
 

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/<image_type>/<name>-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/<image_type>/<name>-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 = (
+            builds_dir = BUILDS_BASE_DIR / ImageType.docker.value
-                Path(os.path.expanduser("./.llama/distributions"))
+            if args.output_dir:
-                / ImageType.docker.value
+                builds_dir = Path(output_dir)
-            )
+            os.makedirs(builds_dir, exist_ok=True)
-            build_config_file = build_dir / f"{args.config}-build.yaml"
-
-            os.makedirs(build_dir, exist_ok=True)
 
             script = pkg_resources.resource_filename(
                 "llama_toolchain", "core/configure_container.sh"
             )
-            script_args = [
+            script_args = [script, docker_image, str(builds_dir)]
-                script,
-                args.config,
-                str(build_config_file),
-            ]
 
             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:
-            if return_code != 0 or not build_config_file.exists():
                 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",
+        )

llama_toolchain/configs/distributions/generate_distribution_specs.py

@@ -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)

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 <<EOF
 
 EOF
 
-add_to_docker "ADD $build_file_path ./build.yaml"
+add_to_docker "ADD $build_file_path ./llamastack-build.yaml"
 
 printf "Dockerfile created successfully in $TEMP_DIR/Dockerfile"
 cat $TEMP_DIR/Dockerfile
@@ -107,10 +108,10 @@ if [ -n "$LLAMA_MODELS_DIR" ]; then
   mounts="$mounts -v $(readlink -f $LLAMA_MODELS_DIR):$models_mount"
 fi
 set -x
-podman build --network host -t $image_name -f "$TEMP_DIR/Dockerfile" "$REPO_DIR" $mounts
+$DOCKER_BINARY build $DOCKER_OPTS -t $image_name -f "$TEMP_DIR/Dockerfile" "$REPO_DIR" $mounts
 set +x
 
 echo "You can run it with: podman run -p 8000:8000 $image_name"
 
 echo "Checking image builds..."
-podman run -it llamastack-local-docker-example cat build.yaml
+podman run -it $image_name cat llamastack-build.yaml

llama_toolchain/core/configure_container.sh

@@ -21,6 +21,8 @@ if [ $# -lt 2 ]; then
 fi
 
 docker_image="$1"
-build_file_path="$2"
+host_build_dir="$2"
+container_build_dir="/app/builds"
 
-podman run -it $docker_image cat build.yaml >> ./$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

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(
+    description: Optional[str] = Field(
-        default="Use code from `llama_toolchain` itself to serve all llama stack APIs",
+        default="",
         description="Description of the distribution",
     )
     docker_image: Optional[str] = None

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.