diff --git a/llama_stack/distribution/build.py b/llama_stack/distribution/build.py
index b8d35ccdc..950338730 100644
--- a/llama_stack/distribution/build.py
+++ b/llama_stack/distribution/build.py
@@ -119,12 +119,15 @@ def build_image(
     normal_deps += SERVER_DEPENDENCIES
 
     if build_config.image_type == ImageType.container.value:
+        if not template_name:
+            raise ValueError("template_name is required for container builds")
+
         script = str(
             importlib.resources.files("llama_stack") / "distribution/build_container.sh"
         )
         args = [
             script,
-            image_name,
+            template_name,
             container_image,
             str(build_file_path),
             str(BUILDS_BASE_DIR / ImageType.container.value),
diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index c7b6211f7..91c1dd1a6 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -12,9 +12,10 @@ TEST_PYPI_VERSION=${TEST_PYPI_VERSION:-}
 PYPI_VERSION=${PYPI_VERSION:-}
 BUILD_PLATFORM=${BUILD_PLATFORM:-}
 
-if [ "$#" -lt 4 ]; then
-  echo "Usage: $0 <build_name> <container_base> <pip_dependencies> [<special_pip_deps>]" >&2
-  echo "Example: $0 my-fastapi-app python:3.9-slim 'fastapi uvicorn' " >&2
+if [ "$#" -lt 5 ]; then
+  # This only works for templates
+  echo "Usage: $0 <template_name> <container_base> <pip_dependencies> <host_build_dir> [<special_pip_deps>]" >&2
+  echo "Example: $0 fireworks python:3.9-slim 'fastapi uvicorn' /path/to/build/dir" >&2
   exit 1
 fi
 
@@ -22,7 +23,7 @@ special_pip_deps="$6"
 
 set -euo pipefail
 
-build_name="$1"
+template_name="$1"
 container_base=$2
 build_file_path=$3
 host_build_dir=$4
@@ -151,7 +152,7 @@ add_to_container << 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", "--template", "$build_name"]
+ENTRYPOINT ["python", "-m", "llama_stack.distribution.server.server", "--template", "$template_name"]
 
 EOF
 
@@ -183,6 +184,7 @@ else
 fi
 
 # Add version tag to image name
+build_name="distribution-$template_name"
 image_tag="$build_name:$version_tag"
 
 # Detect platform architecture