diff --git a/tests/integration-tests/configs/develop.yaml b/tests/integration-tests/configs/develop.yaml index b38406dddf..52aae2e9b2 100644 --- a/tests/integration-tests/configs/develop.yaml +++ b/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: diff --git a/tests/integration-tests/configs/released.yaml b/tests/integration-tests/configs/released.yaml index ae02e04747..e7a16039ed 100644 --- a/tests/integration-tests/configs/released.yaml +++ b/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: diff --git a/tests/integration-tests/tests/basic/test_essential_features.py b/tests/integration-tests/tests/basic/test_essential_features.py index cbe128f99a..22630a0315 100644 --- a/tests/integration-tests/tests/basic/test_essential_features.py +++ b/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( diff --git a/tests/integration-tests/tests/common/utils.py b/tests/integration-tests/tests/common/utils.py index 7bda9f3dbd..a998374c94 100644 --- a/tests/integration-tests/tests/common/utils.py +++ b/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 /. 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 = { diff --git a/tests/integration-tests/tests/patching/test_patching.py b/tests/integration-tests/tests/patching/test_patching.py index f07b9765d9..b654d9bcf1 100644 --- a/tests/integration-tests/tests/patching/test_patching.py +++ b/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): diff --git a/tests/integration-tests/tests/tutorials/__init__.py b/tests/integration-tests/tests/tutorials/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software.py b/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software.py new file mode 100644 index 0000000000..ca69e4dfb1 --- /dev/null +++ b/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software.py @@ -0,0 +1,273 @@ +# Copyright 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# A copy of the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "LICENSE.txt" file accompanying this file. +# This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. +# See the License for the specific language governing permissions and limitations under the License. +import logging +import time + +import boto3 +import pytest +import yaml +from assertpy import assert_that, soft_assertions +from jinja2 import DebugUndefined +from jinja2.sandbox import SandboxedEnvironment +from remote_command_executor import RemoteCommandExecutor +from retrying import retry +from time_utils import minutes, seconds +from utils import generate_stack_name + +from tests.common.assertions import assert_head_node_is_running +from tests.common.utils import ( + generate_random_string, + get_installed_parallelcluster_base_version, + retrieve_latest_ami, + run_gpu_workload, +) + +# Instance types: build on a GPU instance so the NVIDIA driver installation is exercised on GPU +# hardware; the head node validates that the upgraded AMI boots on non-GPU instances too. The GPU +# compute node type comes from the `instance` dimension. +BUILD_INSTANCE_TYPE = "g4dn.2xlarge" +HEAD_NODE_INSTANCE_TYPE = "c5.xlarge" + +QUEUE_NAME = "q1" +COMPUTE_RESOURCE_NAME = "cr1" + +# Software versions installed by the component. The test injects these values into the AWSTOE +# constants of the component document (placeholder markers in update-nvidia.yaml) and +# asserts the same versions on the cluster nodes. +NVIDIA_DRIVER_VERSION = "595.71.05" +CUDA_VERSION = "13.2.2" +CUDA_SAMPLES_VERSION = "13.3" +# Driver version pinned into the CUDA runfile filename by NVIDIA for this CUDA release. It is +# fixed per CUDA version and independent of the driver installed by the component. +CUDA_RELEASE_NVIDIA_VERSION = "595.71.05" +# nvcc reports the major.minor release of the CUDA toolkit. +CUDA_RELEASE = ".".join(CUDA_VERSION.split(".")[:2]) + +# NVLSM is versioned independently of the driver and is not pinned in the component: the component +# installs the version bundled in the NVIDIA driver local repository for NVIDIA_DRIVER_VERSION +# (2025.10.14 is the version bundled in the 580.173.02 driver local repository). +NVLSM_BUNDLED_VERSION = "2025.10.14" + +# Packages whose version must match the NVIDIA driver version exactly. +DRIVER_ALIGNED_PACKAGES = ["nvidia-fabricmanager", "nvidia-imex", "libnvsdm"] + + +@pytest.fixture() +def nvidia_stack_component(region, request): + """Manage the creation/deletion of the EC2 Image Builder component with the NVIDIA upgrade procedure.""" + imagebuilder_client = boto3.client("imagebuilder", region_name=region) + component_arn = None + + def _create_component(component_document): + nonlocal component_arn + # Random suffix to avoid name+version clashes across concurrent test runs in the same account/region. + component_name = f"update-nvidia-{generate_random_string()}" + logging.info("Creating Image Builder component %s", component_name) + response = imagebuilder_client.create_component( + name=component_name, + semanticVersion="1.0.0", + platform="Linux", + data=component_document, + ) + component_arn = response["componentBuildVersionArn"] + logging.info("Created Image Builder component %s", component_arn) + return component_arn + + yield _create_component + + if component_arn and not request.config.getoption("no_delete"): + logging.info("Deleting Image Builder component %s", component_arn) + imagebuilder_client.delete_component(componentBuildVersionArn=component_arn) + logging.info("Deleted Image Builder component %s", component_arn) + elif component_arn: + logging.warning("Skipping deletion of Image Builder component %s because --no-delete is set", component_arn) + + +@pytest.mark.usefixtures("scheduler") +def test_upgrade_nvidia_software( + region, + os, + instance, + architecture, + test_datadir, + pcluster_config_reader, + images_factory, + clusters_factory, + scheduler_commands_factory, + nvidia_stack_component, + request, +): + """ + Validate the tutorial procedure to upgrade the NVIDIA software stack via a custom AMI. + + Steps: + 1. Create an EC2 Image Builder component with the NVIDIA upgrade procedure (driver + CUDA + runfiles, NVLink stack from the NVIDIA driver local repository), injecting the expected + software versions into the AWSTOE constants of the component document. + 2. Build a custom AMI with `pcluster build-image` using the official vanilla OS AMI as + parent and the component from step 1. + 3. Wait for the AMI produced by the build to be available in EC2. + 4. Create a cluster using the custom AMI, with a static GPU compute node. + 5. Assert the NVIDIA software versions (driver, CUDA, Fabric Manager, IMEX, libnvsdm, NVLSM) + on the GPU compute node are the ones installed by the component. + 6. Run a GPU workload through the scheduler, reusing the shared CUDA samples job script + (tests/common/data/gpu_job.sh), and assert it succeeds. + 7. Teardown is managed by the fixtures (cluster, image and Image Builder component). + """ + # Step 1: render the component document, injecting the software versions into its AWSTOE + # constants, and create the Image Builder component from it. + logging.info("Expecting NVIDIA driver %s and CUDA release %s", NVIDIA_DRIVER_VERSION, CUDA_RELEASE) + component_document = _render_component_document(test_datadir, architecture) + component_arn = nvidia_stack_component(component_document) + + # Step 2: build the custom AMI with pcluster build-image, using the official vanilla OS AMI + # as parent image. + parent_image = retrieve_latest_ami(region, os, ami_type="official", architecture=architecture) + image_id = generate_stack_name("integ-tests-upgrade-nvidia", request.config.getoption("stackname_suffix")) + image_config = pcluster_config_reader( + config_file="image.config.yaml", + parent_image=parent_image, + component_arn=component_arn, + build_instance_type=BUILD_INSTANCE_TYPE, + ) + image = images_factory(image_id, image_config, region) + _wait_for_build_image_complete(image) + + # Step 3: the build produced an AMI; wait for it to be available in EC2. + assert_that(image.ec2_image_id).described_as("EC2 AMI id from the image build").is_not_none() + _wait_for_ami_available(region, image.ec2_image_id) + + # Step 4: create a cluster using the custom AMI. + cluster_config = pcluster_config_reader( + custom_ami=image.ec2_image_id, + head_node_instance_type=HEAD_NODE_INSTANCE_TYPE, + queue_name=QUEUE_NAME, + compute_resource_name=COMPUTE_RESOURCE_NAME, + ) + cluster = clusters_factory(cluster_config) + assert_head_node_is_running(region, cluster) + + # Step 5: assert the NVIDIA software versions on the GPU compute node. + compute_node_ip = cluster.get_compute_nodes_private_ip(QUEUE_NAME, COMPUTE_RESOURCE_NAME)[0] + compute_remote_command_executor = RemoteCommandExecutor(cluster, compute_node_ip=compute_node_ip) + _assert_nvidia_stack_versions(compute_remote_command_executor, NVIDIA_DRIVER_VERSION, CUDA_RELEASE) + + # Step 6: run a GPU workload through the scheduler and assert it succeeds. + remote_command_executor = RemoteCommandExecutor(cluster) + scheduler_commands = scheduler_commands_factory(remote_command_executor) + run_gpu_workload(scheduler_commands, partition=QUEUE_NAME) + + # Step 7: teardown is managed by clusters_factory, images_factory and nvidia_stack_component. + + +def _render_component_document(test_datadir, architecture): + """ + Render the component document template, injecting the software versions into its constants. + + The template is the tutorial component document with lowercase jinja variables as the values + of its AWSTOE constants section; every other line is identical to the tutorial document. + Injecting the values here keeps the test as the single source of truth for the expected + versions, instead of parsing them back out of the component. + + The jinja environment uses DebugUndefined so that the uppercase AWSTOE constant references + used by the document steps ("{{ NVIDIA_DRIVER_VERSION }}"), being undefined jinja variables, + are re-emitted verbatim instead of blanked: they must reach Image Builder untouched, since + AWSTOE resolves them at build time. + """ + expected_constants = { + "NVIDIA_DRIVER_VERSION": NVIDIA_DRIVER_VERSION, + # The NVIDIA installers use "aarch64" for arm64. + "ARCH": {"x86_64": "x86_64", "arm64": "aarch64"}[architecture], + "CUDA_VERSION": CUDA_VERSION, + "CUDA_SAMPLES_VERSION": CUDA_SAMPLES_VERSION, + "CUDA_RELEASE_NVIDIA_VERSION": CUDA_RELEASE_NVIDIA_VERSION, + } + template = SandboxedEnvironment(undefined=DebugUndefined).from_string( + (test_datadir / "update-nvidia.yaml").read_text() + ) + document = template.render(**{name.lower(): value for name, value in expected_constants.items()}) + + # Guard: the constants of the rendered document must carry exactly the injected values + # (catches blanked or misnamed jinja variables). + rendered_constants = { + name: attrs["value"] for item in yaml.safe_load(document).get("constants", []) for name, attrs in item.items() + } + assert_that(rendered_constants).described_as("rendered component constants").is_equal_to(expected_constants) + return document + + +def _wait_for_build_image_complete(image): + """Wait for the image build to complete and assert it succeeded.""" + logging.info("Waiting for build of image %s to complete", image.image_id) + while image.image_status.endswith("_IN_PROGRESS"): # e.g. BUILD_IN_PROGRESS + time.sleep(300) + logging.info(image.describe()) + if image.image_status != "BUILD_COMPLETE": + _log_recent_image_build_events(image) + assert_that(image.image_status).is_equal_to("BUILD_COMPLETE") + + +def _log_recent_image_build_events(image): + """Log the last lines of the image build log to ease troubleshooting of a failed build.""" + log_stream_name = f"{get_installed_parallelcluster_base_version()}/1" + nlines = 200 + try: + log_events = image.get_log_events(log_stream_name, start_from_head=False, query="events[*]", limit=nlines) + log_messages = [event["message"] for event in log_events] + logging.error( + "Image build failed for %s, the last %d lines of the log are:\n%s", + image.image_id, + nlines, + "\n".join(log_messages), + ) + except Exception as e: # noqa: BLE001 + logging.error("Could not retrieve build log events for image %s: %s", image.image_id, e) + + +@retry( + retry_on_result=lambda state: state != "available", + wait_fixed=seconds(30), + stop_max_delay=minutes(15), +) +def _wait_for_ami_available(region, ec2_image_id): + """Wait for the AMI produced by the build to reach the "available" state in EC2.""" + images = boto3.client("ec2", region_name=region).describe_images(ImageIds=[ec2_image_id]).get("Images") + state = images[0]["State"] if images else None + logging.info("AMI %s state: %s", ec2_image_id, state) + return state + + +def _assert_nvidia_stack_versions(remote_command_executor, driver_version, cuda_release): + """Assert driver, CUDA and NVLink stack versions on the node (must run on a GPU node).""" + + def _run(command): + return remote_command_executor.run_remote_command(command).stdout.strip() + + with soft_assertions(): + # NVIDIA driver: both the installed kernel module and the one loaded by the running GPU. + assert_that(_run("modinfo -F version nvidia")).described_as("installed nvidia kernel module").is_equal_to( + driver_version + ) + assert_that(_run("nvidia-smi --query-gpu=driver_version --format=csv,noheader")).described_as( + "driver loaded by nvidia-smi" + ).is_equal_to(driver_version) + + # CUDA toolkit. + assert_that(_run("/usr/local/cuda/bin/nvcc --version")).described_as("nvcc release").contains( + f"release {cuda_release}" + ) + + # NVLink software stack: Fabric Manager, IMEX and libnvsdm must match the driver version + # exactly; NVLSM is installed at the version bundled in the driver local repository. + for package in DRIVER_ALIGNED_PACKAGES: + assert_that(_run(f"rpm -q --qf '%{{VERSION}}' {package}")).described_as(package).is_equal_to(driver_version) + assert_that(_run("rpm -q --qf '%{VERSION}' nvlsm")).described_as("nvlsm").is_equal_to(NVLSM_BUNDLED_VERSION) diff --git a/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software/test_upgrade_nvidia_software/image.config.yaml b/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software/test_upgrade_nvidia_software/image.config.yaml new file mode 100644 index 0000000000..740688b11b --- /dev/null +++ b/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software/test_upgrade_nvidia_software/image.config.yaml @@ -0,0 +1,8 @@ +Build: + InstanceType: {{ build_instance_type }} + ParentImage: {{ parent_image }} + Components: + - Type: arn + Value: {{ component_arn }} +DevSettings: + TerminateInstanceOnFailure: True diff --git a/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software/test_upgrade_nvidia_software/pcluster.config.yaml b/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software/test_upgrade_nvidia_software/pcluster.config.yaml new file mode 100644 index 0000000000..d8c9ef515f --- /dev/null +++ b/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software/test_upgrade_nvidia_software/pcluster.config.yaml @@ -0,0 +1,22 @@ +Image: + Os: {{ os }} + CustomAmi: {{ custom_ami }} +HeadNode: + InstanceType: {{ head_node_instance_type }} + Networking: + SubnetId: {{ public_subnet_id }} + Ssh: + KeyName: {{ key_name }} +Scheduling: + Scheduler: slurm + SlurmQueues: + - Name: {{ queue_name }} + Networking: + SubnetIds: + - {{ private_subnet_id }} + ComputeResources: + - Name: {{ compute_resource_name }} + InstanceType: {{ instance }} + # Keep one static GPU node up so the NVIDIA stack versions can be checked directly on it. + MinCount: 1 + MaxCount: 1 diff --git a/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software/test_upgrade_nvidia_software/update-nvidia.yaml b/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software/test_upgrade_nvidia_software/update-nvidia.yaml new file mode 100644 index 0000000000..84fed34b50 --- /dev/null +++ b/tests/integration-tests/tests/tutorials/test_upgrade_nvidia_software/test_upgrade_nvidia_software/update-nvidia.yaml @@ -0,0 +1,206 @@ +name: update-nvidia +description: >- + Update the NVIDIA software stack: driver, CUDA, Fabric Manager, IMEX and NVLSM. +schemaVersion: 1.0 + +# Single source of truth for every step in this document: upgrading the software versions or +# changing the architecture (use "aarch64" for arm64) only requires editing these values. +constants: + - NVIDIA_DRIVER_VERSION: + type: string + value: "{{ nvidia_driver_version }}" + - ARCH: + type: string + value: "{{ arch }}" + - CUDA_VERSION: + type: string + value: "{{ cuda_version }}" + - CUDA_SAMPLES_VERSION: + type: string + value: "{{ cuda_samples_version }}" + # Driver version pinned into the CUDA runfile filename by NVIDIA for this CUDA release. It is + # fixed per CUDA version and independent of the driver installed by this component. + # CUDA 13.3.1 ships as cuda_13.3.1_610.43.02_linux.run. + - CUDA_RELEASE_NVIDIA_VERSION: + type: string + value: "{{ cuda_release_nvidia_version }}" + +phases: + - name: build + steps: + - name: InstallDriver + action: ExecuteBash + inputs: + commands: + - | + #!/bin/bash + set -ex + + # Values injected from the document constants. + NVIDIA_DRIVER_VERSION="{{ NVIDIA_DRIVER_VERSION }}" + ARCH="{{ ARCH }}" + + # Create temporary directory + TMP_DIR="/pcluster-tmp/$(date +"%Y-%m-%dT%H-%M-%S")" + + COMPILER_PATH="/usr/bin/gcc" + export CC="${COMPILER_PATH}" + + NVIDIA_RUNFILE="NVIDIA-Linux-${ARCH}-${NVIDIA_DRIVER_VERSION}.run" + wget -P "${TMP_DIR}" "https://us.download.nvidia.com/tesla/${NVIDIA_DRIVER_VERSION}/${NVIDIA_RUNFILE}" + chmod +x "${TMP_DIR}/${NVIDIA_RUNFILE}" + "${TMP_DIR}/${NVIDIA_RUNFILE}" --silent --dkms --disable-nouveau -m="kernel-open" + + # Cleanup + rm -rf "${TMP_DIR}" + + - name: InstallNVLinkStack + action: ExecuteBash + inputs: + commands: + - | + #!/bin/bash + set -ex + + # -------------------------------------------------------------------------- + # Keep the NVLink software stack aligned with the NVIDIA driver. + # + # Upgrading the driver alone is not enough. NVIDIA Fabric Manager and IMEX + # must match the driver version EXACTLY; a mismatched service refuses to + # start, which breaks multi-GPU (NVSwitch: p4d/p4de, p5/p5e/p5en, p6-b200) + # and multi-node NVLink (GB200 NVL: p6e-gb200) communication even when + # "nvidia-smi" on a single GPU looks healthy. + # + # The packages are installed from the NVIDIA driver local repository + # (https://docs.nvidia.com/datacenter/tesla/driver-installation-guide/), + # which bundles the NVLink software stack validated for this driver + # release: Fabric Manager, IMEX and libnvsdm at the exact driver version, + # plus a compatible NVLSM (NVLink Subnet Manager, versioned independently + # of the driver). + # + # NVIDIA_DRIVER_VERSION and ARCH are injected from the document constants, + # so they are guaranteed to match the values used by the InstallDriver step. + # -------------------------------------------------------------------------- + NVIDIA_DRIVER_VERSION="{{ NVIDIA_DRIVER_VERSION }}" + ARCH="{{ ARCH }}" + + # Create temporary directory + TMP_DIR="/pcluster-tmp/$(date +"%Y-%m-%dT%H-%M-%S")" + + # Download the NVIDIA driver local repository for the target OS/arch and + # register it with dnf: installing the wrapper RPM only drops the repo + # definition, GPG keyring and bundled RPMs under /var, it does not install + # any driver package by itself. Adapt "amzn2023" for other distributions + # (e.g. rhel8, rhel9); the architecture comes from the ARCH document constant. + LOCAL_REPO_PKG="nvidia-driver-local-repo-amzn2023-${NVIDIA_DRIVER_VERSION}" + LOCAL_REPO_RPM="${LOCAL_REPO_PKG}-1.0-1.${ARCH}.rpm" + wget -P "${TMP_DIR}" "https://developer.download.nvidia.com/compute/nvidia-driver/${NVIDIA_DRIVER_VERSION}/local_installers/${LOCAL_REPO_RPM}" + dnf install -y "${TMP_DIR}/${LOCAL_REPO_RPM}" + + # ParallelCluster AMIs pin the NVLink stack with dnf versionlock at build + # time. Stale locks filter the new versions out of the dnf transaction + # ("All matches were filtered out by exclude filtering") and silently keep + # nvlsm at the locked version, so clear them before installing. + dnf versionlock delete nvidia-fabricmanager nvidia-imex libnvsdm nvlsm || true + + # Install/upgrade the NVLink software stack from the local repo. Fabric + # Manager, IMEX and libnvsdm are pinned to the exact driver version; NVLSM + # is installed at the version bundled in this driver's local repo; + # infiniband-diags and libibumad (from the OS repos) are runtime + # dependencies of NVLSM. --best is required on Amazon Linux 2023: dnf.conf + # ships best=False, which turns "dnf install " into a no-op when any + # version is already installed, silently skipping the upgrade of unpinned + # packages such as nvlsm. + dnf install -y --best \ + "nvidia-fabricmanager-${NVIDIA_DRIVER_VERSION}-1*" \ + "nvidia-imex-${NVIDIA_DRIVER_VERSION}-1*" \ + "libnvsdm-${NVIDIA_DRIVER_VERSION}-1*" \ + nvlsm \ + infiniband-diags libibumad + + # Re-apply the version lock so a later "dnf upgrade" cannot drift these + # packages away from the driver version. + dnf versionlock add nvidia-fabricmanager nvidia-imex libnvsdm nvlsm + + # Drop the local repository once the packages are installed: the wrapper + # package keeps a full copy of the bundled RPMs under /var, which would + # otherwise bloat the image. Removing it does not touch the packages + # installed from it. + dnf remove -y "${LOCAL_REPO_PKG}" + + # Cleanup + rm -rf "${TMP_DIR}" + + - name: InstallCuda + action: ExecuteBash + inputs: + commands: + - | + #!/bin/bash + set -ex + + # Values injected from the document constants. + CUDA_VERSION="{{ CUDA_VERSION }}" + CUDA_SAMPLES_VERSION="{{ CUDA_SAMPLES_VERSION }}" + CUDA_RELEASE_NVIDIA_VERSION="{{ CUDA_RELEASE_NVIDIA_VERSION }}" + + # Create temporary directory + TMP_DIR="/pcluster-tmp/$(date +"%Y-%m-%dT%H-%M-%S")" + + CUDA_RUNFILE="cuda_${CUDA_VERSION}_${CUDA_RELEASE_NVIDIA_VERSION}_linux.run" + wget -P "${TMP_DIR}" "https://developer.download.nvidia.com/compute/cuda/${CUDA_VERSION}/local_installers/${CUDA_RUNFILE}" + chmod +x "${TMP_DIR}/${CUDA_RUNFILE}" + CUDA_TMP_INSTALL_DIR="${TMP_DIR}/cuda-install" + mkdir -p "${CUDA_TMP_INSTALL_DIR}" + "${TMP_DIR}/${CUDA_RUNFILE}" --silent --toolkit --samples --tmpdir="${CUDA_TMP_INSTALL_DIR}" + + CUDA_SAMPLES_ARCHIVE="v${CUDA_SAMPLES_VERSION}.tar.gz" + wget -P "${TMP_DIR}" "https://github.com/NVIDIA/cuda-samples/archive/refs/tags/v${CUDA_SAMPLES_VERSION}.tar.gz" + tar xf "${TMP_DIR}/${CUDA_SAMPLES_ARCHIVE}" --directory "/usr/local/" + + # Cleanup + rm -rf "${TMP_DIR}" + + ## Add CUDA to PATH + CUDA_PATH="/usr/local/cuda" + echo "export PATH=${CUDA_PATH}/bin:\${PATH}" > /etc/profile.d/pcluster_cuda.sh + echo "export LD_LIBRARY_PATH=${CUDA_PATH}/lib64:\${LD_LIBRARY_PATH}" >> /etc/profile.d/pcluster_cuda.sh + chmod +x /etc/profile.d/pcluster_cuda.sh + + - name: Validation + action: ExecuteBash + inputs: + commands: + - | + #!/bin/bash + set -ex + + # Injected from the document constant, matching the installed driver version. + NVIDIA_DRIVER_VERSION="{{ NVIDIA_DRIVER_VERSION }}" + + ## Driver and CUDA + source /etc/profile.d/pcluster_cuda.sh + ls -l /usr/local + which nvcc + nvcc --version + which nvidia-smi + nvidia-smi + + ## NVLink stack alignment: Fabric Manager, IMEX and libnvsdm must match + ## the driver version exactly. + for pkg in nvidia-fabricmanager nvidia-imex libnvsdm; do + installed="$(rpm -q --qf '%{VERSION}' "${pkg}")" + echo "${pkg} version: ${installed} (driver: ${NVIDIA_DRIVER_VERSION})" + if [ "${installed}" != "${NVIDIA_DRIVER_VERSION}" ]; then + echo "ERROR: ${pkg} ${installed} does not match driver ${NVIDIA_DRIVER_VERSION}" >&2 + exit 1 + fi + done + + ## NVLSM is versioned independently of the driver. + rpm -q nvlsm + + ## Binary version checks (non-fatal: they require the driver to be loaded, + ## which may not be the case on the build instance). + nv-fabricmanager --version || true + nvidia-imex -v || true