Update override_commit in workspace.bzl (#2221) #8068
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: Build Enzyme-JAX | |
| on: | |
| push: | |
| branches: | |
| - main | |
| paths: &paths | |
| - ".github/scripts/process_benchmark_results.py" | |
| - ".github/workflows/build.yml" | |
| - "**/BUILD" | |
| - "**/WORKSPACE" | |
| - "**/*.bzl" | |
| - "builddeps/**" | |
| - "patches/**" | |
| - "src/**" | |
| - "test/**" | |
| - "third_party/**" | |
| pull_request: | |
| branches: | |
| - main | |
| paths: *paths | |
| merge_group: | |
| concurrency: | |
| # Skip intermediate builds: always. | |
| # Cancel intermediate builds: only if it is a pull request build. | |
| group: ${{ github.workflow }}-${{ github.ref }} | |
| cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} | |
| jobs: | |
| build: | |
| name: Build - ${{ matrix.os }} - python ${{ matrix.python }} | |
| runs-on: ${{ matrix.os }} | |
| container: | |
| image: ${{ (contains(matrix.os, 'linux') && 'ghcr.io/enzymead/reactant-docker-images@sha256:cd45d851f5ea544f88d042eafefa53d948c229fffcab6189019324e3b02b505a' ) || '' }} | |
| # container: | |
| # image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest | |
| # container: | |
| # image: ${{ (contains(matrix.runner, 'a100') && 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda12.8-cudnn9.8:latest') || | |
| # (contains(matrix.runner, 'linux-x86') && 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest') || | |
| # (contains(matrix.runner, 'linux-arm64') && 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest') }} | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| os: | |
| - linux-x86-n2-32 | |
| - linux-arm64-c4a-32 | |
| - linux-x86-ct6e-180-4tpu | |
| - linux-x86-a2-48-a100-4gpu | |
| - macOS-latest | |
| python: | |
| - "3.11" | |
| env: | |
| HERMETIC_PYTHON_VERSION: ${{ matrix.python }} | |
| steps: | |
| - name: Check GPUs | |
| if: ${{ contains(matrix.os, 'a100') }} | |
| run: | | |
| nvidia-smi | |
| - name: Set BAZEL_FLAGS | |
| shell: bash | |
| run: | | |
| set -e | |
| BAZEL_FLAGS="" | |
| if [ ${{ runner.os }} == macOS ]; then | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --define using_clang=true" | |
| fi | |
| if [ ${{ runner.arch }} == ARM64 ] && [ ${{ runner.os }} == Linux ]; then | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --linkopt=-fuse-ld=lld" | |
| fi | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_TOPOLOGY" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_WORKER_ID" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_SKIP_MDS_QUERY=true" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_TOPOLOGY_WRAP" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_CHIPS_PER_HOST_BOUNDS" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_ACCELERATOR_TYPE" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_RUNTIME_METRICS_PORTS" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_TOPOLOGY_ALT" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_HOST_BOUNDS" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=TPU_WORKER_HOSTNAMES" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=CHIPS_PER_HOST_BOUNDS" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=HOST_BOUNDS" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=ALT=false" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=WRAP" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=VBAR_CONTROL_SERVICE_URL" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --repo_env=CC=`which clang`" | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --define=using_clang=true" | |
| if [[ "${{ matrix.os }}" == *"tpu"* ]]; then | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=JAX_PLATFORMS=tpu" | |
| elif [[ "${{ matrix.os }}" == *"a100"* ]]; then | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=JAX_PLATFORMS=cuda" | |
| else | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --test_env=JAX_PLATFORMS=cpu" | |
| fi | |
| BAZEL_FLAGS="${BAZEL_FLAGS} --config=public_cache_push" | |
| echo "BAZEL_FLAGS=${BAZEL_FLAGS}" >> "${GITHUB_ENV}" | |
| - uses: actions/checkout@v6 | |
| - uses: bazel-contrib/setup-bazel@0.16.0 | |
| name: Set up Bazel | |
| with: | |
| # Avoid downloading Bazel every time. | |
| bazelisk-cache: true | |
| # Store build cache per workflow only for macOS. | |
| disk-cache: ${{ runner.os == 'macOS' && format('build-{0}-{1}-{2}', github.workflow, matrix.os, matrix.python) || false }} | |
| # Do not share repository cache between workflows. | |
| repository-cache: false | |
| bazelisk-version: 1.x | |
| - name: Update JAX version | |
| run: | | |
| if [[ "${{ matrix.os }}" == *"tpu"* ]]; then | |
| sed -i.bak0 "s/jax\[cuda12\]/jax[tpu]/g" builddeps/requirements.in | |
| sed -i.bak0 "s/pypi_jax_cuda12_plugin/pypi_libtpu/g" test/BUILD | |
| elif [[ "${{ matrix.os }}" == *"a100"* ]]; then | |
| sed -i.bak0 "s/jax\[cuda12\]/jax[cuda13]/g" builddeps/requirements.in | |
| sed -i.bak0 "s/pypi_jax_cuda12_plugin/pypi_jax_cuda13_plugin/g" test/BUILD | |
| else | |
| sed -i.bak0 "s/jax\[cuda12\]/jax[cpu]/g" builddeps/requirements.in | |
| sed -i.bak0 "/pypi_jax_cuda12_plugin/d" test/BUILD | |
| fi | |
| cat builddeps/requirements.in | |
| cat test/BUILD | |
| shell: bash | |
| - name: Update Python requirements files | |
| run: | | |
| bazel run $BAZEL_FLAGS --color=yes //builddeps:requirements.update | |
| - name: Build enzymexlamlir-opt | |
| timeout-minutes: 180 | |
| run: | | |
| bazel build $BAZEL_FLAGS --color=yes //:enzymexlamlir-opt --test_output=errors | |
| - name: Upload enzymexlamlir-opt | |
| uses: actions/upload-artifact@v7 | |
| with: | |
| name: enzymexlamlir-opt-${{ matrix.os }}-${{ matrix.python }} | |
| path: "bazel-bin/enzymexlamlir-opt" | |
| retention-days: 90 | |
| - name: Run lit tests | |
| id: lit_tests | |
| continue-on-error: true | |
| run: | | |
| bazel test $BAZEL_FLAGS \ | |
| --color=yes \ | |
| --test_output=errors \ | |
| --experimental_ui_max_stdouterr_bytes=-1 \ | |
| --test_verbose_timeout_warnings \ | |
| //test/lit_tests/... | |
| - name: Run lit tests (dbg) | |
| if: steps.lit_tests.outcome == 'failure' | |
| run: | | |
| bazel test $BAZEL_FLAGS -c dbg \ | |
| --color=yes \ | |
| --test_output=errors \ | |
| --experimental_ui_max_stdouterr_bytes=-1 \ | |
| --test_verbose_timeout_warnings \ | |
| //test/lit_tests/... | |
| - name: Build wheel | |
| timeout-minutes: 240 | |
| run: | | |
| bazel build $BAZEL_FLAGS --color=yes :wheel --test_output=errors | |
| - name: Upload wheel | |
| uses: actions/upload-artifact@v7 | |
| if: ${{ !endsWith(matrix.os, 'arm') }} | |
| with: | |
| name: Enzyme-JAX-Wheel-${{ matrix.os }}-${{ matrix.python }} | |
| path: "bazel-bin/*.whl" | |
| retention-days: 90 | |
| - name: Run tests | |
| id: tests | |
| continue-on-error: true | |
| run: | | |
| bazel test $BAZEL_FLAGS \ | |
| --color=yes \ | |
| --test_output=streamed \ | |
| --experimental_ui_max_stdouterr_bytes=-1 \ | |
| --test_verbose_timeout_warnings \ | |
| //test:python_tests | |
| - name: Run tests (dbg) | |
| if: steps.tests.outcome == 'failure' | |
| run: | | |
| bazel test $BAZEL_FLAGS \ | |
| --color=yes \ | |
| --test_output=streamed \ | |
| --experimental_ui_max_stdouterr_bytes=-1 \ | |
| --test_verbose_timeout_warnings \ | |
| -c dbg \ | |
| //test:python_tests | |
| - name: Upload test results | |
| uses: actions/upload-artifact@v7 | |
| if: ${{ !cancelled() }} | |
| with: | |
| name: Enzyme-JAX-tests-${{ matrix.os }}-${{ matrix.python }} | |
| path: | | |
| bazel-out/**/test.log | |
| retention-days: 90 | |
| - name: Upload benchmark results | |
| uses: actions/upload-artifact@v7 | |
| if: ${{ !cancelled() }} | |
| with: | |
| name: Enzyme-JAX-benchmark-${{ matrix.os }}-${{ matrix.python }} | |
| path: | | |
| bazel-testlogs/test/**/test.outputs/*.zip | |
| retention-days: 90 | |
| - name: Extract opt and debug enzymexlamlir-opt for release | |
| if: ${{ !cancelled() && github.ref == 'refs/heads/main' && matrix.os == 'linux-x86-n2-32' }} | |
| run: | | |
| bazel build $BAZEL_FLAGS --color=yes //:enzymexlamlir-opt @xla//xla/hlo/tools:hlo-translate --test_output=errors | |
| cp bazel-bin/enzymexlamlir-opt enzymexlamlir-opt | |
| cp bazel-bin/external/xla/xla/hlo/tools/hlo-translate hlo-translate | |
| - name: Update Nightly Release | |
| uses: ncipollo/release-action@v1 | |
| if: ${{ !cancelled() && github.ref == 'refs/heads/main' && matrix.os == 'linux-x86-n2-32' }} | |
| with: | |
| tag: nightly | |
| name: Nightly Release | |
| body: "Latest nightly build of Enzyme-JAX." | |
| artifacts: "enzymexlamlir-opt,hlo-translate,bazel-bin/*.whl" | |
| allowUpdates: true | |
| removeArtifacts: true | |
| replacesArtifacts: true | |
| token: ${{ secrets.GITHUB_TOKEN }} | |
| benchmark-summary: | |
| name: Benchmark Summary | |
| runs-on: ubuntu-latest | |
| needs: build | |
| if: ${{ !cancelled() }} | |
| permissions: | |
| contents: write # contents permission to update benchmark contents in gh-pages branch | |
| statuses: read | |
| deployments: write # deployments permission to deploy GitHub pages website | |
| pull-requests: write | |
| steps: | |
| - name: Install Dependencies | |
| run: | | |
| curl -LsSf https://astral.sh/uv/install.sh | sh | |
| - uses: actions/checkout@v6 | |
| - name: Download all benchmark artifacts | |
| uses: actions/download-artifact@v8 | |
| with: | |
| pattern: Enzyme-JAX-benchmark-* | |
| path: benchmark-results | |
| merge-multiple: false | |
| - name: Extract and combine benchmark results | |
| run: | | |
| uv run --no-project .github/scripts/process_benchmark_results.py \ | |
| --artifact_dir=benchmark-results \ | |
| --csv_file=combined_results.csv \ | |
| --json_file=combined_results.json | |
| - name: Upload benchmark results | |
| uses: actions/upload-artifact@v7 | |
| if: ${{ !cancelled() }} | |
| with: | |
| name: Enzyme-JAX-combined-benchmark-results | |
| path: | | |
| combined_results.csv | |
| combined_results.json | |
| retention-days: 90 | |
| - name: Cache and upload benchmark results | |
| uses: benchmark-action/github-action-benchmark@v1 | |
| with: | |
| name: EnzymeJAX Benchmarks | |
| tool: "customSmallerIsBetter" | |
| output-file-path: combined_results.json | |
| benchmark-data-dir-path: "benchmarks" | |
| github-token: ${{ secrets.GITHUB_TOKEN }} | |
| comment-always: false | |
| summary-always: ${{ github.event.pull_request.head.repo.full_name == github.repository }} | |
| alert-threshold: "150%" | |
| fail-on-alert: false | |
| auto-push: ${{ github.event_name != 'pull_request' }} | |
| max-items-in-chart: 50 |