From 3ffe51052a0816201637933f38cc6a052cc0ca78 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Sat, 28 Sep 2024 16:21:50 -0700
Subject: [PATCH] remove dependency on CONDA_PREFIX in CLI

---
 llama_stack/cli/stack/build.py              | 10 +---------
 llama_stack/distribution/build.py           |  1 +
 llama_stack/distribution/build_conda_env.sh | 12 ++++++++----
 3 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py
index 2b5b432c8..528aa290a 100644
--- a/llama_stack/cli/stack/build.py
+++ b/llama_stack/cli/stack/build.py
@@ -100,10 +100,7 @@ class StackBuild(Subcommand):
                 llama_stack_path / "tmp/configs/"
             )
         else:
-            build_dir = (
-                Path(os.getenv("CONDA_PREFIX")).parent
-                / f"llamastack-{build_config.name}"
-            )
+            build_dir = DISTRIBS_BASE_DIR / f"llamastack-{build_config.name}"
 
         os.makedirs(build_dir, exist_ok=True)
         build_file_path = build_dir / f"{build_config.name}-build.yaml"
@@ -116,11 +113,6 @@ class StackBuild(Subcommand):
         if return_code != 0:
             return
 
-        cprint(
-            f"Build spec configuration saved at {str(build_file_path)}",
-            color="blue",
-        )
-
         configure_name = (
             build_config.name
             if build_config.image_type == "conda"
diff --git a/llama_stack/distribution/build.py b/llama_stack/distribution/build.py
index 828311ea8..1047c6418 100644
--- a/llama_stack/distribution/build.py
+++ b/llama_stack/distribution/build.py
@@ -92,6 +92,7 @@ def build_image(build_config: BuildConfig, build_file_path: Path):
         args = [
             script,
             build_config.name,
+            str(build_file_path),
             " ".join(deps),
         ]
 
diff --git a/llama_stack/distribution/build_conda_env.sh b/llama_stack/distribution/build_conda_env.sh
index 65b2a8c0e..10218beff 100755
--- a/llama_stack/distribution/build_conda_env.sh
+++ b/llama_stack/distribution/build_conda_env.sh
@@ -17,9 +17,9 @@ if [ -n "$LLAMA_MODELS_DIR" ]; then
   echo "Using llama-models-dir=$LLAMA_MODELS_DIR"
 fi
 
-if [ "$#" -lt 2 ]; then
-  echo "Usage: $0 <distribution_type> <build_name> <pip_dependencies> [<special_pip_deps>]" >&2
-  echo "Example: $0 <distribution_type> mybuild 'numpy pandas scipy'" >&2
+if [ "$#" -lt 3 ]; then
+  echo "Usage: $0 <distribution_type> <build_name> <build_file_path> <pip_dependencies> [<special_pip_deps>]" >&2
+  echo "Example: $0 <distribution_type> mybuild ./my-stack-run.yaml 'numpy pandas scipy'" >&2
   exit 1
 fi
 
@@ -29,7 +29,8 @@ set -euo pipefail
 
 build_name="$1"
 env_name="llamastack-$build_name"
-pip_dependencies="$2"
+build_file_path="$2"
+pip_dependencies="$3"
 
 # Define color codes
 RED='\033[0;31m'
@@ -123,6 +124,9 @@ ensure_conda_env_python310() {
       done
     fi
   fi
+
+  mv $build_file_path $CONDA_PREFIX/
+  echo "Build spec configuration saved at $CONDA_PREFIX/$build_name-build.yaml"
 }
 
 ensure_conda_env_python310 "$env_name" "$pip_dependencies" "$special_pip_deps"