From c95b34d9044a9c24c25311091b664aa260e2df9b Mon Sep 17 00:00:00 2001
From: remsky <jeremy.braun@ucalgary.ca>
Date: Fri, 31 Jan 2025 01:29:56 -0700
Subject: [PATCH] Refactor Docker workflow to support split jobs

---
 .github/workflows/docker-publish.yml | 113 ++++++++++++++++++++++++---
 1 file changed, 104 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml
index fbd6e3e..6a6eb68 100644
--- a/.github/workflows/docker-publish.yml
+++ b/.github/workflows/docker-publish.yml
@@ -6,20 +6,54 @@ on:
   workflow_dispatch:
 
 jobs:
-  build:
+  prepare-release:
+    runs-on: ubuntu-latest
+    outputs:
+      version: ${{ steps.get-version.outputs.version }}
+      is_prerelease: ${{ steps.check-prerelease.outputs.is_prerelease }}
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+      
+      - name: Get version
+        id: get-version
+        run: echo "version=$(cat VERSION)" >> $GITHUB_OUTPUT
+      
+      - name: Check if prerelease
+        id: check-prerelease
+        run: echo "is_prerelease=${{ contains(github.ref, '-pre') }}" >> $GITHUB_OUTPUT
+
+  build-cpu:
+    needs: prepare-release
     runs-on: ubuntu-latest
     permissions:
-      contents: write
       packages: write
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
 
+      - name: Free disk space
+        run: |
+          echo "Listing current disk space"
+          df -h
+          echo "Cleaning up disk space..."
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache
+          docker system prune -af
+          echo "Disk space after cleanup"
+          df -h
+
       - name: Set up QEMU
         uses: docker/setup-qemu-action@v2
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v2
+        with:
+          driver-opts: |
+            image=moby/buildkit:latest
+            network=host
 
       - name: Log in to GitHub Container Registry
         uses: docker/login-action@v2
@@ -28,17 +62,78 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: Run build script
+      - name: Build and push CPU image
+        env:
+          DOCKER_BUILDKIT: 1
+          BUILDKIT_STEP_LOG_MAX_SIZE: 10485760
         run: |
-          chmod +x docker/build.sh
-          VERSION=$(cat VERSION)
-          docker/build.sh $VERSION
+          docker buildx build --platform linux/amd64,linux/arm64 \
+            -t ghcr.io/remsky/kokoro-fastapi-cpu:${{ needs.prepare-release.outputs.version }} \
+            -t ghcr.io/remsky/kokoro-fastapi-cpu:latest \
+            -f docker/cpu/Dockerfile \
+            --push \
+            --build-arg BUILDKIT_INLINE_CACHE=1 .
 
+  build-gpu:
+    needs: prepare-release
+    runs-on: ubuntu-latest
+    permissions:
+      packages: write
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Free disk space
+        run: |
+          echo "Listing current disk space"
+          df -h
+          echo "Cleaning up disk space..."
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /opt/hostedtoolcache
+          docker system prune -af
+          echo "Disk space after cleanup"
+          df -h
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v2
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+        with:
+          driver-opts: |
+            image=moby/buildkit:latest
+            network=host
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v2
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Build and push GPU image
+        env:
+          DOCKER_BUILDKIT: 1
+          BUILDKIT_STEP_LOG_MAX_SIZE: 10485760
+        run: |
+          docker buildx build --platform linux/amd64,linux/arm64 \
+            -t ghcr.io/remsky/kokoro-fastapi-gpu:${{ needs.prepare-release.outputs.version }} \
+            -t ghcr.io/remsky/kokoro-fastapi-gpu:latest \
+            -f docker/gpu/Dockerfile \
+            --push \
+            --build-arg BUILDKIT_INLINE_CACHE=1 .
+
+  create-release:
+    needs: [prepare-release, build-cpu, build-gpu]
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+    steps:
       - name: Create Release
         uses: softprops/action-gh-release@v1
-        env:
-          IS_PRERELEASE: ${{ contains(github.ref, '-pre') }}
         with:
           generate_release_notes: true
           draft: true
-          prerelease: ${{ contains(github.ref, '-pre') }}
+          prerelease: ${{ needs.prepare-release.outputs.is_prerelease }}