ci-build.yml

name: build whl and uploads

on:
  push:
    branches: [ main ]

  # Allows you to run this workflow manually from the Actions tab
  workflow_dispatch:

jobs:
  build-cuda:
    strategy:
      matrix:
        cuda-version: ["11.8", "12.1"]
    runs-on: windows-2022
    env:
      AZURE_STORAGE_CONNECTION_STRING: ${{ secrets.AZURE_STORAGE_CONNECTION_STRING }}
      BAZEL_PATH: "D:\\bazel.exe"
      TEMP: C:\\Users\\runneradmin\\Temp
      TMP: C:\\Users\\runneradmin\\Temp
      PYTHONUNBUFFERED: '1'

    steps:
    - name: Show user home
      run: ls ~
    - name: Show cpu info
      run: Get-CimInstance Win32_Processor
    - name: Limit cpu
      run: |
        $p = Get-CimInstance Win32_Processor
        if ($p.Name -match "E5-") { throw "CPU is too old!" }
    - name: Show memory info
      run: Get-CimInstance Win32_PhysicalMemory | Format-Table Tag, DeviceLocator, Capacity, Speed
    - name: Configure pagefile
      uses: al-cheb/configure-pagefile-action@v1.2
      with:
        minimum-size: 8GB
        maximum-size: 32GB
        disk-root: "C:"
    # - name: Show disk info
    #   run: Get-Volume -DriveLetter CD | Sort-Object DriveLetter
    - name: Workaround https://github.com/bazelbuild/bazel/issues/18592
      run: rm -Recurse -Force "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\vcpkg" -ErrorAction Continue

    - uses: actions/checkout@v2
      with:
        submodules: true

    - name: Download Bazelisk
      run: curl -k -L https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazelisk-windows-amd64.exe -o $env:BAZEL_PATH

    - name: Install CUDA ${{ matrix.cuda-version }}
      run: |
        curl -k -L https://whls.blob.core.windows.net/ci-files/v${{ matrix.cuda-version }}.7z -o cuda.7z
        7z x cuda.7z -o'D:/CUDA'
        rm cuda.7z
        ls D:/CUDA/v${{ matrix.cuda-version }}

    - uses: actions/cache@v2
      with:
        path: ~\AppData\Local\pip\Cache
        key: ${{ runner.os }}-pip-${{ hashFiles('build-requirements.txt') }}
        restore-keys: |
          ${{ runner.os }}-pip-

    #=============#
    # Python 3.11 #
    #=============#
    - name: py311
      uses: actions/setup-python@v2
      with:
        python-version: "3.11"
    - name: py311 pip install
      run: pip install -r "$env:GITHUB_WORKSPACE/build-requirements.txt"
    - name: py311 build whl and upload
      run: |
        rm -Recurse -Force "$env:GITHUB_WORKSPACE/jax/bazel-dist" -ErrorAction Continue
        cd "$env:GITHUB_WORKSPACE/jax"
        ../build-jaxlib.ps1 cuda -bazel_path $env:BAZEL_PATH -vs_version 2022 -cuda_version '${{ matrix.cuda-version }}' -cuda_prefix 'D:/CUDA' -symlink_python
        az storage blob upload-batch --overwrite -d unstable -s "$env:GITHUB_WORKSPACE/jax/bazel-dist" --pattern '*.whl'
    - uses: actions/upload-artifact@v3
      with:
        name: whls
        path: $env:GITHUB_WORKSPACE/jax/bazel-dist/

    #=============#
    # Python 3.10 #
    #=============#
    - name: py310
      uses: actions/setup-python@v2
      with:
        python-version: "3.10"
    - name: py310 pip install
      run: pip install -r "$env:GITHUB_WORKSPACE/build-requirements.txt"
    - name: py310 build whl and upload
      run: |
        rm -Recurse -Force "$env:GITHUB_WORKSPACE/jax/bazel-dist" -ErrorAction Continue
        cd "$env:GITHUB_WORKSPACE/jax"
        ../build-jaxlib.ps1 cuda -bazel_path $env:BAZEL_PATH -vs_version 2022 -cuda_version '${{ matrix.cuda-version }}' -cuda_prefix 'D:/CUDA' -symlink_python
        az storage blob upload-batch --overwrite -d unstable -s "$env:GITHUB_WORKSPACE/jax/bazel-dist" --pattern '*.whl'
    - uses: actions/upload-artifact@v3
      with:
        name: whls
        path: $env:GITHUB_WORKSPACE/jax/bazel-dist/

    #============#
    # Python 3.9 #
    #============#
    - name: py39
      uses: actions/setup-python@v2
      with:
        python-version: "3.9"
    - name: py39 pip install
      run: pip install -r "$env:GITHUB_WORKSPACE/build-requirements.txt"
    - name: py39 build whl and upload
      run: |
        rm -Recurse -Force "$env:GITHUB_WORKSPACE/jax/bazel-dist" -ErrorAction Continue
        cd "$env:GITHUB_WORKSPACE/jax"
        ../build-jaxlib.ps1 cuda -bazel_path $env:BAZEL_PATH -vs_version 2022 -cuda_version '${{ matrix.cuda-version }}' -cuda_prefix 'D:/CUDA' -symlink_python
        az storage blob upload-batch --overwrite -d unstable -s "$env:GITHUB_WORKSPACE/jax/bazel-dist" --pattern '*.whl'
    - uses: actions/upload-artifact@v3
      with:
        name: whls
        path: $env:GITHUB_WORKSPACE/jax/bazel-dist/

  update-index:
    if: ${{ ! cancelled() }}
    needs:
      - build-cuda
    runs-on: windows-2022
    env:
      AZURE_STORAGE_CONNECTION_STRING: ${{ secrets.AZURE_STORAGE_CONNECTION_STRING }}
    steps:
    - uses: actions/checkout@v2
    - name: py311
      uses: actions/setup-python@v2
      with:
        python-version: "3.11"
    - name: update index.html
      run: |
        cd "$env:GITHUB_WORKSPACE/"
        python ./update_index.py --url_mode absolute unstable > index.html
        az storage blob upload --overwrite -c unstable -f index.html -n index.html --content-type='text/html'
        az storage blob upload --overwrite -c '$web' -f index.html -n 'unstable/index.html' --content-type='text/html'