name: Build Docs

on:
  push:
    branches:
      - main
      - release/*
    tags:
      - v[0-9]+.[0-9]+.[0-9]
      - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
  pull_request:
  workflow_dispatch:

concurrency:
  group: build-docs-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
  cancel-in-progress: true

defaults:
  run:
    shell: bash -l -eo pipefail {0}

jobs:
  build_docs:
    if: github.repository_owner == 'pytorch'
    runs-on: ubuntu-latest
    strategy:
      matrix:
        python-version: ['3.11']
    steps:
      - name: Check out repo
        uses: actions/checkout@v4
      - name: Setup conda env
        uses: conda-incubator/setup-miniconda@v2
        with:
          auto-update-conda: true
          miniconda-version: "latest"
          activate-environment: test
          python-version: ${{ matrix.python-version }}
      - name: Update pip
        run: python -m pip install --upgrade pip
      - name: Install dependencies
        run: |
          python -m pip install torch torchvision torchao
          python -m pip install -e .
          cd docs
          python -m pip install -r requirements.txt
      - name: Build docs
        env:
            TORCHTUNE_VERSION_DOCS: ${{ github.ref }}
        run: |
          cd docs
          make html
      - uses: actions/upload-artifact@v4
        with:
          name: Built-Docs
          path: docs/build/html/

  doc-preview:
    runs-on: [self-hosted, linux.2xlarge]
    needs: build_docs
    if: ${{ github.repository_owner == 'pytorch' && github.event_name == 'pull_request' }}
    steps:
      - name: Checkout
        uses: actions/checkout@v4
      - name: Download artifact
        uses: actions/download-artifact@v4
        with:
          name: Built-Docs
          path: docs
      - name: Add no-index tag
        run: |
          find docs -name "*.html" -print0 | xargs -0 sed -i '/<head>/a \ \ <meta name="robots" content="noindex">';
      - name: Upload docs preview
        uses: seemethere/upload-artifact-s3@v5
        if: ${{ github.event_name == 'pull_request' }}
        with:
          retention-days: 14
          s3-bucket: doc-previews
          if-no-files-found: error
          path: docs
          s3-prefix: pytorch/torchtune/${{ github.event.pull_request.number }}

  upload:
    runs-on: ubuntu-latest
    needs: build_docs
    if: github.repository_owner == 'pytorch' && github.event_name == 'push' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/v'))
    environment: ${{ (github.ref == 'refs/heads/main' || startsWith(github.event.ref, 'refs/tags/v')) && 'docs-push' || '' }}
    steps:
      - name: Checkout
        uses: actions/checkout@v4
        with:
          ref: gh-pages
          persist-credentials: false
      - name: Download artifact
        uses: actions/download-artifact@v4
        with:
          name: Built-Docs
          path: docs
      - name: Add no-index tag
        run: |
          REF_NAME=$(echo "${{ github.ref }}")
          echo "Ref name: ${REF_NAME}"
          if [[ "${{ github.ref }}" == 'refs/heads/main' ]]; then
            find docs -name "*.html" -print0 | xargs -0 sed -i '/<head>/a \ \ <meta name="robots" content="noindex">';
          fi
      - name: Move and commit changes
        env:
          GITHUB_PYTORCHBOT_TOKEN: ${{ secrets.GH_PYTORCHBOT_TOKEN }}
        run: |
          git remote set-url origin https://pytorchbot:${GITHUB_PYTORCHBOT_TOKEN}@github.com/pytorch/torchtune.git
          set -euo pipefail

          # Convert refs/tags/v1.12.0rc3 into 1.12.
          # Adopted from https://github.com/pytorch/pytorch/blob/main/.github/workflows/_docs.yml#L150C11-L155C13
          GITHUB_REF=${{ github.ref }}
          if [[ "${GITHUB_REF}" =~ ^refs/tags/v([0-9]+\.[0-9]+)\.* ]]; then
            TARGET_FOLDER="${BASH_REMATCH[1]}"
          else
            TARGET_FOLDER="main"
          fi

          echo "Target Folder: ${TARGET_FOLDER}"
          mkdir -p "${TARGET_FOLDER}"
          rm -rf "${TARGET_FOLDER}"/*
          mv docs/* "${TARGET_FOLDER}"

          git config user.name 'pytorchbot'
          git config user.email 'soumith+bot@pytorch.org'
          git add "${TARGET_FOLDER}" || true
          git commit -m "auto-generating sphinx docs" || true
          git push -f
