ci: add PyPI nightly wheel publish to nightly workflow

.github/workflows/nightly.yml

@@ -43,31 +43,86 @@ jobs:
       - name: Build CI image
         run: $CT build --network=host --build-arg BASE_IMAGE=${{ matrix.base_image }} -t $IMAGE -f docker/Dockerfile.dev .
 
-      - name: Build wheel
+      - name: Build wheels
         run: |
           BASE_VER=$(git describe --tags --abbrev=0 | sed 's/^v//')
           DATE=$(date +%Y%m%d)
           SHA=$(git rev-parse --short HEAD)
-          VERSION="${BASE_VER}.dev${DATE}+${SHA}"
+          GH_VERSION="${BASE_VER}.dev${DATE}+${SHA}"
+          PYPI_VERSION="${BASE_VER}.dev${DATE}"
 
           $CT run --rm --network=host \
             -v $GITHUB_WORKSPACE:$GITHUB_WORKSPACE -w $GITHUB_WORKSPACE \
             -v ${{ runner.temp }}/dist:/tmp/dist \
             $IMAGE bash -c "
               git config --global --add safe.directory '*'
-              echo 'Building wheel for Python ${{ matrix.python }}, version=${VERSION}'
-              SETUPTOOLS_SCM_PRETEND_VERSION=${VERSION} \
+
+              echo '=== Building gh-pages wheel (amd_mori ${GH_VERSION}) ==='
+              SETUPTOOLS_SCM_PRETEND_VERSION=${GH_VERSION} \
+                python3 -m pip wheel --no-deps -w /tmp/dist .
+
+              echo '=== Building PyPI wheel (amd_mori_nightly ${PYPI_VERSION}) ==='
+              sed -i 's/name = \"amd_mori\"/name = \"amd_mori_nightly\"/' pyproject.toml
+              SETUPTOOLS_SCM_PRETEND_VERSION=${PYPI_VERSION} \
                 python3 -m pip wheel --no-deps -w /tmp/dist .
+              git checkout pyproject.toml
             "
 
           echo "=== produced ==="
           ls -lh ${{ runner.temp }}/dist/
 
-      - name: Upload wheel
+      - name: Retag PyPI wheels (manylinux platform tag)
+        run: |
+          mkdir -p ${{ runner.temp }}/dist-pypi
+          cat > ${{ runner.temp }}/retag.py << 'PYEOF'
+          import zipfile, shutil, os, sys, tempfile
+          src, dst, plat = sys.argv[1], sys.argv[2], sys.argv[3]
+          shutil.copy2(src, dst)
+          with zipfile.ZipFile(dst, 'r') as zin:
+              wheel_file = [n for n in zin.namelist() if n.endswith('/WHEEL')][0]
+              content = zin.read(wheel_file).decode()
+          fixed = content.replace('linux_x86_64', plat)
+          tmpfd, tmppath = tempfile.mkstemp(suffix='.whl')
+          os.close(tmpfd)
+          with zipfile.ZipFile(dst, 'r') as zin, zipfile.ZipFile(tmppath, 'w') as zout:
+              for item in zin.infolist():
+                  data = zin.read(item.filename)
+                  if item.filename == wheel_file:
+                      data = fixed.encode()
+                  zout.writestr(item, data)
+          shutil.move(tmppath, dst)
+          print(f'Retagged: {os.path.basename(dst)}')
+          PYEOF
+          $CT run --rm \
+            -v ${{ runner.temp }}/dist:/tmp/dist:ro \
+            -v ${{ runner.temp }}/dist-pypi:/tmp/dist-pypi \
+            -v ${{ runner.temp }}/retag.py:/tmp/retag.py:ro \
+            $IMAGE bash -c "
+              pip install auditwheel -q
+              for whl in /tmp/dist/amd_mori_nightly-*.whl; do
+                PLAT=\$(auditwheel show \"\$whl\" 2>&1 | grep -o 'manylinux_[0-9_]*_x86_64' | head -1)
+                if [ -z \"\$PLAT\" ]; then PLAT=linux_x86_64; fi
+                NEWNAME=\$(basename \"\$whl\" | sed \"s/linux_x86_64/\$PLAT/\")
+                python3 /tmp/retag.py \"\$whl\" \"/tmp/dist-pypi/\$NEWNAME\" \"\$PLAT\"
+              done
+            "
+          $CT run --rm -v ${{ runner.temp }}/dist-pypi:/tmp/dist-pypi $IMAGE \
+            chown -R $(id -u):$(id -g) /tmp/dist-pypi
+          echo "=== PyPI wheels ==="
+          ls -lh ${{ runner.temp }}/dist-pypi/
+
+      - name: Upload gh-pages wheel
         uses: actions/upload-artifact@v4
         with:
           name: wheel-cp${{ matrix.python }}
-          path: ${{ runner.temp }}/dist/*.whl
+          path: ${{ runner.temp }}/dist/amd_mori-*.whl
+          retention-days: 30
+
+      - name: Upload PyPI wheel
+        uses: actions/upload-artifact@v4
+        with:
+          name: pypi-wheel-cp${{ matrix.python }}
+          path: ${{ runner.temp }}/dist-pypi/amd_mori_nightly-*.whl
           retention-days: 30
 
       - name: Cleanup
@@ -800,3 +855,65 @@ jobs:
           git diff --cached --quiet && echo "No changes" && exit 0
           git commit -m "nightly: promote tested wheels to latest $(date +%Y-%m-%d)"
           git push
+
+  # ── Stage 5: Test PyPI wheel (install + import verification) ───────────────
+  test-pypi-wheel:
+    name: test pypi wheel (py${{ matrix.python }})
+    needs: [promote-latest]
+    if: github.event_name != 'pull_request'
+    runs-on: [self-hosted, MI355X-AINIC-TW]
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - python: "3.10"
+            base_image: rocm/pytorch:rocm7.2.1_ubuntu22.04_py3.10_pytorch_release_2.8.0
+          - python: "3.12"
+            base_image: rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.8.0
+    env:
+      IMAGE: rocm/mori:ci-py${{ matrix.python }}
+    steps:
+      - name: Download PyPI wheel
+        uses: actions/download-artifact@v4
+        with:
+          name: pypi-wheel-cp${{ matrix.python }}
+          path: ${{ runner.temp }}/pypi-wheel
+
+      - name: Build CI image
+        run: $CT build --network=host --build-arg BASE_IMAGE=${{ matrix.base_image }} -t $IMAGE -f docker/Dockerfile.dev .
+
+      - name: Test PyPI wheel
+        run: |
+          $CT run --rm --network=host \
+            --device=/dev/kfd --device=/dev/dri \
+            -v ${{ runner.temp }}/pypi-wheel:/tmp/wheel:ro \
+            $IMAGE bash -c "
+              pip install /tmp/wheel/*.whl
+              python3 -c 'import mori; print(\"amd-mori-nightly \" + mori.__version__ + \" OK\")'
+              python3 -c 'from mori import ops; print(\"mori.ops OK\")'
+            "
+
+  # ── Stage 6: Publish to PyPI ───────────────────────────────────────────────
+  publish-pypi:
+    name: publish to PyPI
+    needs: [test-pypi-wheel]
+    if: github.event_name != 'pull_request'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Download PyPI wheels
+        uses: actions/download-artifact@v4
+        with:
+          pattern: pypi-wheel-*
+          merge-multiple: true
+          path: ./dist
+
+      - name: List wheels
+        run: ls -lh ./dist/
+
+      - name: Upload to PyPI
+        run: |
+          pip install twine "packaging>=24.2"
+          python3 -m twine upload --skip-existing ./dist/*.whl
+        env:
+          TWINE_USERNAME: __token__
+          TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}