From 82df23197b1f84125a4cc75cf856b9a5e3e5fa68 Mon Sep 17 00:00:00 2001
From: David Hewitt <david.hewitt@pydantic.dev>
Date: Tue, 26 Nov 2024 15:27:44 +0000
Subject: [PATCH] fix linux PGO wheel build (#170)

---
 .github/actions/build-pgo-wheel/action.yml | 74 ++++++++++++++++++++++
 .github/workflows/ci.yml                   | 48 ++------------
 2 files changed, 78 insertions(+), 44 deletions(-)
 create mode 100644 .github/actions/build-pgo-wheel/action.yml

diff --git a/.github/actions/build-pgo-wheel/action.yml b/.github/actions/build-pgo-wheel/action.yml
new file mode 100644
index 0000000..824d522
--- /dev/null
+++ 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
index 9770d81..9c5b19b 100644
--- 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/