_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

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@v4
        with:
          ref: ${{ inputs.branch }}
          submodules: recursive
      - name: Setup Python ${{ matrix.python-version }} for Windows
        if: ${{ startsWith( matrix.os, 'windows' ) }}
        uses: actions/setup-python@v5
        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@v4
        with:
          name: torchdata-artifact-${{ matrix.os }}-${{ matrix.python-version }}
          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@v4
        with:
          pattern: torchdata-artifact-*
          merge-multiple: true
      - 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

  build_docs:
    if: |
      always() && inputs.branch != '' &&
      ( needs.wheel_upload.outputs.upload == 'true')
    needs: [get_release_type, wheel_upload]
    runs-on: ubuntu-latest
    steps:
      - name: Setup Python 3.9
        uses: actions/setup-python@v5
        with:
          python-version: 3.9
      - name: Checkout
        uses: actions/checkout@v4
        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.