fix linux PGO wheel build (#170)

pydantic · Nov 26, 2024 · 82df231 · 82df231
1 parent 72fc9ef
commit 82df231
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 44 deletions.
diff --git a/.github/actions/build-pgo-wheel/action.yml b/.github/actions/build-pgo-wheel/action.yml
@@ -0,0 +1,74 @@
+name: Build PGO wheel
+description: Builds a PGO-optimized wheel
+inputs:
+  interpreter:
+    description: 'Interpreter to build the wheel for'
+    required: true
+  rust-toolchain:
+    description: 'Rust toolchain to use'
+    required: true
+outputs:
+  wheel:
+    description: 'Path to the built wheel'
+    value: ${{ steps.find_wheel.outputs.path }}
+runs:
+  using: "composite"
+  steps:
+    - name: prepare profiling directory
+      shell: bash
+      # making this ahead of the compile ensures that the local user can write to this
+      # directory; the maturin action (on linux) runs in docker so would create as root
+      run: mkdir -p ${{ github.workspace }}/profdata
+
+    - name: build initial wheel
+      uses: PyO3/maturin-action@v1
+      with:
+        manylinux: auto
+        args: >
+          --release
+          --out pgo-wheel
+          --interpreter ${{ inputs.interpreter }}
+        rust-toolchain: ${{ inputs.rust-toolchain }}
+        docker-options: -e CI
+        working-directory: crates/jiter-python
+      env:
+        RUSTFLAGS: '-Cprofile-generate=${{ github.workspace }}/profdata'
+
+    - name: detect rust host
+      run: echo RUST_HOST=$(rustc -Vv | grep host | cut -d ' ' -f 2) >> "$GITHUB_ENV"
+      shell: bash
+
+    - name: generate pgo data
+      run: |
+        pip install -U pip
+        pip install -r tests/requirements.txt
+        pip install jiter --no-index --no-deps --find-links pgo-wheel --force-reinstall
+        python bench.py jiter jiter-cache
+        RUST_HOST=$(rustc -Vv | grep host | cut -d ' ' -f 2)
+        rustup run ${{ inputs.rust-toolchain }} bash -c 'echo LLVM_PROFDATA=$RUSTUP_HOME/toolchains/$RUSTUP_TOOLCHAIN/lib/rustlib/$RUST_HOST/bin/llvm-profdata >> "$GITHUB_ENV"'
+      shell: bash
+      working-directory: crates/jiter-python
+
+    - name: merge pgo data
+      run: ${{ env.LLVM_PROFDATA }} merge -o ${{ github.workspace }}/merged.profdata ${{ github.workspace }}/profdata
+      shell: pwsh  # because it handles paths on windows better, and works well enough on unix for this step
+
+    - name: build pgo-optimized wheel
+      uses: PyO3/maturin-action@v1
+      with:
+        manylinux: auto
+        args: >
+          --release
+          --out dist
+          --interpreter ${{ inputs.interpreter }}
+        rust-toolchain: ${{inputs.rust-toolchain}}
+        docker-options: -e CI
+        working-directory: crates/jiter-python
+      env:
+        RUSTFLAGS: '-Cprofile-use=${{ github.workspace }}/merged.profdata'
+
+    - name: find built wheel
+      id: find_wheel
+      run: echo "path=$(ls dist/*.whl)" | tee -a "$GITHUB_OUTPUT"
+      shell: bash
+      working-directory: crates/jiter-python
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -380,8 +380,6 @@ jobs:
         with:
           components: llvm-tools
 
-      - run: rustc --version --verbose
-
       # linux: builds are done inside manylinux docker images, need to help maturin-action
       # find the right path
       #
@@ -401,49 +399,11 @@ jobs:
             '3.13t': 'cp313-cp313t',
           }['${{ matrix.interpreter }}'])")/bin/python >> $GITHUB_OUTPUT
 
-      - name: build initial wheel
-        uses: PyO3/maturin-action@v1
+      - name: build pgo wheel
+        uses: ./.github/actions/build-pgo-wheel
         with:
-          manylinux: auto
-          args: >
-            --release
-            --out pgo-wheel
-            --interpreter ${{ matrix.os == 'linux' && steps.resolve-interpreter.outputs.python-path || steps.setup-python.outputs.python-path }}
-          rust-toolchain: stable
-          working-directory: crates/jiter-python
-        env:
-          RUSTFLAGS: "-Cprofile-generate=${{ github.workspace }}/profdata"
-
-      - name: detect rust host
-        run: echo RUST_HOST=$(rustc -Vv | grep host | cut -d ' ' -f 2) >> "$GITHUB_ENV"
-        shell: bash
-
-      - name: generate pgo data
-        run: |
-          cd crates/jiter-python
-          pip install -U pip
-          pip install -r tests/requirements.txt
-          pip install jiter --no-index --no-deps --find-links pgo-wheel --force-reinstall
-          python bench.py jiter jiter-cache
-          rustup run stable bash -c 'echo LLVM_PROFDATA=$RUSTUP_HOME/toolchains/$RUSTUP_TOOLCHAIN/lib/rustlib/${{ env.RUST_HOST }}/bin/llvm-profdata >> "$GITHUB_ENV"'
-
-      - name: merge pgo data
-        run: |
-          cd crates/jiter-python
-          ${{ env.LLVM_PROFDATA }} merge -o ${{ github.workspace }}/merged.profdata ${{ github.workspace }}/profdata
-
-      - name: build pgo-optimized wheel
-        uses: PyO3/maturin-action@v1
-        with:
-          manylinux: auto
-          args: >
-            --release
-            --out dist
-            --interpreter ${{ matrix.os == 'linux' && steps.resolve-interpreter.outputs.python-path || steps.setup-python.outputs.python-path }}
-          rust-toolchain: stable
-          working-directory: crates/jiter-python
-        env:
-          RUSTFLAGS: "-Cprofile-use=${{ github.workspace }}/merged.profdata"
+          interpreter: ${{ matrix.os == 'linux' && steps.resolve-interpreter.outputs.python-path || steps.setup-python.outputs.python-path }}
+          rust-toolchain: ${{ steps.rust-toolchain.outputs.name }}
 
       - run: ${{ matrix.ls || 'ls -lh' }} crates/jiter-python/dist/