From 649efc2f45852e670f3adcf014306a69500f07c7 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Tue, 5 Nov 2024 11:00:37 -0800
Subject: [PATCH] build->run

---
 llama_stack/cli/stack/build.py     | 19 ++++++++++++++-----
 llama_stack/cli/stack/configure.py |  2 --
 llama_stack/cli/stack/run.py       |  4 +---
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py
index d2dae9319..52cb48537 100644
--- a/llama_stack/cli/stack/build.py
+++ b/llama_stack/cli/stack/build.py
@@ -189,16 +189,24 @@ class StackBuild(Subcommand):
         import yaml
         from termcolor import cprint
 
+        from llama_stack.distribution.build import ImageType
+
         from llama_stack.distribution.utils.serialize import EnumEncoder
 
-        # TODO: we should make the run.yaml file invisible to users by with ENV variables
-        # but keeping it visible and exposed to users for now
-        # generate a default run.yaml file using the build_config
         apis = list(build_config.distribution_spec.providers.keys())
         run_config = StackRunConfig(
             built_at=datetime.now(),
+            docker_image=(
+                build_config.name
+                if build_config.image_type == ImageType.docker.value
+                else None
+            ),
             image_name=build_config.name,
-            conda_env=build_config.name,
+            conda_env=(
+                build_config.name
+                if build_config.image_type == ImageType.conda.value
+                else None
+            ),
             apis=apis,
             providers={},
         )
@@ -223,7 +231,9 @@ class StackBuild(Subcommand):
                 p_spec.config = config_type()
                 run_config.providers[api].append(p_spec)
 
+        os.makedirs(build_dir, exist_ok=True)
         run_config_file = build_dir / f"{build_config.name}-run.yaml"
+
         with open(run_config_file, "w") as f:
             to_write = json.loads(json.dumps(run_config.model_dump(), cls=EnumEncoder))
             f.write(yaml.dump(to_write, sort_keys=False))
@@ -240,7 +250,6 @@ class StackBuild(Subcommand):
         import os
 
         import yaml
-        from termcolor import cprint
 
         from llama_stack.distribution.build import build_image, ImageType
         from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR
diff --git a/llama_stack/cli/stack/configure.py b/llama_stack/cli/stack/configure.py
index e5c6fb61b..7f3026930 100644
--- a/llama_stack/cli/stack/configure.py
+++ b/llama_stack/cli/stack/configure.py
@@ -7,8 +7,6 @@
 import argparse
 
 from llama_stack.cli.subcommand import Subcommand
-from llama_stack.distribution.utils.config_dirs import BUILDS_BASE_DIR
-from llama_stack.distribution.datatypes import *  # noqa: F403
 
 
 class StackConfigure(Subcommand):
diff --git a/llama_stack/cli/stack/run.py b/llama_stack/cli/stack/run.py
index dd4247e4b..9cd0314f3 100644
--- a/llama_stack/cli/stack/run.py
+++ b/llama_stack/cli/stack/run.py
@@ -45,7 +45,6 @@ class StackRun(Subcommand):
 
         import pkg_resources
         import yaml
-        from termcolor import cprint
 
         from llama_stack.distribution.build import ImageType
         from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
@@ -71,11 +70,10 @@ class StackRun(Subcommand):
 
         if not config_file.exists():
             self.parser.error(
-                f"File {str(config_file)} does not exist. Please run `llama stack build` and `llama stack configure <name>` to generate a run.yaml file"
+                f"File {str(config_file)} does not exist. Please run `llama stack build` to generate a run.yaml file"
             )
             return
 
-        cprint(f"Using config `{config_file}`", "green")
         with open(config_file, "r") as f:
             config_dict = yaml.safe_load(config_file.read_text())
             config = parse_and_maybe_upgrade_config(config_dict)