switch over to psutil, more robust system detection

12fcbc0d4c04f3c5d67fd17aaad271d79619eceb · 2024-07-17 00:53:46 -0700 · Alex Cheema

Files touched

M exo/topology/device_capabilities.py

Diff

commit 12fcbc0d4c04f3c5d67fd17aaad271d79619eceb
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Wed Jul 17 00:53:46 2024 -0700

    switch over to psutil, more robust system detection
---
 exo/topology/device_capabilities.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/exo/topology/device_capabilities.py b/exo/topology/device_capabilities.py
index 8c9787a1..ba5cef69 100644
--- a/exo/topology/device_capabilities.py
+++ b/exo/topology/device_capabilities.py
@@ -1,7 +1,7 @@
 from exo import DEBUG
 from dataclasses import dataclass
 import subprocess
-import platform
+import psutil
 
 @dataclass
 class DeviceCapabilities:
@@ -10,15 +10,12 @@ class DeviceCapabilities:
     memory: int
 
 def device_capabilities() -> DeviceCapabilities:
-    system = platform.system()
-    if system == 'Darwin':
+    if psutil.MACOS:
         return mac_device_capabilities()
-    elif system == 'Linux':
+    elif psutil.LINUX:
         return linux_device_capabilities()
-    # elif system == 'Windows':
-    #     return windows_device_capabilities()
     else:
-        return DeviceCapabilities(model="Unknown Model", chip="Unknown Chip", memory=0)
+        return DeviceCapabilities(model=f"Unknown Device", chip=f"Unknown Chip", memory=psutil.virtual_memory().total // 2**20)
 
 def mac_device_capabilities() -> DeviceCapabilities:
     # Fetch the model of the Mac using system_profiler

← 7545e060 fix: syntax error in requirements.txt · back to Exo · added setup py bfaeccc7 →