mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 22:26:31 +00:00
b21f8e320f
add tensorrt engine
85 lines
3.2 KiB
Cython
85 lines
3.2 KiB
Cython
import re
|
|
import subprocess
|
|
import psutil
|
|
|
|
cdef class HardwareInfo:
|
|
def __init__(self, str cpu, str gpu, str memory, str mac_address):
|
|
self.cpu = cpu
|
|
self.gpu = gpu
|
|
self.memory = memory
|
|
self.mac_address = mac_address
|
|
|
|
cdef 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}'
|
|
|
|
cdef class HardwareService:
|
|
"""Handles hardware information retrieval and hash generation."""
|
|
|
|
def __init__(self):
|
|
try:
|
|
res = subprocess.check_output("ver", shell=True).decode('utf-8')
|
|
if "Microsoft Windows" in res:
|
|
self.is_windows = True
|
|
else:
|
|
self.is_windows = False
|
|
except Exception:
|
|
print('Error during os type checking')
|
|
self.is_windows = False
|
|
|
|
cdef get_mac_address(self, 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
|
|
|
|
@staticmethod
|
|
cdef has_nvidia_gpu():
|
|
try:
|
|
output = subprocess.check_output(['nvidia-smi']).decode()
|
|
match = re.search(r'CUDA Version:\s*([\d.]+)', output)
|
|
if match:
|
|
return float(match.group(1)) > 11
|
|
return False
|
|
except Exception as e:
|
|
print(e)
|
|
return False
|
|
|
|
cdef HardwareInfo get_hardware_info(self):
|
|
if self.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}' && \""
|
|
)
|
|
# in case of subprocess error do:
|
|
# cdef bytes os_command_bytes = os_command.encode('utf-8')
|
|
# and use os_command_bytes
|
|
result = subprocess.check_output(os_command, shell=True).decode('utf-8')
|
|
lines = [line.strip() for line in result.splitlines() if line.strip()]
|
|
|
|
cdef str cpu = lines[0].replace("Name=", "").replace(" ", " ")
|
|
cdef str gpu = lines[1].replace("Name=", "").replace(" ", " ")
|
|
cdef str memory = lines[2].replace("TotalVisibleMemorySize=", "").replace(" ", " ")
|
|
cdef str mac_address = self.get_mac_address()
|
|
|
|
return HardwareInfo(cpu, gpu, memory, mac_address)
|