#!/usr/bin/env python3
# User prompt: 帮我写一个每隔内跑满gpu的pytorch code，不要占用太多gpu内存

import torch
import time
import argparse

def activate_gpu(duration=10, interval=5, matrix_size=1024, device='cuda:0'):
    """
    周期性地激活GPU，进行密集计算但不占用太多内存
    
    Args:
        duration: 每次激活持续的秒数
        interval: 激活之间的间隔秒数
        matrix_size: 矩阵大小（越大计算量越大，但也会占用更多内存）
        device: GPU设备
    """
    if not torch.cuda.is_available():
        print("错误：CUDA不可用")
        return
    
    device = torch.device(device)
    print(f"使用设备: {device}")
    print(f"GPU名称: {torch.cuda.get_device_name(device)}")
    print(f"激活参数: 持续{duration}秒, 间隔{interval}秒, 矩阵大小{matrix_size}x{matrix_size}")
    print("-" * 60)
    
    iteration = 0
    
    try:
        while True:
            iteration += 1
            print(f"\n[迭代 {iteration}] 开始激活GPU...")
            
            # 记录开始时间
            start_time = time.time()
            
            # 创建两个较小的矩阵在GPU上
            a = torch.randn(matrix_size, matrix_size, device=device)
            b = torch.randn(matrix_size, matrix_size, device=device)
            
            # 打印内存使用情况
            mem_allocated = torch.cuda.memory_allocated(device) / 1024**2
            mem_reserved = torch.cuda.memory_reserved(device) / 1024**2
            print(f"  GPU内存: 已分配 {mem_allocated:.2f} MB, 已保留 {mem_reserved:.2f} MB")
            
            # 持续进行矩阵运算直到达到duration
            ops_count = 0
            while time.time() - start_time < duration:
                # 执行大量矩阵乘法来跑满GPU
                c = torch.matmul(a, b)
                c = torch.matmul(c, a)
                c = torch.matmul(c, b)
                
                # 添加一些其他操作增加计算复杂度
                c = torch.relu(c)
                c = torch.sigmoid(c)
                c = c + a * b
                
                ops_count += 1
                
                # 偶尔更新显示
                if ops_count % 100 == 0:
                    elapsed = time.time() - start_time
                    print(f"  运行中... {elapsed:.1f}s / {duration}s, 完成 {ops_count} 次运算")
            
            # 清理GPU内存
            del a, b, c
            torch.cuda.empty_cache()
            
            elapsed = time.time() - start_time
            print(f"  完成！总共 {ops_count} 次运算，耗时 {elapsed:.2f} 秒")
            
            # 显示清理后的内存状态
            mem_allocated = torch.cuda.memory_allocated(device) / 1024**2
            print(f"  清理后GPU内存: {mem_allocated:.2f} MB")
            
            # 等待间隔
            print(f"  休息 {interval} 秒...")
            time.sleep(interval)
            
    except KeyboardInterrupt:
        print("\n\n收到中断信号，正在停止...")
        torch.cuda.empty_cache()
        print("已清理GPU内存并退出")

def main():
    parser = argparse.ArgumentParser(description='周期性激活GPU进行计算')
    parser.add_argument('--duration', type=int, default=10, 
                        help='每次激活持续的秒数 (默认: 10)')
    parser.add_argument('--interval', type=int, default=5, 
                        help='激活之间的间隔秒数 (默认: 5)')
    parser.add_argument('--matrix-size', type=int, default=1024, 
                        help='矩阵大小，影响计算量和内存占用 (默认: 1024)')
    parser.add_argument('--device', type=str, default='cuda:0', 
                        help='GPU设备 (默认: cuda:0)')
    
    args = parser.parse_args()
    
    activate_gpu(
        duration=args.duration,
        interval=args.interval,
        matrix_size=args.matrix_size,
        device=args.device
    )

if __name__ == '__main__':
    main()

