Adding GPU feature and some dependent tools

lisa/features/__init__.py

@@ -1,7 +1,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
+from .gpu import Gpu
 from .serial_console import SerialConsole
 from .startstop import StartStop
 
-__all__ = ["SerialConsole", "StartStop"]
+__all__ = ["Gpu", "SerialConsole", "StartStop"]

lisa/features/gpu.py

@@ -0,0 +1,133 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+import re
+from enum import Enum
+from typing import Any, cast
+
+from lisa.feature import Feature
+from lisa.operating_system import Linux, Redhat, Ubuntu
+from lisa.tools import Uname, Wget
+from lisa.util import LisaException
+
+FEATURE_NAME_GPU = "Gpu"
+
+ComputeSDK = Enum(
+    "ComputeSDK",
+    [
+        # GRID Driver
+        "GRID",
+        # CUDA Driver
+        "CUDA",
+    ],
+)
+
+
+class Gpu(Feature):
+    def __init__(self, node: Any, platform: Any) -> None:
+        super().__init__(node, platform)
+        self._log = self._node.log
+
+    @classmethod
+    def name(cls) -> str:
+        return FEATURE_NAME_GPU
+
+    def _install_grid_driver(self, version: str) -> None:
+        self._log.info("Starting GRID driver installation")
+
+    def _install_cuda_driver(self, version: str = "10.1.105-1") -> None:
+        self._log.info("Starting CUDA driver installation")
+        cuda_repo = ""
+        linux_os: Linux = cast(Linux, self._node.os)
+
+        # CUDA driver installation for redhat distros
+        if isinstance(self._node.os, Redhat):
+            cuda_repo_pkg = f"cuda-repo-rhel7-{version}.x86_64.rpm"
+            cuda_repo = (
+                "http://developer.download.nvidia.com/"
+                f"compute/cuda/repos/rhel7/x86_64/{cuda_repo_pkg}"
+            )
+            linux_os = Redhat(self._node)
+
+        # CUDA driver installation for Ubuntu distros
+        elif isinstance(self._node.os, Ubuntu):
+            release_version = self._node.os._os_version.release
+            release = re.sub("[^0-9]+", "", release_version)
+            cuda_repo_pkg = f"cuda-repo-ubuntu{release}_{version}_amd64.deb"
+            cuda_repo = (
+                "http://developer.download.nvidia.com/compute/"
+                f"cuda/repos/ubuntu{release}/x86_64/{cuda_repo_pkg}"
+            )
+            linux_os = Ubuntu(self._node)
+
+        else:
+            raise LisaException(
+                f"Distro {self._node.os.__class__.__name__}"
+                "not supported to install CUDA driver."
+            )
+
+        wget_tool = self._node.tools[Wget]
+        # download the cuda driver at /tmp/
+        wget_tool.get(cuda_repo, "/tmp/", cuda_repo_pkg)
+        # install the cuda driver rpm
+        install_result = linux_os.install_packages(
+            f"/tmp/{cuda_repo_pkg}", signed=False
+        )
+        if install_result.exit_code != 0:
+            raise LisaException(
+                f"Failed to install {cuda_repo_pkg}. stdout: {install_result.stdout}"
+            )
+        else:
+            self._log.info("Sucessfully installed cuda-drivers")
+
+    def install_gpu_dep(self) -> None:
+        uname_tool = self._node.tools[Uname]
+        uname_ver = uname_tool.get_linux_information().uname_version
+
+        # install dependency libraries for redhat and CentOS
+        if isinstance(self._node.os, Redhat):
+            # install the kernel-devel and kernel-header packages
+            package_name = f"kernel-devel-{uname_ver} kernel-headers-{uname_ver}"
+            install_result = self._node.os.install_packages(package_name)
+            if install_result.exit_code != 0:
+                raise LisaException(
+                    f"Failed to install {package_name}."
+                    f" stdout: {install_result.stdout}"
+                )
+            # mesa-libEGL install/update is require to avoid a conflict between
+            # libraries - bugzilla.redhat 1584740
+            package_name = "mesa-libGL mesa-libEGL libglvnd-devel"
+            install_result = self._node.os.install_packages(package_name)
+            if install_result.exit_code != 0:
+                raise LisaException(
+                    f"Failed to install {package_name}."
+                    f" stdout: {install_result.stdout}"
+                )
+            # install dkms
+            package_name = "dkms"
+            install_result = self._node.os.install_packages(package_name, signed=False)
+            if install_result.exit_code != 0:
+                raise LisaException(
+                    f"Failed to install {package_name}. stdout: {install_result.stdout}"
+                )
+
+        # install dependency libraraies for Ubuntu
+        elif isinstance(self._node.os, Ubuntu):
+            package_name = (
+                f"build-essential libelf-dev linux-tools-{uname_ver}"
+                f" linux-cloud-tools-{uname_ver} python libglvnd-dev ubuntu-desktop"
+            )
+            install_result = self._node.os.install_packages(package_name)
+            if install_result.exit_code != 0:
+                raise LisaException(
+                    f"Failed to install {package_name}."
+                    f" stdout: {install_result.stdout}"
+                )
+
+    def install_compute_sdk(self, driver: ComputeSDK, version: str) -> None:
+        if driver == ComputeSDK.GRID:
+            self._install_grid_driver(version)
+        elif driver == ComputeSDK.CUDA:
+            self._install_cuda_driver(version)
+        else:
+            raise LisaException("No valid driver SDK name provided to install.")