Refractor (#57)

* use lower case to check platform * add windows darwin solaris freebsd * update if else branch * update intention * fix os name check * init item * remove memory_total func * assert base attrs are valid * init features as empty * check not none * fix memory_total is int * add try catch for getting FreePhysicalMemory failed * add seperator for _run * fix parsing commands * auto remove consecutive \n * [CI] don't copy back * [CI] remove .git folder * disable type hint
ModelCloud · Jan 7, 2025 · e8bc2e7 · e8bc2e7
1 parent d1ebab4
commit e8bc2e7
Show file tree

Hide file tree

Showing 9 changed files with 110 additions and 93 deletions.
diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
@@ -43,9 +43,13 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
+      - name: clean dir
+        run: rm -rf .git
+
       - name: Test in FreeBSD
         uses: vmactions/freebsd-vm@v1
         with:
+          copyback: false
           prepare: |
             env ASSUME_ALWAYS_YES=yes pkg install -y python py311-pip
           run: |
@@ -60,9 +64,13 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
+      - name: clean dir
+        run: rm -rf .git
+
       - name: Test in Solaris
         uses: vmactions/solaris-vm@v1
         with:
+          copyback: false
           run: |
             python -V
             python -m venv venv

diff --git a/device_smi/apple.py b/device_smi/apple.py
@@ -12,12 +12,11 @@ def __init__(self, cls, index):
 
         args = ["system_profiler", "SPDisplaysDataType"]
 
-        result = _run(args=args)
+        result = _run(args=args, seperator="\n")
 
-        output = result.split("\n")
         model = ""
         vendor = ""
-        for o in output:
+        for o in result:
             if "Chipset Model" in o:
                 model = o.split(":")[1].replace("Apple", "").strip()
             if "Vendor" in o:

diff --git a/device_smi/base.py b/device_smi/base.py
@@ -1,5 +1,7 @@
 import subprocess
+import re
 from abc import abstractmethod
+from typing import Optional, Callable
 
 
 class BaseDevice:
@@ -67,7 +69,7 @@ def __repr__(self):
         return self.__str__()
 
 
-def _run(args, line_start: str = None) -> str:
+def _run(args, line_start: Optional[str] = None, seperator: str=None): # -> str | list[str] disable type hint, because solaris test is using python 3.7 which doesn't support | usage
     result = subprocess.run(
         args,
         stdout=subprocess.PIPE,
@@ -79,6 +81,9 @@ def _run(args, line_start: str = None) -> str:
         raise RuntimeError(result.stderr)
 
     result = result.stdout.strip()
+    result = re.sub(r'\n+', '\n', result) # remove consecutive \n
     if line_start:
         return " ".join([line for line in result.splitlines() if line.strip().startswith(line_start)])
+    if seperator:
+        return [l.strip() for l in result.split(seperator)]
     return result
diff --git a/device_smi/cpu.py b/device_smi/cpu.py
@@ -17,7 +17,41 @@ def __init__(self, cls):
         vendor = "Unknown vendor"
         flags = set()
 
-        if os.name == 'posix':
+        if platform.system().lower() == "windows":
+            command_result = _run(["wmic", "cpu", "get", "manufacturer,name,numberofcores,numberoflogicalprocessors", "/format:csv"]).strip()
+            result = command_result.split("\n")[1].split(",")
+
+            cpu_count = command_result.count('\n')
+            model = result[2].strip()
+            cpu_cores = int(result[3])
+            cpu_threads = int(result[4])
+            vendor = result[1].strip()
+
+            command_result = _run(["wmic", "os", "get", "TotalVisibleMemorySize", "/Value", "/format:csv"]).strip()
+            result = command_result.split("\n")[1].split(",")
+
+            mem_total = int(result[1])
+        elif platform.system().lower() == 'darwin':
+            model = (_run(["sysctl", "-n", "machdep.cpu.brand_string"]).replace("Apple", "").strip())
+            try:
+                vendor = (_run(["sysctl", "-n", "machdep.cpu.vendor"]))
+            except BaseException:
+                vendor = "apple"
+
+            sysctl_info = self.to_dict(_run(["sysctl", "-a"]))
+            cpu_count = 1
+            cpu_cores = int(sysctl_info["hw.physicalcpu"])
+            cpu_threads = int(sysctl_info["hw.logicalcpu"])
+
+            mem_total = int(_run(["sysctl", "-n", "hw.memsize"]))
+
+            try:
+                features = sysctl_info["machdep.cpu.features"].splitlines()
+            except Exception:
+                features = []
+
+            flags = set(features)
+        elif os.name == 'posix':
             try:
                 with open("/proc/cpuinfo", "r") as f:
                     lines = f.readlines()
@@ -26,69 +60,28 @@ def __init__(self, cls):
                             flags.update(line.strip().split(":")[1].split())
                         if line.startswith("model name"):
                             model = line.split(":")[1].strip()
-
-
                         elif line.startswith("vendor_id"):
                             vendor = line.split(":")[1].strip()
             except FileNotFoundError:
-                if platform.system() == "Darwin":
-                    model = (
-                        _run(["sysctl", "-n", "machdep.cpu.brand_string"])
-                        .replace("Apple", "")
-                        .strip()
-                    )
-                    try:
-                        vendor = (_run(["sysctl", "-n", "machdep.cpu.vendor"]))
-                    except BaseException:
-                        vendor = "apple"
-                else:
-                    model = platform.processor()
-                    vendor = platform.uname().system
-            if platform.system() == "Darwin":
-                sysctl_info = self.to_dict(_run(["sysctl", "-a"]))
-                cpu_count = 1
-                cpu_cores = int(sysctl_info["hw.physicalcpu"])
-                cpu_threads = int(sysctl_info["hw.logicalcpu"])
-
-                mem_total = int(_run(["sysctl", "-n", "hw.memsize"]))
+                model = platform.processor()
+                vendor = platform.uname().system
 
-                try:
-                    features = sysctl_info["machdep.cpu.features"].splitlines()
-                except Exception:
-                    # machdep.cpu.features is not available on arm arch
-                    features = []
-
-                flags = set(features)
-            else:
-                cpu_info = self.to_dict(_run(['lscpu']))
+            cpu_info = self.to_dict(_run(['lscpu']))
 
-                cpu_count = int(cpu_info["Socket(s)"])
-                cpu_cores_per_socket = int(cpu_info["Core(s) per socket"])
-                cpu_cores = cpu_count * cpu_cores_per_socket
-                cpu_threads = int(cpu_info["CPU(s)"])
+            cpu_count = int(cpu_info["Socket(s)"])
+            cpu_cores_per_socket = int(cpu_info["Core(s) per socket"])
+            cpu_cores = cpu_count * cpu_cores_per_socket
+            cpu_threads = int(cpu_info["CPU(s)"])
 
-                with open("/proc/meminfo", "r") as f:
-                    lines = f.readlines()
-                    mem_total = 0
-                    for line in lines:
-                        if line.startswith("MemTotal:"):
-                            mem_total = int(line.split()[1]) * 1024
-                            break
+            with open("/proc/meminfo", "r") as f:
+                lines = f.readlines()
+                mem_total = 0
+                for line in lines:
+                    if line.startswith("MemTotal:"):
+                        mem_total = int(line.split()[1]) * 1024
+                        break
         else:
-            if platform.system().lower() == "windows":
-                command_result = _run(["wmic", "cpu", "get", "manufacturer,name,numberofcores,numberoflogicalprocessors", "/format:csv"]).strip()
-                command_result = re.sub(r'\n+', '\n', command_result)  # windows uses \n\n
-                result = command_result.split("\n")[1].split(",")
-                cpu_count = command_result.count('\n')
-                model = result[2].strip()
-                cpu_cores = int(result[3])
-                cpu_threads = int(result[4])
-                vendor = result[1].strip()
-
-                command_result = _run(["wmic", "os", "get", "TotalVisibleMemorySize", "/Value", "/format:csv"]).strip()
-                command_result = re.sub(r'\n+', '\n', command_result)
-                result = command_result.split("\n")[1].split(",")
-                mem_total = int(result[1])
+            print("not support")
 
         model = " ".join(i for i in model.lower().split() if not any(x in i for x in ["ghz", "cpu", "(r)", "(tm)", "intel", "amd", "core", "processor", "@"]))
         cls.model = model
@@ -107,7 +100,7 @@ def __init__(self, cls):
 
     def _utilization(self):
         # check if is macOS
-        if platform.system() == "Darwin":
+        if platform.system().lower() == "darwin":
             output = _run(["top", "-l", "1", "-stats", "cpu"])
 
             # CPU usage: 7.61% user, 15.23% sys, 77.15% idle
@@ -133,7 +126,6 @@ def metrics(self):
         if platform.system().lower() == "windows":
             if platform.system().lower() == "windows":
                 command_result = _run(["wmic", "cpu", "get", "loadpercentage"]).strip()
-                command_result = re.sub(r'\n+', '\n', command_result)
                 try:
                     result = command_result.split("\n")[1].split(",")
                     utilization = int(result[0])
@@ -143,11 +135,15 @@ def metrics(self):
                     print(f"------------")
                     raise e
 
-                command_result = _run(["wmic", "os", "get", "FreePhysicalMemory"]).strip()
-                command_result = re.sub(r'\n+', '\n', command_result)
-                result = command_result.split("\n")[1].split(",")
-                memory_used = int(result[0])
-
+                try:
+                    command_result = _run(["wmic", "os", "get", "FreePhysicalMemory"]).strip()
+                    result = command_result.split("\n")[1].split(",")
+                    memory_used = int(result[0])
+                except BaseException as e:
+                    print(f"error occurred, command_result: ")
+                    print(f"{command_result}")
+                    print(f"------------")
+                    raise e
                 return CPUMetrics(
                     memory_used=memory_used,  # bytes
                     memory_process=0,  # bytes
@@ -165,12 +161,12 @@ def metrics(self):
         if total_diff <= 0:
             utilization = 0
         else:
-            if platform.system() == "Darwin":
+            if platform.system().lower() == "darwin":
                 utilization = idle_time_2 - idle_time_1
             else:
                 utilization = (1 - (idle_diff / total_diff)) * 100
 
-        if platform.system() == "Darwin":
+        if platform.system().lower() == "darwin":
             available_mem = _run(["vm_stat"]).replace(".", "").lower()
 
             result = self.to_dict(available_mem)
@@ -180,7 +176,7 @@ def metrics(self):
 
             free_pages = int(result["pages free"])
 
-            mem_free = free_pages *  page_size
+            mem_free = free_pages * page_size
         else:
             with open("/proc/meminfo", "r") as f:
                 lines = f.readlines()
@@ -193,7 +189,7 @@ def metrics(self):
         memory_used = self.memory_total - mem_free
 
         process_id = os.getpid()
-        if platform.system() == "Darwin":
+        if platform.system().lower() == "darwin":
             result = _run(["ps", "-p", str(process_id), "-o", "rss="])
             memory_current_process = int(result) * 1024
         else:

diff --git a/device_smi/device.py b/device_smi/device.py
@@ -19,11 +19,25 @@
 
 class Device:
     def __init__(self, device):
+        # init attribute first to avoid IDE not attr warning
+        # CPU/GPU Device
+        self.memory_total = None
+        self.type = None
+        self.features = []
+        self.vendor = None
+        self.model = None
+        # OS Device
+        self.arch = None
+        self.version = None
+        self.name = None
         if HAS_TORCH and isinstance(device, torch.device):
             device_type = device.type.lower()
             device_index = device.index
         elif f"{device}".lower() == "os":
             self.device = OSDevice(self)
+            assert self.arch
+            assert self.version
+            assert self.name
             return
         else:
             d = f"{device}".lower()
@@ -44,7 +58,7 @@ def __init__(self, device):
             or device_type == "xpu"
             or re.match(r"(gpu|cuda|xpu):\d+", device_type)
         ):
-            if platform.system() == "Darwin":
+            if platform.system().lower() == "darwin":
                 if platform.machine() == 'x86_64':
                     raise Exception("Not supported for macOS on Intel chips.")
 
@@ -68,6 +82,12 @@ def __init__(self, device):
         else:
             raise Exception(f"The device {device_type} is not supported")
 
+        assert self.memory_total
+        assert self.type
+        assert self.features is not None
+        assert self.vendor
+        assert self.model
+
     def info(self):
         warnings.warn(
             "info() method is deprecated and will be removed in next release.",
@@ -76,9 +96,6 @@ def info(self):
         )
         return self
 
-    def memory_total(self):
-        return self.memory_total
-
     def memory_used(self) -> int:
         return self.device.metrics().memory_used
 

diff --git a/device_smi/intel.py b/device_smi/intel.py
@@ -55,13 +55,12 @@ def metrics(self):
                 "-m", "0,18",
                 "-n", "1"
             ]
-            result = _run(args=args)
+            output = _run(args=args, seperator="\n")[-1]
 
             # xpu-smi dump -d 0 -m 0,1,2 -i 1 -n 5
             # Timestamp, DeviceId, GPU Utilization (%), GPU Power (W), GPU Frequency (MHz)
             # 06:14:46.000,    0, 0.00, 14.61,    0
 
-            output = result.split("\n")[-1]
             memory_used = output.split(",")[-1].strip()
             utilization = output.split(",")[-2].strip()
             if utilization.lower() == "n/a":

diff --git a/device_smi/nvidia.py b/device_smi/nvidia.py
@@ -27,13 +27,12 @@ def __init__(self, cls, index):
                 "--format=csv,noheader,nounits",
             ]
 
-            result = _run(args=args)
+            result = _run(args=args,seperator="\n")
 
-            output = result.strip().split("\n")[0]
-            model, total_memory, pci_bus_id, pcie_gen, pcie_width, driver = (output.split(", "))
+            model, total_memory, pci_bus_id, pcie_gen, pcie_width, driver = (result[0].split(", "))
 
-            result = _run(args=["nvidia-smi", "-q", "-i", f"{self.gpu_id}"])
-            firmware = " ".join([line.split(":", 1)[1].strip() for line in result.splitlines() if "VBIOS" in line])
+            result = _run(args=["nvidia-smi", "-q", "-i", f"{self.gpu_id}"],seperator="\n")
+            firmware = " ".join([line.split(":", 1)[1].strip() for line in result if "VBIOS" in line])
 
             if model.lower().startswith("nvidia"):
                 model = model[len("nvidia"):]
@@ -67,11 +66,8 @@ def _get_gpu_id(self):
 
     def metrics(self):
         try:
-            args = ["nvidia-smi", f"--id={self.gpu_id}", "--query-gpu=memory.used,utilization.gpu", "--format=csv,noheader,nounits", ]
-            result = _run(args=args)
-
-            output = result.split("\n")[0]
-            used_memory, utilization = output.split(", ")
+            args = ["nvidia-smi", f"--id={self.gpu_id}", "--query-gpu=memory.used,utilization.gpu", "--format=csv,noheader,nounits" ]
+            used_memory, utilization = _run(args=args, seperator="\n")[0].split(", ")
 
             return NvidiaGPUMetrics(
                 memory_used=int(used_memory) * 1024 * 1024,  # bytes

diff --git a/device_smi/os.py b/device_smi/os.py
@@ -29,11 +29,8 @@ def __init__(self, cls):
             cls.version = release_info["productversion"]
             cls.kernel, cls.arch = _run(["uname", "-mr"]).lower().split()
         elif platform.system().lower() == "windows":
-            command_result = _run(["wmic", "os", "get", "caption", "/format:csv"]).strip()
-            command_result = re.sub(r'\n+', '\n', command_result)  # windows uses \n\n
-            result = command_result.split("\n")[1].split(",")
+            cls.version = _run(["wmic", "os", "get", "caption", "/format:csv"], seperator="\n")[1].split(",")[1].lower().removeprefix("microsoft windows").strip()
             cls.name = "windows"
-            cls.version = result[1].lower().removeprefix("microsoft windows").strip()
             cls.arch = os.environ.get("PROCESSOR_ARCHITECTURE").lower()
 
             cls.kernel = _run(["cmd", "/c", "ver"])

diff --git a/tests/gpu.py b/tests/gpu.py
@@ -12,4 +12,4 @@
     assert dev.gpu.driver
     assert dev.gpu.firmware
 assert dev.model
-assert dev.memory_total > 10, f"wrong memory size: {dev.memory_total}"
+assert dev.memory_total > 10, f"wrong memory size: {dev.memory_total()}"