_build_test_upload.yml

name: Build, Test and Upload Wheel

on:
  workflow_call:
    inputs:
      branch:
        required: true
        type: string
      pre_dev_release:
        required: true
        type: boolean
      pytorch_version:
        required: true
        type: string
      do-upload:
        required: false
        default: true
        type: boolean
    secrets:
      PYPI_TOKEN:
        required: false
      CONDA_PYTORCHBOT_TOKEN:
        required: false
      CONDA_TEST_PYTORCHBOT_TOKEN:
        required: false
      CONDA_NIGHTLY_PYTORCHBOT_TOKEN:
        required: false

permissions:
  id-token: write
  contents: write

jobs:
  get_release_type:
    runs-on: ubuntu-latest
    outputs:
      type: ${{ steps.get_release_type.outputs.type }}
    steps:
      - name: Get Release Type
        run: |
          if [[ "${{ inputs.branch }}" == v* ]] && [[ ${{ inputs.pre_dev_release }} == false ]]; then
            RELEASE_TYPE=official
          elif [[ "${{ inputs.branch }}" == release/* ]] && [[ ${{ inputs.pre_dev_release }} == true ]]; then
            RELEASE_TYPE=test
          else
            if [[ "${{ github.base_ref }}" == release/* ]]; then
              RELEASE_TYPE=test
            else
              RELEASE_TYPE=nightly
            fi
          fi
          echo "Release Type: $RELEASE_TYPE"
          echo "type=$RELEASE_TYPE" >> $GITHUB_OUTPUT
        id: get_release_type

  wheel_build_test:
    needs: get_release_type
    runs-on: ${{ matrix.os }}
    container: ${{ startsWith( matrix.os, 'ubuntu' ) && 'pytorch/manylinux-cpu' || null }}
    strategy:
      fail-fast: false
      matrix:
        os:
          - windows-latest
        python-version:
          - 3.9
          - "3.10"
          - "3.11"
          - "3.12"
          - "3.13"
          - pure
        exclude:
          - os: windows-latest
            python-version: pure
    steps:
      - name: Checkout Source Repository
        uses: actions/checkout@v3
        with:
          ref: ${{ inputs.branch }}
          submodules: recursive
      - name: Setup Python ${{ matrix.python-version }} for Windows
        if: ${{ startsWith( matrix.os, 'windows' ) }}
        uses: actions/setup-python@v4
        with:
          python-version: ${{ matrix.python-version }}
      - name: Add temp runner environment variables
        if: ${{ startsWith( matrix.os, 'macos' ) }}
        shell: bash -l {0}
        run: |
          echo "MINICONDA_INSTALL_PATH_MACOS=${RUNNER_TEMP}/miniconda" >> "${GITHUB_ENV}"
      - name: Setup msbuild on Windows
        if: startsWith( matrix.os, 'windows' )
        uses: microsoft/setup-msbuild@v1.1
      - name: Set up Visual Studio shell
        if: startsWith( matrix.os, 'windows' )
        uses: egor-tensin/vs-shell@v2
        with:
          arch: x64
      - name: Install Build Dependency
        shell: bash -l {0}
        env:
          PYTHON_VERSION: ${{ matrix.python-version }}
        run: |
          pip install cmake ninja setuptools
          echo "/home/runner/.local/bin" >> $GITHUB_PATH
      - name: Install PyTorch and Build TorchData Wheel
        shell: bash -l {0}
        env:
          PYTHON_VERSION: ${{ matrix.python-version }}
          PYTORCH_VERSION: ${{ inputs.pytorch_version }}
          BUILD_S3: 0
        run: |
          set -ex
          packaging/build_wheel.sh
      - name: Validate TorchData Wheel
        shell: bash -l {0}
        env:
          PYTHON_VERSION: ${{ matrix.python-version }}
        run: |
          pip3 install pkginfo
          for pkg in dist/torchdata*.whl; do
              echo "PkgInfo of $pkg:"
              pkginfo $pkg
              if ${{ startsWith( matrix.os, 'ubuntu' ) && matrix.python-version != 'pure' }}; then
                auditwheel show $pkg
                auditwheel repair $pkg --plat manylinux2014_x86_64 -w wheelhouse
              fi
          done
      - name: Install TorchData Wheel
        shell: bash -l {0}
        env:
          PYTHON_VERSION: ${{ matrix.python-version }}
        run: |
          pip3 install dist/torchdata*.whl
      - name: Run Smoke Tests
        shell: bash -l {0}
        env:
          PYTHON_VERSION: ${{ matrix.python-version }}
        run: |
          if ${{ matrix.python-version == 'pure' }}; then
            python test/smoke_test/smoke_test.py --no-s3
          else
            python test/smoke_test/smoke_test.py
          fi

      - name: Upload Wheels to Github
        if: always()
        uses: actions/upload-artifact@v3
        with:
          name: torchdata-artifact
          path: dist/torchdata*.whl

  wheel_upload:
    if: always() && inputs.branch != '' && inputs.do-upload == true
    needs: [get_release_type, wheel_build_test]
    runs-on: ubuntu-latest
    environment: pytorchbot-env
    outputs:
      upload: ${{ steps.trigger_upload.outputs.value }}
    steps:
      - name: Configure aws credentials (pytorch account)
        if: ${{ needs.get_release_type.outputs.type == 'nightly' }}
        uses: aws-actions/configure-aws-credentials@v3
        with:
          role-to-assume: arn:aws:iam::749337293305:role/gha_workflow_nightly_build_wheels
          aws-region: us-east-1

      - name: Configure aws credentials (pytorch account)
        if: ${{ needs.get_release_type.outputs.type == 'test' }}
        uses: aws-actions/configure-aws-credentials@v3
        with:
          role-to-assume: arn:aws:iam::749337293305:role/gha_workflow_test_build_wheels
          aws-region: us-east-1
      - name: Download Artifacts from Github
        continue-on-error: true
        uses: actions/download-artifact@v3
        with:
          name: torchdata-artifact
      - name: Determine if Wheel Uploading is needed
        run: |
          upload=false
          for txt in torchdata*.whl; do
            upload=true
            break
          done
          echo "value=$upload" >> $GITHUB_OUTPUT
        id: trigger_upload
      - name: Display All TorchData Wheels
        if: steps.trigger_upload.outputs.value == 'true'
        run: ls -lh torchdata*.whl
      - name: Upload Wheels to S3 Storage
        if: steps.trigger_upload.outputs.value == 'true'
        run: |
          if [[ ${{ inputs.branch }} == 'main' ]]; then
            S3_PATH=s3://pytorch/whl/nightly/
          else
            S3_PATH=s3://pytorch/whl/test/
          fi
          pip3 install awscli
          set -x
          for pkg in torchdata*.whl; do
            aws s3 cp "$pkg" "$S3_PATH" --acl public-read
          done
      - name: Upload Official Wheels to PYPI
        if: |
          steps.trigger_upload.outputs.value == 'true' &&
          needs.get_release_type.outputs.type == 'official'
        env:
          PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
        run: |
          pip3 install twine
          python -m twine upload \
            --username __token__ \
            --password "$PYPI_TOKEN" \
            torchdata*.whl

  conda_build_test:
    needs: get_release_type
    runs-on: ${{ matrix.os }}
    container: ${{ startsWith( matrix.os, 'ubuntu' ) && 'pytorch/manylinux-cpu' || null }}
    strategy:
      fail-fast: false
      matrix:
        os:
          - windows-latest
        python-version:
          - 3.9
          - "3.10"
          - "3.11"
          - "3.12"
          - "3.13"
    steps:
      - name: Checkout Source Repository
        uses: actions/checkout@v3
        with:
          ref: ${{ inputs.branch }}
          submodules: recursive
      - name: Determine if build AWSSDK
        shell: bash -l {0}
        run: |
          BUILD_S3=0
          echo "value=$BUILD_S3" >> $GITHUB_OUTPUT
        id: build_s3
      - name: Add temp runner environment variables
        shell: bash -l {0}
        run: |
          echo "MINICONDA_INSTALL_PATH=${RUNNER_TEMP}/miniconda" >> "${GITHUB_ENV}"
          echo "CONDA_ENV_PATH=${RUNNER_TEMP}/conda_build_env" >> "${GITHUB_ENV}"
      - name: Create Conda Env on Windows
        if: ${{ startsWith( matrix.os, 'windows' ) }}
        uses: conda-incubator/setup-miniconda@v2
        with:
          python-version: ${{ matrix.python-version }}
          activate-environment: ${{ env.CONDA_ENV_PATH }}
      - name: Build TorchData for Conda
        shell: bash -l {0}
        env:
          PYTHON_VERSION: ${{ matrix.python-version }}
          PYTORCH_VERSION: ${{ inputs.pytorch_version }}
          BUILD_S3: ${{ steps.build_s3.outputs.value }}
        run: |
          set -ex
          if ${{ ! startsWith( matrix.os, 'windows' ) }}; then
            source "${MINICONDA_INSTALL_PATH}/etc/profile.d/conda.sh"
          fi
          conda activate "${CONDA_ENV_PATH}"
          if [[ "${PYTHON_VERSION}" = "3.11" ]]; then
            conda install -yq conda-build -c malfet -c conda-forge
            export CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c malfet"
          else
            conda install -yq conda-build
          fi
          packaging/build_conda.sh
          # conda index ./conda-bld
      - name: Upload Conda Package to Github
        if: always()
        uses: actions/upload-artifact@v3
        with:
          name: torchdata-artifact
          path: conda-bld/*/torchdata-*.tar.bz2

  conda_upload:
    if: always() && inputs.branch != '' && inputs.do-upload == true
    needs: [get_release_type, conda_build_test]
    runs-on: ubuntu-latest
    container: continuumio/miniconda3
    environment: pytorchbot-env
    outputs:
      upload: ${{ steps.trigger_upload.outputs.value }}
    steps:
      - name: Download Artifacts from Github
        continue-on-error: true
        uses: actions/download-artifact@v3
        with:
          name: torchdata-artifact
      - name: Determine if Conda Uploading is needed
        run: |
          upload=false
          for pkg in ./*/torchdata-*.tar.bz2; do
            upload=true
            break
          done
          echo "value=$upload" >> $GITHUB_OUTPUT
        id: trigger_upload
      - name: Display All TorchData Conda Package
        if: steps.trigger_upload.outputs.value == 'true'
        run: ls -lh ./*/torchdata-*.tar.bz2
      - name: Upload Packages to Conda
        if: steps.trigger_upload.outputs.value == 'true'
        shell: bash -l {0}
        env:
          CONDA_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_PYTORCHBOT_TOKEN }}
          CONDA_NIGHTLY_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_NIGHTLY_PYTORCHBOT_TOKEN }}
          CONDA_TEST_PYTORCHBOT_TOKEN: ${{ secrets.CONDA_TEST_PYTORCHBOT_TOKEN }}
        run: |
          conda create -y --name conda_upload_env
          conda activate conda_upload_env

          conda install -yq anaconda-client
          conda install -c conda-forge -yq jq

          if [[ ${{ needs.get_release_type.outputs.type }} == 'official' ]]; then
            CONDA_CHANNEL=pytorch
            CONDA_TOKEN=${CONDA_PYTORCHBOT_TOKEN}
          elif [[ ${{ needs.get_release_type.outputs.type }} == 'test' ]]; then
            CONDA_CHANNEL=pytorch-${{ needs.get_release_type.outputs.type }}
            CONDA_TOKEN=${CONDA_TEST_PYTORCHBOT_TOKEN}
          else
            CONDA_CHANNEL=pytorch-${{ needs.get_release_type.outputs.type }}
            CONDA_TOKEN=${CONDA_NIGHTLY_PYTORCHBOT_TOKEN}
          fi

          if [[ ${{ needs.get_release_type.outputs.type }} == 'nightly' ]]; then
            # Loop over all platforms [win-64, osx-64, linux-64]
            for subdir in $(find . -type f -name '*torchdata*.tar.bz2' | sed -r 's|/[^/]+$||' | uniq | cut -f2 -d/) ; do
              version=""
              # Find existing conda packages on pytorch nightly
              for val in $(conda search --json torchdata --channel=$CONDA_CHANNEL --subdir=$subdir | jq -r '.[][] | .version, .build'); do
                if [[ -z $version ]]; then
                  version=$val
                else
                  build=$val
                  # Check if a new conda package built based on the same Python version and platform exists
                  for new_package in ./$subdir/torchdata-*-$build.tar.bz2; do
                    if [[ -f "$new_package" ]]; then
                      echo "Removing $CONDA_CHANNEL/torchdata/$version/$subdir/torchdata-$version-$build.tar.bz2"
                      anaconda -t "${CONDA_TOKEN}" remove -f "$CONDA_CHANNEL/torchdata/$version/$subdir/torchdata-$version-$build.tar.bz2"
                      break
                    fi
                  done
                  version=""
                  build=""
                fi
              done
              # Upload new conda packages per subdir
              anaconda -t "${CONDA_TOKEN}" upload ./$subdir/torchdata-*.tar.bz2 -u "$CONDA_CHANNEL" --label main --no-progress --force
            done
          else
            anaconda -t "${CONDA_TOKEN}" upload ./*/torchdata-*.tar.bz2 -u "$CONDA_CHANNEL" --label main --no-progress --force
          fi

  build_docs:
    if: |
      always() && inputs.branch != '' &&
      ( needs.wheel_upload.outputs.upload == 'true' || needs.conda_upload.outputs.upload == 'true' )
    needs: [get_release_type, wheel_upload, conda_upload]
    runs-on: ubuntu-latest
    steps:
      - name: Setup Python 3.9
        uses: actions/setup-python@v4
        with:
          python-version: 3.9
      - name: Checkout
        uses: actions/checkout@v3
        with:
          ref: ${{ inputs.branch }}
          submodules: recursive
      - name: Install Dependencies
        run: |
          echo `python3 --version`
          sudo apt-get install -y python-setuptools
          python3 -m pip install --upgrade pip
          python3 -m pip install setuptools
          python3 -m pip install matplotlib
          sudo apt-get install -y yarn
      - name: Install PyTorch & TorchData
        run: |
          pip3 install numpy
          # Add version requirement to PyTorch except nightly release
          if [[ -z "${{ inputs.pytorch_version }}" ]]; then
            PYTORCH_VERSION=torch
          else
            PYTORCH_VERSION=torch==${{ inputs.pytorch_version }}
          fi

          PIP_CHANNEL=${{ needs.get_release_type.outputs.type }}
          if [[ $PIP_CHANNEL == 'official' ]]; then
            pip3 install "$PYTORCH_VERSION" -f https://download.pytorch.org/whl/torch_stable.html
          else
            pip3 install --pre "$PYTORCH_VERSION" --index-url "https://download.pytorch.org/whl/$PIP_CHANNEL/cpu"
          fi

          pip3 install -r requirements.txt
          pip3 install .
      - name: Check env
        run: echo `which spinx-build`
      - name: Build the docset
        run: |
          cd ./docs
          sudo apt-get install -y graphviz
          pip3 install -r requirements.txt
          make html
          cd ..
      - name: Export Target Folder
        run: |
          TARGET_FOLDER=${{ inputs.branch }}
          if [[ $TARGET_FOLDER == release/* ]]; then
            TARGET_FOLDER=${TARGET_FOLDER:8}
          elif [[ $TARGET_FOLDER == tags/* ]]; then
            TARGET_FOLDER=${TARGET_FOLDER:5}
          elif [[ $TARGET_FOLDER == v* ]] && [[ ${{ inputs.pre_dev_release }} == false ]]; then
            if [[ $TARGET_FOLDER == v*.*.* ]]; then
              TARGET_FOLDER=${TARGET_FOLDER%.*}
            fi
            TARGET_FOLDER=${TARGET_FOLDER:1}
          fi
          echo "value=$TARGET_FOLDER" >> $GITHUB_OUTPUT
        id: target_folder
      - name: Deploy
        uses: JamesIves/github-pages-deploy-action@v4.4.1
        with:
          token: ${{ secrets.GITHUB_TOKEN }}
          branch: gh-pages # The branch the action should deploy to.
          folder: docs/build/html # The folder the action should deploy.
          target-folder: ${{ steps.target_folder.outputs.value }} # The destination folder the action should deploy to.