addftool 0.2.8__py3-none-any.whl → 0.2.9__py3-none-any.whl

addftool/blob.py

@@ -76,6 +76,14 @@ def install_main(args):
     # if has root permission, run install script
     # else, print install script
 
+    # make sure wget is installed
+    if not check_package_installed("wget"):
+        print("wget is not installed, installing wget")
+        command = "apt-get install wget -y"
+        if args.sudo:
+            command = "sudo " + command
+        execute_command(command, script_writer)
+
     print("Get ubuntu version: ", ubuntu_version)
     command = f"wget https://packages.microsoft.com/config/ubuntu/{ubuntu_version}/packages-microsoft-prod.deb -O /tmp/packages-microsoft-prod.deb"
     print("Install packages-microsoft-prod.deb")

addftool/sleep.py

@@ -2,6 +2,7 @@ import time
 import subprocess
 import sys
 import multiprocessing as mp
+import re
 
 try:
     import torch
@@ -9,9 +10,61 @@ except ImportError:
     print("PyTorch is not installed. Please install it to run this script.")
     sys.exit(1)
 
+try:
+    import triton
+    import triton.runtime.driver
+except ImportError:
+    print("Triton is not installed. Will try to detect GPU type using command line tools.")
+    triton = None
+
+def is_cuda():
+    """使用 triton 检测是否是 CUDA 环境"""
+    try:
+        if triton is None:
+            return None
+        return triton.runtime.driver.active.get_current_target().backend == "cuda"
+    except:
+        return None
+
+def get_gpu_type():
+    """检测GPU类型（NVIDIA/CUDA或AMD/ROCm）"""
+    # 首先尝试使用 triton 检测
+    cuda_detected = is_cuda()
+    if cuda_detected is True:
+        return "nvidia"
+    elif cuda_detected is False:
+        return "amd"
+    
+    # 如果 triton 检测失败，回退到命令行检测
+    try:
+        # 尝试检测NVIDIA GPU
+        result = subprocess.run("nvidia-smi", shell=True, capture_output=True, text=True)
+        if result.returncode == 0:
+            return "nvidia"
+        
+        # 尝试检测AMD GPU
+        result = subprocess.run("rocm-smi", shell=True, capture_output=True, text=True)
+        if result.returncode == 0:
+            return "amd"
+        
+        return None
+    except:
+        return None
 
 def get_gpu_stats(device_id):
-    """获取指定GPU的利用率和显存使用情况"""
+    """获取指定GPU的利用率和显存使用情况（支持NVIDIA和AMD）"""
+    gpu_type = get_gpu_type()
+    
+    if gpu_type == "nvidia":
+        return get_nvidia_gpu_stats(device_id)
+    elif gpu_type == "amd":
+        return get_amd_gpu_stats(device_id)
+    else:
+        print("No supported GPU found (neither NVIDIA nor AMD)")
+        return None, None
+
+def get_nvidia_gpu_stats(device_id):
+    """获取NVIDIA GPU的统计信息"""
     try:
         cmd = f"nvidia-smi --id={device_id} --query-gpu=utilization.gpu,memory.used --format=csv,noheader,nounits"
         result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
@@ -31,9 +84,75 @@ def get_gpu_stats(device_id):
         
         return None, None
     except Exception as e:
-        print(f"Error getting GPU stats for device {device_id}: {e}")
+        print(f"Error getting NVIDIA GPU stats for device {device_id}: {e}")
         return None, None
 
+def get_amd_gpu_stats(device_id):
+    """获取AMD GPU的统计信息"""
+    try:
+        # 获取GPU利用率和显存使用情况
+        cmd = f"rocm-smi -d {device_id} --showuse --showmemuse"
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+        
+        if result.returncode != 0:
+            # 尝试备用命令
+            cmd = f"rocm-smi -d {device_id}"
+            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+            if result.returncode != 0:
+                print(f"Error running rocm-smi for GPU {device_id}")
+                return None, None
+        
+        gpu_util = None
+        memory_used = None
+        
+        # 解析输出
+        output = result.stdout
+        lines = output.split('\n')
+        
+        for line in lines:
+            # 查找GPU利用率
+            if 'GPU use' in line or '%' in line:
+                # 匹配百分比
+                match = re.search(r'(\d+)%', line)
+                if match:
+                    gpu_util = int(match.group(1))
+            
+            # 查找显存使用（MB）
+            if 'vram' in line.lower() or 'memory' in line.lower() or 'MB' in line:
+                # 匹配MB数值，格式可能是 "1024 MB" 或 "1024MB"
+                match = re.search(r'(\d+)\s*MB', line, re.IGNORECASE)
+                if match:
+                    memory_used = int(match.group(1))
+        
+        # 如果仍然无法获取利用率，设置为0（假设空闲）
+        if gpu_util is None:
+            gpu_util = 0
+        
+        # 如果仍然无法获取内存使用，设置为0
+        if memory_used is None:
+            memory_used = 0
+        
+        return gpu_util, memory_used
+        
+    except Exception as e:
+        print(f"Error getting AMD GPU stats for device {device_id}: {e}")
+        return None, None
+
+# 在程序启动时检测GPU类型
+try:
+    GPU_TYPE = get_gpu_type()
+    if GPU_TYPE:
+        cuda_status = is_cuda()
+        if cuda_status is not None:
+            print(f"Detected {GPU_TYPE.upper()} GPU environment (triton backend: {'cuda' if cuda_status else 'hip'})")
+        else:
+            print(f"Detected {GPU_TYPE.upper()} GPU environment")
+    else:
+        print("No supported GPU environment detected")
+except:
+    GPU_TYPE = None
+    print("Failed to detect GPU environment")
+
 def check_gpu_occupied(device_id, util_threshold=20, memory_threshold=2048):
     """检查GPU是否被其他进程占用
     

{addftool-0.2.8.dist-info → addftool-0.2.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: addftool
-Version: 0.2.8
+Version: 0.2.9
 Requires-Dist: cryptography
 Requires-Dist: requests
 Requires-Dist: PyYAML

{addftool-0.2.8.dist-info → addftool-0.2.9.dist-info}/RECORD

@@ -1,8 +1,8 @@
 addftool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 addftool/addf_portal.py,sha256=vc8opPzValNFPwJne5C5LbZvgcJ0eMBJSWDSiM23OPM,1274
-addftool/blob.py,sha256=y1HZaDBUNeXicVytvwpRXwufvvrgxR33ruBlYpxnSa4,9453
+addftool/blob.py,sha256=vyjJHlQZuGrpEiaRF-Bdmow_TMnhXtXEGL31qA5Hb-g,9742
 addftool/broadcast_folder.py,sha256=GQBuSL8Ch537V_fSBHesWyqT3KRYry68pbYOKy2bDj4,19619
-addftool/sleep.py,sha256=FA1fTUI47eQq-9nBtXElkS7SZMunP_5tLiIBuFNSM6w,7823
+addftool/sleep.py,sha256=Y6gAJb2Ho4qvbd52-UBmwt7Rgv2HpJAcBnWO23asaY8,11787
 addftool/sync.py,sha256=ZpYxbM8uiPFrV7ODmOaM7asVPCWaxBixA-arVc-1kfs,14045
 addftool/tool.py,sha256=FmxRY3-pP0_Z0zCUAngjmEMmPUruMftg_iUlB1t2TnQ,2001
 addftool/util.py,sha256=zlNLu8Be8cGIpNRqBw8_0q7nFxWlsJ9cToN62ohjdXE,2335
@@ -13,8 +13,8 @@ addftool/deploy/vscode_server.py,sha256=tLtSvlcK2fEOaw6udWt8dNELVhwv9F59hF5DJJ-1
 addftool/process/__init__.py,sha256=Dze8OrcyjQlAbPrjE_h8bMi8W4b3OJyZOjTucPrkJvM,3721
 addftool/process/utils.py,sha256=JldxnwanLJOgxaPgmCJh7SeBRaaj5rFxWWxh1hpsvbA,2609
 addftool/ssh/__init__.py,sha256=h5_rCO0A6q2Yw9vFguQZZp_ApAJsT1dcnKnbKKZ0cDM,4409
-addftool-0.2.8.dist-info/METADATA,sha256=rxu5Oy4lH7lQF99Z8gzz5QuoGxnZ739h0OBNhr_0NA0,221
-addftool-0.2.8.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
-addftool-0.2.8.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.2.8.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.2.8.dist-info/RECORD,,
+addftool-0.2.9.dist-info/METADATA,sha256=lWVPrPjbHKv10pKm1w4HPdccUy_122llzDR1WQ6c2mg,221
+addftool-0.2.9.dist-info/WHEEL,sha256=iAkIy5fosb7FzIOwONchHf19Qu7_1wCWyFNR5gu9nU0,91
+addftool-0.2.9.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
+addftool-0.2.9.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
+addftool-0.2.9.dist-info/RECORD,,