Add HardwareInfo

CHANGELOG.md

@@ -5,6 +5,7 @@
 - Added `reset` method to `TickTock`
 - Removed `TickTock.remove_outliers`
 - Moved `is_windows` to static `OS` attribute, which also contains `is_mac` and `is_linux`
+- Added `HardwareInfo`, which contains a bunch of system hardware information
 
 ### Bug fixes
 

pelutils/__init__.py

@@ -6,9 +6,12 @@
 import ctypes
 import os
 import random
+import subprocess
 import sys
 
+import cpuinfo
 import git
+import psutil
 import numpy as np
 try:
     import torch
@@ -225,6 +228,23 @@ def except_keys(d: dict[_T, Any], except_keys: Iterable[_T]) -> dict[_T, Any]:
     except_keys = set(except_keys)
     return { kw: v for kw, v in d.items() if kw not in except_keys }
 
+class HardwareInfo:
+
+    # Name of the CPU
+    cpu     = cpuinfo.get_cpu_info()["brand_raw"]
+    # How many CPU sockets there are on the system
+    # Only works on Linux, otherwise None
+    sockets = int(subprocess.check_output(
+                  'cat /proc/cpuinfo | grep "physical id" | sort -u | wc -l', shell=True
+              )) if OS.is_linux else None
+    # Total threads available across all CPU sockets
+    threads = os.cpu_count()
+    # Total system memory in bytes
+    memory  = psutil.virtual_memory().total
+    # Available gpu
+    # Requires torch, otherwise None
+    gpu     = torch.cuda.get_device_name() if _has_torch and torch.cuda.is_available() else None
+
 # To allow imports directly from utils
 # Placed down here to prevent issues with circular imports
 from .__version__ import __version__

setup.py

@@ -16,6 +16,7 @@
     "click>=7.0.0",
     "python-rapidjson>=1.5",
     "py-cpuinfo>=8.0.0",
+    "psutil>=5.8.0",
 ]
 requirements_ds = [
     "torch>=1.7.0",

tests/test_init.py

@@ -11,7 +11,7 @@
 
 from pelutils import EnvVars, UnsupportedOS, reverse_line_iterator, except_keys,\
     split_path, binary_search, raises, thousands_seperators, OS, c_ptr,\
-    get_timestamp, get_timestamp_for_files
+    get_timestamp, get_timestamp_for_files, HardwareInfo
 from pelutils.tests import UnitTestCollection
 
 class TestInit(UnitTestCollection):
@@ -98,7 +98,7 @@ def test_is_windows(self):
 
     def test_is_mac(self):
         # What a dumb test
-        assert OS.is_mac == platform.platform().startswith("Darwin")
+        assert OS.is_mac == (platform.platform().startswith("Darwin") or platform.platform().startswith("macOS"))
 
     def test_is_linux(self):
         # What a dumb test
@@ -159,3 +159,16 @@ def test_get_timestamp(self):
             ts1 = get_timestamp_for_files(with_date=date)
             assert len(ts0[:-4]) == len(ts1)
             assert ts1 == ts0[:-4].replace(" ", "_").replace(":", "-")
+
+    def test_hardware_info(self):
+        assert isinstance(HardwareInfo.cpu, str) and len(HardwareInfo.cpu) > 0
+        if OS.is_linux:
+            assert isinstance(HardwareInfo.sockets, int) and HardwareInfo.sockets >= 1
+        else:
+            assert HardwareInfo.sockets is None
+        assert isinstance(HardwareInfo.threads, int) and HardwareInfo.threads > 0
+        assert isinstance(HardwareInfo.memory, int) and HardwareInfo.memory > 0
+        if torch.cuda.is_available():
+            assert isinstance(HardwareInfo.gpu, str) and len(HardwareInfo.gpu) > 0
+        else:
+            assert HardwareInfo.gpu is None