Split wheel upload and PyPI publish into separate optimized jobs so all matrix runners build wheels and a unified Linux job publishes them to PyPI

Author: vkarampudi

.github/workflows/build.yml

@@ -72,13 +72,11 @@ jobs:
       shell: bash -l {0}
       run: bazel run //:build_pip_package
 
-    - name: Publish Package to PyPI
-      if: github.event_name == 'push' && (startsWith(github.ref, 'refs/tags/v') || startsWith(github.ref, 'refs/heads/r'))
-      uses: pypa/gh-action-pypi-publish@release/v1
+    - name: Upload Wheel Artifact
+      uses: actions/upload-artifact@v4
       with:
-        password: ${{ secrets.PYPI_API_TOKEN }}
-        packages-dir: dist/
-        skip-existing: true
+        name: wheels-${{ matrix.os }}-${{ matrix.python-version }}
+        path: dist/*.whl
 
   docker-serving-build:
     name: Build Docker Serving Image
@@ -94,3 +92,23 @@ jobs:
       run: |
         cd struct2tensor/tools/tf_serving_docker
         docker build -t struct2tensor-serving:latest .
+
+  pypi-publish:
+    name: Publish Wheels to PyPI
+    needs: build
+    if: github.event_name == 'push' && (startsWith(github.ref, 'refs/tags/v') || startsWith(github.ref, 'refs/heads/r'))
+    runs-on: ubuntu-latest
+    steps:
+    - name: Download all wheels
+      uses: actions/download-artifact@v4
+      with:
+        pattern: wheels-*
+        path: dist/
+        merge-multiple: true
+
+    - name: Publish Package to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+      with:
+        password: ${{ secrets.PYPI_API_TOKEN }}
+        packages-dir: dist/
+        skip-existing: true