Replace netifaces (unmaintained,outdated) with scapy + add dependencies for previous fixes

5f06aa2759a16cc32617bace5d09eda0f6ad6396 · 2025-01-17 05:37:01 -0500 · Sandesh Bharadwaj
Files touched

M exo/helpers.py
M setup.py
Diff

commit 5f06aa2759a16cc32617bace5d09eda0f6ad6396
Author: Sandesh Bharadwaj <sndshvnktsh@gmail.com>
Date:   Fri Jan 17 05:37:01 2025 -0500

    Replace netifaces (unmaintained,outdated) with scapy + add dependencies for previous fixes
---
 exo/helpers.py | 17 ++++++++++-------
 setup.py       | 41 +++++++++++++++++++++++++++++++++++------
 2 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/exo/helpers.py b/exo/helpers.py
index b8445706..be29cdc5 100644
--- a/exo/helpers.py
+++ b/exo/helpers.py
@@ -7,7 +7,8 @@ import random
 import platform
 import psutil
 import uuid
-import netifaces
+from scapy.all import get_if_addr, get_if_list
+import re
 import subprocess
 from pathlib import Path
 import tempfile
@@ -231,12 +232,14 @@ def pretty_print_bytes_per_second(bytes_per_second: int) -> str:
 def get_all_ip_addresses_and_interfaces():
   try:
     ip_addresses = []
-    for interface in netifaces.interfaces():
-      ifaddresses = netifaces.ifaddresses(interface)
-      if netifaces.AF_INET in ifaddresses:
-        for link in ifaddresses[netifaces.AF_INET]:
-          ip = link['addr']
-          ip_addresses.append((ip, interface))
+    for interface in get_if_list():
+      ip = get_if_addr(interface)
+      # Include all addresses, including loopback
+      # Filter out link-local addresses
+      if not ip.startswith('169.254.') and not ip.startswith('0.0.'):
+        # Remove "\\Device\\NPF_" prefix from interface name
+        simplified_interface = re.sub(r'^\\Device\\NPF_', '', interface)
+        ip_addresses.append((ip, simplified_interface))
     return list(set(ip_addresses))
   except:
     if DEBUG >= 1: print("Failed to get all IP addresses. Defaulting to localhost.")
diff --git a/setup.py b/setup.py
index 6b59a9ac..98d80857 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,6 @@
 import sys
 import platform
+import subprocess
 
 from setuptools import find_packages, setup
 
@@ -11,7 +12,6 @@ install_requires = [
   "grpcio==1.68.0",
   "grpcio-tools==1.68.0",
   "Jinja2==3.1.4",
-  "netifaces==0.11.0",
   "numpy==2.0.0",
   "nuitka==2.5.1",
   "nvidia-ml-py==12.560.30",
@@ -23,6 +23,7 @@ install_requires = [
   "pydantic==2.9.2",
   "requests==2.32.3",
   "rich==13.7.1",
+  "scapy==2.6.1",
   "tenacity==9.0.0",
   "tqdm==4.66.4",
   "transformers==4.46.3",
@@ -31,19 +32,47 @@ install_requires = [
 ]
 
 extras_require = {
-  "formatting": [
-    "yapf==0.40.2",
-  ],
-  "apple_silicon": [
+  "formatting": ["yapf==0.40.2",], "apple_silicon": [
     "mlx==0.20.0",
     "mlx-lm==0.19.3",
-  ],
+  ], "windows": ["pywin32==308",], "nvidia-gpu": ["nvidia-ml-py==12.560.30",], "amd-gpu": ["pyrsmi==0.2.0"]
 }
 
 # Check if running on macOS with Apple Silicon
 if sys.platform.startswith("darwin") and platform.machine() == "arm64":
   install_requires.extend(extras_require["apple_silicon"])
 
+# Check if running Windows
+if sys.platform.startswith("win32"):
+  install_requires.extend(extras_require["windows"])
+
+
+def _add_gpu_requires():
+  global install_requires
+  # Add Nvidia-GPU
+  try:
+    out = subprocess.run(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], shell=True, text=True, capture_output=True, check=False)
+    if out.returncode == 0:
+      install_requires.extend(extras_require["nvidia-gpu"])
+  except subprocess.CalledProcessError:
+    pass
+
+  # Add AMD-GPU
+  # This will mostly work only on Linux, amd/rocm-smi is not yet supported on Windows
+  try:
+    out = subprocess.run(['amd-smi', 'list', '--csv'], shell=True, text=True, capture_output=True, check=False)
+    if out.returncode == 0:
+      install_requires.extend(extras_require["amd-gpu"])
+  except:
+    out = subprocess.run(['rocm-smi', 'list', '--csv'], shell=True, text=True, capture_output=True, check=False)
+    if out.returncode == 0:
+      install_requires.extend(extras_require["amd-gpu"])
+  finally:
+    pass
+
+
+_add_gpu_requires()
+
 setup(
   name="exo",
   version="0.0.1",
← 349b5344 Minor fix for Shard typing · back to Exo · Formatting b9eccedc →