name: Create Legacy Ckpts

on:
  workflow_dispatch:
    inputs:
      push_to_s3:
        description: 'Push generated checkpoints to S3.'
        default: false
        required: false
        type: boolean
      pl_version:
        description: 'PL version to generate a legacy checkpoint with. If not provided, uses PL from source.'
        default: ''
        required: false
        type: string
      create_pr:
        description: 'Create a PR to enable testing PL with a new checkpoint.'
        default: false
        required: false
        type: boolean
  workflow_call:
    inputs:
      push_to_s3:
        description: 'Push generated checkpoints to S3.'
        default: false
        required: false
        type: boolean
      pl_version:
        description: 'PL version to generate a legacy checkpoint with. If not provided, uses PL from source.'
        default: ''
        required: false
        type: string
      create_pr:
        description: 'Create a PR to enable testing PL with a new checkpoint.'
        default: false
        required: false
        type: boolean
    secrets:
      AWS_ACCESS_KEY_ID:
        required: true
      AWS_SECRET_KEY_ID:
        required: true

jobs:
  create-legacy-ckpts:
    runs-on: ubuntu-20.04
    outputs:
      pl-version: ${{ steps.decide-pl-version.outputs.pl-version }}
    defaults:
      run:
        working-directory: tests/legacy
    steps:
    - uses: actions/checkout@v3

    - uses: actions/setup-python@v4
      with:
        # Python version here needs to be supported by all PL versions listed in back-compatible-versions.txt.
        python-version: 3.8

    - name: Configure AWS credentials
      uses: aws-actions/configure-aws-credentials@v1
      with:
        aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
        aws-secret-access-key: ${{ secrets.AWS_SECRET_KEY_ID }}
        aws-region: us-east-1

    - name: Install PL from source
      working-directory: ./
      env:
        PACKAGE_NAME: pytorch
      run: |
        pip install -e . -f https://download.pytorch.org/whl/cpu/torch_stable.html
        pip list
      if: inputs.pl_version == ''

    - name: Pull legacy checkpoints
      working-directory: ./
      run: bash .actions/pull_legacy_checkpoints.sh

    - name: Decide PL version to create a PR with
      id: decide-pl-version
      run: |
        python -c "import pytorch_lightning as pl; print(f'pl-version={pl.__version__}')" >> $GITHUB_OUTPUT || echo pl-version='' >> $GITHUB_OUTPUT

    - name: Generate checkpoints
      run: |
        bash generate_checkpoints.sh ${{ inputs.pl_version }}

    - name: Upload checkpoints to GitHub Actions artifact
      uses: actions/upload-artifact@v3
      with:
        name: checkpoints-${{ github.sha }}
        path: tests/legacy/checkpoints/

    - name: Upload checkpoints to S3 (dryrun)
      run: |
        aws s3 sync --dryrun checkpoints/ s3://pl-public-data/legacy/checkpoints/
        zip -r checkpoints.zip checkpoints
        aws s3 cp --dryrun checkpoints.zip s3://pl-public-data/legacy/ --acl public-read

    - name: Upload checkpoints to S3
      run: |
        aws s3 sync checkpoints/ s3://pl-public-data/legacy/checkpoints/
        zip -r checkpoints.zip checkpoints
        aws s3 cp checkpoints.zip s3://pl-public-data/legacy/ --acl public-read
      if: inputs.push_to_s3

  adding-ckpt-test:
    runs-on: ubuntu-20.04
    if: inputs.create_pr
    needs: create-legacy-ckpts
    steps:
    - uses: actions/checkout@v3
      with:
        ref: master

    - name: Append a new version to legacy checkpoint list
      run: echo ${{ needs.create-legacy-ckpts.outputs.pl-version }} >> tests/legacy/back-compatible-versions.txt

    - name: Create Pull Request
      uses: peter-evans/create-pull-request@v4
      with:
        title: Adding test for legacy checkpiont created with ${{ needs.create-legacy-ckpts.outputs.pl-version }}
        delete-branch: true
        labels: |
          tests
          pl