Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 28 additions & 13 deletions .github/workflows/tritonbench.yml
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,8 @@ jobs:
matrix: ${{ fromJson(needs.set-parameters.outputs.benchmark_matrix) }}
fail-fast: false
runs-on: ${{ matrix.runner }}
env:
TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN: ${{ secrets.TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN }}
environment: pytorch-x-vllm
permissions:
id-token: write
Expand Down Expand Up @@ -125,6 +127,14 @@ jobs:
fi
echo "DEVICE_TYPE=$DEVICE_TYPE" >> $GITHUB_ENV

- name: Install dependencies
shell: bash
working-directory: triton-benchmarks/tritonbench
run: |
set -eux

pip install -r .ci/upload/requirements.txt

- name: Setup CUDA GPU_FLAG for docker run
if: env.DEVICE_NAME == 'cuda'
run: |
Expand All @@ -144,13 +154,15 @@ jobs:
exit 1
fi

DOCKER_IMAGE="meta-pytorch/tritonbench:${IMAGE_SUFFIX}"
CONDA_ENV="triton-main"
DOCKER_IMAGE="ghcr.io/meta-pytorch/tritonbench:${IMAGE_SUFFIX}"
echo "DOCKER_IMAGE=$DOCKER_IMAGE" >> "$GITHUB_ENV"
echo "CONDA_ENV=triton-main" >> "$GITHUB_ENV"
echo "CONDA_ENV=$CONDA_ENV" >> "$GITHUB_ENV"
echo "Using docker image: $DOCKER_IMAGE "
echo "Using conda env: $CONDA_ENV "

- name: Run TritonBench benchmark
working-directory: triton-benchmarks/tritonbench
run: |
set -eux

Expand All @@ -170,12 +182,14 @@ jobs:
)

docker exec -t -w /tmp/workspace "${container_name}" bash -c " \
set -eux && cd /workspace/tritonbench &&
bash .ci/tritonbench/run-benchmark.sh ${{ matrix.BENCHMARKS }} --conda-env ${{ env.CONDA_ENV }} "

docker exec -t -w /tmp/workspace "${container_name}" bash -c " \
set -eux && cd /workspace/tritonbench && mv .benchmarks /tmp/workspace/triton-benchmarks/tritonbench/results
"
set -eux && cd /workspace/tritonbench && \
bash .ci/tritonbench/run-benchmark.sh ${{ matrix.BENCHMARKS }} --conda-env ${{ env.CONDA_ENV }} && \
sudo mv /workspace/tritonbench/.benchmarks /tmp/workspace/triton-benchmarks/tritonbench/results-${{ env.CONDA_ENV }} "

# post-process result.json
latest_result_json=$(find ./results-${CONDA_ENV} -name "result.json" | sort -r | head -n 1)
python3 ./.ci/test_infra/oss_ci_benchmark_v3.py --json ${latest_result_json} \
--add-github-env --output ${latest_result_json}

- name: Authenticate with AWS
# AWS CUDA runners already have access to the bucket via its runner IAM role
Expand All @@ -191,22 +205,23 @@ jobs:
- uses: actions/upload-artifact@v4
if: always()
with:
name: tritonbench-results
path: triton-benchmarks/tritonbench/results
name: tritonbench-${{ matrix.runner }}-benchmark-single
path: triton-benchmarks/tritonbench/results-${{ env.CONDA_ENV }}
retention-days: 30

- name: Upload result to Scribe
if: ${{ env.TRITONBENCH_SCRIBE_GRAPHQL_ACCESS_TOKEN != '' }}
working-directory: triton-benchmarks/tritonbench
run: |
latest_result_json=$(find ./results/${TRITONBENCH_SIDE_A_ENV} -name "result.json" | sort -r | head -n 1)
latest_result_json=$(find ./results-${CONDA_ENV} -name "result.json" | sort -r | head -n 1)
python3 ./.ci/upload/scribe.py --json ${latest_result_json}

- name: Rewrite Tritonbench result json to ClickHouse style
working-directory: triton-benchmarks/tritonbench
run: |
latest_result_json=$(find ./results/${TRITONBENCH_SIDE_A_ENV} -name "result.json" | sort -r | head -n 1)
latest_result_json=$(find ./results-${CONDA_ENV} -name "result.json" | sort -r | head -n 1)
python3 ./.ci/test_infra/oss_ci_benchmark_v3.py --json ${latest_result_json} \
--output clickhouse-results/result-${TRITONBENCH_SIDE_A_ENV}.json
--output clickhouse-results/result-${CONDA_ENV}.json

- name: Upload result to ClickHouse
uses: pytorch/test-infra/.github/actions/upload-benchmark-results@main
Expand Down