Skip to content

Fix docker tag for nightly vit tests #1138

Fix docker tag for nightly vit tests

Fix docker tag for nightly vit tests #1138

Workflow file for this run

name: CI
on:
pull_request:
paths-ignore:
- '**.md'
workflow_dispatch:
inputs:
CUDA_IMAGE:
type: string
description: 'Base CUDA image, e.g. nvidia/cuda:X.Y.Z-devel-ubuntu22.04'
required: false
default: 'latest'
SRC_JAX:
description: 'JAX source: <repo>#<branch|tag|commit>'
type: string
required: true
default: 'https://github.com/google/jax.git#main'
SRC_XLA:
description: 'XLA source: <repo>#<branch|tag|commit>'
type: string
required: true
default: 'https://github.com/openxla/xla.git#main'
SRC_TE:
description: 'TE source: <repo>#<branch|tag|commit>'
type: string
required: true
# TODO: This is a temporary pinning of TE as the API in TE no longer matches the TE patch
# This should be reverted to main ASAP
default: 'https://github.com/NVIDIA/TransformerEngine.git#7976bd003fcf084dd068069b92a9a79b1743316a'
SRC_T5X:
description: 'T5X source: <repo>#<branch|tag|commit>'
type: string
required: true
default: 'https://github.com/google-research/t5x.git#main'
SRC_PAXML:
description: 'Paxml source: <repo>#<branch|tag|commit>'
type: string
required: true
default: 'https://github.com/google/paxml.git#main'
SRC_PRAXIS:
description: 'Praxis source: <repo>#<branch|tag|commit>'
type: string
required: true
default: 'https://github.com/google/praxis.git#main'
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
permissions:
contents: read # to fetch code
actions: write # to cancel previous workflows
packages: write # to upload container
jobs:
amd64:
uses: ./.github/workflows/_ci.yaml
with:
ARCHITECTURE: amd64
CUDA_IMAGE: ${{ inputs.CUDA_IMAGE || 'latest' }}
SRC_JAX: ${{ inputs.SRC_JAX || 'https://github.com/google/jax.git#main' }}
SRC_XLA: ${{ inputs.SRC_XLA || 'https://github.com/openxla/xla.git#main'}}
SRC_TE: ${{ inputs.SRC_TE || 'https://github.com/NVIDIA/TransformerEngine.git#7976bd003fcf084dd068069b92a9a79b1743316a'}}
SRC_T5X: ${{ inputs.SRC_T5X || 'https://github.com/google-research/t5x.git#main'}}
SRC_PAXML: ${{ inputs.SRC_PAXML || 'https://github.com/google/paxml.git#main'}}
SRC_PRAXIS: ${{ inputs.SRC_PRAXIS || 'https://github.com/google/praxis.git#main'}}
secrets: inherit
arm64:
uses: ./.github/workflows/_ci.yaml
with:
ARCHITECTURE: arm64
CUDA_IMAGE: ${{ inputs.CUDA_IMAGE || 'latest' }}
SRC_JAX: ${{ inputs.SRC_JAX || 'https://github.com/google/jax.git#main' }}
SRC_XLA: ${{ inputs.SRC_XLA || 'https://github.com/openxla/xla.git#main'}}
SRC_TE: ${{ inputs.SRC_TE || 'https://github.com/NVIDIA/TransformerEngine.git#7976bd003fcf084dd068069b92a9a79b1743316a'}}
SRC_T5X: ${{ inputs.SRC_T5X || 'https://github.com/google-research/t5x.git#main'}}
SRC_PAXML: ${{ inputs.SRC_PAXML || 'https://github.com/google/paxml.git#main'}}
SRC_PRAXIS: ${{ inputs.SRC_PRAXIS || 'https://github.com/google/praxis.git#main'}}
secrets: inherit
finalize:
needs: [amd64, arm64]
if: always()
uses: ./.github/workflows/_finalize.yaml
with:
PUBLISH_BADGE: false
secrets: inherit