From 6cd3e4183f6574c7aa7171a6ea5bf66568e51ec0 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Mon, 30 Sep 2024 08:49:25 -0700
Subject: [PATCH] bake run.yaml inside docker, simplify run

---
 llama_stack/cli/stack/build.py              | 14 +++++---------
 llama_stack/cli/stack/configure.py          | 15 +++++----------
 llama_stack/distribution/build_container.sh | 11 +++++++++--
 3 files changed, 19 insertions(+), 21 deletions(-)

diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py
index 31cf991be..fc2341f5f 100644
--- a/llama_stack/cli/stack/build.py
+++ b/llama_stack/cli/stack/build.py
@@ -113,15 +113,11 @@ class StackBuild(Subcommand):
         if return_code != 0:
             return
 
-        configure_name = (
-            build_config.name
-            if build_config.image_type == "conda"
-            else (f"llamastack-{build_config.name}")
-        )
-        cprint(
-            f"You can now run `llama stack configure {configure_name}`",
-            color="green",
-        )
+        if build_config.image_type == ImageType.conda.value:
+            cprint(
+                f"You can now run `llama stack configure {build_config.name}`",
+                color="green",
+            )
 
     def _run_template_list_cmd(self, args: argparse.Namespace) -> None:
         import json
diff --git a/llama_stack/cli/stack/configure.py b/llama_stack/cli/stack/configure.py
index e8105b7e0..815ab2f38 100644
--- a/llama_stack/cli/stack/configure.py
+++ b/llama_stack/cli/stack/configure.py
@@ -112,12 +112,6 @@ class StackConfigure(Subcommand):
             )
             return
 
-        build_name = docker_image.removeprefix("llamastack-")
-        saved_file = str(builds_dir / f"{build_name}-run.yaml")
-        cprint(
-            f"YAML configuration has been written to {saved_file}. You can now run `llama stack run {saved_file}`",
-            color="green",
-        )
         return
 
     def _configure_llama_distribution(
@@ -173,7 +167,8 @@ class StackConfigure(Subcommand):
             color="blue",
         )
 
-        cprint(
-            f"You can now run `llama stack run {image_name} --port PORT`",
-            color="green",
-        )
+        if build_config.image_type == "conda":
+            cprint(
+                f"You can now run `llama stack run {image_name} --port PORT`",
+                color="green",
+            )
diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index 970da804e..2af46dd96 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -29,9 +29,12 @@ SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
 REPO_DIR=$(dirname $(dirname "$SCRIPT_DIR"))
 DOCKER_BINARY=${DOCKER_BINARY:-docker}
 DOCKER_OPTS=${DOCKER_OPTS:-}
+REPO_CONFIGS_DIR="$REPO_DIR/tmp/configs"
 
 TEMP_DIR=$(mktemp -d)
 
+llama stack configure $build_file_path --output-dir $REPO_CONFIGS_DIR
+
 add_to_docker() {
   local input
   output_file="$TEMP_DIR/Dockerfile"
@@ -99,11 +102,12 @@ add_to_docker <<EOF
 # This would be good in production but for debugging flexibility lets not add it right now
 # We need a more solid production ready entrypoint.sh anyway
 #
-# ENTRYPOINT ["python", "-m", "llama_stack.distribution.server.server"]
+CMD ["python", "-m", "llama_stack.distribution.server.server", "./llamastack-run.yaml"]
 
 EOF
 
 add_to_docker "ADD tmp/configs/$(basename "$build_file_path") ./llamastack-build.yaml"
+add_to_docker "ADD tmp/configs/$build_name-run.yaml ./llamastack-run.yaml"
 
 printf "Dockerfile created successfully in $TEMP_DIR/Dockerfile"
 cat $TEMP_DIR/Dockerfile
@@ -124,7 +128,10 @@ set -x
 $DOCKER_BINARY build $DOCKER_OPTS -t $image_name -f "$TEMP_DIR/Dockerfile" "$REPO_DIR" $mounts
 set +x
 
-echo "You can run it with: podman run -p 8000:8000 $image_name"
+echo "Success! You can run it with: $DOCKER_BINARY run -p 8000:8000 $image_name"
 
 echo "Checking image builds..."
 $DOCKER_BINARY run $DOCKER_OPTS -it $image_name cat llamastack-build.yaml
+
+echo "Checking image run..."
+$DOCKER_BINARY run $DOCKER_OPTS -it $image_name cat llamastack-run.yaml