[Test] Add end2end test to validate tutotial for the upgrade of nvidia software stack.

tests/integration-tests/configs/develop.yaml

@@ -805,6 +805,14 @@ test-suites:
           schedulers: ["slurm"]
           oss: ["ubuntu2404"]
           instances: ["trn1.32xlarge"]
+  tutorials:
+    test_upgrade_nvidia_software.py::test_upgrade_nvidia_software:
+      dimensions:
+        # The NVIDIA upgrade component targets Amazon Linux 2023 x86_64 only.
+        - regions: [ {{ g4dn_2xlarge_CAPACITY_RESERVATION_2_INSTANCES_2_HOURS_NOPG_alinux2023 }} ]
+          instances: [ "g4dn.2xlarge" ]
+          oss: [ "alinux2023" ]
+          schedulers: [ "slurm" ]
   update:
     test_update.py::test_update_slurm:
       dimensions:

tests/integration-tests/configs/released.yaml

@@ -292,6 +292,14 @@ test-suites:
           instances: {{ common.INSTANCES_DEFAULT_X86 }}
           oss: ["alinux2023"]
           schedulers: ["slurm"]
+  tutorials:
+    test_upgrade_nvidia_software.py::test_upgrade_nvidia_software:
+      dimensions:
+        # The NVIDIA upgrade component targets Amazon Linux 2023 x86_64 only.
+        - regions: [ {{ g4dn_2xlarge_CAPACITY_RESERVATION_2_INSTANCES_2_HOURS_NOPG_alinux2023 }} ]
+          instances: [ "g4dn.2xlarge" ]
+          oss: [ "alinux2023" ]
+          schedulers: [ "slurm" ]
   update:
     test_update.py::test_update_slurm:
       dimensions:

tests/integration-tests/tests/basic/test_essential_features.py

@@ -27,7 +27,7 @@
     wait_instance_replaced_or_terminating,
 )
 from tests.common.mpi_common import _test_mpi
-from tests.common.utils import GPU_JOB_SCRIPT, fetch_instance_slots, run_system_analyzer
+from tests.common.utils import fetch_instance_slots, run_gpu_workload, run_system_analyzer
 
 
 def test_essential_features(
@@ -92,7 +92,7 @@ def test_essential_features(
         cluster, region, instance, scheduler, default_threads_per_core, request, scheduler_commands_factory
     )
 
-    _test_gpu_workload(cluster, scheduler_commands_factory, test_datadir)
+    _test_gpu_workload(cluster, scheduler_commands_factory)
 
 
 def _test_mpi_job(
@@ -334,27 +334,11 @@ def _test_custom_bootstrap_scripts_args_quotes(cluster):
     )
 
 
-def _test_gpu_workload(cluster, scheduler_commands_factory, test_datadir):
-    """Submit a Slurm job that builds and runs CUDA samples on a GPU compute node."""
+def _test_gpu_workload(cluster, scheduler_commands_factory):
+    """Submit Slurm jobs that build and run CUDA samples on a GPU compute node."""
     remote_command_executor = RemoteCommandExecutor(cluster)
     scheduler_commands = scheduler_commands_factory(remote_command_executor)
-
-    samples = ["1_Utilities/deviceQuery", "4_CUDA_Libraries/matrixMulCUBLAS"]
-    job_ids = []
-    for sample in samples:
-        logging.info("Submitting CUDA sample job for %s", sample)
-        result = scheduler_commands.submit_script(
-            str(GPU_JOB_SCRIPT),
-            script_args=[sample],
-            partition="gpu",
-            nodes=1,
-            slots=1,
-        )
-        job_ids.append(scheduler_commands.assert_job_submitted(result.stdout))
-
-    for job_id in job_ids:
-        scheduler_commands.wait_job_completed(job_id, timeout=20)
-        scheduler_commands.assert_job_succeeded(job_id)
+    run_gpu_workload(scheduler_commands, partition="gpu")
 
 
 def _test_disable_hyperthreading(

tests/integration-tests/tests/common/utils.py

@@ -46,6 +46,46 @@
 # compute node. Used by multiple tests to validate GPU workloads.
 GPU_JOB_SCRIPT = pathlib.Path(__file__).parent / "data/gpu_job.sh"
 
+# Default CUDA samples run by run_gpu_workload: deviceQuery validates driver/GPU visibility,
+# matrixMulCUBLAS exercises the CUDA libraries.
+GPU_WORKLOAD_SAMPLES = ["1_Utilities/deviceQuery", "4_CUDA_Libraries/matrixMulCUBLAS"]
+
+
+def run_gpu_workload(scheduler_commands, partition=None, samples=None, timeout=20):
+    """
+    Build and run CUDA samples on a GPU compute node and assert they succeed.
+
+    Submits one job per sample through the scheduler using the shared GPU job script
+    (tests/common/data/gpu_job.sh), which compiles the sample from the /usr/local/cuda-samples-*
+    tree installed on the AMI and runs it on a GPU node. All jobs are submitted upfront, then
+    awaited and asserted successful.
+
+    :param scheduler_commands: SchedulerCommands instance used to submit and check the jobs.
+    :param partition: optional scheduler partition to submit the jobs to.
+    :param samples: CUDA samples to run, as <category>/<sample>. Defaults to GPU_WORKLOAD_SAMPLES.
+    :param timeout: per-job completion timeout, in minutes.
+    :return: the list of completed job ids.
+    """
+    if samples is None:
+        samples = GPU_WORKLOAD_SAMPLES
+    job_ids = []
+    for sample in samples:
+        logging.info("Submitting CUDA sample job for %s", sample)
+        result = scheduler_commands.submit_script(
+            str(GPU_JOB_SCRIPT),
+            script_args=[sample],
+            partition=partition,
+            nodes=1,
+            slots=1,
+        )
+        job_ids.append(scheduler_commands.assert_job_submitted(result.stdout))
+
+    for job_id in job_ids:
+        scheduler_commands.wait_job_completed(job_id, timeout=timeout)
+        scheduler_commands.assert_job_succeeded(job_id)
+    return job_ids
+
+
 RHEL_OWNERS = ["309956199498", "841258680906", "219670896067"]
 
 OS_TO_OFFICIAL_AMI_NAME_OWNER_MAP = {

tests/integration-tests/tests/patching/test_patching.py

@@ -21,11 +21,11 @@
 from tests.common.osu_common import PRIVATE_OSES
 from tests.common.utils import (
     COMPUTE_NODE,
-    GPU_JOB_SCRIPT,
     LOGIN_NODE,
     NODE_TYPES,
     reboot_head_node,
     retrieve_cluster_head_node_ami,
+    run_gpu_workload,
     wait_node_reachable,
 )
 
@@ -231,17 +231,8 @@ def _run_gpu_workload(cluster, scheduler_commands_factory, use_login_node):
     logging.info("Submitting GPU validation job from the %s", source)
     remote_command_executor = RemoteCommandExecutor(cluster, use_login_node=use_login_node)
     scheduler_commands = scheduler_commands_factory(remote_command_executor)
-    result = scheduler_commands.submit_script(
-        str(GPU_JOB_SCRIPT),
-        script_args=["1_Utilities/deviceQuery"],
-        partition="q1",
-        nodes=1,
-        slots=1,
-    )
-    job_id = scheduler_commands.assert_job_submitted(result.stdout)
-    scheduler_commands.wait_job_completed(job_id, timeout=20)
-    scheduler_commands.assert_job_succeeded(job_id)
-    logging.info("GPU validation job %s submitted from the %s succeeded", job_id, source)
+    job_ids = run_gpu_workload(scheduler_commands, partition="q1", samples=["1_Utilities/deviceQuery"])
+    logging.info("GPU validation job %s submitted from the %s succeeded", job_ids[0], source)
 
 
 def _trigger_lazy_kernel_modules(cluster, scheduler_commands_factory):