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 = (
-                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",
+        )