name: PR Test

on:
  push:
    branches: [main]
    paths:
      - "fastvideo/**/*.py"
      - ".github/workflows/pr-test.yml"
  pull_request:
    branches: [main]
    types: [opened, ready_for_review, synchronize, reopened]
    paths:
      - "fastvideo/**/*.py"
      - ".github/workflows/pr-test.yml"
  workflow_dispatch:
    inputs:
      custom_image:
        description: "Custom image from this repository (default: fastvideo-dev:latest)"
        required: false
        default: "fastvideo-dev:latest"
        type: string
      run_encoder_test:
        description: "Run encoder-test"
        required: false
        default: false
        type: boolean
      run_vae_test:
        description: "Run vae-test"
        required: false
        default: false
        type: boolean
      run_transformer_test:
        description: "Run transformer-test"
        required: false
        default: false
        type: boolean
      run_ssim_test:
        description: "Run ssim-test"
        required: false
        default: false
        type: boolean
      run_training_test:
        description: "Run training-test"
        required: false
        default: false
        type: boolean

env:
  PYTHONUNBUFFERED: "1"

concurrency:
  group: pr-test-${{ github.ref }}
  cancel-in-progress: true

jobs:
  pre-commit:
    uses: ./.github/workflows/pre-commit.yml

  change-filter:
    runs-on: ubuntu-latest
    needs: pre-commit
    if: ${{ github.event.pull_request.draft == false || github.event_name == 'workflow_dispatch' }}
    outputs:
      encoder-test: ${{ steps.filter.outputs.encoder-test }}
      vae-test: ${{ steps.filter.outputs.vae-test }}
      transformer-test: ${{ steps.filter.outputs.transformer-test }}
    steps:
      - uses: actions/checkout@v4
      - uses: dorny/paths-filter@v3
        id: filter
        with:
          filters: |
            encoder-test:
              - 'fastvideo/v1/models/encoders/**'
              - 'fastvideo/v1/models/loaders/**'
              - 'fastvideo/v1/tests/encoders/**'
            vae-test:
              - 'fastvideo/v1/models/vaes/**'
              - 'fastvideo/v1/models/loaders/**'
              - 'fastvideo/v1/tests/vaes/**'
            transformer-test:
              - 'fastvideo/v1/models/dits/**'
              - 'fastvideo/v1/models/loaders/**'
              - 'fastvideo/v1/tests/transformers/**'
              - 'fastvideo/v1/layers/**'
              - 'fastvideo/v1/attention/**'

  encoder-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.encoder-test == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_encoder_test == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "encoder-test"
      gpu_type: "NVIDIA A40"
      gpu_count: 1
      volume_size: 100
      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

  vae-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.vae-test == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_vae_test == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "vae-test"
      gpu_type: "NVIDIA A40"
      gpu_count: 1
      volume_size: 100
      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

  transformer-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.transformer-test == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_transformer_test == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "transformer-test"
      gpu_type: "NVIDIA L40S"
      gpu_count: 1
      volume_size: 100
      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
      
  ssim-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
    strategy:
      fail-fast: false
      matrix:
        python-version: [
          # {version: "3.10", tag: "latest"},
          # {version: "3.11", tag: "py3.11-latest"},
          {version: "3.12", tag: "py3.12-latest"}
        ]
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "ssim-test-py${{ matrix.python-version.version }}"
      gpu_type: "NVIDIA A40"
      gpu_count: 2
      volume_size: 200
      disk_size: 200
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
      timeout_minutes: 60
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
  
  training-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "training-test"
      gpu_type: "NVIDIA A40"
      gpu_count: 4
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/training -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

  runpod-cleanup:
    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
    if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
    runs-on: ubuntu-latest
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: "3.10"

      - name: Install dependencies
        run: pip install requests

      - name: Cleanup all RunPod instances
        env:
          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12"]'
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        run: python .github/scripts/runpod_cleanup.py