chore: make cprint write to stderr (#2250)

Also do sys.exit(1) in case of errors
This commit is contained in:
raghotham 2025-05-24 23:39:57 -07:00 committed by GitHub
parent c25bd0ad58
commit 5a422e236c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 81 additions and 44 deletions

View file

@ -79,6 +79,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
cprint(
f"Could not find template {args.template}. Please run `llama stack build --list-templates` to check out the available templates",
color="red",
file=sys.stderr,
)
sys.exit(1)
build_config = available_templates[args.template]
@ -88,6 +89,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
cprint(
f"Please specify a image-type ({' | '.join(e.value for e in ImageType)}) for {args.template}",
color="red",
file=sys.stderr,
)
sys.exit(1)
elif args.providers:
@ -97,6 +99,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
cprint(
"Could not parse `--providers`. Please ensure the list is in the format api1=provider1,api2=provider2",
color="red",
file=sys.stderr,
)
sys.exit(1)
api, provider = api_provider.split("=")
@ -105,6 +108,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
cprint(
f"{api} is not a valid API.",
color="red",
file=sys.stderr,
)
sys.exit(1)
if provider in providers_for_api:
@ -113,6 +117,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
cprint(
f"{provider} is not a valid provider for the {api} API.",
color="red",
file=sys.stderr,
)
sys.exit(1)
distribution_spec = DistributionSpec(
@ -123,6 +128,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
cprint(
f"Please specify a image-type (container | conda | venv) for {args.template}",
color="red",
file=sys.stderr,
)
sys.exit(1)
@ -151,12 +157,14 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
cprint(
f"No current conda environment detected or specified, will create a new conda environment with the name `llamastack-{name}`",
color="yellow",
file=sys.stderr,
)
image_name = f"llamastack-{name}"
else:
cprint(
f"Using conda environment {image_name}",
color="green",
file=sys.stderr,
)
else:
image_name = f"llamastack-{name}"
@ -169,9 +177,10 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
""",
),
color="green",
file=sys.stderr,
)
print("Tip: use <TAB> to see options for the providers.\n")
cprint("Tip: use <TAB> to see options for the providers.\n", color="green", file=sys.stderr)
providers = dict()
for api, providers_for_api in get_provider_registry().items():
@ -213,6 +222,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
cprint(
f"Could not parse config file {args.config}: {e}",
color="red",
file=sys.stderr,
)
sys.exit(1)
@ -239,14 +249,17 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
cprint(
f"Error building stack: {exc}",
color="red",
file=sys.stderr,
)
cprint("Stack trace:", color="red")
cprint("Stack trace:", color="red", file=sys.stderr)
traceback.print_exc()
sys.exit(1)
if run_config is None:
cprint(
"Run config path is empty",
color="red",
file=sys.stderr,
)
sys.exit(1)
@ -304,6 +317,7 @@ def _generate_run_config(
cprint(
f"Failed to import provider {provider_type} for API {api} - assuming it's external, skipping",
color="yellow",
file=sys.stderr,
)
# Set config_type to None to avoid UnboundLocalError
config_type = None
@ -331,10 +345,7 @@ def _generate_run_config(
# For non-container builds, the run.yaml is generated at the very end of the build process so it
# makes sense to display this message
if build_config.image_type != LlamaStackImageType.CONTAINER.value:
cprint(
f"You can now run your stack with `llama stack run {run_config_file}`",
color="green",
)
cprint(f"You can now run your stack with `llama stack run {run_config_file}`", color="green", file=sys.stderr)
return run_config_file
@ -372,7 +383,7 @@ def _run_stack_build_command_from_build_config(
# Generate the run.yaml so it can be included in the container image with the proper entrypoint
# Only do this if we're building a container image and we're not using a template
if build_config.image_type == LlamaStackImageType.CONTAINER.value and not template_name and config_path:
cprint("Generating run.yaml file", color="green")
cprint("Generating run.yaml file", color="yellow", file=sys.stderr)
run_config_file = _generate_run_config(build_config, build_dir, image_name)
with open(build_file_path, "w") as f:
@ -396,11 +407,13 @@ def _run_stack_build_command_from_build_config(
run_config_file = build_dir / f"{template_name}-run.yaml"
shutil.copy(path, run_config_file)
cprint("Build Successful!", color="green")
cprint("You can find the newly-built template here: " + colored(template_path, "light_blue"))
cprint("Build Successful!", color="green", file=sys.stderr)
cprint(f"You can find the newly-built template here: {template_path}", color="light_blue", file=sys.stderr)
cprint(
"You can run the new Llama Stack distro via: "
+ colored(f"llama stack run {template_path} --image-type {build_config.image_type}", "light_blue")
+ colored(f"llama stack run {template_path} --image-type {build_config.image_type}", "light_blue"),
color="green",
file=sys.stderr,
)
return template_path
else: