GPU case: refactor Ubuntu Debian Installation

Author: adityagesh

lisa/microsoft/testsuites/cvm/cvm_attestation_tool.py

@@ -89,7 +89,7 @@ def _install(self) -> bool:
 
         git.clone(self.repo, root_path)
 
-        posix_os._install_package_from_url(
+        posix_os.install_package_from_url(
             self.deb_file, package_name="azguestattestation1.deb"
         )
         self.node.execute(

lisa/microsoft/testsuites/gpu/gpusuite.py

@@ -23,7 +23,6 @@
     Debian,
     Linux,
     Oracle,
-    Suse,
     Ubuntu,
     Windows,
 )
@@ -33,6 +32,7 @@
 from lisa.tools.python import PythonVenv
 from lisa.util import UnsupportedOperationException, get_matched_str
 
+
 _cudnn_location = (
     "https://developer.download.nvidia.com/compute/redist/cudnn/"
     "v7.5.0/cudnn-10.0-linux-x64-v7.5.0.56.tgz"
@@ -77,7 +77,7 @@ def before_case(self, log: Logger, **kwargs: Any) -> None:
         timeout=TIMEOUT,
         requirement=simple_requirement(
             supported_features=[GpuEnabled(), SerialConsole, AzureExtension],
-            unsupported_os=[Suse, AlmaLinux, Oracle],
+            unsupported_os=[AlmaLinux, Oracle],
         ),
         priority=1,
     )
@@ -143,7 +143,7 @@ def verify_max_gpu_provision(self, node: Node, log: Logger) -> None:
         timeout=TIMEOUT,
         requirement=simple_requirement(
             supported_features=[GpuEnabled(), SerialConsole, AzureExtension],
-            unsupported_os=[Suse, AlmaLinux, Oracle],
+            unsupported_os=[AlmaLinux, Oracle],
         ),
         priority=2,
     )
@@ -178,7 +178,7 @@ def verify_gpu_extension_installation(
         timeout=TIMEOUT,
         requirement=simple_requirement(
             supported_features=[GpuEnabled()],
-            unsupported_os=[Suse, AlmaLinux, Oracle],
+            unsupported_os=[AlmaLinux, Oracle],
         ),
         priority=2,
     )
@@ -254,7 +254,7 @@ def verify_gpu_rescind_validation(
         priority=3,
         requirement=simple_requirement(
             supported_features=[GpuEnabled()],
-            unsupported_os=[Suse, AlmaLinux, Oracle],
+            unsupported_os=[AlmaLinux, Oracle],
         ),
     )
     def verify_gpu_cuda_with_pytorch(

lisa/operating_system.py

@@ -687,7 +687,7 @@ def _get_package_list(
             self._initialize_package_installation()
         return package_names
 
-    def _install_package_from_url(
+    def install_package_from_url(
         self,
         package_url: str,
         package_name: str = "",

lisa/tools/gpu_drivers.py

@@ -15,6 +15,7 @@
     Oracle,
     Posix,
     Redhat,
+    Suse,
     Ubuntu,
 )
 
@@ -303,6 +304,10 @@ class NvidiaCudaDriver(GpuDriver):
     """
 
     DEFAULT_CUDA_VERSION = "10.1.243-1"
+    DEFAULT_CUDA_PACKAGE = "cuda-drivers"
+    NVIDIA_CUDA_REPO_BASE_URL = (
+        "https://developer.download.nvidia.com/compute/cuda/repos"
+    )
 
     @classmethod
     def type_name(cls) -> str:
@@ -406,10 +411,11 @@ def _install_dependencies(self) -> None:
             if release == "7":
                 assert isinstance(self.node.os, Posix)
                 self._log.debug("Installing vulkan-filesystem for CentOS 7")
-                self.node.os._install_package_from_url(
+                package_url = (
                     "https://vault.centos.org/centos/7/os/x86_64/Packages/"
                     "vulkan-filesystem-1.1.97.0-1.el7.noarch.rpm"
                 )
+                self.node.os.install_package_from_url(package_url, signed=False)
 
     def _install_driver(self) -> None:
         """Install CUDA driver based on OS"""
@@ -436,8 +442,8 @@ def _install_cuda_redhat(self) -> None:
 
         # Add CUDA repository
         self.node.os.add_repository(
-            f"http://developer.download.nvidia.com/compute/cuda/"
-            f"repos/rhel{release}/x86_64/cuda-rhel{release}.repo"
+            f"{self.NVIDIA_CUDA_REPO_BASE_URL}/"
+            f"rhel{release}/x86_64/cuda-rhel{release}.repo"
         )
 
         # Install CUDA packages
@@ -450,108 +456,93 @@ def _install_cuda_redhat(self) -> None:
 
     def _install_cuda_ubuntu(self) -> None:
         """Install CUDA driver on Ubuntu"""
-        self._log.debug("Installing CUDA driver for Ubuntu")
-
-        assert isinstance(self.node.os, Posix), "CUDA installation requires Posix OS"
-
-        cuda_package_name = "cuda-drivers"
-        cuda_drivers_package_pattern = re.compile(
-            r"^cuda-drivers-(\d+)/.*$", re.MULTILINE
-        )
-
         os_info = self.node.os.information
-        release = re.sub("[^0-9]+", "", os_info.release)
 
-        # Handle unsupported releases by using closest supported version
-        if release in ["2110", "2104"]:
-            release = "2004"
-        if release in ["2210", "2304", "2310"]:
-            release = "2204"
+        # NVIDIA only provides CUDA repos for LTS releases (even major version, XX.04)
+        # For non-LTS releases or odd major versions, use the previous LTS release
+        # e.g., 21.10 -> 20.04, 22.10 -> 22.04, 23.04 -> 22.04, 23.10 -> 22.04
+        major_version = os_info.version.major
+        minor_version = os_info.version.minor
 
-        # Install CUDA public GPG key
-        cuda_keyring = "cuda-keyring_1.1-1_all.deb"
-        self.node.tools[Wget].get(
-            f"https://developer.download.nvidia.com/compute/cuda/repos/"
-            f"ubuntu{release}/x86_64/{cuda_keyring}"
-        )
-        self.node.execute(
-            f"dpkg -i {cuda_keyring}",
-            sudo=True,
-            cwd=self.node.get_working_path(),
-        )
+        # If odd major version (e.g., 21, 23), use previous even major version
+        if major_version % 2 == 1:
+            major_version = major_version - 1
+            release = f"{major_version}04"
+            self._log.debug(
+                f"Using previous LTS release {release} for CUDA repository "
+                f"(original: {os_info.release}, odd major version)"
+            )
+        elif minor_version != 4:
+            # Even major but non-LTS release (not XX.04), use corresponding LTS
+            release = f"{major_version}04"
+            self._log.debug(
+                f"Using LTS release {release} for CUDA repository "
+                f"(original: {os_info.release})"
+            )
+        else:
+            # LTS release (even major, XX.04)
+            release = f"{major_version}{minor_version:0>2}"
 
         # For Ubuntu 16.04, use legacy installation method
         if release == "1604":
-            cuda_repo_pkg = (
-                f"cuda-repo-ubuntu{release}_" f"{self.DEFAULT_CUDA_VERSION}_amd64.deb"
-            )
-            cuda_repo = (
-                f"http://developer.download.nvidia.com/compute/cuda/repos/"
-                f"ubuntu{release}/x86_64/{cuda_repo_pkg}"
+            self._log.debug("Installing CUDA driver for Ubuntu 16.04 (legacy method)")
+
+            assert isinstance(self.node.os, Ubuntu), "Ubuntu installation expected"
+
+            # Install CUDA keyring
+            cuda_keyring_url = (
+                f"{self.NVIDIA_CUDA_REPO_BASE_URL}/"
+                f"ubuntu{release}/x86_64/cuda-keyring_1.1-1_all.deb"
             )
-            self.node.os._install_package_from_url(
-                cuda_repo, package_name="cuda-drivers.deb", signed=False
-            )
-        else:
-            # Modern Ubuntu versions
-            self.node.tools[Wget].get(
-                f"https://developer.download.nvidia.com/compute/cuda/repos/"
-                f"ubuntu{release}/x86_64/cuda-ubuntu{release}.pin",
-                "/etc/apt/preferences.d",
-                "cuda-repository-pin-600",
-                sudo=True,
-                overwrite=False,
+            self.node.os.install_package_from_url(
+                cuda_keyring_url,
+                package_name="cuda-keyring.deb",
+                signed=False,
             )
 
-            # Add CUDA repository
-            repo_entry = (
-                f"deb http://developer.download.nvidia.com/compute/cuda/repos/"
-                f"ubuntu{release}/x86_64/ /"
+            # Install CUDA repository package
+            cuda_repo_url = (
+                f"{self.NVIDIA_CUDA_REPO_BASE_URL}/"
+                f"ubuntu{release}/x86_64/cuda-repo-ubuntu{release}_"
+                f"{self.DEFAULT_CUDA_VERSION}_amd64.deb"
             )
-            self.node.execute(
-                f'add-apt-repository -y "{repo_entry}"',
-                sudo=True,
-                expected_exit_code=0,
-                expected_exit_code_failure_message=f"failed to add repo {repo_entry}",
+            self.node.os.install_package_from_url(
+                cuda_repo_url,
+                package_name="cuda-drivers.deb",
+                signed=False,
             )
-
-            # Find available CUDA driver versions
-            result = self.node.execute(f"apt search {cuda_package_name}", sudo=True)
-            available_versions = cuda_drivers_package_pattern.findall(result.stdout)
-
-            if available_versions:
-                # Sort versions and select the highest one
-                highest_version = max(available_versions, key=int)
-                package_name = f"{cuda_package_name}-{highest_version}"
-            else:
-                raise MissingPackagesException([cuda_package_name])
-
-            self.node.os.install_packages(package_name)
+        else:
+            # Modern Ubuntu versions use the same method as Debian
+            self._install_cuda_debian_based(f"ubuntu{release}")
 
         self._log.info("Successfully installed CUDA driver for Ubuntu")
 
     def _install_cuda_debian(self) -> None:
         """Install CUDA driver on Debian"""
-        self._log.debug("Installing CUDA driver for Debian")
+        os_info = self.node.os.information
+        major_version = str(os_info.version.major)
+        release = f"debian{major_version}"
+        self._install_cuda_debian_based(release)
+        self._log.info("Successfully installed CUDA driver for Debian")
+
+    def _install_cuda_debian_based(self, release: str) -> None:
+        """
+        Shared installation method for Debian-based distributions.
+        Supports both Debian and Ubuntu.
+        """
+        self._log.debug(f"Installing CUDA driver for {release}")
 
         assert isinstance(self.node.os, Posix), "CUDA installation requires Posix OS"
 
-        cuda_package_name = "cuda-drivers"
+        cuda_package_name = self.DEFAULT_CUDA_PACKAGE
         cuda_drivers_package_pattern = re.compile(
-            r"^cuda-drivers-(\d+)\s", re.MULTILINE
+            rf"^{cuda_package_name}-(\d+)\s", re.MULTILINE
         )
 
-        os_info = self.node.os.information
-        major_version = str(os_info.version.major)
-
-        # Debian CUDA repos follow pattern: debian{major_version}
-        release = f"debian{major_version}"
-
         # Install CUDA public GPG key
         cuda_keyring = "cuda-keyring_1.1-1_all.deb"
         self.node.tools[Wget].get(
-            f"https://developer.download.nvidia.com/compute/cuda/repos/"
-            f"{release}/x86_64/{cuda_keyring}"
+            f"{self.NVIDIA_CUDA_REPO_BASE_URL}/{release}/x86_64/{cuda_keyring}"
         )
         self.node.execute(
             f"dpkg -i {cuda_keyring}",
@@ -562,7 +553,7 @@ def _install_cuda_debian(self) -> None:
         # Add CUDA repository
         repo_entry = (
             f"deb [signed-by=/usr/share/keyrings/cuda-archive-keyring.gpg] "
-            f"https://developer.download.nvidia.com/compute/cuda/repos/"
+            f"{self.NVIDIA_CUDA_REPO_BASE_URL}/"
             f"{release}/x86_64/ /"
         )
 
@@ -593,8 +584,6 @@ def _install_cuda_debian(self) -> None:
 
         self.node.os.install_packages(package_name)
 
-        self._log.info("Successfully installed CUDA driver for Debian")
-
     def _install_cuda_mariner(self) -> None:
         """Install CUDA driver on CBL-Mariner"""
         self._log.debug("Installing CUDA driver for CBL-Mariner")