Create class for gpu_smi, move nvidia smi

Author: adityagesh

lisa/tools/__init__.py

@@ -15,7 +15,6 @@
     YumConfigManager,
 )
 
-from .amdsmi import AmdSmi
 from .aria import Aria
 from .b4 import B4
 from .blkid import Blkid
@@ -56,6 +55,7 @@
     NvidiaCudaDriver,
     NvidiaGridDriver,
 )
+from .gpu_smi import AmdSmi, GpuSmi, NvidiaSmi
 from .grub_config import GrubConfig
 from .hibernation_setup import HibernationSetup
 from .hostname import Hostname
@@ -99,7 +99,6 @@
 from .ntp import Ntp
 from .ntpstat import Ntpstat
 from .ntttcp import Ntttcp
-from .nvidiasmi import NvidiaSmi
 from .nvmecli import Nvmecli
 from .openssl import OpenSSL
 from .parted import Parted
@@ -185,6 +184,7 @@
     "Git",
     "GpuDriver",
     "GpuDriverInstaller",
+    "GpuSmi",
     "GrubConfig",
     "Ip",
     "IpInfo",

lisa/tools/gpu_drivers.py

@@ -4,7 +4,7 @@
 import re
 from abc import abstractmethod
 from pathlib import PurePosixPath
-from typing import TYPE_CHECKING, Any, List, Optional, Type, Union
+from typing import Any, List, Optional, Type
 
 from lisa.base_tools import Sed, Uname, Wget
 from lisa.executable import Tool
@@ -16,21 +16,17 @@
     Redhat,
     Ubuntu,
 )
-from lisa.tools.amdsmi import AmdSmi
 
 # Import tools directly from their modules to avoid circular import.
 # lisa.tools.__init__.py imports from this file (gpu_drivers.py), so we cannot
 # import from lisa.tools package directly. Instead, import from individual modules.
 from lisa.tools.df import Df
 from lisa.tools.echo import Echo
+from lisa.tools.gpu_smi import GpuSmi
 from lisa.tools.mkdir import Mkdir
 from lisa.tools.modprobe import Modprobe
-from lisa.tools.nvidiasmi import NvidiaSmi
 from lisa.util import LisaException, MissingPackagesException, SkippedException
 
-if TYPE_CHECKING:
-    from lisa.node import Node
-
 
 class GpuDriver(Tool):
     """
@@ -47,7 +43,7 @@ class GpuDriver(Tool):
     """
 
     _driver_class: Type["GpuDriverInstaller"]
-    _smi_class: Type[Union[AmdSmi, NvidiaSmi]]
+    _smi_class: Type[GpuSmi]
 
     @property
     def command(self) -> str:
@@ -56,9 +52,8 @@ def command(self) -> str:
     @classmethod
     def create(
         cls,
-        node: "Node",
+        node: Any,
         *args: Any,
-        driver_class: Type["GpuDriverInstaller"],
         **kwargs: Any,
     ) -> "GpuDriver":
         """
@@ -72,6 +67,11 @@ def create(
         Returns:
             GpuDriver instance configured for the specified driver
         """
+        driver_class: Type["GpuDriverInstaller"] = kwargs.pop("driver_class", None)
+        assert driver_class is not None, (
+            "driver_class parameter is required when creating GpuDriver. "
+            "Use node.tools.create(GpuDriver, driver_class=AmdGpuDriver) or similar."
+        )
 
         instance = cls(node)
         instance._driver_class = driver_class
@@ -84,7 +84,7 @@ def get_gpu_count(self) -> int:
         """
         Get GPU count using the appropriate monitoring tool.
         """
-        smi_tool: Union[AmdSmi, NvidiaSmi] = self.node.tools[self._smi_class]
+        smi_tool: GpuSmi = self.node.tools[self._smi_class]
         return smi_tool.get_gpu_count()
 
     def install_driver(self) -> None:
@@ -102,7 +102,7 @@ class GpuDriverInstaller(Tool):
 
     @classmethod
     @abstractmethod
-    def smi(cls) -> Type[Union[AmdSmi, NvidiaSmi]]:
+    def smi(cls) -> Type[GpuSmi]:
         """Return the monitoring tool class for this driver"""
         raise NotImplementedError
 
@@ -185,8 +185,10 @@ class NvidiaGridDriver(GpuDriverInstaller):
     }
 
     @classmethod
-    def smi(cls) -> Type[NvidiaSmi]:
+    def smi(cls) -> Type[GpuSmi]:
         """Return the monitoring tool class for NVIDIA GRID driver"""
+        from lisa.tools.gpu_smi import NvidiaSmi
+
         return NvidiaSmi
 
     @property
@@ -303,8 +305,10 @@ class NvidiaCudaDriver(GpuDriverInstaller):
     DEFAULT_CUDA_VERSION = "10.1.243-1"
 
     @classmethod
-    def smi(cls) -> Type[NvidiaSmi]:
+    def smi(cls) -> Type[GpuSmi]:
         """Return the monitoring tool class for NVIDIA CUDA driver"""
+        from lisa.tools.gpu_smi import NvidiaSmi
+
         return NvidiaSmi
 
     @property
@@ -552,8 +556,10 @@ class AmdGpuDriver(GpuDriverInstaller):
     ROCM_BUILD = "70001"
 
     @classmethod
-    def smi(cls) -> Type[AmdSmi]:
+    def smi(cls) -> Type[GpuSmi]:
         """Return the monitoring tool class for AMD GPU driver"""
+        from lisa.tools.gpu_smi import AmdSmi
+
         return AmdSmi
 
     @property

lisa/tools/amdsmi.py renamed to lisa/tools/gpu_smi.py

@@ -2,13 +2,62 @@
 # Licensed under the MIT license.
 
 import re
+from abc import abstractmethod
 from typing import List, Type
 
 from lisa.executable import Tool
 from lisa.util import LisaException, find_groups_in_lines
 
 
-class AmdSmi(Tool):
+class GpuSmi(Tool):
+    """
+    Base class for GPU monitoring tools (nvidia-smi, amd-smi, etc.).
+    """
+
+    @abstractmethod
+    def get_gpu_count(self) -> int:
+        """Get the number of GPUs detected by the SMI tool"""
+        raise NotImplementedError
+
+
+class NvidiaSmi(GpuSmi):
+    # tuple of gpu device names and their device id pattern
+    # e.g. Tesla GPU device has device id "47505500-0001-0000-3130-444531303244"
+    # A10-4Q device id "56475055-0002-0000-3130-444532323336"
+    gpu_devices = (
+        ("Tesla", "47505500", 0),
+        ("A100", "44450000", 6),
+        ("H100", "44453233", 0),
+        ("A10-4Q", "56475055", 0),
+        ("A10-8Q", "3e810200", 0),
+        ("GB200", "42333130", 0),
+    )
+
+    @property
+    def command(self) -> str:
+        return "nvidia-smi"
+
+    @property
+    def can_install(self) -> bool:
+        return False
+
+    def get_gpu_count(self) -> int:
+        result = self.run("-L")
+        if result.exit_code != 0 or (result.exit_code == 0 and result.stdout == ""):
+            result = self.run("-L", sudo=True)
+            if result.exit_code != 0 or (result.exit_code == 0 and result.stdout == ""):
+                raise LisaException(
+                    f"nvidia-smi command exited with exit_code {result.exit_code}"
+                )
+        gpu_types = [x[0] for x in self.gpu_devices]
+        device_count = 0
+        for gpu_type in gpu_types:
+            device_count += result.stdout.count(gpu_type)
+
+        return device_count
+
+
+class AmdSmi(GpuSmi):
     # Pattern to match GPU entries in amd-smi list output
     # Example:
     # GPU: 0
@@ -25,7 +74,7 @@ def command(self) -> str:
 
     @property
     def can_install(self) -> bool:
-        return False
+        return True
 
     @property
     def dependencies(self) -> List[Type[Tool]]: