ai-training/hardware_service.py

import os
import subprocess

import psutil


class HardwareInfo:
    def __init__(self, cpu, gpu, memory, mac_address):
        self.cpu = cpu
        self.gpu = gpu
        self.memory = memory
        self.mac_address = mac_address

    def to_json_object(self):
        return {
            "CPU": self.cpu,
            "GPU": self.gpu,
            "MacAddress": self.mac_address,
            "Memory": self.memory
        }

    def __str__(self):
        return f'CPU: {self.cpu}. GPU: {self.gpu}. Memory: {self.memory}. MAC Address: {self.mac_address}'


def get_mac_address(interface="Ethernet"):
    addresses = psutil.net_if_addrs()
    for interface_name, interface_info in addresses.items():
        if interface_name == interface:
            for addr in interface_info:
                if addr.family == psutil.AF_LINK:
                    return addr.address.replace('-', '')
    return None


def get_hardware_info():
    is_windows = os.name == 'nt'
    res = subprocess.check_output("ver", shell=True).decode('utf-8')
    if "Microsoft Windows" in res:
        is_windows = True
    else:
        is_windows = False

    if is_windows:
        os_command = (
            "powershell -Command \""
            "Get-CimInstance -ClassName Win32_Processor | Select-Object -ExpandProperty Name | Write-Output; "
            "Get-CimInstance -ClassName Win32_VideoController | Select-Object -ExpandProperty Name | Write-Output; "
            "Get-CimInstance -ClassName Win32_OperatingSystem | Select-Object -ExpandProperty TotalVisibleMemorySize | Write-Output"
            "\""
        )
    else:
        os_command = (
            "/bin/bash -c \" lscpu | grep 'Model name:' | cut -d':' -f2 && "
            "lspci | grep VGA | cut -d':' -f3 && "
            "free -g | grep Mem: | awk '{print $2}' && \""
        )
    result = subprocess.check_output(os_command, shell=True).decode('utf-8')
    lines = [line.strip() for line in result.splitlines() if line.strip()]

    cpu = lines[0].replace("Name=", "").replace("  ", " ")
    gpu = lines[1].replace("Name=", "").replace("  ", " ")
    memory = lines[2].replace("TotalVisibleMemorySize=", "").replace("  ", " ")
    mac_address = get_mac_address()

    return HardwareInfo(cpu, gpu, memory, mac_address)