import os
import time
import torch
import multiprocessing as mp
from pynvml import *

# 配置参数
TARGET_GPUS = [4,5,6,7]
TARGET_MEMORY_MB = 75000  # 目标占用显存 75GB
IDLE_THRESHOLD = 0.99     # 20%
IDLE_DURATION = 10       # 秒
CHECK_INTERVAL = 5        

def get_gpu_utilization(gpu_id):
    handle = nvmlDeviceGetHandleByIndex(gpu_id)
    util = nvmlDeviceGetUtilizationRates(handle)
    return util.gpu / 100.0

def occupy_task(gpu_id, target_mb):
    """执行矩阵乘法占卡任务"""
    # 限制该进程只能看到指定的显卡
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    device = torch.device("cuda:0")
    
    # 计算 N: 75000MB / 4 bytes (float32) = 19.6亿个元素
    # 每个矩阵占一半空间，即 9.8亿个元素。 sqrt(9.8e8) ≈ 31300
    n = 60000 
    
    print(f"[GPU {gpu_id}] 正在分配显存并启动运算...")
    try:
        # 分配两个大矩阵，总计约占 73-75GB
        a = torch.randn(n, n, device=device)
        b = torch.randn(n, n, device=device)
        print(f"[GPU {gpu_id}] 显存分配完成，开始循环矩阵乘法。")
        
        while True:
            # 持续计算确保利用率
            _ = torch.matmul(a, b)
    except Exception as e:
        print(f"[GPU {gpu_id}] 任务异常: {e}")

def monitor_and_run(gpu_id):
    nvmlInit()
    idle_start_time = None
    is_running = False
    p_occupy = None

    print(f"[GPU {gpu_id}] 监控进程已启动...")

    try:
        while True:
            util = get_gpu_utilization(gpu_id)
            
            # 只有在没有运行占卡任务时，才去判断是否空闲
            if not is_running:
                if util < IDLE_THRESHOLD:
                    if idle_start_time is None:
                        idle_start_time = time.time()
                    
                    elapsed = time.time() - idle_start_time
                    if elapsed >= IDLE_DURATION:
                        print(f"[GPU {gpu_id}] 已持续空闲，启动占卡。")
                        p_occupy = mp.Process(target=occupy_task, args=(gpu_id, TARGET_MEMORY_MB))
                        p_occupy.start()
                        is_running = True
                else:
                    idle_start_time = None
            else:
                # 任务运行中，我们不再根据 util > 20% 来关闭它
                # 如果你想手动停止，请直接 Ctrl+C
                pass

            time.sleep(CHECK_INTERVAL)
    except KeyboardInterrupt:
        if p_occupy: p_occupy.terminate()

if __name__ == "__main__":
    mp.set_start_method('spawn', force=True)
    processes = []
    for gid in TARGET_GPUS:
        p = mp.Process(target=monitor_and_run, args=(gid,))
        p.start()
        processes.append(p)
    
    for p in processes:
        p.join()