From 6a8c2ae1df5b2c7115c12ff7483811c466077568 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Sat, 28 Sep 2024 16:46:47 -0700
Subject: [PATCH] [CLI] remove dependency on CONDA_PREFIX in CLI (#144)

* remove dependency on CONDA_PREFIX in CLI

* lint

* typo

* more robust
---
 llama_stack/cli/stack/build.py              | 10 +--------
 llama_stack/cli/stack/configure.py          | 23 ++++++++++++++-------
 llama_stack/distribution/build.py           |  1 +
 llama_stack/distribution/build_conda_env.sh | 12 +++++++----
 4 files changed, 26 insertions(+), 20 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/cli/stack/configure.py b/llama_stack/cli/stack/configure.py
index 5b1fbba86..e8105b7e0 100644
--- a/llama_stack/cli/stack/configure.py
+++ b/llama_stack/cli/stack/configure.py
@@ -65,18 +65,27 @@ class StackConfigure(Subcommand):
             f"Could not find {build_config_file}. Trying conda build name instead...",
             color="green",
         )
-        if os.getenv("CONDA_PREFIX"):
+        if os.getenv("CONDA_PREFIX", ""):
             conda_dir = (
                 Path(os.getenv("CONDA_PREFIX")).parent / f"llamastack-{args.config}"
             )
-            build_config_file = Path(conda_dir) / f"{args.config}-build.yaml"
+        else:
+            cprint(
+                "Cannot find CONDA_PREFIX. Trying default conda path ~/.conda/envs...",
+                color="green",
+            )
+            conda_dir = (
+                Path(os.path.expanduser("~/.conda/envs")) / f"llamastack-{args.config}"
+            )
 
-            if build_config_file.exists():
-                with open(build_config_file, "r") as f:
-                    build_config = BuildConfig(**yaml.safe_load(f))
+        build_config_file = Path(conda_dir) / f"{args.config}-build.yaml"
 
-                self._configure_llama_distribution(build_config, args.output_dir)
-                return
+        if build_config_file.exists():
+            with open(build_config_file, "r") as f:
+                build_config = BuildConfig(**yaml.safe_load(f))
+
+            self._configure_llama_distribution(build_config, args.output_dir)
+            return
 
         # if we get here, we need to try to find the docker image
         cprint(
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..2a5205f79 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-build.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"