name: PR Test - JIT Kernel

on:
  workflow_call:
    inputs:
      jit_kernel:
        required: true
        type: string
      sgl_kernel:
        required: true
        type: string
      rust_ext_artifact:
        description: 'Artifact of prebuilt Rust extension modules, from rust-ext-build. Empty, or a download that fails, falls back to the cache; a miss there compiles during install.'
        type: string
        default: ''
      runs_on_map:
        required: true
        type: string
      git_ref:
        required: false
        type: string
        default: ''
      test_parallel_dispatch:
        required: false
        type: string
        default: 'false'
      skip_pr_test_health_check:
        required: false
        type: boolean
        default: false

# Workflow-level env is NOT inherited from the caller in reusable workflows (verified by CI test).
# The github context (including github.event_name) IS inherited from the caller.
env:
  SGLANG_IS_IN_CI: true
  NCCL_NVLS_ENABLE: "0"
  SGLANG_CUDA_COREDUMP: "1"
  SGLANG_JIT_DEEPGEMM_FAST_WARMUP: true
  PR_TEST_BYPASS_MAINTENANCE_ON_MAIN: ${{ github.ref == 'refs/heads/main' && 'true' || 'false' }}
  SKIP_PR_TEST_HEALTH_CHECK: ${{ inputs.skip_pr_test_health_check == true && 'true' || 'false' }}

jobs:
  # `name:` reproduces the job names this workflow published before the table.
  jit-kernel-test:
    # The job name is the suite it runs, with the stage prefix swapped for this
    # workflow's, so a row declares a (suite, runner) pair and nothing else.
    name: jit-kernel-${{ matrix.suite }}${{ matrix.shard }}
    # Runs whenever call-jit-kernel-tests dispatches this workflow. That caller is the
    # single gate (PR jit_kernel changes, or scheduled/parallel-dispatch full runs), so
    # the sub-jobs no longer re-exclude schedule/parallel-dispatch here.
    strategy:
      fail-fast: false
      matrix:
        include:
          - suite: unit-test-1-gpu-large
            runner_config: 1-gpu-large
            shard: " (0)"
            suite_args: --auto-partition-id 0 --auto-partition-size 2 --fork-worker-batch-size 20
          - suite: unit-test-1-gpu-large
            runner_config: 1-gpu-large
            shard: " (1)"
            suite_args: --auto-partition-id 1 --auto-partition-size 2 --fork-worker-batch-size 20
          - suite: unit-test-1-gpu-small
            runner_config: 1-gpu-small
          - suite: unit-test-4-gpu-b200
            runner_config: 4-gpu-b200
          - suite: unit-test-8-gpu-h200
            runner_config: 8-gpu-h200
            # Alone among these, it has never run the health check.
            skip_health_check: true
          - suite: benchmark-test-1-gpu-large
            runner_config: 1-gpu-large
    runs-on: ${{ fromJson(inputs.runs_on_map)[matrix.runner_config] }}
    timeout-minutes: 240
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.git_ref || github.sha }}

      - uses: ./.github/actions/check-pr-test-health
        if: ${{ !matrix.skip_health_check }}

      - uses: ./.github/actions/check-maintenance

      - name: Cleanup
        if: inputs.sgl_kernel == 'true'
        run: |
          ls -alh python/sglang/kernels/aot/dist || true
          rm -rf python/sglang/kernels/aot/dist/* || true

      - name: Download artifacts
        if: inputs.sgl_kernel == 'true'
        uses: actions/download-artifact@v4
        with:
          path: python/sglang/kernels/aot/dist/
          merge-multiple: true
          pattern: wheel-python3.10-cuda13.0

      - uses: ./.github/actions/download-rust-ext
        with:
          artifact_name: ${{ inputs.rust_ext_artifact }}

      - name: Install dependencies
        timeout-minutes: 20
        run: |
          CUSTOM_BUILD_SGL_KERNEL=${{ inputs.sgl_kernel }} bash scripts/ci/cuda/ci_install_dependency.sh diffusion

      - name: Run test
        timeout-minutes: 60
        run: |
          cd test/
          python3 run_suite.py --hw cuda --suite base-b-kernel-${{ matrix.suite }} ${{ matrix.suite_args }}