From 0a70d42fb9918de34a81a09f085682cb5fcd500f Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Tue, 13 Jun 2023 16:22:27 -0400
Subject: [PATCH] Windows CI improvements:

* upload wheel artifacts to the github workflow.
* use XLA from github head in CI builds.
* update release builds to match the CI build configuration.
---
 .github/workflows/wheel_win_x64.yml | 43 ++++++-----------------------
 .github/workflows/windows_ci.yml    | 10 +++++--
 2 files changed, 17 insertions(+), 36 deletions(-)

diff --git a/.github/workflows/wheel_win_x64.yml b/.github/workflows/wheel_win_x64.yml
index c44e0e03a729..e6707255ae32 100644
--- a/.github/workflows/wheel_win_x64.yml
+++ b/.github/workflows/wheel_win_x64.yml
@@ -1,9 +1,8 @@
 name: Wheel::Windows::x86_64
 on:
   workflow_dispatch: # allows triggering the workflow run manually
-#  pull_request:
-#  release:
-#    types: [published]
+  release:
+    types: [published]
 
 env:
   DISTUTILS_USE_SDK: 1
@@ -14,10 +13,10 @@ jobs:
     strategy:
       fail-fast: true
       matrix:
-        os: [windows-2019]
+        os: [win-2019-16core]
         arch: [AMD64]
         pyver: ['3.9', '3.10', '3.11']
-    name: ${{ matrix.os }} jaxlib Wheel-builder
+    name: ${{ matrix.os }} ${{ matrix.pyver }} jaxlib wheel build
     runs-on: ${{ matrix.os }}
 
     steps:
@@ -39,30 +38,7 @@ jobs:
         run: |
           python -m pip install -r build/test-requirements.txt
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
-          python.exe build\build.py `
-            --noenable_tpu `
-            --noenable_remote_tpu `
-            --noenable_mkl_dnn `
-            --target_cpu_features release `
-            --noenable_plugin_device `
-            --noremote_build `
-            --noenable_cuda `
-            --configure_only 
-          bazel build --jobs=1 `
-            --cxxopt="/Zm2000" `
-            --verbose_failures `
-            --jobs=1 `
-            --discard_analysis_cache `
-            --notrack_incremental_state `
-            --nokeep_state_after_build `
-            --experimental_local_memory_estimate="true" `
-            --local_ram_resources=2048 `
-            //jaxlib
-          bazel run `
-            --verbose_failures=true `
-            //build:build_wheel -- `
-            --output_path=$(Join-Path -Path (Get-Location) -ChildPath "\tmp") `
-            --cpu="AMD64"
+          python.exe build\build.py
 
       - name: Run tests
         env:
@@ -74,8 +50,7 @@ jobs:
           echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
           pytest -n auto --tb=short tests examples
 
-      - uses: actions/upload-artifact@v2
-        if: ${{ github.event_name == 'release' }}
-        with:
-          name: jaxlib_wheels
-          path: ${{ github.workspace }}\dist\*.whl
+      - uses: actions/upload-artifact@v3
+        name: jaxlib_wheels
+        path: ${{ github.workspace }}\dist\*.whl
+        retention-days: 5
\ No newline at end of file
diff --git a/.github/workflows/windows_ci.yml b/.github/workflows/windows_ci.yml
index 49cb5562e96f..64756249291f 100644
--- a/.github/workflows/windows_ci.yml
+++ b/.github/workflows/windows_ci.yml
@@ -16,7 +16,7 @@ jobs:
         os: [win-2019-16core]
         arch: [AMD64]
         pyver: ['3.9']
-    name: ${{ matrix.os }} jaxlib Wheel-builder
+    name: ${{ matrix.os }} CI build
     runs-on: ${{ matrix.os }}
 
     steps:
@@ -36,9 +36,10 @@ jobs:
         env:
           BAZEL_VC: "C:\\Program Files (x86)\\Microsoft Visual Studio\\2019\\Enterprise\\VC"
         run: |
+          git clone https://github.com/openxla/xla.git
           python -m pip install -r build/test-requirements.txt
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
-          python.exe build\build.py
+          python.exe build\build.py --override_repository=xla=${{ github.workspace }}\xla
 
       - name: Run tests
         env:
@@ -49,3 +50,8 @@ jobs:
           python -m pip install --no-index --find-links ${{ github.workspace }}\dist jaxlib
           echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
           pytest -n auto --tb=short tests examples
+
+      - uses: actions/upload-artifact@v3
+        name: jaxlib_wheels
+        path: ${{ github.workspace }}\dist\*.whl
+        retention-days: 5
\ No newline at end of file