81 lines
2.6 KiB
Bash
81 lines
2.6 KiB
Bash
|
#!/usr/bin/env bash
|
||
|
|
|
||
|
|
set -euo pipefail
|
||
|
|
|
||
|
|
die() {
|
||
|
|
local message="$1"
|
||
|
|
local code="${2:-2}"
|
||
|
|
echo "${message}" >&2
|
||
|
|
exit "${code}"
|
||
|
|
}
|
||
|
|
|
||
|
|
usage() {
|
||
|
|
cat <<EOF
|
||
|
|
Usage: $0 [bench-yaml-path]
|
||
|
|
|
||
|
|
Parse ci/bench.yaml and emit a GitHub Actions strategy matrix JSON object:
|
||
|
|
{"include":[...]}
|
||
|
|
|
||
|
|
Each include entry maps one enabled GPU to a benchmark workflow invocation.
|
||
|
|
EOF
|
||
|
|
}
|
||
|
|
|
||
|
|
if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
|
||
|
|
usage
|
||
|
|
exit 0
|
||
|
|
fi
|
||
|
|
|
||
|
|
bench_yaml_path="${1:-ci/bench.yaml}"
|
||
|
|
if [[ ! -f "${bench_yaml_path}" ]]; then
|
||
|
|
die "Benchmark config file not found: ${bench_yaml_path}"
|
||
|
|
fi
|
||
|
|
|
||
|
|
command -v "yq" >/dev/null 2>&1 || die "'yq' is required to parse ${bench_yaml_path}."
|
||
|
|
command -v "jq" >/dev/null 2>&1 || die "'jq' is required to build the dispatch matrix."
|
||
|
|
|
||
|
|
if ! bench_cfg_json="$(yq -o=json '.benchmarks // {}' "${bench_yaml_path}" 2>&1)"; then
|
||
|
|
die "Failed to parse ${bench_yaml_path} as YAML: ${bench_cfg_json}"
|
||
|
|
fi
|
||
|
|
|
||
|
|
# Extract CUB and Python filter arrays (default to empty arrays).
|
||
|
|
cub_filters_json="$(jq -c '.filters.cub // []' <<<"${bench_cfg_json}")"
|
||
|
|
python_filters_json="$(jq -c '.filters.python // []' <<<"${bench_cfg_json}")"
|
||
|
|
|
||
|
|
has_cub_filters="$(jq -e 'type == "array" and length > 0 and all(.[]; type == "string")' <<<"${cub_filters_json}" >/dev/null 2>&1 && echo true || echo false)"
|
||
|
|
has_python_filters="$(jq -e 'type == "array" and length > 0 and all(.[]; type == "string")' <<<"${python_filters_json}" >/dev/null 2>&1 && echo true || echo false)"
|
||
|
|
|
||
|
|
if [[ "${has_cub_filters}" != "true" && "${has_python_filters}" != "true" ]]; then
|
||
|
|
die "${bench_yaml_path} must define at least one string entry in benchmarks.filters.cub or benchmarks.filters.python."
|
||
|
|
fi
|
||
|
|
|
||
|
|
cub_filters_arg=""
|
||
|
|
if [[ "${has_cub_filters}" == "true" ]]; then
|
||
|
|
cub_filters_arg="$(jq -r '.filters.cub | map(@sh) | join(" ")' <<<"${bench_cfg_json}")"
|
||
|
|
fi
|
||
|
|
|
||
|
|
python_filters_arg=""
|
||
|
|
if [[ "${has_python_filters}" == "true" ]]; then
|
||
|
|
python_filters_arg="$(jq -r '.filters.python | map(@sh) | join(" ")' <<<"${bench_cfg_json}")"
|
||
|
|
fi
|
||
|
|
|
||
|
|
jq -cn \
|
||
|
|
--argjson cfg "${bench_cfg_json}" \
|
||
|
|
--arg cub_filters "${cub_filters_arg}" \
|
||
|
|
--arg python_filters "${python_filters_arg}" \
|
||
|
|
'{
|
||
|
|
"include": [
|
||
|
|
($cfg.gpus // [])[] as $gpu
|
||
|
|
| {
|
||
|
|
"gpu": $gpu,
|
||
|
|
"launch_args": ($cfg.launch_args // ""),
|
||
|
|
"arch": ($cfg.arch // "native"),
|
||
|
|
"base_ref": ($cfg.base_ref // "origin/main"),
|
||
|
|
"test_ref": ($cfg.test_ref // "HEAD"),
|
||
|
|
"cub_filters": $cub_filters,
|
||
|
|
"python_filters": $python_filters,
|
||
|
|
"nvbench_args": ($cfg.nvbench_args // ""),
|
||
|
|
"nvbench_compare_args": ($cfg.nvbench_compare_args // "")
|
||
|
|
}
|
||
|
|
]
|
||
|
|
}'
|