xformers/.github/workflows/wheels.yml

name: wheels

on:
  pull_request:
    paths:
      - "packaging/compute_wheel_version.sh"
      - ".github/workflows/wheel*"
      - ".github/actions/setup-windows-runner/action.yml"
      - "setup.py"
      - "requirements*.txt"
  push:
    branches:
      - main
    tags:
      - "v[0-9]+*"

jobs:
  target_determinator:
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.set-matrix.outputs.matrix }}
    steps:
    - id: set-matrix
      shell: python
      run: |
        import os
        import json
        environ = os.environ

        PY_VERSIONS = ['3.8', '3.9', '3.10', '3.11', '3.12']
        if os.environ["GITHUB_EVENT_NAME"] == "pull_request":
          print("pull-request: only keeping 2 python versions")
          PY_VERSIONS = [PY_VERSIONS[0], PY_VERSIONS[-1]]

        include = []
        for os in ['8-core-ubuntu', 'windows-8-core']:
          for python in PY_VERSIONS:
            for torch_version in ['2.3.1']:
              for cuda_short_version in ["118", "121"]:
                include.append(dict(
                  os=os,
                  python=python,
                  torch_version=torch_version,
                  cuda_short_version=cuda_short_version,
                ))
                print(include[-1])
        matrix = {'include': include}
        print(json.dumps(matrix))
        with open(environ["GITHUB_OUTPUT"], "a") as fd:
          fd.write("matrix="+json.dumps(matrix))
  build:
    needs: target_determinator
    strategy:
      fail-fast: false
      matrix: ${{ fromJson(needs.target_determinator.outputs.matrix) }}

    uses: ./.github/workflows/wheels_build.yml
    if: github.repository == 'facebookresearch/xformers' || github.event_name == 'pull_request'
    with:
      os: ${{ matrix.os }}
      python: ${{ matrix.python }}
      torch_version: ${{ matrix.torch_version }}
      cuda_short_version: ${{ matrix.cuda_short_version }}

  upload_pip:
    needs: build
    uses: ./.github/workflows/wheels_upload_pip.yml
    with:
      twine_username: __token__
      filter: "*torch2.3.1+cu121*"
      execute: ${{ github.repository == 'facebookresearch/xformers' && github.event_name != 'pull_request' }}
    secrets:
      twine_password: ${{ secrets.PYPI_TOKEN }}

  upload_pt_cu118:
    needs: build
    uses: ./.github/workflows/wheels_upload_s3.yml
    with:
      aws_role: "arn:aws:iam::749337293305:role/pytorch_bot_uploader_role"
      s3_path: s3://pytorch/whl/cu118/
      aws_s3_cp_extra_args: --acl public-read
      filter: "*torch2.3.1+cu118*"
      execute: ${{ github.repository == 'facebookresearch/xformers' && github.ref_type == 'tag' }}

  upload_pt_cu121:
    needs: build
    uses: ./.github/workflows/wheels_upload_s3.yml
    with:
      aws_role: "arn:aws:iam::749337293305:role/pytorch_bot_uploader_role"
      s3_path: s3://pytorch/whl/cu121/
      aws_s3_cp_extra_args: --acl public-read
      filter: "*torch2.3.1+cu121*"
      execute: ${{ github.repository == 'facebookresearch/xformers' && github.ref_type == 'tag' }}