File size: 2,833 Bytes
f0fc238 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | 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() |