Skip to content

Support bounded host layer staging for pool reshard sends #574

Support bounded host layer staging for pool reshard sends

Support bounded host layer staging for pool reshard sends #574

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
name: Postsubmit Benchmarks
on:
# Required for internal services to invoke workflow runs.
# Required inputs: benchmark_filter, environment_filter, custom_metadata.
workflow_dispatch:
inputs:
benchmark_filter:
description: "Regex to filter by benchmark name (e.g., 'h2d_d2h_gating')."
required: false
type: string
default: ""
environment_filter:
description: "Regex to filter by environment configuration ID (e.g., 'tpu-v5e-single-node')."
required: false
type: string
default: ""
custom_metadata:
description: "JSON string of metadata to append to the benchmark result. Overwrites static metadata from the registry but is overwritten by dynamic metadata from the workload."
required: false
type: string
default: "{}"
experiment_ref:
description: "Git ref (SHA, branch, or tag) of the user repository to run benchmarks against."
required: false
type: string
default: ""
push:
branches:
- main
permissions:
contents: read
concurrency:
# Group by workflow name and branch to ensure runs execute sequentially.
# This prevents parallel postsubmit runs from clashing and ensures accurate
# reporting history.
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: false
jobs:
run_benchmarks:
name: Run Single Host Benchmarks
uses: google-ml-infra/bap/.github/workflows/run-benchmarks.yaml@82682aceae058a272d67c5e8bb4e8eba2da7b7c2 # main
permissions:
contents: read
pull-requests: write
with:
registry_file: "tpu_sync/benchmarks/benchmark_registry.pbtxt"
bap_ref: "82682aceae058a272d67c5e8bb4e8eba2da7b7c2"
runner: "linux-x86-n2-8"
publish_metrics: true
pub_sub_gcp_project_id: "ml-oss-benchmarking-production"
pub_sub_gcp_topic_id: "public-results-prod"
benchmark_filter: ${{ inputs.benchmark_filter || '' }}
environment_filter: ${{ inputs.environment_filter || '' }}
custom_metadata: ${{ inputs.custom_metadata || '{}' }}
experiment_ref: ${{ inputs.experiment_ref || '' }}
run_multihost_benchmarks:
name: Run Multi-Host Benchmarks
uses: google-ml-infra/bap/.github/workflows/run-benchmarks.yaml@82682aceae058a272d67c5e8bb4e8eba2da7b7c2 # main
permissions:
contents: read
pull-requests: write
with:
registry_file: "tpu_sync/benchmarks/benchmark_registry_h2h_multihost_gate.pbtxt"
bap_ref: "82682aceae058a272d67c5e8bb4e8eba2da7b7c2"
runner: "linux-x86-n2-8"
publish_metrics: true
pub_sub_gcp_project_id: "ml-oss-benchmarking-production"
pub_sub_gcp_topic_id: "public-results-prod"
benchmark_filter: ${{ inputs.benchmark_filter || '' }}
environment_filter: ${{ inputs.environment_filter || '' }}
custom_metadata: ${{ inputs.custom_metadata || '{}' }}
experiment_ref: ${{ inputs.experiment_ref || '' }}