Skip to content

Commit acfb071

Browse files
juliagmt-googletensorflower-gardener
juliagmt-google
authored and
tensorflower-gardener
committed
[XLA:benchmarks] Add postsubmit workflow to run benchmarks on B200 and updates relevant scripts and configs
PiperOrigin-RevId: 765284686
1 parent bf044cf commit acfb071

4 files changed

Lines changed: 250 additions & 10 deletions

File tree

‎third_party/xla/.github/workflows/postsubmit_benchmark.yml‎

Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
1+
# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ============================================================================
15+
# .github/workflows/postsubmit_benchmarks.yml
16+
name: Postsubmit - Run Benchmarks
17+
18+
permissions:
19+
contents: read
20+
21+
on:
22+
workflow_dispatch:
23+
inputs:
24+
halt-for-connection:
25+
description: 'Should this workflow run wait for a remote connection?'
26+
type: choice
27+
required: true
28+
default: 'no'
29+
options:
30+
- 'yes'
31+
- 'no'
32+
pull_request: # TODO(juliagmt): update after testing
33+
branches:
34+
- main
35+
36+
# Cancel in-progress runs for the same branch if a new commit is pushed
37+
concurrency:
38+
group: ${{ github.workflow }}-${{ github.ref }}
39+
cancel-in-progress: true
40+
41+
jobs:
42+
# =================================================================
43+
# Job 1: Generate the matrix specifically for POSTSUBMIT benchmarks
44+
# =================================================================
45+
generate_matrix:
46+
name: Generate Postsubmit Matrix
47+
# Condition: Run if manually dispatched OR if it's a push event to the main branch.
48+
# TODO(juliagmt): update after testing.
49+
# if: |
50+
# github.event_name == 'workflow_dispatch' ||
51+
# (github.event_name == 'pull_request' && github.ref == 'refs/heads/main')
52+
uses: ./.github/workflows/generate_benchmark_matrix.yml
53+
with:
54+
workflow_type: 'POSTSUBMIT'
55+
registry_file: 'xla/tools/benchmarks/registries/default_registry.yml'
56+
checkout_ref: ${{ github.sha }} # On push/dispatch to main, github.sha is the commit SHA
57+
58+
run_benchmarks:
59+
name: Run Benchmark (${{ matrix.benchmark_entry.config_id }}) # config_id will indicate the workflow type, e.g., '_postsubmit'
60+
needs: generate_matrix
61+
# Also check if the generate_matrix job itself was skipped due to its `if` condition
62+
if: success() && needs.generate_matrix.result != 'skipped' && needs.generate_matrix.outputs.matrix_include_json != '[]' && needs.generate_matrix.outputs.matrix_include_json != ''
63+
64+
strategy:
65+
fail-fast: false # Continue running other benchmarks even if one fails
66+
matrix:
67+
benchmark_entry: ${{ fromJson(needs.generate_matrix.outputs.matrix_include_json || '[]') }}
68+
69+
runs-on: ${{ matrix.benchmark_entry.runner_label }}
70+
container: ${{ matrix.benchmark_entry.container_image }}
71+
72+
defaults:
73+
run:
74+
shell: bash
75+
76+
timeout-minutes: 60
77+
78+
env:
79+
# Variables from the matrix
80+
BENCHMARK_NAME: ${{ matrix.benchmark_entry.benchmark_name }}
81+
CONFIG_ID: ${{ matrix.benchmark_entry.config_id }}
82+
RUNNER_LABEL: ${{ matrix.benchmark_entry.runner_label }}
83+
CONTAINER_IMAGE: ${{ matrix.benchmark_entry.container_image }}
84+
ARTIFACT_LOCATION: ${{ matrix.benchmark_entry.artifact_location }}
85+
IS_GCS_ARTIFACT: ${{ matrix.benchmark_entry.is_gcs_artifact }}
86+
INPUT_FORMAT: ${{ matrix.benchmark_entry.input_format }}
87+
XLA_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.xla_compilation_flags) }}
88+
RUNTIME_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.runtime_flags) }}
89+
TARGET_METRICS_JSON: ${{ toJson(matrix.benchmark_entry.target_metrics) }}
90+
TOPOLOGY_JSON: ${{ toJson(matrix.benchmark_entry.topology) }}
91+
HARDWARE_CATEGORY: ${{ matrix.benchmark_entry.hardware_category }} # e.g., NVIDIA_GPU_B200
92+
93+
# Workflow context variables
94+
CHECKOUT_REF: ${{ github.sha }} # For postsubmit, use the commit SHA from the push/dispatch event
95+
COMMIT_SHA: ${{ github.sha }} # For postsubmit, use the commit SHA
96+
WORKFLOW_RUN_ID: ${{ github.run_id }}
97+
98+
# Paths and script locations
99+
OUTPUT_DIR_NAME: benchmark_output_${{ matrix.benchmark_entry.config_id }} # config_id includes workflow type
100+
SCRIPT_DIR_RELATIVE: .github/workflows/benchmarks
101+
BASELINE_YAML_FILE_RELATIVE: xla/tools/benchmarks/baseline/postsubmit_baseline.yml
102+
COMPARISON_SCRIPT_RELATIVE: .github/workflows/benchmarks/compare_with_baseline.py
103+
104+
steps:
105+
- name: "Wait For Connection"
106+
uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
107+
with:
108+
halt-dispatch-input: ${{ inputs.halt-for-connection }}
109+
110+
- name: Print Job Info & Set Full Paths in ENV
111+
run: |
112+
# Resolve full paths based on GITHUB_WORKSPACE and relative paths defined in env
113+
FULL_OUTPUT_DIR_PATH="${GITHUB_WORKSPACE}/${OUTPUT_DIR_NAME}"
114+
RESOLVED_SCRIPT_DIR_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}"
115+
RESOLVED_BASELINE_YAML_PATH="${GITHUB_WORKSPACE}/${BASELINE_YAML_FILE_RELATIVE}" # Will use postsubmit_baseline.yml
116+
RESOLVED_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${COMPARISON_SCRIPT_RELATIVE}"
117+
RESOLVED_RUN_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}/run_comparison.sh"
118+
119+
echo "--- Benchmark Job Info (Postsubmit) ---"
120+
echo "Config ID (from matrix): $CONFIG_ID"
121+
echo "Benchmark Name: $BENCHMARK_NAME"
122+
echo "Runner Label: $RUNNER_LABEL"
123+
echo "Hardware Category: $HARDWARE_CATEGORY"
124+
echo "Output Directory Name (relative): $OUTPUT_DIR_NAME"
125+
echo "Full Output Directory Path: $FULL_OUTPUT_DIR_PATH"
126+
echo "Full Script Directory Path: $RESOLVED_SCRIPT_DIR_PATH"
127+
echo "Full Baseline YAML Path: $RESOLVED_BASELINE_YAML_PATH" # This will point to postsubmit_baseline.yml
128+
echo "Full Comparison Script Path: $RESOLVED_COMPARISON_SCRIPT_PATH"
129+
echo "Full Run Comparison Script Path: $RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
130+
echo "GITHUB_WORKSPACE: ${GITHUB_WORKSPACE}"
131+
echo "Current PWD: $(pwd)"
132+
echo "---------------------------------------"
133+
134+
# Make resolved paths available to subsequent steps via GITHUB_ENV
135+
echo "RESOLVED_OUTPUT_DIR=$FULL_OUTPUT_DIR_PATH" >> $GITHUB_ENV
136+
echo "RESOLVED_SCRIPT_DIR=$RESOLVED_SCRIPT_DIR_PATH" >> $GITHUB_ENV
137+
echo "RESOLVED_BASELINE_YAML=$RESOLVED_BASELINE_YAML_PATH" >> $GITHUB_ENV
138+
echo "RESOLVED_COMPARISON_SCRIPT=$RESOLVED_COMPARISON_SCRIPT_PATH" >> $GITHUB_ENV
139+
echo "RESOLVED_RUN_COMPARISON_SCRIPT=$RESOLVED_RUN_COMPARISON_SCRIPT_PATH" >> $GITHUB_ENV
140+
141+
# Create the output directory; scripts will assume it exists
142+
mkdir -p "$FULL_OUTPUT_DIR_PATH"
143+
if [ ! -d "$FULL_OUTPUT_DIR_PATH" ]; then
144+
echo "::error::Failed to create output directory: $FULL_OUTPUT_DIR_PATH"
145+
exit 1
146+
fi
147+
148+
- name: Checkout OpenXLA Repository
149+
uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
150+
with:
151+
ref: ${{ env.CHECKOUT_REF }}
152+
153+
- name: Build Binaries
154+
id: build_binaries
155+
run: |
156+
bash "${RESOLVED_SCRIPT_DIR}/build_binaries.sh"
157+
env:
158+
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
159+
160+
- name: Prepare Benchmark Artifact
161+
id: prep_artifact
162+
run: |
163+
bash "${RESOLVED_SCRIPT_DIR}/prepare_artifact.sh"
164+
env:
165+
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
166+
167+
- name: Run Benchmark and Generate Stats
168+
id: run_hlo
169+
env:
170+
RUNNER_BINARY: "${{ steps.build_binaries.outputs.runner_binary }}"
171+
STATS_BINARY: "${{ steps.build_binaries.outputs.stats_binary }}"
172+
DEVICE_TYPE_FLAG: "${{ steps.build_binaries.outputs.device_type_flag }}"
173+
LOCAL_ARTIFACT_PATH: "${{ steps.prep_artifact.outputs.artifact_local_path }}"
174+
OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
175+
# Other job-level env vars are automatically inherited
176+
run: |
177+
# The run_benchmark.sh script is expected to create $RESOLVED_OUTPUT_DIR/results.json
178+
bash "${RESOLVED_SCRIPT_DIR}/run_benchmark.sh"
179+
180+
# Print the content of results.json for debugging purposes
181+
RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
182+
echo "--- Content of results.json (from workflow) ---"
183+
if [ -f "$RESULTS_JSON_FILE_PATH" ]; then
184+
if command -v jq &> /dev/null && jq '.' "$RESULTS_JSON_FILE_PATH" > /dev/null 2>&1; then
185+
jq '.' "$RESULTS_JSON_FILE_PATH"
186+
else
187+
echo "results.json may not be valid JSON or jq failed, printing with cat:"
188+
cat "$RESULTS_JSON_FILE_PATH"
189+
fi
190+
elif [ -f "${RESULTS_JSON_FILE_PATH}.txt" ]; then # Fallback for non-JSON or renamed output
191+
echo "results.json not found, printing fallback .txt file:"
192+
cat "${RESULTS_JSON_FILE_PATH}.txt"
193+
else
194+
echo "::warning::Neither results.json nor results.json.txt found at $RESULTS_JSON_FILE_PATH"
195+
fi
196+
echo "---------------------------------------------"
197+
198+
- name: Compare Results to Baseline
199+
if: (steps.run_hlo.outcome == 'success' || (steps.run_hlo.outcome == 'failure' && (hashFiles(format('{0}/results.json', env.RESOLVED_OUTPUT_DIR)) != '' || hashFiles(format('{0}/results.json.txt', env.RESOLVED_OUTPUT_DIR)) != ''))) && env.skip_comparison_due_to_yaml_install_failure != 'true'
200+
run: |
201+
echo "Starting baseline comparison (against postsubmit baseline)..."
202+
bash "$RESOLVED_RUN_COMPARISON_SCRIPT" # This script calls the Python comparison script
203+
echo "Baseline comparison finished."
204+
echo "---------------------------------------------"
205+
206+
- name: Upload Benchmark Artifacts
207+
if: always()
208+
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
209+
with:
210+
name: results-${{ env.CONFIG_ID }}
211+
path: ${{ env.RESOLVED_OUTPUT_DIR }}
212+
retention-days: 14
213+
if-no-files-found: error

‎third_party/xla/xla/tools/benchmarks/baseline/postsubmit_baseline.yml‎

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,21 +27,41 @@
2727
},
2828
"gemma3_1b_flax_call_x86_1h1d_postsubmit": { # config_id
2929
"CPU_TIME": {
30-
"baseline_ms": 2500,
30+
"baseline_ms": 10000,
3131
"threshold": 0.30 # Allow 30% regression max
3232
},
3333
"WALL_TIME": {
34-
"baseline_ms": 800,
34+
"baseline_ms": 3000,
3535
"threshold": 0.30 # Allow 30% regression max
3636
}
3737
},
38-
"gemma3_1b_flax_call_b200_1h1d_postsubmit": { # config_id
38+
"gemma2_2b_keras_jax_b200_1h1d_postsubmit": { # config_id
3939
"GPU_DEVICE_TIME": {
40-
"baseline_ms": 10,
40+
"baseline_ms": 100,
4141
"threshold": 0.30 # Allow 30% regression max
4242
},
4343
"GPU_DEVICE_MEMCPY_TIME": {
44-
"baseline_ms": 0.1,
44+
"baseline_ms": 0.085,
45+
"threshold": 0.30 # Allow 30% regression max
46+
}
47+
},
48+
"gemma2_2b_keras_jax_x86_1h1d_postsubmit": { # config_id
49+
"CPU_TIME": {
50+
"baseline_ms": 42000,
51+
"threshold": 0.30 # Allow 30% regression max
52+
},
53+
"WALL_TIME": {
54+
"baseline_ms": 14000,
55+
"threshold": 0.30 # Allow 30% regression max
56+
}
57+
},
58+
"gemma2_2b_keras_jax_l4_1h1d_postsubmit": {
59+
"GPU_DEVICE_TIME": {
60+
"baseline_ms": 1000,
61+
"threshold": 0.30 # Allow 30% regression max
62+
},
63+
"GPU_DEVICE_MEMCPY_TIME": {
64+
"baseline_ms": 0.05,
4565
"threshold": 0.30 # Allow 30% regression max
4666
}
4767
}

‎third_party/xla/xla/tools/benchmarks/baseline/presubmit_baseline.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,11 @@
2727
},
2828
"gemma3_1b_flax_call_x86_1h1d_presubmit": { # config_id
2929
"CPU_TIME": {
30-
"baseline_ms": 2500,
30+
"baseline_ms": 10000,
3131
"threshold": 0.30 # Allow 30% regression max
3232
},
3333
"WALL_TIME": {
34-
"baseline_ms": 800,
34+
"baseline_ms": 3000,
3535
"threshold": 0.30 # Allow 30% regression max
3636
}
3737
},

‎third_party/xla/xla/tools/benchmarks/registries/default_registry.yml‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
benchmarks: [
1919
{
2020
name: "gemma3_1b_flax_call"
21-
description: "Benchmarks Gemma3 1b in Flax."
21+
description: "Gemma3 1b in Flax."
2222
owner: "juliagmt-google@"
2323
input_artifact: {
2424
input_format: HLO_TEXT
@@ -47,13 +47,20 @@ benchmarks: [
4747
{
4848
name: "gemma2_2b_keras_jax"
4949
description: "Gemma2 2B in Keras."
50-
owner: "company-A@"
50+
owner: "juliagmt-google@"
5151
input_artifact: {
5252
input_format: HLO_TEXT
5353
artifact_gcs_bucket_path: "https://storage.googleapis.com/xla-benchmarking-temp/gemma2_2b_keras_jax.hlo"
5454
}
5555
model_source_info: ["Gemma2 2B"]
5656
hardware_execution_configs: [{
57+
hardware_category: GPU_L4
58+
topology: { num_hosts: 1, num_devices_per_host: 1, multi_host: false, multi_device: false }
59+
target_metrics: [GPU_DEVICE_TIME, GPU_DEVICE_MEMCPY_TIME]
60+
workflow_type: [PRESUBMIT, POSTSUBMIT]
61+
runtime_flags: ["--num_repeats=5"]
62+
},
63+
{
5764
hardware_category: GPU_B200
5865
topology: { num_hosts: 1, num_devices_per_host: 1, multi_host: false, multi_device: false }
5966
target_metrics: [GPU_DEVICE_TIME, GPU_DEVICE_MEMCPY_TIME]
@@ -64,7 +71,7 @@ benchmarks: [
6471
hardware_category: CPU_X86
6572
topology: { num_hosts: 1, num_devices_per_host: 1, multi_host: false, multi_device: false }
6673
target_metrics: [CPU_TIME]
67-
workflow_type: [POSTSUBMIT]
74+
workflow_type: [PRESUBMIT, POSTSUBMIT]
6875
runtime_flags: ["--num_repeats=5"]
6976
}]
7077
update_frequency_policy: QUARTERLY

0 commit comments

Comments
 (0)