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

env:
  legacy_dir: tests/legacy

defaults:
  run:
    shell: bash

jobs:
  create-legacy-ckpts:
    runs-on: ubuntu-20.04
    outputs:
      pl-version: ${{ steps.decide-version.outputs.pl-version }}
    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@v2
      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
      env:
        PACKAGE_NAME: pytorch
        FREEZE_REQUIREMENTS: 1
      run: |
        pip install . -f https://download.pytorch.org/whl/cpu/torch_stable.html
        pip list
      if: inputs.pl_version == ''

    - name: Install PL version
      run: |
        pip install "pytorch-lightning==${{ inputs.pl_version }}" \
          -f https://download.pytorch.org/whl/cpu/torch_stable.html
        pip list
      if: inputs.pl_version != ''

    - name: Adjust tests -> PL
      if: ${{ matrix.pkg-name != 'lightning' }}
      run: |
        pip install -q -r .actions/requirements.txt
        python .actions/assistant.py copy_replace_imports --source_dir="./tests" \
          --source_import="lightning.fabric,lightning.pytorch" \
          --target_import="lightning_fabric,pytorch_lightning"

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

    - name: Decide PL version to create a PR with
      id: decide-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
      working-directory: ${{ env.legacy_dir }}
      run: |
        bash generate_checkpoints.sh ${{ inputs.pl_version }}

    - name: Keep artifact
      id: keep-artifact
      run: python -c "print('DAYS=' + str(30 if '${{ github.event_name }}'.startswith('pull_request') else 0))" >> $GITHUB_OUTPUT

    - name: Upload checkpoints to GitHub Actions artifact
      uses: actions/upload-artifact@v3
      with:
        name: checkpoints-${{ github.sha }}
        path: ${{ env.legacy_dir }}/checkpoints/
        retention-days: ${{ steps.keep-artifact.outputs.DAYS }}

    - name: Upload checkpoints to S3 (dryrun)
      working-directory: ${{ env.legacy_dir }}
      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
      working-directory: ${{ env.legacy_dir }}
      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

  add-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 }} >> ${{ env.legacy_dir }}/back-compatible-versions.txt

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