nightly-rosetta-t5x-build-test.yaml

name: Nightly Rosetta T5x build and test
run-name: Nightly Rosetta T5x build and test (${{ github.event_name == 'workflow_run' && format('nightly {0}', github.event.workflow_run.created_at) || github.event_name }})

on:
  workflow_run:
    workflows: [Nightly T5X build]
    types: [completed]
    branches: [main]
  workflow_dispatch:
    inputs:
      BASE_IMAGE:
        type: string
        description: 'T5x image built by NVIDIA/JAX-Toolbox'
        default: 'ghcr.io/nvidia/upstream-t5x:latest'
        required: true
      PUBLISH:
        type: boolean
        description: Publish dated images and update the 'latest' tag?
        default: false
        required: false

env:
  BASE_LIBRARY: t5x
  DOCKER_REGISTRY: ghcr.io/nvidia

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 }}
      BASE_LIBRARY: ${{ steps.base-metadata.outputs.BASE_LIBRARY }}
      BASE_IMAGE: ${{ steps.base-metadata.outputs.BASE_IMAGE }}
      PUBLISH: ${{ steps.if-publish.outputs.PUBLISH }}
    steps:
      - name: Check if the triggering workflow failed
        id: if-upstream-failed
        shell: bash -x -e {0}
        run: |
          echo "UPSTREAM_FAILED=${{ github.event.workflow_run.conclusion != 'success' }}" >> $GITHUB_OUTPUT

      - name: Cancel workflow if upstream workflow did not success
        if: ${{ steps.if-upstream-failed.outputs.UPSTREAM_FAILED == 'true' }}
        uses: styfle/cancel-workflow-action@0.12.0

      - name: Determine if the resulting container should be 'published'
        id: if-publish
        shell: bash -x -e {0}
        run:
          # A container should be published if:
          # 1) the workflow is triggered by workflow_dispatch and the PUBLISH input is true, or
          # 2) the workflow is triggered by workflow_run (i.e., a nightly build)
          echo "PUBLISH=${{ github.event_name == 'workflow_run' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH) }}" >> $GITHUB_OUTPUT

      - 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: Set base library and image
        id: base-metadata
        shell: bash -x -e {0}
        run: |
          if [[ -z "${{ inputs.BASE_IMAGE }}" ]]; then
            BASE_IMAGE=${{ env.DOCKER_REGISTRY }}/upstream-${{ env.BASE_LIBRARY }}:latest
          else
            BASE_IMAGE=${{ inputs.BASE_IMAGE }}
          fi
          
          echo "BASE_LIBRARY=${{ env.BASE_LIBRARY }}" >> $GITHUB_OUTPUT
          echo "BASE_IMAGE=${BASE_IMAGE}" >> $GITHUB_OUTPUT
          
  amd64:
    needs: metadata
    uses: ./.github/workflows/_build_rosetta.yaml
    with:
      ARCHITECTURE: amd64
      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
      BASE_LIBRARY: ${{ needs.metadata.outputs.BASE_LIBRARY }}
      BASE_IMAGE: ${{ needs.metadata.outputs.BASE_IMAGE }}
    secrets: inherit
      
  arm64:
    needs: metadata
    runs-on: ubuntu-22.04
    outputs:
      DOCKER_TAG_MEALKIT=''
    steps:
      - name: Generate placeholder warning
        shell: bash -x -e {0}
        run: |
          echo "WARNING: arm64 build is not yet supported"
    
  public-mealkit:
    needs: [metadata, amd64, arm64]
    if: needs.metadata.output.PUBLISH == 'true' 
    uses: ./.github/workflows/_publish_container.yaml
    with: |
        SOURCE_IMAGE: |
        ${{ needs.amd64.outputs.DOCKER_TAG_MEALKIT }}
        ${{ needs.arm64.outputs.DOCKER_TAG_MEALKIT }}
      TARGET_IMAGE: upstream-pax
      TARGET_TAGS: |
        type=raw,value=mealkit,priority=500
        type=raw,value=mealkit-${{ needs.metadata.outputs.BUILD_DATE }},priority=500  

  publish-final:
    needs: [metadata, amd64, arm64]
    if: needs.metadata.outputs.PUBLISH == 'true'
    uses: ./.github/workflows/_publish_container.yaml
    with:
      SOURCE_IMAGE: |
        ${{ needs.amd64.outputs.DOCKER_TAG_FINAL }}
        ${{ needs.arm64.outputs.DOCKER_TAG_FINAL }}
      TARGET_IMAGE: upstream-t5x
      TARGET_TAGS: |
        type=raw,value=latest,priority=1000
        type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }},priority=900  

  test-unit:
    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
    needs: [metadata, amd64, arm64]
    uses: ./.github/workflows/_test_rosetta.yaml
    with:
      ROSETTA_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
    secrets: inherit

  test-t5x:
    needs: [metadata, amd64, arm64]
    uses: ./.github/workflows/_test_t5x_rosetta.yaml
    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
    with:
      T5X_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
    secrets: inherit

  test-vit:
    needs: [metadata, amd64, arm64]
    uses: ./.github/workflows/_test_vit.yaml
    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
    with:
      ROSETTA_T5X_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
    secrets: inherit

  publish-t5x:
    needs: [metadata, test-t5x, test-vit]
    uses: ./.github/workflows/_publish_t5x_results.yaml
    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
    with:
      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
      EXPERIMENT_SUBDIR: ROSETTA_T5X
    secrets: inherit

  publish-test:
    needs: [metadata, test-unit, test-t5x, test-vit]
    uses: ./.github/workflows/_publish_badge.yaml
    if: ( always() )
    secrets: inherit
    with:
      ENDPOINT_FILENAME: 'rosetta-t5x-overall-test-status.json'
      PUBLISH: ${{ github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true' }}
      SCRIPT: |
        UNIT_STATUS=${{ needs.test-unit.outputs.TEST_STATUS }}
        T5X_STATUS=${{ needs.test-t5x.outputs.TEST_STATUS }}
        VIT_STATUS=${{ needs.test-vit.outputs.TEST_STATUS }}

        echo "LABEL='Tests'" >> $GITHUB_OUTPUT

        if [[ ${{ needs.build.result }} == "success" ]]; then
          if [[ $UNIT_STATUS == "success" ]] && [[ $T5X_STATUS == "success" ]] && [[ $VIT_STATUS == "success" ]]; then
            COLOR=brightgreen
            MESSAGE="Unit passed / MGMN passed"
          elif [[ $UNIT_STATUS == "success" ]]; then 
            COLOR=yellow
            MESSAGE="Unit passed / MGMN failed"
          elif [[ $T5X_STATUS == "success" ]] && [[ $VIT_STATUS == "success" ]]; then
            COLOR=yellow
            MESSAGE="Unit failed / MGMN passed"
          else
            COLOR=red
            MESSAGE="Unit failed / MGMN failed"
          fi
        else
          MESSAGE="n/a"
          COLOR="red"
        fi

        echo "MESSAGE='${MESSAGE}'" >> $GITHUB_OUTPUT
        echo "COLOR='${COLOR}'" >> $GITHUB_OUTPUT

  finalize:
    if: always()
    needs: [metadata, amd64, arm64]
    uses: ./.github/workflows/_finalize.yaml
    with:
      PUBLISH_BADGE: ${{ needs.metadata.outputs.PUBLISH == 'true' }}
    secrets: inherit