Blame for gpu_utils/gpu_stress.py

'''
运行:
python gpu_utils/gpu_stress.py --target_utilization 0.8 --duration 4 --num_gpus 8



'''


import torch
import torch.nn as nn
import torch.nn.functional as F
import time
import threading
import argparse
import numpy as np
from typing import List, Dict
import psutil
import os
from concurrent.futures import ThreadPoolExecutor
import multiprocessing as mp

class H100StressTest:
    def __init__(self, target_utilization: float = 0.5, num_gpus: int = None):
        self.target_utilization = target_utilization
        self.running = False
        self.stats = {
            'operations': 0,
            'start_time': None,
            'peak_memory': 0,
            'avg_utilization': 0
        }
        
        # 检查GPU可用性并确定使用的GPU数量
        if torch.cuda.is_available():
            available_gpus = torch.cuda.device_count()
            if num_gpus is None:
                self.num_gpus = available_gpus
            else:
                self.num_gpus = min(num_gpus, available_gpus)
            
            self.devices = [torch.device(f"cuda:{i}") for i in range(self.num_gpus)]
            
            print(f"检测到 {available_gpus} 个GPU，使用 {self.num_gpus} 个GPU")
            for i in range(self.num_gpus):
                print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
                print(f"GPU {i} 内存: {torch.cuda.get_device_properties(i).total_memory / 1024**3:.1f} GB")
                print(f"GPU {i} 计算能力: {torch.cuda.get_device_properties(i).major}.{torch.cuda.get_device_properties(i).minor}")
        else:
            raise RuntimeError("CUDA不可用，无法进行GPU压力测试")
    
    def create_large_transformer_model(self, hidden_size: int = 4096, num_layers: int = 8, device: torch.device = None) -> nn.Module:
        """创建大型Transformer模型以充分利用H100"""
        class LargeTransformer(nn.Module):
            def __init__(self, hidden_size, num_layers):
                super().__init__()
                self.hidden_size = hidden_size
                self.num_layers = num_layers
                
                # 多层Transformer编码器
                encoder_layer = nn.TransformerEncoderLayer(
                    d_model=hidden_size,
                    nhead=16,  # 多头注意力
                    dim_feedforward=hidden_size * 4,  # 前馈网络
                    dropout=0.1,
                    batch_first=True,
                    norm_first=True
                )
                self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
                
                # 额外的线性层
                self.input_projection = nn.Linear(hidden_size, hidden_size)
                self.output_projection = nn.Linear(hidden_size, hidden_size)
                self.layer_norm = nn.LayerNorm(hidden_size)
                
            def forward(self, x):
                x = self.input_projection(x)
                x = self.transformer(x)
                x = self.layer_norm(x)
                x = self.output_projection(x)
                return x
        
        if device is None:
            device = self.devices[0]
        model = LargeTransformer(hidden_size, num_layers).to(device)
        return model.half()  # 使用半精度以节省内存并提高性能
    
    def create_convolutional_model(self, device: torch.device = None) -> nn.Module:
        """创建卷积神经网络模型"""
        class ConvNet(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
                self.conv2 = nn.Conv2d(64, 128, 3, padding=1)
                self.conv3 = nn.Conv2d(128, 256, 3, padding=1)
                self.conv4 = nn.Conv2d(256, 512, 3, padding=1)
                self.conv5 = nn.Conv2d(512, 1024, 3, padding=1)
                
                self.pool = nn.MaxPool2d(2, 2)
                self.adaptive_pool = nn.AdaptiveAvgPool2d((1, 1))
                
                self.fc1 = nn.Linear(1024, 2048)
                self.fc2 = nn.Linear(2048, 4096)
                self.fc3 = nn.Linear(4096, 2048)
                self.fc4 = nn.Linear(2048, 1024)
                
                self.dropout = nn.Dropout(0.1)
                
            def forward(self, x):
                x = F.relu(self.conv1(x))
                x = self.pool(x)
                x = F.relu(self.conv2(x))
                x = self.pool(x)
                x = F.relu(self.conv3(x))
                x = self.pool(x)
                x = F.relu(self.conv4(x))
                x = self.pool(x)
                x = F.relu(self.conv5(x))
                x = self.adaptive_pool(x)
                
                x = x.view(x.size(0), -1)
                x = F.relu(self.fc1(x))
                x = self.dropout(x)
                x = F.relu(self.fc2(x))
                x = self.dropout(x)
                x = F.relu(self.fc3(x))
                x = self.dropout(x)
                x = self.fc4(x)
                return x
        
        if device is None:
            device = self.devices[0]
        model = ConvNet().to(device)
        return model.half()
    
    def matrix_operations_intensive(self, size: int = 8192, duration: float = 10.0, device: torch.device = None):
        """高强度矩阵运算"""
        if device is None:
            device = self.devices[0]
            
        # 创建大矩阵
        a = torch.randn(size, size, device=device, dtype=torch.float16)
        b = torch.randn(size, size, device=device, dtype=torch.float16)
        c = torch.randn(size, size, device=device, dtype=torch.float16)
        
        start_time = time.time()
        operations = 0
        
        while time.time() - start_time < duration:
            # 矩阵乘法链
            result = torch.mm(a, b)
            result = torch.mm(result, c)
            result = torch.mm(result, a.transpose(0, 1))
            
            # 矩阵求逆（小矩阵）
            if size <= 2048:
                try:
                    inv_result = torch.inverse(result[:min(1024, size), :min(1024, size)])
                    result = torch.mm(result, inv_result)
                except:
                    pass
            
            # 矩阵加法
            result = result + b
            result = result + c
            
            # 矩阵转置
            result = result.transpose(0, 1)
            
            operations += 1
            
            # 每100次操作同步一次
            if operations % 100 == 0:
                torch.cuda.synchronize(device)
        
        return operations
    
    def transformer_training_loop(self, model, optimizer, criterion, batch_size: int = 16, seq_length: int = 2048, device: torch.device = None):
        """Transformer训练循环"""
        if device is None:
            device = self.devices[0]
            
        hidden_size = model.hidden_size
        
        # 创建输入数据
        x = torch.randn(batch_size, seq_length, hidden_size, device=device, dtype=torch.float16)
        y = torch.randn(batch_size, seq_length, hidden_size, device=device, dtype=torch.float16)
        
        # 前向传播
        outputs = model(x)
        loss = criterion(outputs, y)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        return loss.item()
    
    def conv_training_loop(self, model, optimizer, criterion, batch_size: int = 32, device: torch.device = None):
        """卷积网络训练循环"""
        if device is None:
            device = self.devices[0]
            
        # 创建输入数据
        x = torch.randn(batch_size, 3, 224, 224, device=device, dtype=torch.float16)
        y = torch.randn(batch_size, 1024, device=device, dtype=torch.float16)
        
        # 前向传播
        outputs = model(x)
        loss = criterion(outputs, y)
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        return loss.item()
    
    def monitor_gpu_stats(self):
        """监控所有GPU统计信息"""
        while self.running:
            try:
                total_allocated = 0
                total_cached = 0
                gpu_stats = []
                
                # 获取所有GPU的内存使用情况
                for i in range(self.num_gpus):
                    allocated = torch.cuda.memory_allocated(i) / 1024**3
                    cached = torch.cuda.memory_reserved(i) / 1024**3
                    total_allocated += allocated
                    total_cached += cached
                    gpu_stats.append((allocated, cached))
                
                # 更新峰值内存
                if total_allocated > self.stats['peak_memory']:
                    self.stats['peak_memory'] = total_allocated
                
                # 每10秒打印一次状态
                if self.stats['start_time'] and int(time.time() - self.stats['start_time']) % 10 == 0:
                    elapsed = time.time() - self.stats['start_time']
                    ops_per_sec = self.stats['operations'] / elapsed if elapsed > 0 else 0
                    
                    print(f"[{elapsed:.0f}s] 操作数: {self.stats['operations']}, "
                          f"操作/秒: {ops_per_sec:.1f}")
                    print(f"总GPU内存: {total_allocated:.2f}GB/{total_cached:.2f}GB, "
                          f"峰值内存: {self.stats['peak_memory']:.2f}GB")
                    
                    # 显示每个GPU的详细状态
                    for i, (allocated, cached) in enumerate(gpu_stats):
                        print(f"  GPU {i}: {allocated:.2f}GB/{cached:.2f}GB")
                
                time.sleep(1)
            except Exception as e:
                print(f"监控错误: {e}")
                time.sleep(5)
    
    def gpu_worker(self, gpu_id: int, duration_seconds: int, shared_stats: dict, stats_lock: threading.Lock):
        """单个GPU的工作线程"""
        device = self.devices[gpu_id]
        print(f"GPU {gpu_id} 开始工作...")
        
        # 为每个GPU创建独立的模型
        transformer_model = self.create_large_transformer_model(
            hidden_size=4096, num_layers=12, device=device
        )
        transformer_optimizer = torch.optim.AdamW(
            transformer_model.parameters(), lr=1e-4, weight_decay=0.01
        )
        transformer_criterion = nn.MSELoss()
        
        conv_model = self.create_convolutional_model(device=device)
        conv_optimizer = torch.optim.AdamW(
            conv_model.parameters(), lr=1e-4, weight_decay=0.01
        )
        conv_criterion = nn.MSELoss()
        
        # 预热GPU
        for _ in range(5):
            self.matrix_operations_intensive(size=4096, duration=1.0, device=device)
        
        start_time = time.time()
        local_operations = 0
        cycle_count = 0
        
        try:
            while time.time() - start_time < duration_seconds:
                cycle_count += 1
                current_time = time.time() - start_time
                
                # 每30秒轮换不同类型的操作
                if int(current_time) % 30 < 10:
                    # 矩阵运算
                    ops = self.matrix_operations_intensive(size=8192, duration=8, device=device)
                    local_operations += ops
                    
                elif int(current_time) % 30 < 20:
                    # Transformer训练
                    for _ in range(5):
                        loss = self.transformer_training_loop(
                            transformer_model, transformer_optimizer, transformer_criterion,
                            batch_size=8, seq_length=2048, device=device
                        )
                        local_operations += 1
                    
                else:
                    # 卷积网络训练
                    for _ in range(10):
                        loss = self.conv_training_loop(
                            conv_model, conv_optimizer, conv_criterion, 
                            batch_size=16, device=device
                        )
                        local_operations += 1
                
                # 每100个周期清理一次缓存
                if cycle_count % 100 == 0:
                    torch.cuda.empty_cache()
                
                # 短暂休息以控制利用率
                time.sleep(0.1)
                
        except Exception as e:
            print(f"GPU {gpu_id} 发生错误: {e}")
        finally:
            # 更新共享统计信息
            with stats_lock:
                shared_stats['operations'] += local_operations
            print(f"GPU {gpu_id} 完成工作，操作数: {local_operations}")

    def run_intensive_stress_test(self, duration_hours: float = 2.0):
        """运行多GPU高强度压力测试"""
        duration_seconds = int(duration_hours * 3600)
        print(f"开始多GPU高强度压力测试")
        print(f"使用GPU数量: {self.num_gpus}")
        print(f"目标利用率: {self.target_utilization*100:.1f}%")
        print(f"持续时间: {duration_hours:.1f}小时 ({duration_seconds}秒)")
        
        # 启动监控线程
        self.running = True
        self.stats['start_time'] = time.time()
        monitor_thread = threading.Thread(target=self.monitor_gpu_stats)
        monitor_thread.daemon = True
        monitor_thread.start()
        
        # 创建共享统计信息和锁
        shared_stats = {'operations': 0}
        stats_lock = threading.Lock()
        
        # 使用线程池并行运行所有GPU
        with ThreadPoolExecutor(max_workers=self.num_gpus) as executor:
            # 为每个GPU提交任务
            futures = []
            for gpu_id in range(self.num_gpus):
                future = executor.submit(
                    self.gpu_worker, gpu_id, duration_seconds, shared_stats, stats_lock
                )
                futures.append(future)
            
            try:
                # 等待所有任务完成
                for future in futures:
                    future.result()
            except KeyboardInterrupt:
                print("\n用户中断测试")
                # 取消所有未完成的任务
                for future in futures:
                    future.cancel()
        
        self.running = False
        elapsed = time.time() - self.stats['start_time']
        print(f"\n多GPU压力测试完成!")
        print(f"总运行时间: {elapsed:.1f}秒 ({elapsed/3600:.2f}小时)")
        print(f"总操作数: {shared_stats['operations']}")
        print(f"平均操作/秒: {shared_stats['operations']/elapsed:.1f}")
        print(f"峰值GPU内存使用: {self.stats['peak_memory']:.2f}GB")
    
    def get_gpu_info(self):
        """获取所有GPU的详细信息"""
        if torch.cuda.is_available():
            print("=== 多GPU 详细信息 ===")
            print(f"GPU数量: {torch.cuda.device_count()}")
            print(f"使用GPU数量: {self.num_gpus}")
            
            total_memory = 0
            for i in range(self.num_gpus):
                print(f"\n--- GPU {i} ---")
                print(f"GPU名称: {torch.cuda.get_device_name(i)}")
                
                props = torch.cuda.get_device_properties(i)
                gpu_memory = props.total_memory / 1024**3
                total_memory += gpu_memory
                
                print(f"总内存: {gpu_memory:.1f} GB")
                print(f"多处理器数量: {props.multi_processor_count}")
                print(f"计算能力: {props.major}.{props.minor}")
                
                # 安全获取可能不存在的属性
                try:
                    print(f"最大线程数/块: {props.max_threads_per_block}")
                except AttributeError:
                    print("最大线程数/块: 信息不可用")
                
                try:
                    print(f"最大网格尺寸: {props.max_grid_size}")
                except AttributeError:
                    print("最大网格尺寸: 信息不可用")
                
                try:
                    print(f"共享内存/块: {props.shared_memory_per_block / 1024:.0f} KB")
                except AttributeError:
                    print("共享内存/块: 信息不可用")
                
                # 获取当前内存使用情况
                allocated = torch.cuda.memory_allocated(i) / 1024**3
                cached = torch.cuda.memory_reserved(i) / 1024**3
                print(f"已分配内存: {allocated:.2f} GB")
                print(f"缓存内存: {cached:.2f} GB")
                print(f"可用内存: {gpu_memory - cached:.2f} GB")
            
            print(f"\n总GPU内存: {total_memory:.1f} GB")
        else:
            print("CUDA不可用")

def main():
    parser = argparse.ArgumentParser(description='多GPU高强度压力测试工具')
    parser.add_argument('--duration', type=float, default=2.0,
                       help='测试持续时间（小时）')
    parser.add_argument('--target_utilization', type=float, default=0.5,
                       help='目标GPU利用率（0-1）')
    parser.add_argument('--num_gpus', type=int, default=None,
                       help='使用的GPU数量（默认使用所有可用GPU）')
    parser.add_argument('--info_only', action='store_true',
                       help='只显示GPU信息，不进行测试')
    
    args = parser.parse_args()
    
    try:
        # 创建压力测试实例
        stress_test = H100StressTest(
            target_utilization=args.target_utilization,
            num_gpus=args.num_gpus
        )
        
        # 显示GPU信息
        stress_test.get_gpu_info()
        
        if args.info_only:
            return
        
        # 运行压力测试
        stress_test.run_intensive_stress_test(duration_hours=args.duration)
        
    except Exception as e:
        print(f"错误: {e}")
        return 1

if __name__ == "__main__":
    exit(main())