pax-cuda-121.yaml

name: Nightly Containers on CUDA 12.1
run-name: Nightly Containers on CUDA 12.1 (${{ github.event_name == 'workflow_run' && format('nightly {0}', github.event.workflow_run.created_at) || github.event_name }})

on:
  schedule:
    - cron: '30 9 * * *'  # Pacific Time 01:30 AM in UTC
  workflow_dispatch:
    inputs:
      PUBLISH:
        type: boolean
        description: Publish dated images and update the 'latest' tag?
        default: false
        required: false

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:

  metadata:
    runs-on: ubuntu-22.04
    outputs:
      BUILD_DATE: ${{ steps.date.outputs.BUILD_DATE }}
    steps:
      - 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

  build-base:
    needs: metadata
    uses: ./.github/workflows/_build_base.yaml
    with:
      BASE_IMAGE: 'nvidia/cuda:12.1.1-devel-ubuntu22.04'
      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
    secrets: inherit

  build-jax:
    needs: [metadata, build-base]
    uses: ./.github/workflows/_build_jax.yaml
    with:
      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-base.outputs.DOCKER_TAGS }}
    secrets: inherit

  build-pax:
    needs: [metadata, build-jax]
    uses: ./.github/workflows/_build_pax.yaml
    with:
      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAGS }}
    secrets: inherit

  build-rosetta-pax:
    uses: ./.github/workflows/_build_rosetta.yaml
    needs: [metadata, build-pax]
    with:
      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAGS }}
      BASE_LIBRARY: pax
      PLATFORMS: '["amd64"]'
    secrets: inherit

  build-t5x:
    needs: [metadata, build-jax]
    uses: ./.github/workflows/_build_t5x.yaml
    with:
      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAGS }}
    secrets: inherit

  build-rosetta-t5x:
    uses: ./.github/workflows/_build_rosetta.yaml
    needs: [metadata, build-t5x]
    with:
      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-t5x.outputs.DOCKER_TAGS }}
      BASE_LIBRARY: t5x
      PLATFORMS: '["amd64"]'
    secrets: inherit

  build-summary:
    needs: [build-base, build-jax, build-t5x, build-rosetta-t5x, build-pax, build-rosetta-pax]
    if: always()
    runs-on: ubuntu-22.04
    steps:
      - name: Generate job summary for container build
        shell: bash -x -e {0}
        run: |
          cat > $GITHUB_STEP_SUMMARY << EOF
          # Images created          

          | Image        | Link                                               |
          | ------------ | -------------------------------------------------- |
          | Base         | ${{ needs.build-base.outputs.DOCKER_TAGS }}        |
          | JAX          | ${{ needs.build-jax.outputs.DOCKER_TAGS }}         |
          | T5X          | ${{ needs.build-t5x.outputs.DOCKER_TAGS }}         |
          | ROSETTA(T5X) | ${{ needs.build-rosetta-t5x.outputs.DOCKER_TAGS }} |
          | PAX          | ${{ needs.build-pax.outputs.DOCKER_TAGS }}         |
          | ROSETTA(pax) | ${{ needs.build-rosetta-pax.outputs.DOCKER_TAGS }} |
          EOF

  publish-upstream-pax:
    if: github.event_name == 'schedule' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH)
    needs: [metadata, build-pax]
    uses: ./.github/workflows/_publish_container.yaml
    secrets: inherit
    with:
      SOURCE_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAGS }}
      TARGET_IMAGE: upstream-pax
      TARGET_TAGS: |
        type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }}-cuda-12.1,priority=900

  publish-pax:
    if: github.event_name == 'schedule' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH)
    needs: [metadata, build-rosetta-pax]
    uses: ./.github/workflows/_publish_container.yaml
    secrets: inherit
    with:
      SOURCE_IMAGE: ${{ needs.build-rosetta-pax.outputs.DOCKER_TAGS }}
      TARGET_IMAGE: pax
      TARGET_TAGS: |
        type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }}-cuda-12.1,priority=900

  publish-upstream-t5x:
    if: github.event_name == 'schedule' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH)
    needs: [metadata, build-t5x]
    uses: ./.github/workflows/_publish_container.yaml
    secrets: inherit
    with:
      SOURCE_IMAGE: ${{ needs.build-t5x.outputs.DOCKER_TAGS }}
      TARGET_IMAGE: upstream-t5x
      TARGET_TAGS: |
        type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }}-cuda-12.1,priority=900

  publish-t5x:
    if: github.event_name == 'schedule' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH)
    needs: [metadata, build-rosetta-t5x]
    uses: ./.github/workflows/_publish_container.yaml
    secrets: inherit
    with:
      SOURCE_IMAGE: ${{ needs.build-rosetta-t5x.outputs.DOCKER_TAGS }}
      TARGET_IMAGE: t5x
      TARGET_TAGS: |
        type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }}-cuda-12.1,priority=900

  test-jax:
    needs: build-jax
    uses: ./.github/workflows/_test_jax.yaml
    with:
      JAX_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAGS }}
    secrets: inherit

  test-pax:
    needs: build-pax
    uses: ./.github/workflows/_test_pax.yaml
    with:
      PAX_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAGS }}
    secrets: inherit

  test-t5x:
    needs: build-t5x
    uses: ./.github/workflows/_test_t5x.yaml
    with:
      T5X_IMAGE: ${{ needs.build-t5x.outputs.DOCKER_TAGS }}
    secrets: inherit

  # TODO(terry): This is missing the rosetta tests which can only be added
  # After a fix for the TB log collision is pushed.

  finalize:
    if: always()
    # TODO: use dynamic matrix to make dependencies self-updating
    needs: [build-summary, test-jax, test-pax]
    uses: ./.github/workflows/_finalize.yaml
    with:
      PUBLISH_BADGE: false
    secrets: inherit