diff --git a/.gitignore b/.gitignore
index db254fa7f..a4fac0e8b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,4 @@ dist
 *.egg-info
 dev_requirements.txt
 build
+.llama
diff --git a/llama_toolchain/cli/stack/build.py b/llama_toolchain/cli/stack/build.py
index 0785ee43d..d69f35408 100644
--- a/llama_toolchain/cli/stack/build.py
+++ b/llama_toolchain/cli/stack/build.py
@@ -8,6 +8,8 @@ import argparse
 
 from llama_toolchain.cli.subcommand import Subcommand
 from llama_toolchain.core.datatypes import *  # noqa: F403
+from pathlib import Path
+
 import yaml
 
 
@@ -47,10 +49,16 @@ class StackBuild(Subcommand):
         from llama_toolchain.core.package import ApiInput, build_package, ImageType
         from termcolor import cprint
 
-        build_package(build_config)
-
         # save build.yaml spec for building same distribution again
-        build_dir = DISTRIBS_BASE_DIR / build_config.image_type
+        if build_config.image_type == ImageType.docker.value:
+            # TODO (xiyan): docker needs build file to be in the llama-stack repo dir
+            build_dir = (
+                Path(os.path.expanduser("./.llama/distributions"))
+                / build_config.image_type
+            )
+        else:
+            build_dir = DISTRIBS_BASE_DIR / build_config.image_type
+
         os.makedirs(build_dir, exist_ok=True)
         build_file_path = build_dir / f"{build_config.name}-build.yaml"
 
@@ -58,6 +66,8 @@ class StackBuild(Subcommand):
             to_write = json.loads(json.dumps(build_config.dict(), cls=EnumEncoder))
             f.write(yaml.dump(to_write, sort_keys=False))
 
+        build_package(build_config, build_file_path)
+
         cprint(
             f"Build spec configuration saved at {str(build_file_path)}",
             color="green",
diff --git a/llama_toolchain/core/build_container.sh b/llama_toolchain/core/build_container.sh
index e297885c8..ce7e3438f 100755
--- a/llama_toolchain/core/build_container.sh
+++ b/llama_toolchain/core/build_container.sh
@@ -4,17 +4,17 @@ LLAMA_MODELS_DIR=${LLAMA_MODELS_DIR:-}
 LLAMA_TOOLCHAIN_DIR=${LLAMA_TOOLCHAIN_DIR:-}
 TEST_PYPI_VERSION=${TEST_PYPI_VERSION:-}
 
-if [ "$#" -ne 3 ]; then
+if [ "$#" -ne 4 ]; then
   echo "Usage: $0 <build_name> <docker_base> <pip_dependencies>
   echo "Example: $0 my-fastapi-app python:3.9-slim 'fastapi uvicorn'
   exit 1
 fi
 
-# distribution_type=$1
 build_name="$1"
 image_name="llamastack-$build_name"
 docker_base=$2
-pip_dependencies=$3
+build_file_path=$3
+pip_dependencies=$4
 
 # Define color codes
 RED='\033[0;31m'
@@ -25,6 +25,7 @@ set -euo pipefail
 
 SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
 REPO_DIR=$(dirname $(dirname "$SCRIPT_DIR"))
+BUILD_DIR=$(dirname $(dirname "$build_file_path"))
 
 TEMP_DIR=$(mktemp -d)
 
@@ -92,6 +93,8 @@ add_to_docker <<EOF
 
 EOF
 
+add_to_docker "ADD $build_file_path ./build.yaml"
+
 printf "Dockerfile created successfully in $TEMP_DIR/Dockerfile"
 cat $TEMP_DIR/Dockerfile
 printf "\n"
@@ -108,3 +111,6 @@ podman build --network host -t $image_name -f "$TEMP_DIR/Dockerfile" "$REPO_DIR"
 set +x
 
 echo "You can run it with: podman run -p 8000:8000 $image_name"
+
+echo "Checking image builds..."
+podman run -it llamastack-local-docker-example cat build.yaml
diff --git a/llama_toolchain/core/package.py b/llama_toolchain/core/package.py
index 7ea0265e4..7987384e2 100644
--- a/llama_toolchain/core/package.py
+++ b/llama_toolchain/core/package.py
@@ -21,6 +21,8 @@ from pydantic import BaseModel
 from termcolor import cprint
 
 from llama_toolchain.core.datatypes import *  # noqa: F403
+from pathlib import Path
+
 from llama_toolchain.core.distribution import api_providers, SERVER_DEPENDENCIES
 
 
@@ -39,7 +41,7 @@ class ApiInput(BaseModel):
     provider: str
 
 
-def build_package(build_config: BuildConfig):
+def build_package(build_config: BuildConfig, build_file_path: Path):
     package_deps = Dependencies(
         docker_image=build_config.distribution_spec.docker_image or "python:3.10-slim",
         pip_packages=SERVER_DEPENDENCIES,
@@ -67,6 +69,7 @@ def build_package(build_config: BuildConfig):
             script,
             build_config.name,
             package_deps.docker_image,
+            str(build_file_path),
             " ".join(package_deps.pip_packages),
         ]
     else: