[CLI] remove dependency on CONDA_PREFIX in CLI (meta-llama#144)

yanxi0830 · web-flow · commit 6a8c2ae1df5b · 2024-09-28T16:46:47.000-07:00
* remove dependency on CONDA_PREFIX in CLI

* lint

* typo

* more robust
diff --git a/llama_stack/cli/stack/build.py b/llama_stack/cli/stack/build.py
@@ -100,10 +100,7 @@ def _run_stack_build_command_from_build_config(
                 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 @@ def _run_stack_build_command_from_build_config(
         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
@@ -65,18 +65,27 @@ def _run_stack_configure_cmd(self, args: argparse.Namespace) -> None:
             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
@@ -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
@@ -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"

Original file line number	Diff line number	Diff line change
`@@ -92,6 +92,7 @@ def build_image(build_config: BuildConfig, build_file_path: Path):`
`92`	`92`	`args = [`
`93`	`93`	`script,`
`94`	`94`	`build_config.name,`
	`95`	`+ str(build_file_path),`
`95`	`96`	`" ".join(deps),`
`96`	`97`	`]`
`97`	`98`