rocm_jax/.github/workflows/wheel_win_x64.yml
Nitin Srinivasan 4b4f2f9cb9 Use uv to install Python packages
PiperOrigin-RevId: 730499307
2025-02-24 10:13:39 -08:00

65 lines
2.1 KiB
YAML

name: Wheel build - Windows CPU x86_64
on:
workflow_dispatch: # allows triggering the workflow run manually
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
cancel-in-progress: true
env:
DISTUTILS_USE_SDK: 1
MSSdk: 1
jobs:
win-wheels:
strategy:
fail-fast: false # Don't stop all wheel builds if one has a test failure.
matrix:
os: [windows-2019-32core]
arch: [AMD64]
pyver: ['3.10', '3.11', '3.12', '3.13']
name: ${{ matrix.os }} ${{ matrix.pyver }} jaxlib wheel build
runs-on: ${{ matrix.os }}
steps:
- name: Install LLVM/Clang
run: choco install llvm --version=18.1.4 --yes --no-progress --allow-downgrade
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- uses: actions/setup-python@42375524e23c412d93fb67b49958b491fce71c38 # v5.4.0
with:
python-version: ${{ matrix.pyver }}
cache: 'pip'
- name: Build wheels
env:
BAZEL_VC: "C:\\Program Files (x86)\\Microsoft Visual Studio\\2019\\Enterprise\\VC"
JAXLIB_RELEASE: true
run: |
python -m pip install uv~=0.5.30
python -m uv pip install -r build/test-requirements.txt \
--upgrade numpy==2.0.0 scipy==1.13.1
"C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
python.exe build\build.py build --wheels=jaxlib `
--bazel_options=--color=yes `
--bazel_options=--config=win_clang `
--verbose
- uses: actions/upload-artifact@6f51ac03b9356f520e9adb1b1b7802705f340c2b # v4.5.0
with:
name: wheels-${{ matrix.os }}-${{ matrix.pyver }}
path: ${{ github.workspace }}\dist\*.whl
retention-days: 5
- name: Run tests
env:
JAX_ENABLE_CHECKS: true
JAX_SKIP_SLOW_TESTS: true
PY_COLORS: 1
run: |
python -m uv pip install --find-links ${{ github.workspace }}\dist jaxlib \
-e ${{ github.workspace }}
echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
pytest -n auto --tb=short tests examples