feat: Publish arm64 images as well

.github/workflows/publish-dockers.yaml

@@ -11,9 +11,14 @@ on:
         description: 'Optional comma-separated templates to publish'
         required: false
         type: string
+      dockerhub_organization:
+        description: 'DockerHub organization'
+        required: false
+        type: string
+        default: llamastack
 
 jobs:
-  publish-docker-images:
+  publish-docker-images-amd64:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
@@ -23,3 +28,34 @@ jobs:
         templates: ${{ inputs.templates }}
         dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
         dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+        dockerhub_organization: ${{ inputs.dockerhub_organization }}
+        arch: amd64
+
+  publish-docker-images-arm64:
+    runs-on: ubuntu-24.04-arm
+    steps:
+    - uses: actions/checkout@v4
+    - uses: ./actions/publish-dockers
+      with:
+        version: ${{ inputs.version }}
+        templates: ${{ inputs.templates }}
+        dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+        dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+        dockerhub_organization: ${{ inputs.dockerhub_organization }}
+        continue_on_error: true
+        arch: arm64
+
+  publish-docker-manifest-lists:
+    needs:
+    - publish-docker-images-amd64
+    - publish-docker-images-arm64
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: ./actions/publish-docker-manifest-lists
+      with:
+        version: ${{ inputs.version }}
+        templates: ${{ inputs.templates }}
+        dockerhub_username: ${{ secrets.DOCKERHUB_USERNAME }}
+        dockerhub_token: ${{ secrets.DOCKERHUB_TOKEN }}
+        dockerhub_organization: ${{ inputs.dockerhub_organization }}

actions/publish-docker-manifest-lists/action.yaml

@@ -0,0 +1,50 @@
+name: 'Publish Docker Manifest Lists'
+description: 'Publish Docker manifest lists; architecture-specific images must already be published'
+inputs:
+  version:
+    description: 'Version number (e.g. 0.1.1rc2, 0.1.1.dev20250201)'
+    required: true
+    type: string
+  templates:
+    description: 'Optional comma-separated templates to publish'
+    required: false
+    type: string
+  dockerhub_username:
+    description: 'DockerHub username'
+    required: true
+    type: string
+  dockerhub_token:
+    description: 'DockerHub token'
+    required: true
+    type: string
+  dockerhub_organization:
+    description: 'DockerHub organization'
+    required: false
+    type: string
+    default: llamastack
+runs:
+  using: 'composite'
+  steps:
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.11'
+    - uses: astral-sh/setup-uv@v6
+      with:
+        python-version: '3.11'
+        enable-cache: false
+
+    - uses: docker/setup-buildx-action@v3
+    - uses: docker/login-action@v3
+      with:
+        username: ${{ inputs.dockerhub_username }}
+        password: ${{ inputs.dockerhub_token }}
+
+    - name: Publish manifest lists
+      shell: bash
+      env:
+        VERSION: ${{ inputs.version }}
+        TEMPLATES: ${{ inputs.templates }}
+        DOCKERHUB_ORGANIZATION: ${{ inputs.dockerhub_organization }}
+      run: |
+        chmod +x ${{ github.action_path }}/main.sh
+        ${{ github.action_path }}/main.sh

actions/publish-docker-manifest-lists/main.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+set -euo pipefail
+
+THIS_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
+source $THIS_DIR/../publish-dockers/common.sh
+
+set -x
+
+push_manifest_list() {
+  template=$1
+
+  echo "Pushing multi-arch manifest list for template $template"
+  if [ "$PYPI_SOURCE" = "testpypi" ]; then
+    manifests="$DOCKERHUB_ORGANIZATION/distribution-$template:test-${VERSION}-amd64"
+    has_arm64_image=$(docker buildx imagetools inspect $DOCKERHUB_ORGANIZATION/distribution-$template:test-${VERSION}-arm64 > /dev/null 2>&1; echo $?)
+    if [ "$has_arm64_image" -eq 0 ]; then
+      manifests="$manifests $DOCKERHUB_ORGANIZATION/distribution-$template:test-${VERSION}-arm64"
+    fi
+    docker buildx imagetools create \
+      -t $DOCKERHUB_ORGANIZATION/distribution-$template:test-${VERSION} \
+      $manifests
+  else
+    manifests="$DOCKERHUB_ORGANIZATION/distribution-$template:${VERSION}-amd64"
+    has_arm64_image=$(docker buildx imagetools inspect $DOCKERHUB_ORGANIZATION/distribution-$template:${VERSION}-arm64 > /dev/null 2>&1; echo $?)
+    if [ "$has_arm64_image" -eq 0 ]; then
+      manifests="$manifests $DOCKERHUB_ORGANIZATION/distribution-$template:${VERSION}-arm64"
+    fi
+    docker buildx imagetools create \
+      -t $DOCKERHUB_ORGANIZATION/distribution-$template:${VERSION} \
+      $manifests
+    docker buildx imagetools create \
+      -t $DOCKERHUB_ORGANIZATION/distribution-$template:latest \
+      $manifests
+  fi
+}
+
+for template in "${TEMPLATES[@]}"; do
+  push_manifest_list $template
+done
+
+echo "Done"

actions/publish-dockers/action.yaml

@@ -17,6 +17,21 @@ inputs:
     description: 'DockerHub token'
     required: true
     type: string
+  dockerhub_organization:
+    description: 'DockerHub organization'
+    required: false
+    type: string
+    default: llamastack
+  continue_on_error:
+    description: 'Whether to continue or exit when individual templates fail to build'
+    required: false
+    type: string
+    default: false
+  arch:
+    description: Architecture we're building for (amd64 or arm64)
+    required: true
+    type: string
+    default: amd64
 runs:
   using: 'composite'
   steps:
@@ -39,6 +54,9 @@ runs:
       env:
         VERSION: ${{ inputs.version }}
         TEMPLATES: ${{ inputs.templates }}
+        DOCKERHUB_ORGANIZATION: ${{ inputs.dockerhub_organization }}
+        CONTINUE_ON_ERROR: ${{ inputs.continue_on_error }}
+        ARCH: ${{ inputs.arch }}
       run: |
         chmod +x ${{ github.action_path }}/main.sh
         ${{ github.action_path }}/main.sh

actions/publish-dockers/common.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+
+# This file is shared between the publish-dockers and
+# publish-docker-manifest-lists actions. It should only contain content
+# that is common to both actions.
+
+if [ -z "$VERSION" ]; then
+  echo "You must set the VERSION environment variable" >&2
+  exit 1
+fi
+
+DOCKERHUB_ORGANIZATION=${DOCKERHUB_ORGANIZATION:-llamastack}
+TEMPLATES=${TEMPLATES:-}
+
+if [ -z "$TEMPLATES" ]; then
+  TEMPLATES=(starter tgi meta-reference-gpu postgres-demo)
+else
+  TEMPLATES=(${TEMPLATES//,/ })
+fi
+
+release_exists() {
+  local source=$1
+  releases=$(curl -s https://${source}.org/pypi/llama-stack/json | jq -r '.releases | keys[]')
+  for release in $releases; do
+    if [ x"$release" = x"$VERSION" ]; then
+      return 0
+    fi
+  done
+  return 1
+}
+
+if release_exists "test.pypi"; then
+  echo "Version $VERSION found in test.pypi"
+  PYPI_SOURCE="testpypi"
+elif release_exists "pypi"; then
+  echo "Version $VERSION found in pypi"
+  PYPI_SOURCE="pypi"
+else
+  echo "Version $VERSION not found in either test.pypi or pypi" >&2
+  exit 1
+fi

actions/publish-dockers/main.sh

@@ -1,34 +1,15 @@
 #!/bin/bash
 
-if [ -z "$VERSION" ]; then
-  echo "You must set the VERSION environment variable" >&2
-  exit 1
-fi
-TEMPLATES=${TEMPLATES:-}
-
 set -euo pipefail
 
-release_exists() {
-  local source=$1
-  releases=$(curl -s https://${source}.org/pypi/llama-stack/json | jq -r '.releases | keys[]')
-  for release in $releases; do
-    if [ x"$release" = x"$VERSION" ]; then
-      return 0
-    fi
-  done
-  return 1
-}
+THIS_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
+source $THIS_DIR/common.sh
+
+CONTINUE_ON_ERROR=${CONTINUE_ON_ERROR:-}
+ARCH=${ARCH:-amd64}
 
-if release_exists "test.pypi"; then
-  echo "Version $VERSION found in test.pypi"
-  PYPI_SOURCE="testpypi"
-elif release_exists "pypi"; then
-  echo "Version $VERSION found in pypi"
-  PYPI_SOURCE="pypi"
-else
-  echo "Version $VERSION not found in either test.pypi or pypi" >&2
-  exit 1
-fi
+# `llama stack build` uses the `BUILD_PLATFORM` as the architecture to build
+export BUILD_PLATFORM="linux/$ARCH"
 
 set -x
 TMPDIR=$(mktemp -d)
@@ -44,35 +25,49 @@ uv pip install --index-url https://test.pypi.org/simple/ \
 which llama
 llama stack list-apis
 
+last_build_error="false"
+handle_build_error() {
+  template=$1
+  last_build_error="true"
+
+  echo "Error building template $template" >&2
+  if [ "$CONTINUE_ON_ERROR" = "true" ]; then
+    echo "Continuing on error" >&2
+  else
+    echo "Stopping on error" >&2
+    exit 1
+  fi
+}
+
 build_and_push_docker() {
   template=$1
+  last_build_error="false"
 
-  echo "Building and pushing docker for template $template"
+  echo "Building docker image for template $template and platform $BUILD_PLATFORM"
   if [ "$PYPI_SOURCE" = "testpypi" ]; then
-    TEST_PYPI_VERSION=${VERSION} llama stack build --template $template --image-type container
+    TEST_PYPI_VERSION=${VERSION} llama stack build --template $template --image-type container || handle_build_error $template
   else
-    PYPI_VERSION=${VERSION} llama stack build --template $template --image-type container
+    PYPI_VERSION=${VERSION} llama stack build --template $template --image-type container || handle_build_error $template
   fi
   docker images
 
-  echo "Pushing docker image"
+  if [ "$last_build_error" = "true" ]; then
+    echo "Skipping push for template $template because of build error"
+    return
+  fi
+
+  echo "Pushing docker image for template $template and platform $BUILD_PLATFORM"
   if [ "$PYPI_SOURCE" = "testpypi" ]; then
-    docker tag distribution-$template:test-${VERSION} llamastack/distribution-$template:test-${VERSION}
-    docker push llamastack/distribution-$template:test-${VERSION}
+    docker tag distribution-$template:test-${VERSION} $DOCKERHUB_ORGANIZATION/distribution-$template:test-${VERSION}-${ARCH}
+    docker push $DOCKERHUB_ORGANIZATION/distribution-$template:test-${VERSION}-${ARCH}
   else
-    docker tag distribution-$template:${VERSION} llamastack/distribution-$template:${VERSION}
-    docker tag distribution-$template:${VERSION} llamastack/distribution-$template:latest
-    docker push llamastack/distribution-$template:${VERSION}
-    docker push llamastack/distribution-$template:latest
+    docker tag distribution-$template:${VERSION} $DOCKERHUB_ORGANIZATION/distribution-$template:${VERSION}-${ARCH}
+    docker tag distribution-$template:${VERSION} $DOCKERHUB_ORGANIZATION/distribution-$template:latest-${ARCH}
+    docker push $DOCKERHUB_ORGANIZATION/distribution-$template:${VERSION}-${ARCH}
+    docker push $DOCKERHUB_ORGANIZATION/distribution-$template:latest-${ARCH}
   fi
 }
 
-if [ -z "$TEMPLATES" ]; then
-  TEMPLATES=(starter tgi meta-reference-gpu postgres-demo)
-else
-  TEMPLATES=(${TEMPLATES//,/ })
-fi
-
 for template in "${TEMPLATES[@]}"; do
   build_and_push_docker $template
 done