Fix stack start

This commit is contained in:
Ashwin Bharambe 2024-08-30 15:03:23 -07:00
parent 5172d9a79d
commit d99c06fce8

View file

@ -13,6 +13,7 @@ import yaml
from llama_toolchain.cli.subcommand import Subcommand from llama_toolchain.cli.subcommand import Subcommand
from llama_toolchain.core.datatypes import * # noqa: F403 from llama_toolchain.core.datatypes import * # noqa: F403
from llama_toolchain.common.config_dirs import BUILDS_BASE_DIR
class StackStart(Subcommand): class StackStart(Subcommand):
@ -28,10 +29,23 @@ class StackStart(Subcommand):
self.parser.set_defaults(func=self._run_stack_start_cmd) self.parser.set_defaults(func=self._run_stack_start_cmd)
def _add_arguments(self): def _add_arguments(self):
from llama_toolchain.core.package import BuildType
self.parser.add_argument( self.parser.add_argument(
"yaml_config", "distribution",
type=str, type=str,
help="Yaml config containing the API build configuration", help="Distribution whose build you want to start",
)
self.parser.add_argument(
"--build-name",
type=str,
help="Name of the API build you want to start",
)
self.parser.add_argument(
"--build-type",
type=str,
default="conda_env",
choices=[v.value for v in BuildType],
) )
self.parser.add_argument( self.parser.add_argument(
"--port", "--port",
@ -48,13 +62,16 @@ class StackStart(Subcommand):
def _run_stack_start_cmd(self, args: argparse.Namespace) -> None: def _run_stack_start_cmd(self, args: argparse.Namespace) -> None:
from llama_toolchain.common.exec import run_with_pty from llama_toolchain.common.exec import run_with_pty
from llama_toolchain.core.package import BuildType
config_file = Path(args.yaml_config) if args.build_name.endswith(".yaml"):
if not config_file.exists(): path = args.build_name
self.parser.error( else:
f"Could not find {config_file}. Please run `llama stack build` first" build_type = BuildType(args.build_type)
) build_dir = BUILDS_BASE_DIR / args.distribution / build_type.descriptor()
return path = build_dir / f"{args.build_name}.yaml"
config_file = Path(path)
with open(config_file, "r") as f: with open(config_file, "r") as f:
config = PackageConfig(**yaml.safe_load(f)) config = PackageConfig(**yaml.safe_load(f))