diff --git a/.github/workflows/publish-to-docker.yml b/.github/workflows/publish-to-docker.yml
index f63f52cbd..1010041b7 100644
--- a/.github/workflows/publish-to-docker.yml
+++ b/.github/workflows/publish-to-docker.yml
@@ -11,6 +11,10 @@ on:
 jobs:
   build-and-push:
     runs-on: ubuntu-latest
+    env:
+      TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
+      FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }}
+      TAVILY_SEARCH_API_KEY: ${{ secrets.TAVILY_SEARCH_API_KEY }}
     permissions:
       contents: read
       packages: write
@@ -32,7 +36,7 @@ jobs:
         id: version
         run: |
           if [ "${{ github.event_name }}" = "push" ]; then
-            echo "VERSION=0.0.63.dev20250114" >> $GITHUB_OUTPUT
+            echo "VERSION=0.0.63.dev51206766" >> $GITHUB_OUTPUT
           else
             echo "VERSION=${{ inputs.version }}" >> $GITHUB_OUTPUT
           fi
@@ -85,6 +89,41 @@ jobs:
         run: |
           docker images
 
+      # TODO (xiyan): make the following 2 steps into a matrix and test all templates other than fireworks
+      - name: Start up built docker image
+        run: |
+          cd distributions/fireworks
+          if [ "$PYPI_SOURCE" = "testpypi" ]; then
+            sed -i 's|image: llamastack/distribution-fireworks|image: distribution-fireworks:test-${{ steps.version.outputs.version }}|' ./compose.yaml
+          else
+            sed -i 's|image: llamastack/distribution-fireworks|image: distribution-fireworks:${{ steps.version.outputs.version }}|' ./compose.yaml
+          fi
+          docker compose up -d
+          cd ..
+          # Wait for the container to start
+          timeout=300
+          while ! curl -s -f http://localhost:8321/v1/version > /dev/null && [ $timeout -gt 0 ]; do
+            echo "Waiting for endpoint to be available..."
+            sleep 5
+            timeout=$((timeout - 5))
+          done
+
+          if [ $timeout -le 0 ]; then
+            echo "Timeout waiting for endpoint to become available"
+            exit 1
+          fi
+
+      - name: Run simple models list test on docker server
+        run: |
+          curl http://localhost:8321/v1/models
+
+      # TODO (xiyan): figure out why client cannot find server but curl works
+      # - name: Run pytest on docker server
+      #   run: |
+      #     pip install pytest pytest-md-report
+      #     export LLAMA_STACK_BASE_URL="http://localhost:8321"
+      #     LLAMA_STACK_BASE_URL="http://localhost:8321" pytest -v tests/client-sdk/inference/test_inference.py --md-report --md-report-verbose=1
+
       - name: Push to dockerhub
         run: |
           TEMPLATES=("ollama" "bedrock" "remote-vllm" "fireworks" "together" "tgi" "meta-reference-gpu")
diff --git a/distributions/fireworks/compose.yaml b/distributions/fireworks/compose.yaml
index 4b53fcf00..84b8491e4 100644
--- a/distributions/fireworks/compose.yaml
+++ b/distributions/fireworks/compose.yaml
@@ -1,13 +1,11 @@
 services:
   llamastack:
     image: llamastack/distribution-fireworks
-    network_mode: "host"
-    volumes:
-      - ~/.llama:/root/.llama
-      - ./run.yaml:/root/llamastack-run-fireworks.yaml
     ports:
       - "8321:8321"
-    entrypoint: bash -c "python -m llama_stack.distribution.server.server --yaml_config /root/llamastack-run-fireworks.yaml"
+    environment:
+      - FIREWORKS_API_KEY=${FIREWORKS_API_KEY}
+    entrypoint: bash -c "python -m llama_stack.distribution.server.server --template fireworks"
     deploy:
       restart_policy:
         condition: on-failure
diff --git a/distributions/together/compose.yaml b/distributions/together/compose.yaml
index c7251d0a7..f66ee69f9 100644
--- a/distributions/together/compose.yaml
+++ b/distributions/together/compose.yaml
@@ -1,13 +1,11 @@
 services:
   llamastack:
     image: llamastack/distribution-together
-    network_mode: "host"
-    volumes:
-      - ~/.llama:/root/.llama
-      - ./run.yaml:/root/llamastack-run-together.yaml
     ports:
       - "8321:8321"
-    entrypoint: bash -c "python -m llama_stack.distribution.server.server --yaml_config /root/llamastack-run-together.yaml"
+    environment:
+      - TOGETHER_API_KEY=${TOGETHER_API_KEY}
+    entrypoint: bash -c "python -m llama_stack.distribution.server.server --template together"
     deploy:
       restart_policy:
         condition: on-failure
diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index 4c2425004..c7b6211f7 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -23,7 +23,6 @@ special_pip_deps="$6"
 set -euo pipefail
 
 build_name="$1"
-image_name="distribution-$build_name"
 container_base=$2
 build_file_path=$3
 host_build_dir=$4
@@ -184,7 +183,7 @@ else
 fi
 
 # Add version tag to image name
-image_tag="$image_name:$version_tag"
+image_tag="$build_name:$version_tag"
 
 # Detect platform architecture
 ARCH=$(uname -m)