Test CI scripts and workflows #295
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 JAX Artifacts | |
| on: | |
| pull_request: | |
| branches: | |
| - main | |
| workflow_dispatch: | |
| inputs: | |
| runner: | |
| description: "Which runner should the workflow run on?" | |
| type: string | |
| required: true | |
| default: "linux-x86-n2-16" | |
| artifact: | |
| description: "Which JAX artifact to build?" | |
| type: string | |
| required: true | |
| default: "jaxlib" | |
| python: | |
| description: "Which python version should the artifact be built for?" | |
| type: string | |
| required: true | |
| default: "3.12" | |
| clone_main_xla: | |
| description: "Should latest XLA be used? (1 to enable, 0 to disable)" | |
| type: string | |
| required: false | |
| default: "0" | |
| halt-for-connection: | |
| description: 'Should this workflow run wait for a remote connection?' | |
| type: choice | |
| required: true | |
| default: 'no' | |
| options: | |
| - 'yes' | |
| - 'no' | |
| workflow_call: | |
| inputs: | |
| runner: | |
| description: "Which runner should the workflow run on?" | |
| type: string | |
| required: true | |
| default: "linux-x86-n2-16" | |
| artifact: | |
| description: "Which JAX artifact to build?" | |
| type: string | |
| required: true | |
| default: "jaxlib" | |
| python: | |
| description: "Which python version should the artifact be built for?" | |
| type: string | |
| required: true | |
| default: "3.12" | |
| clone_main_xla: | |
| description: "Should latest XLA be used? (1 to enable, 0 to disable)" | |
| type: string | |
| required: false | |
| default: "0" | |
| upload_artifacts: | |
| description: "Should the artifacts be uploaded to a GCS bucket?" | |
| required: false | |
| default: false | |
| type: boolean | |
| upload_destination_prefix: | |
| description: "GCS location prefix to where the artifacts should be uploaded" | |
| required: false | |
| default: '${{ github.workflow }}/${{ github.run_number }}/${{ github.run_attempt }}' | |
| type: string | |
| jobs: | |
| build_artifacts: | |
| defaults: | |
| run: | |
| # Explicitly set the shell to bash to override the default Windows environment, i.e, cmd. | |
| shell: bash | |
| runs-on: ${{ inputs.runner }} | |
| container: ${{ (contains(inputs.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') || | |
| (contains(inputs.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') || | |
| (contains(inputs.runner, 'windows-x86') && null) }} | |
| env: | |
| JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}" | |
| JAXCI_CLONE_MAIN_XLA: "${{ inputs.clone_main_xla }}" | |
| name: Build ${{ inputs.artifact }} (${{ inputs.runner }}, Python ${{ inputs.python }}, Clone main XLA=${{ inputs.clone_main_xla }}) | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: Enable RBE if building on Linux x86 or Windows x86 | |
| if: contains(inputs.runner, 'linux-x86') || contains(inputs.runner, 'windows-x86') | |
| run: echo "JAXCI_BUILD_ARTIFACT_WITH_RBE=1" >> $GITHUB_ENV | |
| # Halt for testing | |
| - name: Wait For Connection | |
| uses: google-ml-infra/actions/ci_connection@main | |
| with: | |
| halt-dispatch-input: ${{ inputs.halt-for-connection }} | |
| - name: Build ${{ inputs.artifact }} | |
| run: ./ci/build_artifacts.sh "${{ inputs.artifact }}" | |
| - name: Set PLATFORM env var for use in upload destination | |
| run: | | |
| os=$(uname -s | awk '{print tolower($0)}') | |
| arch=$(uname -m) | |
| # Adjust name for Windows | |
| if [[ $os =~ "msys_nt" ]]; then | |
| os="windows" | |
| fi | |
| echo "PLATFORM=${os}_${arch}" >> $GITHUB_ENV | |
| - name: Upload artifacts to GCS bucket (non-Windows) | |
| if: | | |
| inputs.upload_artifacts && ${{ !contains(inputs.runner, 'windows-x86') }} | |
| run: gsutil -m cp -r $(pwd)/dist/*.whl gs://general-ml-ci-transient/jax-github-actions/"${{ inputs.upload_destination_prefix }}"/$PLATFORM/python${JAXCI_HERMETIC_PYTHON_VERSION}/ | |
| - name: Upload artifacts to GCS bucket (Windows) | |
| if: | | |
| inputs.upload_artifacts && ${{ contains(inputs.runner, 'windows-x86') }} | |
| shell: cmd | |
| run: gsutil -m cp -r dist/*.whl gs://general-ml-ci-transient/jax-github-actions/"${{ inputs.upload_destination_prefix }}"/$PLATFORM/python${JAXCI_HERMETIC_PYTHON_VERSION}/ | |