lemonade-sdk 8.1.8__py3-none-any.whl → 8.1.10__py3-none-any.whl

lemonade/common/system_info.py

@@ -1,9 +1,11 @@
 from abc import ABC, abstractmethod
 import importlib.metadata
+import logging
 import platform
 import re
 import subprocess
 import ctypes
+import glob
 from .inference_engines import detect_inference_engines
 
 # AMD GPU classification keywords - shared across all OS implementations
@@ -19,6 +21,28 @@ AMD_DISCRETE_GPU_KEYWORDS = [
     "fury",
 ]
 
+# NVIDIA GPU classification keywords - shared across all OS implementations
+# NVIDIA GPUs are typically discrete by default, but we include keywords for clarity
+NVIDIA_DISCRETE_GPU_KEYWORDS = [
+    "geforce",
+    "rtx",
+    "gtx",
+    "quadro",
+    "tesla",
+    "titan",
+    "a100",
+    "a40",
+    "a30",
+    "a10",
+    "a6000",
+    "a5000",
+    "a4000",
+    "a2000",
+    "t1000",
+    "t600",
+    "t400",
+]
+
 
 class SystemInfo(ABC):
     """
@@ -51,6 +75,7 @@ class SystemInfo(ABC):
             "cpu": self.get_cpu_device(),
             "amd_igpu": self.get_amd_igpu_device(include_inference_engines=True),
             "amd_dgpu": self.get_amd_dgpu_devices(include_inference_engines=True),
+            "nvidia_dgpu": self.get_nvidia_dgpu_devices(include_inference_engines=True),
             "npu": self.get_npu_device(),
         }
         return device_dict
@@ -82,6 +107,15 @@ class SystemInfo(ABC):
             list: List of AMD dGPU device information.
         """
 
+    @abstractmethod
+    def get_nvidia_dgpu_devices(self, include_inference_engines: bool = False) -> list:
+        """
+        Retrieves NVIDIA discrete GPU device information.
+
+        Returns:
+            list: List of NVIDIA dGPU device information.
+        """
+
     @abstractmethod
     def get_npu_device(self) -> dict:
         """
@@ -173,34 +207,56 @@ class WindowsSystemInfo(SystemInfo):
         Returns:
             list: List of detected GPU info dictionaries
         """
+        logging.debug(f"Starting AMD GPU detection for type: {gpu_type}")
         gpu_devices = []
         try:
             video_controllers = self.connection.Win32_VideoController()
-            for controller in video_controllers:
+            logging.debug(f"Found {len(video_controllers)} video controllers")
+
+            for i, controller in enumerate(video_controllers):
+                logging.debug(
+                    f"Controller {i}: Name='{controller.Name}', "
+                    f"PNPDeviceID='{getattr(controller, 'PNPDeviceID', 'N/A')}'"
+                )
+
                 if (
                     controller.Name
                     and "AMD" in controller.Name
                     and "Radeon" in controller.Name
                 ):
+                    logging.debug(f"Found AMD Radeon GPU: {controller.Name}")
 
                     name_lower = controller.Name.lower()
+                    logging.debug(f"GPU name (lowercase): {name_lower}")
 
                     # Keyword-based classification - simple and reliable
+                    matching_keywords = [
+                        kw for kw in AMD_DISCRETE_GPU_KEYWORDS if kw in name_lower
+                    ]
                     is_discrete_by_name = any(
                         kw in name_lower for kw in AMD_DISCRETE_GPU_KEYWORDS
                     )
                     is_integrated = not is_discrete_by_name
 
+                    logging.debug(f"Matching discrete keywords: {matching_keywords}")
+                    logging.debug(
+                        f"Classified as discrete: {not is_integrated}, integrated: {is_integrated}"
+                    )
+
                     # Filter based on requested type
                     if (gpu_type == "integrated" and is_integrated) or (
                         gpu_type == "discrete" and not is_integrated
                     ):
+                        logging.debug(
+                            f"GPU matches requested type '{gpu_type}', processing..."
+                        )
 
                         device_type = "amd_igpu" if is_integrated else "amd_dgpu"
                         gpu_info = {
                             "name": controller.Name,
                             "available": True,
                         }
+                        logging.debug(f"Created GPU info for {device_type}: {gpu_info}")
 
                         driver_version = self.get_driver_version(
                             "AMD-OpenCL User Mode Driver"
@@ -208,6 +264,21 @@ class WindowsSystemInfo(SystemInfo):
                         gpu_info["driver_version"] = (
                             driver_version if driver_version else "Unknown"
                         )
+                        logging.debug(f"Driver version: {gpu_info['driver_version']}")
+
+                        # Get VRAM information for discrete GPUs
+                        if not is_integrated:  # Only add VRAM for discrete GPUs
+                            # Try dxdiag first (most reliable for dedicated memory)
+                            vram_gb = self._get_gpu_vram_dxdiag_simple(controller.Name)
+
+                            # Fallback to WMI if dxdiag fails
+                            if vram_gb == 0.0:
+                                vram_gb = self._get_gpu_vram_wmi(controller)
+
+                            if vram_gb > 0.0:
+                                gpu_info["vram_gb"] = vram_gb
+                            else:
+                                gpu_info["vram_gb"] = "Unknown"
 
                         if include_inference_engines:
                             gpu_info["inference_engines"] = (
@@ -216,11 +287,26 @@ class WindowsSystemInfo(SystemInfo):
                                 )
                             )
                         gpu_devices.append(gpu_info)
+                        logging.debug(f"Added GPU to devices list: {gpu_info}")
+                    else:
+                        logging.debug(
+                            f"GPU does not match requested type '{gpu_type}', skipping"
+                        )
+                        continue
+                else:
+                    logging.debug(
+                        f"Skipping non-AMD/non-Radeon controller: {controller.Name}"
+                    )
 
         except Exception as e:  # pylint: disable=broad-except
             error_msg = f"AMD {gpu_type} GPU detection failed: {e}"
+            logging.debug(f"Exception in AMD GPU detection: {e}")
             return [{"available": False, "error": error_msg}]
 
+        logging.debug(
+            f"AMD GPU detection completed. Found {len(gpu_devices)} {gpu_type} GPUs: "
+            f"{[gpu.get('name', 'Unknown') for gpu in gpu_devices]}"
+        )
         return gpu_devices
 
     def get_amd_igpu_device(self, include_inference_engines: bool = False) -> dict:
@@ -255,6 +341,67 @@ class WindowsSystemInfo(SystemInfo):
             else [{"available": False, "error": "No AMD discrete GPU found"}]
         )
 
+    def get_nvidia_dgpu_devices(self, include_inference_engines: bool = False) -> list:
+        """
+        Retrieves NVIDIA discrete GPU device information using WMI.
+
+        Returns:
+            list: List of NVIDIA dGPU device information.
+        """
+        gpu_devices = []
+        try:
+            video_controllers = self.connection.Win32_VideoController()
+            for controller in video_controllers:
+                if controller.Name and "NVIDIA" in controller.Name.upper():
+                    name_lower = controller.Name.lower()
+
+                    # Most NVIDIA GPUs are discrete, but we can check keywords for confirmation
+                    is_discrete = (
+                        any(kw in name_lower for kw in NVIDIA_DISCRETE_GPU_KEYWORDS)
+                        or "nvidia" in name_lower
+                    )  # Default to discrete for NVIDIA
+
+                    if is_discrete:
+                        gpu_info = {
+                            "name": controller.Name,
+                            "available": True,
+                        }
+
+                        # Try to get NVIDIA driver version using multiple methods
+                        driver_version = self._get_nvidia_driver_version_windows()
+                        gpu_info["driver_version"] = (
+                            driver_version if driver_version else "Unknown"
+                        )
+
+                        # Get VRAM information
+                        vram_gb = self._get_gpu_vram_wmi(controller)
+                        if vram_gb == 0.0:
+                            # Fallback to nvidia-smi
+                            vram_gb = self._get_nvidia_vram_smi()
+
+                        if vram_gb > 0.0:
+                            gpu_info["vram_gb"] = vram_gb
+                        else:
+                            gpu_info["vram_gb"] = "Unknown"
+
+                        if include_inference_engines:
+                            gpu_info["inference_engines"] = (
+                                self._detect_inference_engines(
+                                    "nvidia_dgpu", controller.Name
+                                )
+                            )
+                        gpu_devices.append(gpu_info)
+
+        except Exception as e:  # pylint: disable=broad-except
+            error_msg = f"NVIDIA discrete GPU detection failed: {e}"
+            return [{"available": False, "error": error_msg}]
+
+        return (
+            gpu_devices
+            if gpu_devices
+            else [{"available": False, "error": "No NVIDIA discrete GPU found"}]
+        )
+
     def get_npu_device(self) -> dict:
         """
         Retrieves NPU device information using existing methods.
@@ -374,6 +521,169 @@ class WindowsSystemInfo(SystemInfo):
             return drivers[0].DriverVersion
         return ""
 
+    def _get_gpu_vram_wmi(self, controller) -> float:
+        """
+        Get GPU VRAM from WMI VideoController.
+
+        Args:
+            controller: WMI Win32_VideoController object
+
+        Returns:
+            float: VRAM in GB, or 0.0 if detection fails
+        """
+        try:
+            if hasattr(controller, "AdapterRAM"):
+                adapter_ram = controller.AdapterRAM
+                if adapter_ram and adapter_ram > 0:
+                    # AdapterRAM is in bytes, convert to GB
+                    vram_bytes = int(adapter_ram)
+                    vram_gb = round(vram_bytes / (1024**3), 1)
+                    return vram_gb
+        except (ValueError, AttributeError):
+            pass
+        return 0.0
+
+    def _get_gpu_vram_dxdiag_simple(self, gpu_name: str) -> float:
+        """
+        Get GPU VRAM using dxdiag, looking specifically for dedicated memory.
+
+        Args:
+            gpu_name: Name of the GPU to look for
+
+        Returns:
+            float: VRAM in GB, or 0.0 if detection fails
+        """
+        try:
+            import tempfile
+            import os
+
+            with tempfile.NamedTemporaryFile(
+                mode="w+", suffix=".txt", delete=False
+            ) as temp_file:
+                temp_path = temp_file.name
+
+            try:
+                subprocess.run(
+                    ["dxdiag", "/t", temp_path],
+                    check=True,
+                    timeout=30,
+                    capture_output=True,
+                )
+
+                with open(temp_path, "r", encoding="utf-8", errors="ignore") as f:
+                    dxdiag_output = f.read()
+
+                lines = dxdiag_output.split("\n")
+                found_gpu = False
+
+                for line in lines:
+                    line = line.strip()
+
+                    # Check if this is our GPU
+                    if "Card name:" in line and gpu_name.lower() in line.lower():
+                        found_gpu = True
+                        continue
+
+                    # Look for dedicated memory line
+                    if found_gpu and "Dedicated Memory:" in line:
+                        memory_match = re.search(
+                            r"(\d+(?:\.\d+)?)\s*MB", line, re.IGNORECASE
+                        )
+                        if memory_match:
+                            vram_mb = float(memory_match.group(1))
+                            vram_gb = round(vram_mb / 1024, 1)
+                            return vram_gb
+
+                    # Reset if we hit another display device
+                    if "Card name:" in line and gpu_name.lower() not in line.lower():
+                        found_gpu = False
+
+            finally:
+                try:
+                    os.unlink(temp_path)
+                except Exception:  # pylint: disable=broad-except
+                    pass
+
+        except Exception:  # pylint: disable=broad-except
+            pass
+
+        return 0.0
+
+    def _get_nvidia_driver_version_windows(self) -> str:
+        """
+        Get NVIDIA driver version on Windows using nvidia-smi and WMI fallback.
+
+        Returns:
+            str: Driver version, or empty string if detection fails
+        """
+        # Primary: Try nvidia-smi command
+        try:
+            output = (
+                subprocess.check_output(
+                    [
+                        "nvidia-smi",
+                        "--query-gpu=driver_version",
+                        "--format=csv,noheader,nounits",
+                    ],
+                    stderr=subprocess.DEVNULL,
+                )
+                .decode()
+                .strip()
+            )
+            if output and output != "N/A":
+                return output.split("\n")[0]
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            pass
+
+        # Fallback: Try WMI Win32_PnPSignedDriver with NVIDIA patterns
+        try:
+            nvidia_patterns = [
+                "NVIDIA GeForce",
+                "NVIDIA RTX",
+                "NVIDIA GTX",
+                "NVIDIA Quadro",
+            ]
+            all_drivers = self.connection.Win32_PnPSignedDriver()
+            for driver in all_drivers:
+                if driver.DeviceName and any(
+                    pattern in driver.DeviceName for pattern in nvidia_patterns
+                ):
+                    if driver.DriverVersion:
+                        return driver.DriverVersion
+        except Exception:  # pylint: disable=broad-except
+            pass
+
+        return ""
+
+    def _get_nvidia_vram_smi(self) -> float:
+        """
+        Get NVIDIA GPU VRAM using nvidia-smi command.
+
+        Returns:
+            float: VRAM in GB, or 0.0 if detection fails
+        """
+        try:
+            output = (
+                subprocess.check_output(
+                    [
+                        "nvidia-smi",
+                        "--query-gpu=memory.total",
+                        "--format=csv,noheader,nounits",
+                    ],
+                    stderr=subprocess.DEVNULL,
+                )
+                .decode()
+                .strip()
+            )
+
+            # nvidia-smi returns memory in MB
+            vram_mb = int(output.split("\n")[0])
+            vram_gb = round(vram_mb / 1024, 1)
+            return vram_gb
+        except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
+            pass
+        return 0.0
+
     @staticmethod
     def get_npu_power_mode() -> str:
         """
@@ -490,6 +800,14 @@ class WSLSystemInfo(SystemInfo):
         """
         return []
 
+    def get_nvidia_dgpu_devices(self, include_inference_engines: bool = False) -> list:
+        """
+        Retrieves NVIDIA discrete GPU device information in WSL environment.
+        """
+        return [
+            {"available": False, "error": "NVIDIA GPU detection not supported in WSL"}
+        ]
+
     def get_npu_device(self) -> dict:
         """
         Retrieves NPU device information in WSL environment.
@@ -625,6 +943,20 @@ class LinuxSystemInfo(SystemInfo):
                             "name": device_name,
                             "available": True,
                         }
+
+                        # Get VRAM information for discrete GPUs
+                        if not is_integrated:  # Only add VRAM for discrete GPUs
+                            vram_gb = self._get_amd_vram_rocm_smi_linux()
+                            if vram_gb == 0.0:
+                                # Fallback to sysfs - extract PCI ID from lspci line
+                                pci_id = line.split()[0] if line else ""
+                                vram_gb = self._get_amd_vram_sysfs(pci_id)
+
+                            if vram_gb > 0.0:
+                                gpu_info["vram_gb"] = vram_gb
+                            else:
+                                gpu_info["vram_gb"] = "Unknown"
+
                         if include_inference_engines:
                             gpu_info["inference_engines"] = (
                                 self._detect_inference_engines(device_type, device_name)
@@ -669,6 +1001,66 @@ class LinuxSystemInfo(SystemInfo):
             else [{"available": False, "error": "No AMD discrete GPU found"}]
         )
 
+    def get_nvidia_dgpu_devices(self, include_inference_engines: bool = False) -> list:
+        """
+        Retrieves NVIDIA discrete GPU device information using lspci.
+
+        Returns:
+            list: List of NVIDIA dGPU device information.
+        """
+        gpu_devices = []
+        try:
+            lspci_output = subprocess.check_output(
+                "lspci | grep -i 'vga\\|3d\\|display'", shell=True
+            ).decode()
+
+            for line in lspci_output.split("\n"):
+                if line.strip() and "NVIDIA" in line.upper():
+                    name_lower = line.lower()
+
+                    # Most NVIDIA GPUs are discrete, check keywords for confirmation
+                    is_discrete = (
+                        any(kw in name_lower for kw in NVIDIA_DISCRETE_GPU_KEYWORDS)
+                        or "nvidia" in name_lower
+                    )  # Default to discrete for NVIDIA
+
+                    if is_discrete:
+                        device_name = line.split(": ")[1] if ": " in line else line
+
+                        gpu_info = {
+                            "name": device_name,
+                            "available": True,
+                        }
+
+                        # Try to get NVIDIA driver version using multiple methods
+                        driver_version = self._get_nvidia_driver_version_linux()
+                        gpu_info["driver_version"] = (
+                            driver_version if driver_version else "Unknown"
+                        )
+
+                        # Get VRAM information
+                        vram_gb = self._get_nvidia_vram_smi_linux()
+                        if vram_gb > 0.0:
+                            gpu_info["vram_gb"] = vram_gb
+
+                        if include_inference_engines:
+                            gpu_info["inference_engines"] = (
+                                self._detect_inference_engines(
+                                    "nvidia_dgpu", device_name
+                                )
+                            )
+                        gpu_devices.append(gpu_info)
+
+        except Exception as e:  # pylint: disable=broad-except
+            error_msg = f"NVIDIA discrete GPU detection failed: {e}"
+            return [{"available": False, "error": error_msg}]
+
+        return (
+            gpu_devices
+            if gpu_devices
+            else [{"available": False, "error": "No NVIDIA discrete GPU found"}]
+        )
+
     def get_npu_device(self) -> dict:
         """
         Retrieves NPU device information (limited support on Linux).
@@ -681,6 +1073,69 @@ class LinuxSystemInfo(SystemInfo):
             "error": "NPU detection not yet implemented for Linux",
         }
 
+    def _get_nvidia_driver_version_linux(self) -> str:
+        """
+        Get NVIDIA driver version on Linux using nvidia-smi and proc fallback.
+
+        Returns:
+            str: Driver version, or empty string if detection fails
+        """
+        # Primary: Try nvidia-smi command
+        try:
+            output = (
+                subprocess.check_output(
+                    "nvidia-smi --query-gpu=driver_version --format=csv,noheader,nounits",
+                    shell=True,
+                    stderr=subprocess.DEVNULL,
+                )
+                .decode()
+                .strip()
+            )
+            if output and output != "N/A":
+                return output.split("\n")[0]
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            pass
+
+        # Fallback: Try /proc/driver/nvidia/version
+        try:
+            with open("/proc/driver/nvidia/version", "r", encoding="utf-8") as f:
+                content = f.read()
+                # Look for version pattern like "NVRM version:
+                #   NVIDIA UNIX x86_64 Kernel Module  470.82.00"
+                match = re.search(r"Kernel Module\s+(\d+\.\d+(?:\.\d+)?)", content)
+                if match:
+                    return match.group(1)
+        except (FileNotFoundError, IOError):
+            pass
+
+        return ""
+
+    def _get_nvidia_vram_smi_linux(self) -> float:
+        """
+        Get NVIDIA GPU VRAM on Linux using nvidia-smi command.
+
+        Returns:
+            float: VRAM in GB, or 0.0 if detection fails
+        """
+        try:
+            output = (
+                subprocess.check_output(
+                    "nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits",
+                    shell=True,
+                    stderr=subprocess.DEVNULL,
+                )
+                .decode()
+                .strip()
+            )
+
+            # nvidia-smi returns memory in MB
+            vram_mb = int(output.split("\n")[0])
+            vram_gb = round(vram_mb / 1024, 1)
+            return vram_gb
+        except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
+            pass
+        return 0.0
+
     @staticmethod
     def get_processor_name() -> str:
         """
@@ -758,6 +1213,109 @@ class LinuxSystemInfo(SystemInfo):
         info_dict["Physical Memory"] = self.get_physical_memory()
         return info_dict
 
+    def _get_nvidia_vram_smi_linux(self) -> float:
+        """
+        Get NVIDIA GPU VRAM using nvidia-smi command on Linux.
+
+        Returns:
+            float: VRAM in GB, or 0.0 if detection fails
+        """
+        try:
+            output = (
+                subprocess.check_output(
+                    [
+                        "nvidia-smi",
+                        "--query-gpu=memory.total",
+                        "--format=csv,noheader,nounits",
+                    ],
+                    stderr=subprocess.DEVNULL,
+                )
+                .decode()
+                .strip()
+            )
+
+            # nvidia-smi returns memory in MB
+            vram_mb = int(output.split("\n")[0])
+            vram_gb = round(vram_mb / 1024, 1)
+            return vram_gb
+        except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
+            pass
+        return 0.0
+
+    def _get_amd_vram_rocm_smi_linux(self) -> float:
+        """
+        Get AMD GPU VRAM using rocm-smi command on Linux.
+
+        Returns:
+            float: VRAM in GB, or 0.0 if detection fails
+        """
+        try:
+            output = (
+                subprocess.check_output(
+                    ["rocm-smi", "--showmeminfo", "vram", "--csv"],
+                    stderr=subprocess.DEVNULL,
+                )
+                .decode()
+                .strip()
+            )
+
+            # Parse CSV output to extract VRAM
+            lines = output.split("\n")
+            for line in lines:
+                if "Total VRAM" in line or "vram" in line.lower():
+                    # Extract numeric value (assuming it's in MB or GB)
+                    numbers = re.findall(r"\d+", line)
+                    if numbers:
+                        vram_value = int(numbers[0])
+                        # Assume MB if value is large, GB if small
+                        if vram_value > 100:  # Likely MB
+                            vram_gb = round(vram_value / 1024, 1)
+                        else:  # Likely GB
+                            vram_gb = float(vram_value)
+                        return vram_gb
+        except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
+            pass
+        return 0.0
+
+    def _get_amd_vram_sysfs(self, pci_id: str) -> float:
+        """
+        Get AMD GPU VRAM using sysfs on Linux.
+
+        Args:
+            pci_id: PCI ID of the GPU (e.g., "0000:01:00.0")
+
+        Returns:
+            float: VRAM in GB, or 0.0 if detection fails
+        """
+        try:
+            # Try different sysfs paths for VRAM information
+            sysfs_paths = [
+                f"/sys/bus/pci/devices/{pci_id}/mem_info_vram_total",
+                "/sys/class/drm/card*/device/mem_info_vram_total",
+            ]
+
+            for path in sysfs_paths:
+                try:
+                    if "*" in path:
+                        # Handle wildcard paths
+                        matching_paths = glob.glob(path)
+                        for match_path in matching_paths:
+                            with open(match_path, "r", encoding="utf-8") as f:
+                                vram_bytes = int(f.read().strip())
+                                vram_gb = round(vram_bytes / (1024**3), 1)
+                                if vram_gb > 0:
+                                    return vram_gb
+                    else:
+                        with open(path, "r", encoding="utf-8") as f:
+                            vram_bytes = int(f.read().strip())
+                            vram_gb = round(vram_bytes / (1024**3), 1)
+                            return vram_gb
+                except (FileNotFoundError, ValueError, PermissionError):
+                    continue
+        except Exception:  # pylint: disable=broad-except
+            pass
+        return 0.0
+
     def _detect_inference_engines(self, device_type: str, device_name: str) -> dict:
         """
         Detect available inference engines for a specific device type.
@@ -803,6 +1361,17 @@ class UnsupportedOSSystemInfo(SystemInfo):
         """
         return []
 
+    def get_nvidia_dgpu_devices(self, include_inference_engines: bool = False) -> list:
+        """
+        Retrieves NVIDIA discrete GPU device information for unsupported OS.
+        """
+        return [
+            {
+                "available": False,
+                "error": "Device detection not supported on this operating system",
+            }
+        ]
+
     def get_npu_device(self) -> dict:
         """
         Retrieves NPU device information for unsupported OS.