Skip to content

Commit

Permalink
Integrated wmi module for accessing system information
Browse files Browse the repository at this point in the history 
Signed-off-by: amd-pworfolk <[email protected]>
  • Loading branch information
@amd-pworfolk
amd-pworfolk committed Dec 6, 2024
1 parent 952f772 commit ce06ffc
Show file tree
Hide file tree
Showing 4 changed files with 413 additions and 206 deletions.
1 change: 1 addition & 0 deletions setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
"fasteners",
"GitPython>=3.1.40",
"psutil",
"wmi",
],
extras_require={
"llm": [
Expand Down
206 changes: 1 addition & 205 deletions src/turnkeyml/common/build.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,8 @@
import logging
import sys
import traceback
import platform
import subprocess
from typing import Dict, Union
import hashlib
import re
import importlib.metadata
import psutil
import yaml
import torch
Expand Down Expand Up @@ -196,7 +192,7 @@ def get_shapes_and_dtypes(inputs: dict):

class Logger:
"""
Redirects stdout to to file (and console if needed)
Redirects stdout to file (and console if needed)
"""

def __init__(
Expand Down Expand Up @@ -265,203 +261,3 @@ def write(self, message):
def flush(self):
# needed for python 3 compatibility.
pass


def get_windows_driver_version(device_name: str) -> str:
"""
Returns the driver version for a given device name. If not found, returns None.
This information is extracted from parsing the output of the powershell Get-WmiObject command.
All drivers and versions can be viewed with the PowerShell command:
> Get-WmiObject Win32_PnPSignedDriver | select DeviceName, Manufacturer, DriverVersion
To find a specific driver version by name:
> Get-WmiObject Win32_PnPSignedDriver | select DeviceName, DriverVersion |
where-object {$_.DeviceName -eq 'NPU Compute Accelerator Device'}
"""

cmd = (
"Get-WmiObject Win32_PnPSignedDriver | select DeviceName, DriverVersion | "
+ "where-object {$_.DeviceName -eq '"
+ device_name
+ "'}"
)
try:
output = subprocess.run(
"powershell " + cmd, capture_output=True, text=True, check=True
).stdout
return output.split("\n")[3].split()[-1]
except Exception: # pylint: disable=broad-except
return ""


def get_system_info():
os_type = platform.system()
info_dict = {}

# Get OS Version
try:
info_dict["OS Version"] = platform.platform()
except Exception as e: # pylint: disable=broad-except
info_dict["Error OS Version"] = str(e)

def get_cim_info(command):
try:
output = subprocess.run(
"powershell " + command, capture_output=True, text=True, check=True
).stdout
return output.strip()
except Exception as e: # pylint: disable=broad-except
return str(e)

if os_type == "Windows":

processor = get_cim_info(
"Get-CimInstance -ClassName Win32_Processor | "
"%{ '{0}###({1} cores, {2} logical processors)' "
"-f $_.Name,$_.NumberOfCores,$_.NumberOfLogicalProcessors }"
)
info_dict["Processor"] = " ".join(
[str.strip() for str in processor.split("###")]
)

info_dict["OEM System"] = get_cim_info(
"Get-CimInstance -ClassName Win32_ComputerSystem | "
"Select-Object -ExpandProperty Model"
)

memory = get_cim_info(
"Get-CimInstance -ClassName Win32_PhysicalMemory | "
"%{ '{0} {1} GB {2} ns' -f $_.Manufacturer,($_.Capacity/1gb),$_.Speed }"
)
info_dict["Physical Memory"] = " + ".join(memory.split("\n"))

info_dict["BIOS Version"] = get_cim_info(
"Get-CimInstance -ClassName Win32_BIOS | "
"Select-Object -ExpandProperty SMBIOSBIOSVersion"
)

info_dict["CPU Max Clock"] = get_cim_info(
"Get-CimInstance -ClassName Win32_Processor | "
"Select-Object -ExpandProperty MaxClockSpeed"
)
info_dict["CPU Max Clock"] += " MHz"

# Get driver versions
device_names = [
"NPU Compute Accelerator Device",
"AMD-OpenCL User Mode Driver",
]
driver_versions = {}
for device_name in device_names:
driver_version = get_windows_driver_version(device_name)
driver_versions[device_name] = (
driver_version if len(driver_version) else "DRIVER NOT FOUND"
)
info_dict["Driver Versions"] = driver_versions

# Get NPU power mode
if "AMD" in info_dict["Processor"]:
try:
out = subprocess.check_output(
[
r"C:\Windows\System32\AMD\xrt-smi.exe",
"examine",
"-r",
"platform",
],
stderr=subprocess.STDOUT,
).decode()
lines = out.splitlines()
modes = [line.split()[-1] for line in lines if "Mode" in line]
if len(modes) > 0:
info_dict["NPU Power Mode"] = modes[0]
except FileNotFoundError:
# xrt-smi not present
pass
except subprocess.CalledProcessError as e:
info_dict["NPU Power Mode ERROR"] = e.output.decode()

# Get windows power setting
try:
out = subprocess.check_output(["powercfg", "/getactivescheme"]).decode()
info_dict["Windows Power Setting"] = re.search(r"\((.*?)\)", out).group(1)
except subprocess.CalledProcessError as e:
info_dict["Windows Power Setting ERROR"] = e.output.decode()

elif os_type == "Linux":
# WSL has to be handled differently compared to native Linux
if "microsoft" in str(platform.release()):
try:
oem_info = (
subprocess.check_output(
'powershell.exe -Command "wmic computersystem get model"',
shell=True,
)
.decode()
.strip()
)
oem_info = (
oem_info.replace("\r", "")
.replace("\n", "")
.split("Model")[-1]
.strip()
)
info_dict["OEM System"] = oem_info
except Exception as e: # pylint: disable=broad-except
info_dict["Error OEM System (WSL)"] = str(e)

else:
# Get OEM System Information
try:
oem_info = (
subprocess.check_output(
"sudo -n dmidecode -s system-product-name",
shell=True,
stderr=subprocess.DEVNULL,
)
.decode()
.strip()
.replace("\n", " ")
)
info_dict["OEM System"] = oem_info
except subprocess.CalledProcessError:
# This catches the case where sudo requires a password
info_dict["OEM System"] = "Unable to get oem info - password required"
except Exception as e: # pylint: disable=broad-except
info_dict["Error OEM System"] = str(e)

# Get CPU Information
try:
cpu_info = subprocess.check_output("lscpu", shell=True).decode()
for line in cpu_info.split("\n"):
if "Model name:" in line:
info_dict["Processor"] = line.split(":")[1].strip()
break
except Exception as e: # pylint: disable=broad-except
info_dict["Error Processor"] = str(e)

# Get Memory Information
try:
mem_info = (
subprocess.check_output("free -m", shell=True)
.decode()
.split("\n")[1]
.split()[1]
)
mem_info_gb = round(int(mem_info) / 1024, 2)
info_dict["Memory Info"] = f"{mem_info_gb} GB"
except Exception as e: # pylint: disable=broad-except
info_dict["Error Memory Info"] = str(e)

else:
info_dict["Error"] = "Unsupported OS"

# Get Python Packages
distributions = importlib.metadata.distributions()
info_dict["Python Packages"] = [
f"{dist.metadata['name']}=={dist.metadata['version']}" for dist in distributions
]

return info_dict
Loading

0 comments on commit ce06ffc

Please sign in to comment.