name: RunPod Test

on:
  workflow_call:
    inputs:
      job_id:
        required: true
        type: string
        description: "Unique identifier for this test job"
      gpu_type:
        required: true
        type: string
        description: "GPU type to use (e.g. NVIDIA A40, NVIDIA L40S)"
      gpu_count:
        required: true
        type: number
        description: "Number of GPUs to use"
      volume_size:
        required: false
        type: number
        default: 20
        description: "Volume size in GB"
      disk_size:
        required: false
        type: number
        default: 20
        description: "Disk size in GB"
      image:
        required: true
        type: string
        description: "Docker image to use"
      test_command:
        required: true
        type: string
        description: "Command to run tests"
      timeout_minutes:
        required: false
        type: number
        default: 30
        description: "Timeout in minutes"
    secrets:
      RUNPOD_API_KEY:
        required: true
      RUNPOD_PRIVATE_KEY:
        required: true

jobs:
  run-test:
    runs-on: ubuntu-latest
    environment: runpod-runners
    steps:
      - name: Checkout code
        uses: actions/checkout@v4

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

      - name: Set up SSH key
        run: |
          mkdir -p ~/.ssh
          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
          chmod 600 ~/.ssh/id_rsa
          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub

      - name: Install dependencies
        run: pip install requests

      - name: Run tests on RunPod
        env:
          JOB_ID: ${{ inputs.job_id }}
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        timeout-minutes: ${{ inputs.timeout_minutes }}
        run: >-
          python .github/scripts/runpod_api.py 
          --gpu-type "${{ inputs.gpu_type }}"
          --gpu-count ${{ inputs.gpu_count }}
          --volume-size ${{ inputs.volume_size }}
          --disk-size ${{ inputs.disk_size }}
          --image "${{ inputs.image }}"
          --test-command "${{ inputs.test_command }}"

      - name: Terminate RunPod Instances
        if: ${{ always() }}
        env:
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
          JOB_ID: ${{ inputs.job_id }}
        run: python .github/scripts/runpod_cleanup.py