Update build-bases workflow

.github/workflows/build-bases.yaml

@@ -4,8 +4,6 @@ on:
   workflow_dispatch:
     branches:
       - main
-      # TODO: Remove once merged to main
-      - cog-base-images
 
 concurrency:
   group: ${{ github.workflow }}
@@ -16,39 +14,43 @@ env:
   REGISTRY_NAMESPACE: replicate/cog
 
 jobs:
+  build-python:
+    name: Build & verify python package
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0 # needed for setuptools_scm version determination
+      - uses: hynek/build-and-inspect-python-package@v2
+
   generate-matrix:
     runs-on: ubuntu-latest
+    needs: build-python
     outputs:
       matrix: ${{ steps.matrix.outputs.matrix }}
     steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-
-      - name: Setup Go Runtime
-        uses: actions/setup-go@v5
+      - uses: actions/checkout@v4
         with:
-          go-version-file: go.mod
-
-      - uses: actions/setup-python@v5
+          fetch-depth: 0 # needed for goreleaser version determination
+      - name: Download pre-built packages
+        uses: actions/download-artifact@v4
         with:
-          python-version: 3.11
-      - name: Install Python dependencies
-        run: |
-          python -m pip install '.[dev]'
-
-      - name: 'Build Base-Image command'
-        run: |
-          make base-image
-
+          name: Packages
+          path: dist
+      - name: Set COG_WHEEL
+        run: echo COG_WHEEL=$(ls dist/*.whl) >>"$GITHUB_ENV"
+      - uses: actions/setup-go@v5
+        with:
+          go-version-file: go.mod
+      - name: Build
+        run: make base-image
       - id: matrix
         run: |
           MATRIX_JSON=$(./base-image generate-matrix)
           echo "matrix=$MATRIX_JSON" >> $GITHUB_OUTPUT
         shell: bash
-
-      - name: 'Check Matrix'
-        run: |
-          jq . <<< '${{ steps.matrix.outputs.matrix }}'
+      - name: Check matrix
+        run: jq . <<< '${{ steps.matrix.outputs.matrix }}'
 
   build:
     name: Build and Push Image
@@ -68,19 +70,26 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-        name: 'Checkout Repository'
-
-      - name: 'Set up Buildx'
+        with:
+          fetch-depth: 0 # needed for goreleaser version determination
+      - name: Set up Buildx
         uses: docker/setup-buildx-action@v3
         with:
           platforms: linux/amd64,linux/arm64
           driver-opts: |
             image=moby/buildkit:v0.13.1
-
-      - name: Setup Go Runtime
-        uses: actions/setup-go@v5
+      - name: Download pre-built packages
+        uses: actions/download-artifact@v4
+        with:
+          name: Packages
+          path: dist
+      - name: Set COG_WHEEL
+        run: echo COG_WHEEL=$(ls dist/*.whl) >>"$GITHUB_ENV"
+      - uses: actions/setup-go@v5
         with:
           go-version-file: go.mod
+      - name: Build
+        run: make base-image
 
       - name: 'Set Python Version'
         run: |
@@ -105,18 +114,6 @@ jobs:
           fi
         shell: bash
 
-      - uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python_version }}
-
-      - name: Install Python dependencies
-        run: |
-          python -m pip install '.[dev]'
-
-      - name: 'Build Base-Image command'
-        run: |
-          make base-image
-
       - name: 'Generate Dockerfile'
         run: |
           ./base-image dockerfile \