CI - Build JAX Artifacts #305
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: | |
| 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 | |
| gcs_upload_uri: | |
| description: "GCS location URI to where the artifacts should be uploaded" | |
| required: false | |
| default: ' gs://general-ml-ci-transient/jax-github-actions/jax-fork/${{ github.workflow }}/${{ github.run_number }}' | |
| 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: Upload artifacts to GCS bucket (non-Windows) | |
| if: >- | |
| ${{ inputs.upload_artifacts && !contains(inputs.runner, 'windows-x86') }} | |
| run: gsutil -m cp -r $(pwd)/dist/*.whl "${{ inputs.gcs_upload_uri }}"/ | |
| - 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 "${{ inputs.gcs_upload_uri }}"/ |