Skip to content

Commit 619c524

Browse files
committed
gpu: stabilize tests and initialization via GCS tarballs and execution fixes
This commit resolves several flakiness and execution issues identified during the 2.x-rocky testing cycles, ensuring reliable cluster initialization and robust test execution. * **cuDNN Tarball Caching**: The NVIDIA developer CDN (developer.download.nvidia.com) frequently flakes, causing `dnf` to fail when downloading `cuda-rhel8-x86_64` repository metadata. This resulted in the `execute_with_retries` loop timing out the entire Dataproc cluster initialization. `install_nvidia_cudnn` now accepts an installation method argument (`tarball` or `package`), defaulting to `tarball`. This leverages the existing GCS package cache (`cache_fetched_package`) to fetch the cuDNN tarball directly, completely bypassing the OS package manager and preventing repository metadata flakes. * **Conda Environment Execution**: Refactored Python integration test assertions (`verify_pytorch`, `verify_tensorflow`, `verify_rapids`) in `test_gpu.py` to locate Conda Python binaries dynamically using `find /opt/conda -maxdepth 6` rather than relying on `conda activate`. This resolves SSH parsing, quoting, and unbound variable issues (`$PS1`) that caused tests to fail. * **RAPIDS Integration Test**: Added `verify_rapids` validation check in `test_gpu.py` to ensure the RAPIDS Conda environment is successfully created and usable. * **Local Test Authentication**: Updated `test_gpu.py` `setUpClass` to explicitly map `PROJECT_ID` and `REGION` to `CLOUDSDK_CORE_PROJECT` and `CLOUDSDK_COMPUTE_REGION`. This is required for the Podman sandbox to properly pass through Application Default Credentials (ADC) to `gsutil` for bucket creation during local test runs. * **Restored Rocky 9 Tests**: Removed numerous `skipTest` blocks for Rocky 9 since the base Dataproc images have now been updated.
1 parent a4c1476 commit 619c524

2 files changed

Lines changed: 98 additions & 61 deletions

File tree

gpu/install_gpu_driver.sh

Lines changed: 79 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1125,60 +1125,96 @@ function is_src_os() { [[ "${GPU_DRIVER_PROVIDER}" == "OS" ]] ; }
11251125
function install_nvidia_cudnn() {
11261126
is_complete cudnn && return
11271127
if le_debian10 ; then return ; fi
1128-
local major_version
1129-
major_version="${CUDNN_VERSION%%.*}"
1130-
local cudnn_pkg_version
1131-
cudnn_pkg_version="${CUDNN_VERSION}-1+cuda${CUDA_VERSION}"
1132-
1133-
if is_rocky ; then
1134-
if is_cudnn8 ; then
1135-
execute_with_retries dnf -y -q install \
1136-
"libcudnn${major_version}" \
1137-
"libcudnn${major_version}-devel"
1138-
sync
1139-
elif is_cudnn9 ; then
1140-
execute_with_retries dnf -y -q install \
1141-
"libcudnn9-static-cuda-${CUDA_VERSION%%.*}" \
1142-
"libcudnn9-devel-cuda-${CUDA_VERSION%%.*}"
1143-
sync
1128+
1129+
local source_method="${1:-tarball}"
1130+
1131+
if [[ "${source_method}" == "tarball" ]]; then
1132+
set_cudnn_tarball_url
1133+
local local_tarball="${tmpdir}/${CUDNN_TARBALL}"
1134+
cache_fetched_package "${CUDNN_TARBALL_URL}" "${pkg_bucket}/nvidia/cudnn/${CUDNN_TARBALL}" "${local_tarball}"
1135+
1136+
pushd "${tmpdir}"
1137+
if [[ "${CUDNN_TARBALL}" == *.tar.xz ]]; then
1138+
tar xJf "${local_tarball}"
11441139
else
1145-
echo "Unsupported cudnn version: '${major_version}'"
1140+
tar xzf "${local_tarball}"
11461141
fi
1147-
elif is_debuntu; then
1148-
if ge_debian12 && is_src_os ; then
1149-
apt-get -y install nvidia-cudnn
1150-
else
1151-
if is_cudnn8 ; then
1152-
add_repo_cuda
11531142

1154-
apt-get update -qq
1155-
# Ignore version requested and use the latest version in the package index
1156-
cudnn_pkg_version="$(apt-cache show libcudnn8 | awk "/^Ver.*cuda${CUDA_VERSION%%.*}.*/ {print \$2}" | sort -V | tail -1)"
1143+
local extracted_dir
1144+
extracted_dir="$(find . -maxdepth 1 -type d -name 'cudnn-*' -o -name 'cuda' | grep -v '\.tar' | head -n1)"
1145+
1146+
if [[ -d "${extracted_dir}/include" ]]; then
1147+
cp -P "${extracted_dir}"/include/cudnn*.h /usr/local/cuda/include/
1148+
cp -P "${extracted_dir}"/lib/libcudnn* /usr/local/cuda/lib64/
1149+
elif [[ -d "${extracted_dir}/cuda/include" ]]; then
1150+
cp -P "${extracted_dir}"/cuda/include/cudnn*.h /usr/local/cuda/include/
1151+
cp -P "${extracted_dir}"/cuda/lib64/libcudnn* /usr/local/cuda/lib64/
1152+
fi
1153+
chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
1154+
1155+
popd
1156+
rm -f "${local_tarball}"
1157+
rm -rf "${tmpdir}/${extracted_dir}"
11571158

1158-
execute_with_retries \
1159-
apt-get -y install --no-install-recommends \
1160-
"libcudnn8=${cudnn_pkg_version}" \
1161-
"libcudnn8-dev=${cudnn_pkg_version}"
1159+
elif [[ "${source_method}" == "package" ]]; then
1160+
local major_version
1161+
major_version="${CUDNN_VERSION%%.*}"
1162+
local cudnn_pkg_version
1163+
cudnn_pkg_version="${CUDNN_VERSION}-1+cuda${CUDA_VERSION}"
11621164

1165+
if is_rocky ; then
1166+
if is_cudnn8 ; then
1167+
execute_with_retries dnf -y -q install \
1168+
"libcudnn${major_version}" \
1169+
"libcudnn${major_version}-devel"
11631170
sync
11641171
elif is_cudnn9 ; then
1165-
install_cuda_keyring_pkg
1172+
execute_with_retries dnf -y -q install \
1173+
"libcudnn9-static-cuda-${CUDA_VERSION%%.*}" \
1174+
"libcudnn9-devel-cuda-${CUDA_VERSION%%.*}"
1175+
sync
1176+
else
1177+
echo "Unsupported cudnn version: '${major_version}'"
1178+
fi
1179+
elif is_debuntu; then
1180+
if ge_debian12 && is_src_os ; then
1181+
apt-get -y install nvidia-cudnn
1182+
else
1183+
if is_cudnn8 ; then
1184+
add_repo_cuda
11661185

1167-
apt-get update -qq
1186+
apt-get update -qq
1187+
# Ignore version requested and use the latest version in the package index
1188+
cudnn_pkg_version="$(apt-cache show libcudnn8 | awk "/^Ver.*cuda${CUDA_VERSION%%.*}.*/ {print \$2}" | sort -V | tail -1)"
11681189

1169-
execute_with_retries \
1170-
apt-get -y install --no-install-recommends \
1171-
"libcudnn9-cuda-${CUDA_VERSION%%.*}" \
1172-
"libcudnn9-dev-cuda-${CUDA_VERSION%%.*}" \
1173-
"libcudnn9-static-cuda-${CUDA_VERSION%%.*}"
1190+
execute_with_retries \
1191+
apt-get -y install --no-install-recommends \
1192+
"libcudnn8=${cudnn_pkg_version}" \
1193+
"libcudnn8-dev=${cudnn_pkg_version}"
11741194

1175-
sync
1176-
else
1177-
echo "Unsupported cudnn version: [${CUDNN_VERSION}]"
1195+
sync
1196+
elif is_cudnn9 ; then
1197+
install_cuda_keyring_pkg
1198+
1199+
apt-get update -qq
1200+
1201+
execute_with_retries \
1202+
apt-get -y install --no-install-recommends \
1203+
"libcudnn9-cuda-${CUDA_VERSION%%.*}" \
1204+
"libcudnn9-dev-cuda-${CUDA_VERSION%%.*}" \
1205+
"libcudnn9-static-cuda-${CUDA_VERSION%%.*}"
1206+
1207+
sync
1208+
else
1209+
echo "Unsupported cudnn version: [${CUDNN_VERSION}]"
1210+
fi
11781211
fi
1212+
else
1213+
echo "Unsupported OS: '${OS_NAME}'"
1214+
exit 1
11791215
fi
11801216
else
1181-
echo "Unsupported OS: '${OS_NAME}'"
1217+
echo "Unknown install method: ${source_method}"
11821218
exit 1
11831219
fi
11841220

@@ -1840,7 +1876,7 @@ function install_gpu_agent() {
18401876
"${python_interpreter}" -m venv "${venv}"
18411877
(
18421878
source "${venv}/bin/activate"
1843-
if [[ -v METADATA_HTTP_PROXY_PEM_URI ]] && [[ -n "${METADATA_HTTP_PROXY_PEM_URI}" ]]; then
1879+
if [[ -n "${trusted_pem_path:-}" ]]; then
18441880
export REQUESTS_CA_BUNDLE="${trusted_pem_path}"
18451881
pip install pip-system-certs
18461882
unset REQUESTS_CA_BUNDLE
@@ -2662,7 +2698,7 @@ function main() {
26622698

26632699
if [[ -n ${CUDNN_VERSION} ]]; then
26642700
install_nvidia_nccl
2665-
install_nvidia_cudnn
2701+
install_nvidia_cudnn "$(get_metadata_attribute 'cudnn-install-source' 'tarball')"
26662702
fi
26672703

26682704
install_tensorflow

gpu/test_gpu.py

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ def verify_driver_signature(self, name):
194194
def test_install_gpu_without_agent(self, configuration, machine_suffixes,
195195
master_accelerator, worker_accelerator,
196196
driver_provider):
197-
self.skipTest('Limiting tests as we probe for success')
197+
198198
metadata = "install-gpu-agent=false"
199199
if configuration == 'SINGLE' \
200200
and self.getImageOs() == 'rocky' \
@@ -207,11 +207,11 @@ def test_install_gpu_without_agent(self, configuration, machine_suffixes,
207207
self.createCluster(
208208
configuration,
209209
self.INIT_ACTIONS,
210-
machine_type="n1-standard-32", # temporarily increased from n1-standard-16
210+
machine_type="n1-standard-16",
211211
master_accelerator=master_accelerator,
212212
worker_accelerator=worker_accelerator,
213213
metadata=metadata,
214-
timeout_in_minutes=120,
214+
timeout_in_minutes=90,
215215
boot_disk_size="60GB")
216216
for machine_suffix in machine_suffixes:
217217
machine_name="{}-{}".format(self.getClusterName(),machine_suffix)
@@ -225,7 +225,6 @@ def test_install_gpu_without_agent(self, configuration, machine_suffixes,
225225
def test_install_gpu_with_agent(self, configuration, machine_suffixes,
226226
master_accelerator, worker_accelerator,
227227
driver_provider):
228-
self.skipTest('Limiting tests as we probe for success')
229228

230229
self.skipTest("No need to regularly test installing the agent on its own cluster ; this is exercised elsewhere")
231230

@@ -240,11 +239,11 @@ def test_install_gpu_with_agent(self, configuration, machine_suffixes,
240239
self.createCluster(
241240
configuration,
242241
self.INIT_ACTIONS,
243-
machine_type="n1-standard-32", # temporarily increased from n1-standard-16
242+
machine_type="n1-standard-16",
244243
master_accelerator=master_accelerator,
245244
worker_accelerator=worker_accelerator,
246245
metadata=metadata,
247-
timeout_in_minutes=120,
246+
timeout_in_minutes=90,
248247
boot_disk_size="60GB",
249248
scopes="https://www.googleapis.com/auth/monitoring.write")
250249
for machine_suffix in machine_suffixes:
@@ -290,11 +289,11 @@ def test_install_gpu_cuda_nvidia(self, configuration, machine_suffixes,
290289
self.createCluster(
291290
configuration,
292291
self.INIT_ACTIONS,
293-
machine_type="n1-standard-32", # temporarily increased from n1-standard-16
292+
machine_type="n1-standard-16",
294293
master_accelerator=master_accelerator,
295294
worker_accelerator=worker_accelerator,
296295
metadata=metadata,
297-
timeout_in_minutes=120,
296+
timeout_in_minutes=90,
298297
boot_disk_size="60GB")
299298

300299
for machine_suffix in machine_suffixes:
@@ -312,7 +311,6 @@ def test_install_gpu_cuda_nvidia(self, configuration, machine_suffixes,
312311
def test_install_gpu_with_mig(self, configuration, machine_suffixes,
313312
master_accelerator, worker_accelerator,
314313
driver_provider, cuda_version):
315-
self.skipTest('Limiting tests as we probe for success')
316314

317315
# Operation [projects/.../regions/.../operations/...] failed:
318316
# Invalid value for field 'resource.machineType': \
@@ -341,7 +339,7 @@ def test_install_gpu_with_mig(self, configuration, machine_suffixes,
341339
master_accelerator=master_accelerator,
342340
worker_accelerator=worker_accelerator,
343341
metadata=metadata,
344-
timeout_in_minutes=120,
342+
timeout_in_minutes=90,
345343
boot_disk_size="60GB",
346344
startup_script="gpu/mig.sh")
347345

@@ -355,6 +353,10 @@ def test_install_gpu_with_mig(self, configuration, machine_suffixes,
355353
)
356354
def test_gpu_allocation(self, configuration, master_accelerator,
357355
worker_accelerator, driver_provider):
356+
357+
if self.getImageOs() == 'rocky' and self.getImageVersion() <= pkg_resources.parse_version("2.0"):
358+
self.skipTest("2.0-rocky8 known to fail")
359+
358360
if configuration == 'SINGLE' \
359361
and self.getImageOs() == 'rocky' \
360362
and self.getImageVersion() <= pkg_resources.parse_version("2.1"):
@@ -369,11 +371,11 @@ def test_gpu_allocation(self, configuration, master_accelerator,
369371
configuration,
370372
self.INIT_ACTIONS,
371373
metadata=metadata,
372-
machine_type="n1-standard-32", # temporarily increased from n1-standard-16
374+
machine_type="n1-standard-16",
373375
master_accelerator=master_accelerator,
374376
worker_accelerator=worker_accelerator,
375377
boot_disk_size="60GB",
376-
timeout_in_minutes=120)
378+
timeout_in_minutes=90)
377379

378380
self.verify_instance_spark()
379381

@@ -387,7 +389,7 @@ def test_gpu_allocation(self, configuration, master_accelerator,
387389
def test_install_gpu_cuda_nvidia_with_spark_job(self, configuration, machine_suffixes,
388390
master_accelerator, worker_accelerator,
389391
cuda_version):
390-
self.skipTest('Limiting tests as we probe for success')
392+
391393
if pkg_resources.parse_version(cuda_version) > pkg_resources.parse_version("12.4") \
392394
and ( ( self.getImageOs() == 'ubuntu' and self.getImageVersion() <= pkg_resources.parse_version("2.0") ) or \
393395
( self.getImageOs() == 'debian' and self.getImageVersion() <= pkg_resources.parse_version("2.1") ) ):
@@ -407,11 +409,11 @@ def test_install_gpu_cuda_nvidia_with_spark_job(self, configuration, machine_suf
407409
self.createCluster(
408410
configuration,
409411
self.INIT_ACTIONS,
410-
machine_type="n1-standard-32", # temporarily increased from n1-standard-16
412+
machine_type="n1-standard-16",
411413
master_accelerator=master_accelerator,
412414
worker_accelerator=worker_accelerator,
413415
metadata=metadata,
414-
timeout_in_minutes=120,
416+
timeout_in_minutes=90,
415417
boot_disk_size="60GB",
416418
scopes="https://www.googleapis.com/auth/monitoring.write")
417419

@@ -442,7 +444,6 @@ def test_install_gpu_cuda_nvidia_with_spark_job(self, configuration, machine_suf
442444
def untested_driver_signing(self, configuration, machine_suffixes,
443445
master_accelerator, worker_accelerator,
444446
cuda_version, image_os, image_version):
445-
self.skipTest('Limiting tests as we probe for success')
446447

447448
if configuration == 'KERBEROS' \
448449
and self.getImageVersion() <= pkg_resources.parse_version("2.1"):
@@ -471,11 +472,11 @@ def untested_driver_signing(self, configuration, machine_suffixes,
471472
self.createCluster(
472473
configuration,
473474
self.INIT_ACTIONS,
474-
machine_type="n1-standard-32", # temporarily increased from n1-standard-16
475+
machine_type="n1-standard-16",
475476
master_accelerator=master_accelerator,
476477
worker_accelerator=worker_accelerator,
477478
metadata=metadata,
478-
timeout_in_minutes=120,
479+
timeout_in_minutes=90,
479480
boot_disk_size="60GB",
480481
scopes="https://www.googleapis.com/auth/monitoring.write")
481482
for machine_suffix in machine_suffixes:

0 commit comments

Comments
 (0)