diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index d0874d99f..81d8811e5 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -140,6 +140,13 @@ if command -v selinuxenabled &>/dev/null && selinuxenabled; then
   DOCKER_OPTS="$DOCKER_OPTS --security-opt label=disable"
 fi
 
+if [ -n "$TEST_PYPI_VERSION" ]; then
+  image_name="$image_name-test-$TEST_PYPI_VERSION"
+else
+  URL="https://pypi.org/pypi/llama-stack/json"
+  image_name="$image_name-$(curl -s $URL | jq -r '.info.version')"
+fi
+
 set -x
 $DOCKER_BINARY build $DOCKER_OPTS -t $image_name -f "$TEMP_DIR/Dockerfile" "$REPO_DIR" $mounts