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"
      - "pyproject.toml"
      - "docker/Dockerfile.python3.12"
      - "csrc/**"
  workflow_dispatch:
    inputs:
      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
      run_training_test_VSA:
        description: "Run training-test-VSA"
        required: false
        default: false
        type: boolean
      run_inference_test_STA:
        description: "Run inference-test-STA"
        required: false
        default: false
        type: boolean
      run_precision_test_STA:
        description: "Run precision-test-STA"
        required: false
        default: false
        type: boolean
      run_precision_test_VSA:
        description: "Run precision-test-VSA"
        required: false
        default: false
        type: boolean
      run_nightly_test:
        description: "Run nightly-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 }}
      training-test: ${{ steps.filter.outputs.training-test }}
      training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }}
      inference-test-STA: ${{ steps.filter.outputs.inference-test-STA }}
      precision-test-STA: ${{ steps.filter.outputs.precision-test-STA }}
      precision-test-VSA: ${{ steps.filter.outputs.precision-test-VSA }}
    steps:
      - uses: actions/checkout@v4
      - uses: dorny/paths-filter@v3
        id: filter
        with:
          filters: |
            # Define reusable path patterns
            common-paths: &common-paths
              - 'pyproject.toml'
              - 'docker/Dockerfile.python3.12'
            sta-kernel-paths: &sta-kernel-paths
              - 'csrc/attn/sliding_tile_attn/**'
              - 'csrc/attn/sliding_tile_attn/tk/**'
              - 'csrc/attn/sliding_tile_attn/setup.py'
              - 'csrc/attn/sliding_tile_attn/config_sta.py'
              - 'csrc/attn/sliding_tile_attn/st_attn.cpp'
            vsa-kernel-paths: &vsa-kernel-paths
              - 'csrc/attn/video_sparse_attn/**'
              - 'csrc/attn/video_sparse_attn/tk/**'
              - 'csrc/attn/video_sparse_attn/setup.py'
              - 'csrc/attn/video_sparse_attn/config_vsa.py'
              - 'csrc/attn/video_sparse_attn/vsa.cpp'
            vsa-paths: &vsa-paths
              - 'fastvideo/**'
              - *common-paths
              - *vsa-kernel-paths
            
            # Actual tests
            encoder-test:
              - 'fastvideo/models/encoders/**'
              - 'fastvideo/models/loader/**'
              - 'fastvideo/tests/encoders/**'
              - *common-paths
            vae-test:
              - 'fastvideo/models/vaes/**'
              - 'fastvideo/models/loader/**'
              - 'fastvideo/tests/vaes/**'
              - *common-paths
            transformer-test:
              - 'fastvideo/models/dits/**'
              - 'fastvideo/models/loader/**'
              - 'fastvideo/tests/transformers/**'
              - 'fastvideo/layers/**'
              - 'fastvideo/attention/**'
              - *common-paths
            training-test:
              - 'fastvideo/**'
              - *common-paths
            training-test-VSA:
              - 'fastvideo/**'
              - *common-paths
              - *vsa-kernel-paths
            inference-test-STA:
              - 'fastvideo/**'
              - *common-paths
              - *sta-kernel-paths
            precision-test-STA:
              - *common-paths
              - *sta-kernel-paths
            precision-test-VSA:
              - *common-paths
              - *vsa-kernel-paths

  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 }}/fastvideo-dev:py3.12-latest"
      test_command: "uv pip install -e .[test] && pytest ./fastvideo/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 }}/fastvideo-dev:py3.12-latest"
      test_command: "uv pip install -e .[test] && pytest ./fastvideo/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 }}/fastvideo-dev:py3.12-latest"
      test_command: "uv pip install -e .[test] && pytest ./fastvideo/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_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: "uv pip install -e .[test] && pytest ./fastvideo/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' && needs.change-filter.outputs.training-test == 'true') || 
      (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 }}/fastvideo-dev:py3.12-latest"
      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/Vanilla -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}

  training-test-VSA:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test-VSA == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test_VSA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "training-test-VSA"
      gpu_type: "NVIDIA H100 NVL"
      gpu_count: 2
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/VSA -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}

  inference-test-STA:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.inference-test-STA == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_inference_test_STA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "inference-test-STA"
      gpu_type: "NVIDIA H100 NVL"
      gpu_count: 2
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/inference/STA -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

  precision-test-STA:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-STA == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_STA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "precision-test-STA"
      gpu_type: "NVIDIA H100 NVL"
      gpu_count: 1
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_sta.py"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

  precision-test-VSA:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-VSA == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_VSA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "precision-test-VSA"
      gpu_type: "NVIDIA H100 NVL"
      gpu_count: 1
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_vsa.py"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

  nightly-test:
    if: >-
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "nightly-test"
      gpu_type: "NVIDIA A40"
      gpu_count: 4
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/nightly/test_e2e_overfit_single_sample.py -vs"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}

  runpod-cleanup:
    # Add other jobs to this list as you create them
    needs: [encoder-test, vae-test, transformer-test, ssim-test, training-test, training-test-VSA, inference-test-STA, precision-test-STA, precision-test-VSA] 
    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", "training-test", "training-test-VSA", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]'
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        run: python .github/scripts/runpod_cleanup.py