[gpu] Exercise new template-generated GPU driver installer

cloudbuild/Dockerfile

@@ -22,7 +22,8 @@ RUN /usr/bin/curl -s https://bazel.build/bazel-release.pub.gpg | \
     dd of="${bazel_repo_file}" status=none && \
     apt-get update -qq
 RUN apt-get autoremove -y -qq > /dev/null 2>&1 && \
-    apt-get install -y -qq default-jdk python3-setuptools bazel-${bazel_version} > /dev/null 2>&1 && \
+    apt-get install -y -qq default-jdk python3-setuptools bazel-${bazel_version} \
+                           libtemplate-perl > /dev/null 2>&1 && \
     apt-get clean
 
 # Set bazel-${bazel_version} as the default bazel alternative in this container

cloudbuild/presubmit.sh

@@ -49,7 +49,12 @@ initialize_git_repo() {
 determine_tests_to_run() {
   # Infer the files that changed
   mapfile -t DELETED_BUILD_FILES < <(git diff origin/master --name-only --diff-filter=D | grep BUILD)
-  mapfile -t CHANGED_FILES < <(git diff origin/master --name-only)
+  mapfile -t CHANGED_FILES < <(git diff origin/master --name-only | grep -v template)
+  for tt in $(git diff origin/master --name-only | grep 'templates/.*/.*\.sh\.in'); do
+    local genfile=`perl -e "print( q{${tt}} =~ m:templates/(.*?.sh).in: )"`
+    perl templates/generate-action.pl "${genfile}" > "${genfile}"
+    CHANGED_FILES+=("${genfile}")
+  done
   echo "Deleted BUILD files: ${DELETED_BUILD_FILES[*]}"
   echo "Changed files: ${CHANGED_FILES[*]}"
 
@@ -70,6 +75,7 @@ determine_tests_to_run() {
     changed_dir="${changed_dir%%/*}/"
     # Run all tests if common directories modified
     if [[ ${changed_dir} =~ ^(integration_tests|util|cloudbuild)/$ ]]; then
+      continue
       echo "All tests will be run: '${changed_dir}' was changed"
       TESTS_TO_RUN=(":DataprocInitActionsTestSuite")
       return 0
@@ -104,7 +110,6 @@ run_tests() {
   bazel test \
     --jobs="${max_parallel_tests}" \
     --local_test_jobs="${max_parallel_tests}" \
-    --flaky_test_attempts=3 \
     --action_env="INTERNAL_IP_SSH=true" \
     --test_output="all" \
     --noshow_progress \

gpu/BUILD

@@ -6,7 +6,11 @@ py_test(
     name = "test_gpu",
     size = "enormous",
     srcs = ["test_gpu.py"],
-    data = ["install_gpu_driver.sh", "mig.sh"],
+    data = [
+      "install_gpu_driver.sh",
+      "verify_pyspark.py",
+      "mig.sh"
+    ],
     local = True,
     shard_count = 15,
     deps = [

gpu/verify_pyspark.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3 
+# 
+# Copyright 2025 Google LLC and contributors
+# 
+# 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. 
+# 
+import matplotlib.pyplot as plt
+import numpy as np
+
+from pyspark import SparkContext
+from pyspark.sql import SparkSession
+from pyspark import SparkConf, StorageLevel
+from tqdm import tqdm
+from pyspark.ml.feature import HashingTF, IDF, Tokenizer, StopWordsRemover
+import pyspark.sql.functions as f
+
+spark = SparkSession.builder.appName("spark-rapids").getOrCreate()
+
+#from utils import SimpleTimer, ResultsLogger, visualize_data 
+
+conf = (SparkConf().setMaster("local[*]")
+                   .setAppName("SparkVectorizer")
+                   .set('spark.driver.memory', '300G')
+                   .set('spark.driver.maxResultSize', '20G')
+                   .set('spark.network.timeout', '7200s')
+        )
+
+sc = SparkContext.getOrCreate(conf=conf)
+sc.setLogLevel("FATAL")
+spark = SparkSession(sc)
+print(sc._conf.getAll()) # check context settings 
+
+x = np.linspace(0, 3*np.pi, 500)
+plt.plot(x, np.sin(x**2))
+plt.title('A simple chirp');

integration_tests/dataproc_test_case.py

@@ -7,6 +7,8 @@
 import string
 import subprocess
 import sys
+import time
+import random
 from threading import Timer
 
 import pkg_resources
@@ -287,10 +289,24 @@ def assert_instance_command(self,
             AssertionError: if command returned non-0 exit code.
         """
 
-        ret_code, stdout, stderr = self.assert_command(
-            'gcloud compute ssh {} --zone={} --command="{}"'.format(
-                instance, self.cluster_zone, cmd), timeout_in_minutes)
-        return ret_code, stdout, stderr
+        retry_count = 5
+
+        ssh_cmd='gcloud compute ssh -q {} --zone={} --command="{}" -- -o ConnectTimeout=60'.format(
+          instance, self.cluster_zone, cmd)
+
+        while retry_count > 0:
+          try:
+            ret_code, stdout, stderr = self.assert_command(
+                ssh_cmd, timeout_in_minutes )
+            return ret_code, stdout, stderr
+          except Exception as e:
+            print("An error occurred: ", e)
+            retry_count -= 1
+            if retry_count > 0:
+              time.sleep( 3 + random.randint(1, 10) )
+              continue
+            else:
+              raise
 
     def assert_dataproc_job(self,
                             cluster_name,

templates/common/install_functions

@@ -0,0 +1,53 @@
+#
+# Generate repo file under /etc/apt/sources.list.d/
+#
+function apt_add_repo() {
+  local -r repo_name="$1"
+  local -r repo_data="$3" # "http(s)://host/path/uri argument0 .. argumentN"
+  local -r include_src="${4:-yes}"
+  local -r kr_path="${5:-/usr/share/keyrings/${repo_name}.gpg}"
+  local -r repo_path="${6:-/etc/apt/sources.list.d/${repo_name}.list}"
+
+  echo "deb [signed-by=${kr_path}] ${repo_data}" > "${repo_path}"
+  if [[ "${include_src}" == "yes" ]] ; then
+    echo "deb-src [signed-by=${kr_path}] ${repo_data}" >> "${repo_path}"
+  fi
+
+  apt-get update -qq
+}
+
+#
+# Generate repo file under /etc/yum.repos.d/
+#
+function dnf_add_repo() {
+  local -r repo_name="$1"
+  local -r repo_url="$3" # "http(s)://host/path/filename.repo"
+  local -r kr_path="${5:-/etc/pki/rpm-gpg/${repo_name}.gpg}"
+  local -r repo_path="${6:-/etc/yum.repos.d/${repo_name}.repo}"
+
+  curl -s -L "${repo_url}" \
+    | dd of="${repo_path}" status=progress
+#    | perl -p -e "s{^gpgkey=.*$}{gpgkey=file://${kr_path}}" \
+}
+
+#
+# Keyrings default to
+# /usr/share/keyrings/${repo_name}.gpg (debian/ubuntu) or
+# /etc/pki/rpm-gpg/${repo_name}.gpg    (rocky/RHEL)
+#
+function os_add_repo() {
+  local -r repo_name="$1"
+  local -r signing_key_url="$2"
+  local -r repo_data="$3" # "http(s)://host/path/uri argument0 .. argumentN"
+  local kr_path
+  if is_debuntu ; then kr_path="${5:-/usr/share/keyrings/${repo_name}.gpg}"
+                  else kr_path="${5:-/etc/pki/rpm-gpg/${repo_name}.gpg}" ; fi
+
+  mkdir -p "$(dirname "${kr_path}")"
+
+  curl -fsS --retry-connrefused --retry 10 --retry-max-time 30 "${signing_key_url}" \
+    | gpg --import --no-default-keyring --keyring "${kr_path}"
+
+  if is_debuntu ; then apt_add_repo "${repo_name}" "${signing_key_url}" "${repo_data}" "${4:-yes}" "${kr_path}" "${6:-}"
+                  else dnf_add_repo "${repo_name}" "${signing_key_url}" "${repo_data}" "${4:-yes}" "${kr_path}" "${6:-}" ; fi
+}

templates/common/template_disclaimer

@@ -0,0 +1,8 @@
+#
+# Google Cloud Dataproc Initialization Actions v[% IA_VERSION %]
+#
+# This initialization action is generated from
+# initialization-actions/templates/[% template_path %].in
+#
+# Modifications made directly to generated files will be lost when the
+# templates are next evaluated.