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

addftool/addf_portal.py

@@ -3,6 +3,7 @@ from addftool.process import add_killer_args, killer_main
 from addftool.sync import add_sync_args, sync_main
 from addftool.deploy import add_deploy_args, deploy_main
 from addftool.broadcast_folder import add_broadcast_folder_args, broadcast_folder_main
+from addftool.sleep import add_sleep_args, sleep_main
 
 from addftool.blob import add_blob_args, blob_main
 
@@ -16,6 +17,7 @@ def get_args():
     add_deploy_args(subparsers)
     add_broadcast_folder_args(subparsers)
     add_blob_args(subparsers)
+    add_sleep_args(subparsers)
 
     return parser.parse_args()
 
@@ -32,6 +34,8 @@ def main():
         broadcast_folder_main(args)
     elif args.command == "blob":
         blob_main(args)
+    elif args.command == "sleep":
+        sleep_main(args)
     else:
         print("Unknown command: ", args.command)
 

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

@@ -0,0 +1,337 @@
+import time
+import subprocess
+import sys
+import multiprocessing as mp
+import re
+
+try:
+    import torch
+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的利用率和显存使用情况（支持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)
+        
+        if result.returncode != 0:
+            print(f"Error running nvidia-smi for GPU {device_id}")
+            return None, None
+        
+        # 解析输出
+        output = result.stdout.strip()
+        if output:
+            parts = output.split(',')
+            if len(parts) == 2:
+                gpu_util = int(parts[0])  # GPU利用率百分比
+                memory_used = int(parts[1])  # 显存使用量(MB)
+                return gpu_util, memory_used
+        
+        return None, None
+    except Exception as 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是否被其他进程占用
+    
+    Args:
+        device_id: GPU设备ID
+        util_threshold: GPU利用率阈值(默认20%)
+        memory_threshold: 显存使用阈值(默认2048MB = 2GB)
+    
+    Returns:
+        bool: True表示GPU被占用，False表示GPU空闲
+    """
+    gpu_util, memory_used = get_gpu_stats(device_id)
+    
+    if gpu_util is None or memory_used is None:
+        # 获取失败时保守处理
+        return True
+    
+    # 判断是否被占用
+    is_occupied = gpu_util > util_threshold or (memory_threshold > 0 and memory_used > memory_threshold)
+    
+    if is_occupied:
+        print(f"GPU {device_id}: Util={gpu_util}%, Memory={memory_used}MB - Occupied")
+    
+    return is_occupied
+
+def check_all_gpus(num_gpus, util_threshold=20, memory_threshold=-1):
+    """检查所有GPU是否被占用"""
+    for device_id in range(num_gpus):
+        if check_gpu_occupied(device_id, util_threshold, memory_threshold):
+            return True, device_id
+    return False, -1
+
+def get_all_gpu_status(num_gpus):
+    """获取所有GPU的状态信息"""
+    print("\nGPU Status:")
+    print("-" * 50)
+    for device_id in range(num_gpus):
+        gpu_util, memory_used = get_gpu_stats(device_id)
+        if gpu_util is not None and memory_used is not None:
+            status = "Available" if (gpu_util <= 20 and memory_used <= 2048) else "Occupied"
+            print(f"GPU {device_id}: Util={gpu_util:3d}%, Memory={memory_used:5d}MB - {status}")
+        else:
+            print(f"GPU {device_id}: Unable to get stats")
+    print("-" * 50)
+
+def matrix_multiply_worker(matrix_size=8192, time_duration=4.0, sleep_duration=1.0, util_threshold=20, memory_threshold=-1):
+    # 获取GPU数量
+    num_gpus = torch.cuda.device_count()
+    if num_gpus == 0:
+        print("No GPUs available!")
+        return
+
+    matrices = {}
+    # print(f"Creating {matrix_size}x{matrix_size} matrices on all GPUs...")
+    for device_id in range(num_gpus):
+        device = torch.device(f'cuda:{device_id}')
+        matrices[device_id] = {
+            'a': torch.randn(matrix_size, matrix_size, device=device, dtype=torch.float32),
+            'b': torch.randn(matrix_size, matrix_size, device=device, dtype=torch.float32)
+        }
+
+    # 主循环
+    while True:
+        try:
+            # 检查所有GPU是否被占用
+            has_occupied_gpu, occupied_gpu = check_all_gpus(num_gpus, util_threshold, memory_threshold)
+            if has_occupied_gpu:
+                break
+
+            start_time = time.time()
+            perform_count = 0
+            while True:
+                # 在所有GPU上同时执行矩阵乘法
+                results = {}
+                for device_id in range(num_gpus):
+                    results[device_id] = torch.matmul(matrices[device_id]['a'], matrices[device_id]['b'])
+                
+                perform_count += 1
+
+                if perform_count % 10 == 0:
+                    for device_id in range(num_gpus):
+                        torch.cuda.synchronize(device_id)
+                
+                torch.cuda.synchronize()  # 确保所有GPU操作完成
+                elapsed_time = time.time() - start_time
+                if elapsed_time > time_duration:
+                    break
+
+            # 清理内存
+            
+            time.sleep(sleep_duration)
+            
+        except KeyboardInterrupt:
+            print("\nKeyboardInterrupt received, stopping...")
+            stop_flag = True
+            exit(0)
+        except Exception as e:
+            print(f"\nError occurred: {e}")
+            # 尝试清理内存
+            try:
+                for device_id in range(num_gpus):
+                    torch.cuda.set_device(device_id)
+                    torch.cuda.empty_cache()
+            except:
+                pass
+            time.sleep(5)
+    
+def sleep_main(args):
+    # 设置多进程启动方法
+    mp.set_start_method('spawn', force=True)
+
+    num_gpus = torch.cuda.device_count()
+    if num_gpus == 0:
+        print("No GPUs available!")
+        exit(1)
+        
+    # 显示初始GPU状态
+    get_all_gpu_status(num_gpus)
+
+    current_process = None
+    
+    # 主循环
+    while True:
+        try:
+            # 检查所有GPU是否被占用
+            has_occupied_gpu, occupied_gpu = check_all_gpus(num_gpus, util_threshold=args.util_threshold, memory_threshold=args.memory_threshold)
+            
+            if has_occupied_gpu:
+                # 休眠60秒
+                print("Holding for 60 seconds...")
+                time.sleep(60)
+                
+            else:
+                # GPU空闲，启动矩阵乘法进程
+                current_process = mp.Process(
+                    target=matrix_multiply_worker,
+                    args=(args.matrix_size, args.time_duration, args.sleep_duration, args.util_threshold, args.memory_threshold), 
+                )
+                current_process.start()
+                current_process.join()
+            
+        except KeyboardInterrupt:
+            print("\nKeyboardInterrupt received, stopping...")
+            stop_flag = True
+            break
+        except Exception as e:
+            print(f"\nError occurred: {e}")
+            time.sleep(5)
+
+    print("\nProgram stopped")
+
+
+def add_sleep_args(subparsers):
+    sleep_parser = subparsers.add_parser('sleep', help='Sleep for a while and check GPU status')
+    add_args(sleep_parser)
+
+
+def add_args(parser):
+    parser.add_argument('--matrix_size', type=int, default=8192, help='Size of the matrices to multiply')
+    parser.add_argument('--time_duration', type=float, default=4.0, help='Duration to perform matrix multiplication')
+    parser.add_argument('--sleep_duration', type=float, default=1.0, help='Duration to sleep between checks')
+
+    parser.add_argument('--util_threshold', type=int, default=20, help='GPU utilization threshold to consider it occupied')
+    parser.add_argument('--memory_threshold', type=int, default=-1, help='Memory usage threshold (in GB) to consider it occupied, set to -1 to disable')
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description='Sleep and check GPU status')
+    add_args(parser)
+    args = parser.parse_args()
+    while True:
+        try:
+            sleep_main(args)
+        except KeyboardInterrupt:
+            print("\nKeyboardInterrupt received, exiting...")
+            sys.exit(0)
+        except Exception as e:
+            print(f"\nUnexpected error: {e}")
+            print("Restarting the program in 5 seconds...")
+            time.sleep(5)
+            continue

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.1
 Name: addftool
-Version: 0.2.7
+Version: 0.2.9
 Requires-Dist: cryptography
 Requires-Dist: requests
 Requires-Dist: PyYAML
@@ -8,3 +8,4 @@ Requires-Dist: psutil
 Requires-Dist: fabric
 Requires-Dist: gevent
 Requires-Dist: parallel-ssh
+

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

@@ -1,7 +1,8 @@
 addftool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-addftool/addf_portal.py,sha256=6XjwGs5m2mRVDWVvCPOiqn1NlxDcGBTQ9Kr_0g5RsJc,1130
-addftool/blob.py,sha256=y1HZaDBUNeXicVytvwpRXwufvvrgxR33ruBlYpxnSa4,9453
+addftool/addf_portal.py,sha256=vc8opPzValNFPwJne5C5LbZvgcJ0eMBJSWDSiM23OPM,1274
+addftool/blob.py,sha256=vyjJHlQZuGrpEiaRF-Bdmow_TMnhXtXEGL31qA5Hb-g,9742
 addftool/broadcast_folder.py,sha256=GQBuSL8Ch537V_fSBHesWyqT3KRYry68pbYOKy2bDj4,19619
+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
@@ -12,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.7.dist-info/METADATA,sha256=7wElFYgZp3OX387bAsiQTwzfrc4pHH8dlK6vPnDULWU,220
-addftool-0.2.7.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
-addftool-0.2.7.dist-info/entry_points.txt,sha256=9lkmuWMInwUAtev8w8poNkNd7iML9Bjd5CBCFVxg2b8,111
-addftool-0.2.7.dist-info/top_level.txt,sha256=jqj56-plrBbyzY0tIxB6wPzjAA8kte4hUlajyyQygN4,9
-addftool-0.2.7.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,,

{addftool-0.2.7.dist-info → addftool-0.2.9.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.0.0)
+Generator: setuptools (75.3.2)
 Root-Is-Purelib: true
 Tag: py3-none-any