diff --git a/.github/workflows/bazel_oneapi_presubmit.yml b/.github/workflows/bazel_oneapi_presubmit.yml new file mode 100644 index 000000000000..668d8747205e --- /dev/null +++ b/.github/workflows/bazel_oneapi_presubmit.yml @@ -0,0 +1,44 @@ +name: Bazel oneAPI + +on: + workflow_dispatch: + inputs: + halt-for-connection: + description: "Should this workflow run wait for a remote connection?" + type: choice + required: true + default: "no" + options: + - "yes" + - "no" + pull_request: + branches: + - main + push: + branches: + - main + - "release/**" + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} + cancel-in-progress: ${{ !contains(github.ref, 'release/') && github.ref != 'main' }} + +permissions: {} + +jobs: + build-oneapi-artifacts: + if: github.repository_owner == 'jax-ml' + uses: ./.github/workflows/build_oneapi_artifacts.yml + strategy: + fail-fast: false + matrix: + runner: ["linux-x86-n4-16"] + artifact: ["jax-oneapi-plugin", "jax-oneapi-pjrt"] + python: ["3.14"] + + name: "Build oneAPI wheels" + with: + runner: ${{ matrix.runner }} + artifact: ${{ matrix.artifact }} + python: ${{ matrix.python }} + clone_main_xla: 0 diff --git a/.github/workflows/build_oneapi_artifacts.yml b/.github/workflows/build_oneapi_artifacts.yml new file mode 100644 index 000000000000..a316596a5e7d --- /dev/null +++ b/.github/workflows/build_oneapi_artifacts.yml @@ -0,0 +1,126 @@ +# CI - Build JAX oneAPI Wheels +# +# This reusable workflow builds the JAX oneAPI plugin and PJRT wheels. +# It can be triggered manually or called from another workflow. + +name: CI - Build JAX oneAPI Wheels + +on: + workflow_dispatch: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: choice + default: "linux-x86-n4-16" + options: + - "linux-x86-n4-16" + + artifact: + description: "Which JAX oneAPI wheel should be built?" + type: choice + default: "jax-oneapi-plugin" + options: + - "jax-oneapi-plugin" + - "jax-oneapi-pjrt" + + python: + description: "Which Python version should the wheel be built for?" + type: choice + default: "3.12" + options: + - "3.12" + - "3.13" + - "3.14" + + clone_main_xla: + description: "Should the latest XLA main branch be used?" + type: choice + default: "0" + options: + - "1" + - "0" + + halt-for-connection: + description: "Should this workflow wait for a remote connection?" + type: choice + default: "no" + options: + - "yes" + - "no" + + workflow_call: + inputs: + runner: + description: "Which runner should the workflow run on?" + type: string + default: "linux-x86-n4-16" + + artifact: + description: "Which JAX oneAPI wheel should be built?" + type: string + default: "jax-oneapi-plugin" + + python: + description: "Which Python version should the wheel be built for?" + type: string + default: "3.12" + + clone_main_xla: + description: "Should the latest XLA main branch be used?" + type: string + default: "0" + + halt-for-connection: + description: "Should this workflow wait for a remote connection?" + type: string + default: "no" + +permissions: + contents: read + +jobs: + build-oneapi-wheel: + name: >- + Build ${{ inputs.artifact }}, + Python ${{ inputs.python }}, + clone main XLA=${{ inputs.clone_main_xla }} + + runs-on: ${{ inputs.runner }} + + container: + image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest # zizmor: ignore[unpinned-images] + + defaults: + run: + shell: bash + + env: + JAXCI_HERMETIC_PYTHON_VERSION: ${{ inputs.python }} + JAXCI_CLONE_MAIN_XLA: ${{ inputs.clone_main_xla }} + JAXCI_BUILD_ARTIFACT_WITH_RBE: "1" + + steps: + - name: Check out JAX + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + + - name: Wait for connection + uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c + with: + halt-dispatch-input: ${{ inputs.halt-for-connection }} + + - name: Build ${{ inputs.artifact }} wheel + timeout-minutes: 60 + env: + ARTIFACT: ${{ inputs.artifact }} + run: ./ci/build_oneapi_artifacts.sh "$ARTIFACT" + + - name: List generated wheels + run: | + echo "Generated wheel artifacts:" + ls -lh dist/*.whl + + - name: Clean up root-owned __pycache__ from actions + if: always() + run: find "$GITHUB_ACTION_PATH/../.." -name __pycache__ -type d -exec rm -rf {} + 2>/dev/null || true diff --git a/ci/build_oneapi_artifacts.sh b/ci/build_oneapi_artifacts.sh index c6d4fd7999ab..37941fa5b50c 100755 --- a/ci/build_oneapi_artifacts.sh +++ b/ci/build_oneapi_artifacts.sh @@ -80,7 +80,7 @@ bazel_remote_cache="" if [[ "$JAXCI_BUILD_ARTIFACT_WITH_RBE" == 1 ]]; then bazelrc_config="rbe_${bazelrc_config}" - bazel_cpu_pool_config="--config=rbe_cpu_pool" + bazel_cpu_pool_config="--bazel_options=--config=rbe_cpu_pool" else bazelrc_config="ci_${bazelrc_config}" bazel_cpu_pool_config=""