pytorch/.github/workflows/vllm-benchmark.yml at main · bab321-AI/pytorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
name: vLLM Benchmark

on:
  workflow_dispatch:
    inputs:
      pytorch_branch:
        description: |
          PyTorch branch (main or refs/pull/PR_NUMBER/head for pull request)
        required: true
        type: string
        default: main
      pytorch_commit:
        description: |
          PyTorch commit (optional, default to use the latest commit from the branch)
        required: false
        type: string
      models:
        description: |
          A comma-separated list of models from pytorch-integration-testing repo (optional, default to run everything)
        required: false
        type: string
        default: ''
      compilation_config:
        description: |
          A JSON string for vLLM --compilation-config, applied to all benchmark tests, e.g. {"cudagraph_mode": "PIECEWISE", "use_inductor_graph_partition": true}
        required: false
        type: string
        default: ''
      runners:
        description: |
          A comma-separated list of runners from .github/scripts/generate_vllm_benchmark_matrix.py to run the benchmark (optional, default to run everything)
        required: true
        type: string
        default: h100,b200
  schedule:
    # Run daily at 5:15 AM PST
    - cron: '15 13 * * *'

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }}
  cancel-in-progress: true

jobs:
  set-parameters:
    if: ${{ !github.event.pull_request.head.repo.fork && github.repository_owner == 'pytorch' }}
    runs-on: linux.c7i.2xlarge
    outputs:
      benchmark_matrix: ${{ steps.set-parameters.outputs.benchmark_matrix }}
      docker_image: ${{ steps.calculate-docker-image.outputs.docker-image }}
      torch_cuda_arch_list: '8.0 8.9 9.0 10.0 12.0'
      build_environment: linux-jammy-cuda13.0-py3.12-gcc11
    steps:
      - uses: pytorch/test-infra/.github/actions/setup-uv@main
        with:
          python-version: "3.12"
          activate-environment: "true"
          ignore-empty-workdir: "true"

      - name: Checkout pytorch-integration-testing repository
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: pytorch/pytorch-integration-testing
          path: pytorch/pytorch-integration-testing
          ref: main

      - name: Set parameters
        working-directory: pytorch/pytorch-integration-testing
        id: set-parameters
        env:
          MODELS: ${{ inputs.models || '' }}
          # Only need CUDA for now, we can add ROCm later if needed
          RUNNERS: ${{ inputs.runners || 'h100,b200' }}
        run: |
          set -eux

          # The generated matrix is grouped by model and runner
          python .github/scripts/generate_vllm_benchmark_matrix.py \
            --benchmark-configs-dir vllm-benchmarks/benchmarks \
            --models "${MODELS}" \
            --runners "${RUNNERS}"

      - name: Checkout PyTorch
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          path: pytorch/pytorch
          ref: ${{ inputs.pytorch_commit || inputs.pytorch_branch }}
          show-progress: false

      - name: Calculate docker image
        id: calculate-docker-image
        uses: pytorch/test-infra/.github/actions/calculate-docker-image@main
        with:
          working-directory: pytorch/pytorch
          docker-image-name: ci-image:pytorch-linux-jammy-cuda13.0-cudnn9-py3.12-gcc11-vllm

  build:
    name: Build PyTorch and vLLM
    needs:
      - set-parameters
    uses: ./.github/workflows/_vllm-build.yml
    with:
      runner: linux.24xlarge.memory
      docker_image: ${{ needs.set-parameters.outputs.docker_image }}
      build_environment: ${{ needs.set-parameters.outputs.build_environment }}
      pytorch_branch: ${{ inputs.pytorch_branch }}
      pytorch_commit: ${{ inputs.pytorch_commit }}
      torch_cuda_arch_list: ${{ needs.set-parameters.outputs.torch_cuda_arch_list }}
    secrets: inherit

  benchmarks:
    name: Run vLLM benchmarks
    needs:
      - set-parameters
      - build
    strategy:
      matrix: ${{ fromJson(needs.set-parameters.outputs.benchmark_matrix) }}
      fail-fast: false
    uses: ./.github/workflows/_vllm-benchmark.yml
    permissions:
      id-token: write
      contents: read
    with:
      runner: ${{ matrix.runner }}
      docker_image: ${{ needs.set-parameters.outputs.docker_image }}
      build_environment: ${{ needs.set-parameters.outputs.build_environment }}
      pytorch_branch: ${{ inputs.pytorch_branch }}
      pytorch_commit: ${{ inputs.pytorch_commit }}
      models: ${{ matrix.models }}
      compilation_config: ${{ inputs.compilation_config }}
      is_nightly: ${{ github.event_name == 'schedule' }}
    secrets: inherit