Update GitHub workflows

Author: raphaelsty

.github/workflows/pypi.yaml

@@ -1,8 +1,3 @@
-# This workflow builds and publishes Python distributions to PyPI.
-# It uses a self-hosted runner and a PyTorch container image
-# to ensure the correct CUDA environment is available for the build.
-# It relies on Trusted Publishing for secure authentication with PyPI.
-
 name: Build and Publish Python 🐍 distributions
 
 on:
@@ -15,7 +10,6 @@ jobs:
     name: Build wheel for Python ${{ matrix.variant.python-version }} on self-hosted
     runs-on: self-hosted
 
-    # The PyTorch container provides the necessary CUDA environment
     container:
       image: pytorch/pytorch:2.3.1-cuda12.1-cudnn8-devel
 
@@ -25,31 +19,31 @@ jobs:
         variant:
           - {
               name: "cuda",
-              python-version: "3.9", # Added Python 3.9
+              python-version: "3.9",
               cuda-version: "12.1",
               compute-caps: "80",
             }
           - {
               name: "cuda",
-              python-version: "3.10", # Existing Python 3.10
+              python-version: "3.10",
               cuda-version: "12.1",
               compute-caps: "80",
             }
           - {
               name: "cuda",
-              python-version: "3.11", # Added Python 3.11
+              python-version: "3.11",
               cuda-version: "12.1",
               compute-caps: "80",
             }
           - {
               name: "cuda",
-              python-version: "3.12", # Added Python 3.12
+              python-version: "3.12",
               cuda-version: "12.1",
               compute-caps: "80",
             }
           - {
               name: "cuda",
-              python-version: "3.13", # Added Python 3.13
+              python-version: "3.13",
               cuda-version: "12.1",
               compute-caps: "80",
             }
@@ -63,6 +57,13 @@ jobs:
         with:
           python-version: ${{ matrix.variant.python-version }}
 
+      - name: Get Python Library Path and Set LD_LIBRARY_PATH
+        run: |
+          PYTHON_LIB_DIR=$(python -c "import sysconfig; print(sysconfig.get_config_var('LIBDIR'))")
+          echo "Found Python library directory: $PYTHON_LIB_DIR"
+          echo "LD_LIBRARY_PATH=$PYTHON_LIB_DIR:${LD_LIBRARY_PATH}" >> $GITHUB_ENV
+        shell: bash
+
       - name: Install build dependencies
         run: |
           apt-get update
@@ -95,7 +96,7 @@ jobs:
 
       - name: Build wheel
         env:
-          LD_LIBRARY_PATH: ${{ env.CUDA_LIB_DIR }}:${LD_LIBRARY_PATH}
+          LD_LIBRARY_PATH: ${{ env.CUDA_LIB_DIR }}:${LD_LIBRARY_PATH} # This will now include Python's lib dir
           CUDA_COMPUTE_CAP: ${{ matrix.variant.compute-caps }}
         run: maturin build --release --out dist
         shell: bash
@@ -143,7 +144,6 @@ jobs:
       - name: Upload wheel artifact
         uses: actions/upload-artifact@v4
         with:
-          # Updated artifact name to include the Python version
           name: wheels-${{ matrix.variant.name }}-cuda${{ matrix.variant.cuda-version }}-py${{ matrix.variant.python-version }}
           path: dist/
 
@@ -152,15 +152,15 @@ jobs:
     needs: build-wheels
     runs-on: ubuntu-latest
     permissions:
-      id-token: write # Required for Trusted Publishing
+      id-token: write
 
     steps:
       - name: Download all wheels
         uses: actions/download-artifact@v4
         with:
-          pattern: wheels-* # This pattern will download all artifacts matching 'wheels-*'
+          pattern: wheels-*
           path: dist
-          merge-multiple: true # Merge all downloaded artifacts into a single 'dist' directory
+          merge-multiple: true
 
       - name: Publish to PyPI
         uses: pypa/gh-action-pypi-publish@v1.12.4