name: Python package

on:
  release:
    types: [published]
  workflow_dispatch:

jobs:
  build-and-publish:
    # do not run in forks
    if: ${{ github.repository_owner == 'OpenRLHF' && (github.event_name == 'release' || github.event_name == 'workflow_dispatch') }}
    name: build wheel and upload
    runs-on: ubuntu-22.04

    strategy:
      matrix:
        python-version: [3.10.14, 3.11.0]
        cuda-version: [12.1]

    steps:
    - uses: actions/checkout@v3

    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v4
      with:
        python-version: ${{ matrix.python-version }}

    - name: Install CUDA ${{ matrix.cuda-version }}
      run: |
        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
        sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
        sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub
        sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /"
        sudo apt-get update
        sudo apt-get -y install cuda-${{ matrix.cuda-version }}

    - name: Set up CUDA environment variables
      run: |
        echo "/usr/local/cuda-${{ matrix.cuda-version }}/lib64" | sudo tee -a /etc/ld.so.conf.d/cuda.conf
        echo "export PATH=/usr/local/cuda-${{ matrix.cuda-version }}/bin:\$PATH" | sudo tee -a /etc/environment
        sudo ldconfig
      shell: bash

    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
        pip install setuptools wheel twine packaging

    - name: Build package
      run: |
        python setup.py bdist_wheel --dist-dir=dist

    - name: Publish package
      env:
        TWINE_USERNAME: __token__
        TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
      run: |
        python -m twine upload dist/*