~Sandbox #530
Workflow file for this run
This file contains 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: "~Sandbox" | |
| on: | |
| workflow_dispatch: | |
| inputs: | |
| ARCHITECTURE: | |
| type: string | |
| required: true | |
| CUDA_IMAGE: | |
| type: string | |
| required: true | |
| default: 'latest' | |
| SRC_JAX: | |
| type: string | |
| required: true | |
| default: 'https://github.com/google/jax.git#main' | |
| SRC_XLA: | |
| type: string | |
| required: true | |
| default: 'https://github.com/openxla/xla.git#main' | |
| SRC_TE: | |
| type: string | |
| required: true | |
| default: 'https://github.com/NVIDIA/TransformerEngine.git#main' | |
| SRC_T5X: | |
| type: string | |
| required: true | |
| default: 'https://github.com/google-research/t5x.git#main' | |
| permissions: | |
| contents: read # to fetch code | |
| actions: write # to cancel previous workflows | |
| packages: write # to upload container | |
| jobs: | |
| metadata: | |
| runs-on: ubuntu-22.04 | |
| outputs: | |
| BUILD_DATE: ${{ steps.date.outputs.BUILD_DATE }} | |
| REPO_JAX: ${{ steps.parse-inputs.outputs.REPO_JAX }} | |
| REF_JAX: ${{ steps.parse-inputs.outputs.REF_JAX }} | |
| REPO_XLA: ${{ steps.parse-inputs.outputs.REPO_XLA }} | |
| REF_XLA: ${{ steps.parse-inputs.outputs.REF_XLA }} | |
| REPO_TE: ${{ steps.parse-inputs.outputs.REPO_TE }} | |
| REF_TE: ${{ steps.parse-inputs.outputs.REF_TE }} | |
| REPO_T5X: ${{ steps.parse-inputs.outputs.REPO_T5X }} | |
| REF_T5X: ${{ steps.parse-inputs.outputs.REF_T5X }} | |
| steps: | |
| - name: Check out the repository under ${GITHUB_WORKSPACE} | |
| uses: actions/checkout@v3 | |
| - name: Set build date | |
| id: date | |
| shell: bash -x -e {0} | |
| run: | | |
| BUILD_DATE=$(TZ='US/Los_Angeles' date '+%Y-%m-%d') | |
| echo "BUILD_DATE=${BUILD_DATE}" >> $GITHUB_OUTPUT | |
| - name: split input "repo#ref" into repo and ref parts | |
| id: parse-inputs | |
| shell: bash -x -e {0} | |
| run: | | |
| source .github/workflows/scripts/parse_git_src.sh | |
| # default values are for `pull_request` event types | |
| parse_git_src JAX "${{ inputs.SRC_JAX }}" | |
| parse_git_src XLA "${{ inputs.SRC_XLA }}" | |
| parse_git_src TE "${{ inputs.SRC_TE }}" | |
| parse_git_src T5X "${{ inputs.SRC_T5X }}" | |
| build-base: | |
| needs: metadata | |
| uses: ./.github/workflows/_build_base.yaml | |
| with: | |
| ARCHITECTURE: ${{ inputs.ARCHITECTURE }} | |
| BASE_IMAGE: ${{ inputs.CUDA_IMAGE || 'latest' }} | |
| BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }} | |
| secrets: inherit | |
| build-jax: | |
| needs: [metadata, build-base] | |
| uses: ./.github/workflows/_build_jax.yaml | |
| with: | |
| ARCHITECTURE: ${{ inputs.ARCHITECTURE }} | |
| BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }} | |
| BASE_IMAGE: ${{ needs.build-base.outputs.DOCKER_TAG }} | |
| REPO_JAX: ${{ needs.metadata.outputs.REPO_JAX }} | |
| REF_JAX: ${{ needs.metadata.outputs.REF_JAX }} | |
| REPO_XLA: ${{ needs.metadata.outputs.REPO_XLA }} | |
| REF_XLA: ${{ needs.metadata.outputs.REF_XLA }} | |
| REPO_TE: ${{ needs.metadata.outputs.REPO_TE }} | |
| REF_TE: ${{ needs.metadata.outputs.REF_TE }} | |
| secrets: inherit | |
| build-t5x: | |
| needs: [metadata, build-jax] | |
| if: inputs.ARCHITECTURE == 'amd64' # T5X arm64 build is wip in PR 252 | |
| uses: ./.github/workflows/_build_t5x.yaml | |
| with: | |
| ARCHITECTURE: amd64 | |
| BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }} | |
| BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }} | |
| REPO_T5X: ${{ needs.metadata.outputs.REPO_T5X }} | |
| REF_T5X: ${{ needs.metadata.outputs.REF_T5X }} | |
| secrets: inherit | |
| test-distribution: | |
| needs: metadata | |
| uses: ./.github/workflows/_test_distribution.yaml | |
| secrets: inherit | |
| test-t5x: | |
| needs: build-t5x | |
| if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a | |
| uses: ./.github/workflows/_test_t5x.yaml | |
| with: | |
| T5X_IMAGE: ${{ needs.build-t5x.outputs.DOCKER_TAG_FINAL }} | |
| secrets: inherit | |
| test-t5x-fmha: | |
| needs: build-t5x | |
| if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a | |
| uses: ./.github/workflows/_test_t5x_fmha.yaml | |
| with: | |
| T5X_IMAGE: ${{ needs.build-t5x.outputs.DOCKER_TAG_FINAL }} | |
| secrets: inherit | |