Add config file based CLI (#60)

* config file for build

* fix build command

* configure script with config

* fix configure script to work with config file

* update build.sh

* update readme

* distribution_type -> distribution

* fix run-config/config-file to config

* move import to inline

* only consume config as argument

* update configure to only consume config

* update readme

* update readme

docs/cli_reference.md

@@ -295,13 +295,18 @@ As you can see above, each “distribution” details the “providers” it is
 Let's imagine you are working with a 8B-Instruct model. The following command will build a package (in the form of a Conda environment) _and_ configure it. As part of the configuration, you will be asked for some inputs (model_id, max_seq_len, etc.) Since we are working with a 8B model, we will name our build `8b-instruct` to help us remember the config.
 
 ```
-llama stack build local --name 8b-instruct
+llama stack build
 ```
 
-Once it runs successfully , you should see some outputs in the form:
+Once it runs, you will be prompted to enter build name and optional arguments, and should see some outputs in the form:
 
 ```
-$ llama stack build local --name 8b-instruct
+$ llama stack build
+Enter value for name (required): 8b-instruct
+Enter value for distribution (default: local) (required): local
+Enter value for api_providers (optional):
+Enter value for image_type (default: conda) (required):
+
 ....
 ....
 Successfully installed cfgv-3.4.0 distlib-0.3.8 identify-2.6.0 libcst-1.4.0 llama_toolchain-0.0.2 moreorless-0.4.0 nodeenv-1.9.1 pre-commit-3.8.0 stdlibs-2024.5.15 toml-0.10.2 tomlkit-0.13.0 trailrunner-1.4.0 ufmt-2.7.0 usort-1.0.8 virtualenv-20.26.3
@@ -312,17 +317,41 @@ Successfully setup conda environment. Configuring build...
 ...
 
 YAML configuration has been written to ~/.llama/builds/local/conda/8b-instruct.yaml
+Target `8b-test` built with configuration at /home/xiyan/.llama/builds/local/conda/8b-test.yaml
+Build spec configuration saved at /home/xiyan/.llama/distributions/local/conda/8b-test-build.yaml
 ```
+
+You can re-build package based on build config
+```
+$ cat ~/.llama/distributions/local/conda/8b-instruct-build.yaml
+name: 8b-instruct
+distribution: local
+api_providers: null
+image_type: conda
+
+$ llama stack build --config ~/.llama/distributions/local/conda/8b-instruct-build.yaml
+
+Successfully setup conda environment. Configuring build...
+
+...
+...
+
+YAML configuration has been written to ~/.llama/builds/local/conda/8b-instruct.yaml
+Target `8b-instruct` built with configuration at ~/.llama/builds/local/conda/8b-instruct.yaml
+Build spec configuration saved at ~/.llama/distributions/local/conda/8b-instruct-build.yaml
+```
+
 ### Step 3.3: Configure a distribution
 
 You can re-configure this distribution by running:
 ```
-llama stack configure local --name 8b-instruct
+llama stack configure ~/.llama/builds/local/conda/8b-instruct.yaml
 ```
 
 Here is an example run of how the CLI will guide you to fill the configuration
+
 ```
-$ llama stack configure local --name 8b-instruct
+$ llama stack configure ~/.llama/builds/local/conda/8b-instruct.yaml
 
 Configuring API: inference (meta-reference)
 Enter value for model (required): Meta-Llama3.1-8B-Instruct
@@ -363,12 +392,12 @@ Now let’s start Llama Stack Distribution Server.
 You need the YAML configuration file which was written out at the end by the `llama stack build` step.
 
 ```
-llama stack run local --name 8b-instruct --port 5000
+llama stack run ~/.llama/builds/local/conda/8b-instruct.yaml --port 5000
 ```
 You should see the Stack server start and print the APIs that it is supporting,
 
 ```
-$ llama stack run local --name 8b-instruct --port 5000
+$ llama stack run ~/.llama/builds/local/conda/8b-instruct.yaml --port 5000
 
 > initializing model parallel with size 1
 > initializing ddp with size 1

llama_toolchain/cli/stack/build.py

@@ -8,6 +8,7 @@ import argparse
 
 from llama_toolchain.cli.subcommand import Subcommand
 from llama_toolchain.core.datatypes import *  # noqa: F403
+import yaml
 
 
 def parse_api_provider_tuples(
@@ -47,55 +48,45 @@ class StackBuild(Subcommand):
         self.parser.set_defaults(func=self._run_stack_build_command)
 
     def _add_arguments(self):
-        from llama_toolchain.core.distribution_registry import available_distribution_specs
+        from llama_toolchain.core.distribution_registry import (
-        from llama_toolchain.core.package import (
+            available_distribution_specs,
-            BuildType,
         )
+        from llama_toolchain.core.package import ImageType
 
         allowed_ids = [d.distribution_type for d in available_distribution_specs()]
         self.parser.add_argument(
-            "distribution",
+            "--config",
             type=str,
-            help="Distribution to build (either \"adhoc\" OR one of: {})".format(allowed_ids),
+            help="Path to a config file to use for the build",
-        )
-        self.parser.add_argument(
-            "api_providers",
-            nargs='?',
-            help="Comma separated list of (api=provider) tuples",
         )
 
-        self.parser.add_argument(
+    def _run_stack_build_command_from_build_config(
-            "--name",
+        self, build_config: BuildConfig
-            type=str,
+    ) -> None:
-            help="Name of the build target (image, conda env)",
+        import json
-            required=True,
+        import os
-        )
-        self.parser.add_argument(
-            "--type",
-            type=str,
-            default="conda_env",
-            choices=[v.value for v in BuildType],
-        )
 
-    def _run_stack_build_command(self, args: argparse.Namespace) -> None:
+        from llama_toolchain.common.config_dirs import DISTRIBS_BASE_DIR
+        from llama_toolchain.common.serialize import EnumEncoder
         from llama_toolchain.core.distribution_registry import resolve_distribution_spec
-        from llama_toolchain.core.package import (
+        from llama_toolchain.core.package import ApiInput, build_package, ImageType
-            ApiInput,
+        from termcolor import cprint
-            BuildType,
-            build_package,
-        )
 
         api_inputs = []
-        if args.distribution == "adhoc":
+        if build_config.distribution == "adhoc":
-            if not args.api_providers:
+            if not build_config.api_providers:
-                self.parser.error("You must specify API providers with (api=provider,...) for building an adhoc distribution")
+                self.parser.error(
+                    "You must specify API providers with (api=provider,...) for building an adhoc distribution"
+                )
                 return
 
-            parsed = parse_api_provider_tuples(args.api_providers, self.parser)
+            parsed = parse_api_provider_tuples(build_config.api_providers, self.parser)
             for api, provider_spec in parsed.items():
                 for dep in provider_spec.api_dependencies:
                     if dep not in parsed:
-                        self.parser.error(f"API {api} needs dependency {dep} provided also")
+                        self.parser.error(
+                            f"API {api} needs dependency {dep} provided also"
+                        )
                         return
 
                 api_inputs.append(
@@ -106,13 +97,17 @@ class StackBuild(Subcommand):
                 )
             docker_image = None
         else:
-            if args.api_providers:
+            if build_config.api_providers:
-                self.parser.error("You cannot specify API providers for pre-registered distributions")
+                self.parser.error(
+                    "You cannot specify API providers for pre-registered distributions"
+                )
                 return
 
-            dist = resolve_distribution_spec(args.distribution)
+            dist = resolve_distribution_spec(build_config.distribution)
             if dist is None:
-                self.parser.error(f"Could not find distribution {args.distribution}")
+                self.parser.error(
+                    f"Could not find distribution {build_config.distribution}"
+                )
                 return
 
             for api, provider_type in dist.providers.items():
@@ -126,8 +121,41 @@ class StackBuild(Subcommand):
 
         build_package(
             api_inputs,
-            build_type=BuildType(args.type),
+            image_type=ImageType(build_config.image_type),
-            name=args.name,
+            name=build_config.name,
-            distribution_type=args.distribution,
+            distribution_type=build_config.distribution,
             docker_image=docker_image,
         )
+
+        # save build.yaml spec for building same distribution again
+        build_dir = (
+            DISTRIBS_BASE_DIR / build_config.distribution / build_config.image_type
+        )
+        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))
+
+        cprint(
+            f"Build spec configuration saved at {str(build_file_path)}",
+            color="green",
+        )
+
+    def _run_stack_build_command(self, args: argparse.Namespace) -> None:
+        from llama_toolchain.common.prompt_for_config import prompt_for_config
+        from llama_toolchain.core.dynamic import instantiate_class_type
+
+        if args.config:
+            with open(args.config, "r") as f:
+                try:
+                    build_config = BuildConfig(**yaml.safe_load(f))
+                except Exception as e:
+                    self.parser.error(f"Could not parse config file {args.config}: {e}")
+                    return
+                self._run_stack_build_command_from_build_config(build_config)
+            return
+
+        build_config = prompt_for_config(BuildConfig, None)
+        self._run_stack_build_command_from_build_config(build_config)

llama_toolchain/cli/stack/configure.py

@@ -9,10 +9,10 @@ import json
 from pathlib import Path
 
 import yaml
-from termcolor import cprint
 
 from llama_toolchain.cli.subcommand import Subcommand
 from llama_toolchain.common.config_dirs import BUILDS_BASE_DIR
+from termcolor import cprint
 from llama_toolchain.core.datatypes import *  # noqa: F403
 
 
@@ -34,38 +34,19 @@ class StackConfigure(Subcommand):
         from llama_toolchain.core.distribution_registry import (
             available_distribution_specs,
         )
-        from llama_toolchain.core.package import BuildType
+        from llama_toolchain.core.package import ImageType
 
         allowed_ids = [d.distribution_type for d in available_distribution_specs()]
         self.parser.add_argument(
-            "distribution",
+            "config",
             type=str,
-            help='Distribution ("adhoc" or one of: {})'.format(allowed_ids),
+            help="Path to the package config file (e.g. ~/.llama/builds/<distribution>/<image_type>/<name>.yaml)",
-        )
-        self.parser.add_argument(
-            "--name",
-            type=str,
-            help="Name of the build",
-            required=True,
-        )
-        self.parser.add_argument(
-            "--type",
-            type=str,
-            default="conda_env",
-            choices=[v.value for v in BuildType],
         )
 
     def _run_stack_configure_cmd(self, args: argparse.Namespace) -> None:
-        from llama_toolchain.core.package import BuildType
+        from llama_toolchain.core.package import ImageType
 
-        build_type = BuildType(args.type)
+        config_file = Path(args.config)
-        name = args.name
-        config_file = (
-            BUILDS_BASE_DIR
-            / args.distribution
-            / build_type.descriptor()
-            / f"{name}.yaml"
-        )
         if not config_file.exists():
             self.parser.error(
                 f"Could not find {config_file}. Please run `llama stack build` first"

llama_toolchain/cli/stack/run.py

@@ -29,24 +29,12 @@ class StackRun(Subcommand):
         self.parser.set_defaults(func=self._run_stack_run_cmd)
 
     def _add_arguments(self):
-        from llama_toolchain.core.package import BuildType
+        from llama_toolchain.core.package import ImageType
 
         self.parser.add_argument(
-            "distribution",
+            "config",
             type=str,
-            help="Distribution whose build you want to start",
+            help="Path to config file to use for the run",
-        )
-        self.parser.add_argument(
-            "--name",
-            type=str,
-            help="Name of the build you want to start",
-            required=True,
-        )
-        self.parser.add_argument(
-            "--type",
-            type=str,
-            default="conda_env",
-            choices=[v.value for v in BuildType],
         )
         self.parser.add_argument(
             "--port",
@@ -63,12 +51,13 @@ class StackRun(Subcommand):
 
     def _run_stack_run_cmd(self, args: argparse.Namespace) -> None:
         from llama_toolchain.common.exec import run_with_pty
-        from llama_toolchain.core.package import BuildType
+        from llama_toolchain.core.package import ImageType
 
-        build_type = BuildType(args.type)
+        if not args.config:
-        build_dir = BUILDS_BASE_DIR / args.distribution / build_type.descriptor()
+            self.parser.error("Must specify a config file to run")
-        path = build_dir / f"{args.name}.yaml"
+            return
 
+        path = args.config
         config_file = Path(path)
 
         if not config_file.exists():

llama_toolchain/core/build_conda_env.sh

@@ -19,7 +19,7 @@ fi
 
 set -euo pipefail
 
-if [ "$#" -ne 3 ]; then
+if [ "$#" -ne 4 ]; then
   echo "Usage: $0 <distribution_type> <build_name> <pip_dependencies>" >&2
   echo "Example: $0 <distribution_type> mybuild 'numpy pandas scipy'" >&2
   exit 1
@@ -28,7 +28,8 @@ fi
 distribution_type="$1"
 build_name="$2"
 env_name="llamastack-$build_name"
-pip_dependencies="$3"
+config_file="$3"
+pip_dependencies="$4"
 
 # Define color codes
 RED='\033[0;31m'
@@ -117,4 +118,4 @@ ensure_conda_env_python310 "$env_name" "$pip_dependencies"
 
 printf "${GREEN}Successfully setup conda environment. Configuring build...${NC}\n"
 
-$CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama stack configure $distribution_type --name "$build_name" --type conda_env
+$CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama stack configure $config_file

llama_toolchain/core/build_container.sh

@@ -4,7 +4,7 @@ LLAMA_MODELS_DIR=${LLAMA_MODELS_DIR:-}
 LLAMA_TOOLCHAIN_DIR=${LLAMA_TOOLCHAIN_DIR:-}
 TEST_PYPI_VERSION=${TEST_PYPI_VERSION:-}
 
-if [ "$#" -ne 4 ]; then
+if [ "$#" -ne 5 ]; then
   echo "Usage: $0 <distribution_type> <build_name> <docker_base> <pip_dependencies>
   echo "Example: $0 distribution_type my-fastapi-app python:3.9-slim 'fastapi uvicorn'
   exit 1
@@ -14,7 +14,8 @@ distribution_type=$1
 build_name="$2"
 image_name="llamastack-$build_name"
 docker_base=$3
-pip_dependencies=$4
+config_file=$4
+pip_dependencies=$5
 
 # Define color codes
 RED='\033[0;31m'
@@ -110,4 +111,4 @@ set +x
 printf "${GREEN}Succesfully setup Podman image. Configuring build...${NC}"
 echo "You can run it with: podman run -p 8000:8000 $image_name"
 
-$CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama stack configure $distribution_type --name "$build_name" --type container
+$CONDA_PREFIX/bin/python3 -m llama_toolchain.cli.llama stack configure $config_file

llama_toolchain/core/datatypes.py

@@ -188,3 +188,19 @@ Provider configurations for each of the APIs provided by this package. This incl
 the dependencies of these providers as well.
 """,
     )
+
+
+@json_schema_type
+class BuildConfig(BaseModel):
+    name: str
+    distribution: str = Field(
+        default="local", description="Type of distribution to build (adhoc | {})"
+    )
+    api_providers: Optional[str] = Field(
+        default_factory=list,
+        description="List of API provider names to build",
+    )
+    image_type: str = Field(
+        default="conda",
+        description="Type of package to build (conda | container)",
+    )

llama_toolchain/core/package.py

@@ -12,24 +12,21 @@ from typing import List, Optional
 
 import pkg_resources
 import yaml
-from pydantic import BaseModel
-
-from termcolor import cprint
 
 from llama_toolchain.common.config_dirs import BUILDS_BASE_DIR
 from llama_toolchain.common.exec import run_with_pty
 from llama_toolchain.common.serialize import EnumEncoder
+from pydantic import BaseModel
+
+from termcolor import cprint
 
 from llama_toolchain.core.datatypes import *  # noqa: F403
 from llama_toolchain.core.distribution import api_providers, SERVER_DEPENDENCIES
 
 
-class BuildType(Enum):
+class ImageType(Enum):
-    container = "container"
+    docker = "docker"
-    conda_env = "conda_env"
+    conda = "conda"
-
-    def descriptor(self) -> str:
-        return "docker" if self == self.container else "conda"
 
 
 class Dependencies(BaseModel):
@@ -44,7 +41,7 @@ class ApiInput(BaseModel):
 
 def build_package(
     api_inputs: List[ApiInput],
-    build_type: BuildType,
+    image_type: ImageType,
     name: str,
     distribution_type: Optional[str] = None,
     docker_image: Optional[str] = None,
@@ -52,7 +49,7 @@ def build_package(
     if not distribution_type:
         distribution_type = "adhoc"
 
-    build_dir = BUILDS_BASE_DIR / distribution_type / build_type.descriptor()
+    build_dir = BUILDS_BASE_DIR / distribution_type / image_type.value
     os.makedirs(build_dir, exist_ok=True)
 
     package_name = name.replace("::", "-")
@@ -106,14 +103,14 @@ def build_package(
         )
 
     c.distribution_type = distribution_type
-    c.docker_image = package_name if build_type == BuildType.container else None
+    c.docker_image = package_name if image_type == ImageType.docker else None
-    c.conda_env = package_name if build_type == BuildType.conda_env else None
+    c.conda_env = package_name if image_type == ImageType.conda else None
 
     with open(package_file, "w") as f:
         to_write = json.loads(json.dumps(c.dict(), cls=EnumEncoder))
         f.write(yaml.dump(to_write, sort_keys=False))
 
-    if build_type == BuildType.container:
+    if image_type == ImageType.docker:
         script = pkg_resources.resource_filename(
             "llama_toolchain", "core/build_container.sh"
         )
@@ -122,6 +119,7 @@ def build_package(
             distribution_type,
             package_name,
             package_deps.docker_image,
+            str(package_file),
             " ".join(package_deps.pip_packages),
         ]
     else:
@@ -132,6 +130,7 @@ def build_package(
             script,
             distribution_type,
             package_name,
+            str(package_file),
             " ".join(package_deps.pip_packages),
         ]