mirror of
https://github.com/NVIDIA/nvbench.git
synced 2026-03-14 20:27:24 +00:00
84 lines
2.2 KiB
Bash
Executable File
84 lines
2.2 KiB
Bash
Executable File
#!/bin/bash
|
|
set -euo pipefail
|
|
|
|
ci_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
|
|
usage="Usage: $0 -py-version <python_version> -cuda-version <cuda_version>"
|
|
|
|
source "$ci_dir/util/python/common_arg_parser.sh"
|
|
|
|
# Parse arguments including CUDA version
|
|
parse_python_args "$@"
|
|
|
|
cuda_version=""
|
|
while [[ $# -gt 0 ]]; do
|
|
case $1 in
|
|
-cuda-version=*)
|
|
cuda_version="${1#*=}"
|
|
shift
|
|
;;
|
|
-cuda-version)
|
|
if [[ $# -lt 2 ]]; then
|
|
echo "Error: -cuda-version requires a value" >&2
|
|
exit 1
|
|
fi
|
|
cuda_version="$2"
|
|
shift 2
|
|
;;
|
|
*)
|
|
shift
|
|
;;
|
|
esac
|
|
done
|
|
|
|
# Check if py_version was provided (this script requires it)
|
|
require_py_version "$usage" || exit 1
|
|
|
|
if [[ -z "$cuda_version" ]]; then
|
|
echo "Error: -cuda-version is required"
|
|
echo "$usage"
|
|
exit 1
|
|
fi
|
|
|
|
# Map cuda_version to full version and set CUDA extra
|
|
if [[ "$cuda_version" == "12" ]]; then
|
|
cuda_full_version="12.9.1"
|
|
cuda_extra="cu12"
|
|
elif [[ "$cuda_version" == "13" ]]; then
|
|
cuda_full_version="13.0.1"
|
|
cuda_extra="cu13"
|
|
else
|
|
echo "Error: Unsupported CUDA version: $cuda_version"
|
|
exit 1
|
|
fi
|
|
|
|
# Use the same rapidsai/ci-wheel images as the build
|
|
readonly devcontainer_version=25.12
|
|
readonly devcontainer_distro=rockylinux8
|
|
|
|
if [[ "$(uname -m)" == "aarch64" ]]; then
|
|
readonly cuda_image=rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda_full_version}-${devcontainer_distro}-py${py_version}-arm64
|
|
else
|
|
readonly cuda_image=rapidsai/ci-wheel:${devcontainer_version}-cuda${cuda_full_version}-${devcontainer_distro}-py${py_version}
|
|
fi
|
|
|
|
echo "::group::🧪 Testing CUDA ${cuda_version} wheel on ${cuda_image}"
|
|
(
|
|
set -x
|
|
docker pull $cuda_image
|
|
docker run --rm -i \
|
|
--workdir /workspace \
|
|
--gpus all \
|
|
--mount type=bind,source=$(pwd),target=/workspace/ \
|
|
--env py_version=${py_version} \
|
|
--env cuda_version=${cuda_version} \
|
|
--env cuda_extra="${cuda_extra}" \
|
|
$cuda_image \
|
|
/workspace/ci/test_cuda_bench_inner.sh
|
|
# Prevent GHA runners from exhausting available storage with leftover images:
|
|
if [[ -n "${GITHUB_ACTIONS:-}" ]]; then
|
|
docker rmi -f $cuda_image
|
|
fi
|
|
)
|
|
echo "::endgroup::"
|