Re-factor build CLI to a subcommand based approach

PiperOrigin-RevId: 691466647
jax-ml · Nov 6, 2024 · 848eed9 · 848eed9
1 parent 37af100
commit 848eed9
Show file tree

Hide file tree

Showing 8 changed files with 804 additions and 521 deletions.
diff --git a/.github/workflows/asan.yaml b/.github/workflows/asan.yaml
@@ -65,9 +65,9 @@ jobs:
         run: |
           source ${GITHUB_WORKSPACE}/venv/bin/activate
           cd jax
-          python build/build.py \
-            --bazel_options=--color=yes \
-            --bazel_options=--copt=-fsanitize=address \
+          python build/build.py jaxlib --verbose \
+            --bazel_build_options='--verbose_failures=true' \
+            --bazel_build_options='--copt=-fsanitize=address' \
             --clang_path=/usr/bin/clang-18
           pip install dist/jaxlib-*.whl
           pip install -e .

diff --git a/.github/workflows/wheel_win_x64.yml b/.github/workflows/wheel_win_x64.yml
@@ -40,10 +40,9 @@ jobs:
           python -m pip install -r build/test-requirements.txt
           python -m pip install --upgrade numpy==2.0.0 scipy==1.13.1
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
-          python.exe build\build.py `
-            --bazel_options=--color=yes `
-            --bazel_options=--config=win_clang `
-            --verbose
+          python.exe build\build.py jaxlib --verbose `
+            --bazel_build_options='--verbose_failures=true' `
+            --bazel_build_options='--config=win_clang'
 
       - uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882  # v4.4.3
         with:

diff --git a/.github/workflows/windows_ci.yml b/.github/workflows/windows_ci.yml
@@ -49,9 +49,9 @@ jobs:
           python -m pip install -r build/test-requirements.txt
           python -m pip install --upgrade numpy==2.0.0 scipy==1.13.1
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
-          python.exe build\build.py `
-            --bazel_options=--color=yes `
-            --bazel_options=--config=win_clang
+          python.exe build\build.py jaxlib --verbose `
+            --bazel_build_options='--verbose_failures=true' `
+            --bazel_build_options='--config=win_clang'
 
       - uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882  # v4.4.3
         with:

diff --git a/.github/workflows/windows_presubmit.yml b/.github/workflows/windows_presubmit.yml
@@ -0,0 +1,70 @@
+name: Presubmit - Windows CPU
+on:
+  # TODO(DO_NOT_SUBMIT): temporary check
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+permissions:
+  contents: read  # to fetch code
+  actions: write  # to cancel previous workflows
+
+env:
+  DISTUTILS_USE_SDK: 1
+  MSSdk: 1
+
+jobs:
+  presubmit-win-wheels:
+    if: ${{ (github.event.action != 'labeled') || (github.event.label.name == 'windows:force-run')}}
+    strategy:
+      fail-fast: true
+      matrix:
+        os: [windows-2019-32core]
+        arch: [AMD64]
+        pyver: ['3.10']
+    name: ${{ matrix.os }} Windows build
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - name: Cancel Previous Runs
+        uses: styfle/cancel-workflow-action@85880fa0301c86cca9da44039ee3bb12d3bedbfa # ratchet: styfle/cancel-workflow-action@0.12.1
+        with:
+          access_token: ${{ github.token }}
+
+      - 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@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.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 -r build/test-requirements.txt
+          python -m pip install --upgrade numpy==2.0.0 scipy==1.13.1
+          "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
+          python.exe build\build.py jaxlib --verbose `
+            --bazel_build_options='--verbose_failures=true' `
+            --bazel_build_options="--config=win_clang" `
+            --clang_path="C:\Program Files\LLVM\bin\clang.exe" `
+            --bazel_build_options='--color=yes'
+
+      - name: Run tests
+        env:
+          JAX_ENABLE_CHECKS: true
+          JAX_SKIP_SLOW_TESTS: true
+          PY_COLORS: 1
+        run: |
+          python -m pip install --find-links ${{ github.workspace }}\dist jaxlib
+          python -m pip install -e ${{ github.workspace }}
+          echo "JAX_ENABLE_CHECKS=$JAX_ENABLE_CHECKS"
+          pytest -n auto --tb=short tests examples