上传文件至 /

getGPUUsedInfo.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python
+
+import pynvml
+import redis
+
+# 大模型进程（端口）与gpu绑定关系：port: [gpu_device_id..]
+gpu={'1001':['g0','g1'],'1002':['g3'],'1003':['g4']}
+
+def getGpuInfo():
+    gpuUsageInfo={}
+    pynvml.nvmlInit()
+    deviceCount = pynvml.nvmlDeviceGetCount()
+    for i in range(deviceCount):
+        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+        memInfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        gpuUtil = pynvml.nvmlDeviceGetUtilizationRates(handle)
+        gpuUsageInfo[f"g{i}"] = {
+            'used': (memInfo.used/memInfo.total)*100,
+            'util': gpuUtil.gpu
+        }
+
+    pynvml.nvmlShutdown()
+    return gpuUsageInfo
+
+def getBanlanceRate(memory,gpu):
+    return memory+gpu
+
+def zsetData(**kwargs):
+    data={}
+    for key in kwargs:
+        memory=0
+        gpu=0
+        for d in kwargs[key]:
+            memory+=gpuUsageInfo[d]['used']
+            gpu+=gpuUsageInfo[d]['util']
+        banlance=int(getBanlanceRate(memory,gpu)/len(kwargs[key]))
+        data[key]=banlance
+    return data
+
+
+if __name__ == '__main__':
+    local="172.20.1.3"
+    pool=redis.ConnectionPool(host="60.204.148.84",port=6379,password='xm!redis123',db=15)
+    r=redis.Redis(connection_pool=pool)
+    pipe=r.pipeline(transaction=True)
+    gpuUsageInfo=getGpuInfo()
+    data=zsetData(**gpu)
+    for key in data:
+        #print(data[key],f"{local}:{key}")
+        r.zadd("danceai",{f"{local}:{key}":data[key]})