python-client-build.yml

# NOTE - this workflow is run on every push to 'main' in order to populate the cache
# If you ever switch this to a non-free runner (e.g. a Namespace runner), be aware that
# this might end up costing us much more than expected.
on:
  workflow_call:

permissions:
  contents: read

env:
  CARGO_INCREMENTAL: 0
  RUST_TOOLCHAIN: "1.90"

jobs:
  linux:
    runs-on: ${{ matrix.platform.runner }}
    if: github.repository == 'tensorzero/tensorzero'
    strategy:
      matrix:
        platform:
          - runner: ubuntu-24.04
            target: x86_64
          - runner: ubuntu-24.04-arm
            target: aarch64
    steps:
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
      - uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38
        with:
          python-version: "3.10"

      - name: Install uv
        uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b
        with:
          version: "0.9.27"

      - name: Install Rust toolchain
        run: |
          for attempt in 1 2 3; do
            if rustup toolchain install $RUST_TOOLCHAIN --profile minimal && rustup default $RUST_TOOLCHAIN; then
              break
            fi
            if [ $attempt -eq 3 ]; then
              echo "Failed to install Rust toolchain after 3 attempts"
              exit 1
            fi
            sleep $((10 * attempt))
          done
        shell: bash

      - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
        with:
          save-if: ${{ github.ref == 'refs/heads/main' }}
      - name: Build wheels
        uses: PyO3/maturin-action@86b9d133d34bc1b40018696f782949dac11bd380
        with:
          working-directory: ./clients/python
          target: ${{ matrix.platform.target }}
          args: --release --out dist -i python3.10
          #sccache: "true"
          manylinux: auto
      - name: Upload wheels
        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4
        with:
          name: wheels-linux-${{ matrix.platform.target }}
          path: ./clients/python/dist
      - name: Check built wheel
        run: |
          ./ci/check-local-wheel.sh ./clients/python/dist/tensorzero*.whl

  musllinux:
    runs-on: ${{ matrix.platform.runner }}
    if: github.repository == 'tensorzero/tensorzero'
    strategy:
      matrix:
        platform:
          - runner: ubuntu-24.04
            target: x86_64
          - runner: ubuntu-24.04-arm
            target: aarch64
    steps:
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
      - name: Install Rust toolchain
        run: |
          for attempt in 1 2 3; do
            if rustup toolchain install $RUST_TOOLCHAIN --profile minimal && rustup default $RUST_TOOLCHAIN; then
              break
            fi
            if [ $attempt -eq 3 ]; then
              echo "Failed to install Rust toolchain after 3 attempts"
              exit 1
            fi
            sleep $((10 * attempt))
          done
        shell: bash
      - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
        with:
          save-if: ${{ github.ref == 'refs/heads/main' }}
      - uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38
        with:
          python-version: "3.10"
      - name: Install uv
        uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b
        with:
          version: "0.9.27"
      - name: Build wheels
        uses: PyO3/maturin-action@86b9d133d34bc1b40018696f782949dac11bd380
        with:
          working-directory: ./clients/python
          target: ${{ matrix.platform.target }}
          args: --release --out dist -i python3.10
          #sccache: "true"
          manylinux: musllinux_1_2
      - name: Upload wheels
        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4
        with:
          name: wheels-musllinux-${{ matrix.platform.target }}
          path: ./clients/python/dist
      - name: Check built musl wheel
        run: |
          ./ci/check-musl-wheel.sh ./clients/python/dist/tensorzero*.whl

  windows:
    runs-on: ${{ matrix.platform.runner }}
    if: github.repository == 'tensorzero/tensorzero'
    strategy:
      matrix:
        platform:
          - runner: windows-latest
            target: x64
    steps:
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
      - name: Install Rust toolchain
        run: |
          for attempt in 1 2 3; do
            if rustup toolchain install $RUST_TOOLCHAIN --profile minimal && rustup default $RUST_TOOLCHAIN; then
              break
            fi
            if [ $attempt -eq 3 ]; then
              echo "Failed to install Rust toolchain after 3 attempts"
              exit 1
            fi
            sleep $((10 * attempt))
          done
        shell: bash
      - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
        with:
          save-if: ${{ github.ref == 'refs/heads/main' }}
      - uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38
        with:
          python-version: "3.10"
          architecture: ${{ matrix.platform.target }}
      - name: Install uv
        uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b
        with:
          version: "0.9.27"
      - name: Build wheels
        uses: PyO3/maturin-action@86b9d133d34bc1b40018696f782949dac11bd380
        with:
          working-directory: ./clients/python
          target: ${{ matrix.platform.target }}
          args: --release --out dist -i python3.10
          #sccache: "true"
      - name: Upload wheels
        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4
        with:
          name: wheels-windows-${{ matrix.platform.target }}
          path: ./clients/python/dist
      - name: Check built wheel
        run: |
          ./ci/check-local-wheel.sh ./clients/python/dist/tensorzero*.whl
        shell: bash

  macos:
    runs-on: ${{ matrix.platform.runner }}
    if: github.repository == 'tensorzero/tensorzero'
    strategy:
      matrix:
        platform:
          # TODO - add back x86_64 support
          - runner: macos-14
            target: aarch64
          # TODO - add this back (without the conflicting artifact upload)
          #- runner: macos-15
          #  target: aarch64
    steps:
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
      - name: Install Rust toolchain
        run: |
          for attempt in 1 2 3; do
            if rustup toolchain install $RUST_TOOLCHAIN --profile minimal && rustup default $RUST_TOOLCHAIN; then
              break
            fi
            if [ $attempt -eq 3 ]; then
              echo "Failed to install Rust toolchain after 3 attempts"
              exit 1
            fi
            sleep $((10 * attempt))
          done
        shell: bash
      - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
        with:
          save-if: ${{ github.ref == 'refs/heads/main' }}
      - uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38
        with:
          python-version: "3.10"

      - name: Install uv
        uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b
        with:
          version: "0.9.27"

      - name: Build wheels
        uses: PyO3/maturin-action@86b9d133d34bc1b40018696f782949dac11bd380
        with:
          working-directory: ./clients/python
          target: ${{ matrix.platform.target }}
          args: --release --out dist -i python3.10
          #sccache: "true"

      - name: Upload wheels
        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4
        with:
          name: wheels-macos-${{ matrix.platform.target }}
          path: ./clients/python/dist
      - name: Check built wheel
        run: |
          ./ci/check-local-wheel.sh ./clients/python/dist/tensorzero*.whl

  sdist:
    runs-on: ubuntu-latest
    if: github.repository == 'tensorzero/tensorzero'
    steps:
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
      - name: Build sdist
        uses: PyO3/maturin-action@86b9d133d34bc1b40018696f782949dac11bd380
        with:
          working-directory: ./clients/python
          command: sdist
          args: --out dist
      - name: Upload sdist
        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4
        with:
          name: wheels-sdist
          path: ./clients/python/dist
      - name: Install uv
        uses: astral-sh/setup-uv@61cb8a9741eeb8a550a1b8544337180c0fc8476b
        with:
          version: "0.9.27"
      - name: Check local sdist
        run: |
          ./ci/check-local-wheel.sh ./clients/python/dist/tensorzero-*.tar.gz
      - name: Run twine on sdist
        run: |
          uv pip install twine
          uv run twine check ./clients/python/dist/tensorzero-*.tar.gz

  check-artifacts:
    needs: [linux, musllinux, windows, macos, sdist]
    runs-on: ubuntu-latest
    if: github.repository == 'tensorzero/tensorzero'
    steps:
      - uses: actions/download-artifact@018cc2cf5baa6db3ef3c5f8a56943fffe632ef53
        with:
          pattern: wheels-*
      - name: Count wheels
        run: |
          ls -lh ./wheels-*/*
          # Count the number of files named 'tensorzero-*.whl', and verify that it equals 5
          # If we add new platforms / python versions, update this number
          ls -lh ./wheels-*/* | grep -c 'tensorzero-.*\.whl' | grep -q 6