()
| 47 | pynvml.nvmlInit() |
| 48 | |
| 49 | # 获取GPU数量 |
| 50 | dev_cnt = pynvml.nvmlDeviceGetCount() |
| 51 | # torch.cuda.device_count() |
| 52 | logging.info(f"Total {dev_cnt} GPU(s) found") |
| 53 | IP_addr = config['ServerConfig']['serverIp_'] |
| 54 | Port = config['ServerConfig']['serverPort_'] |
| 55 | |
| 56 | logging.debug(f"IP_addr: {IP_addr}, Port: {Port}") |
| 57 | r = create_redis_connection() |
| 58 | |
| 59 | # 打印每个GPU的内存信息和利用率 |
| 60 | for i in range(dev_cnt): |
| 61 | gpu_info = {} |
| 62 | # 获取当前GPU的句柄 |
| 63 | handle = pynvml.nvmlDeviceGetHandleByIndex(i) |
| 64 | gpu_info['gpu_id'] = i |
| 65 | gpu_info['IP_addr'] = IP_addr |
| 66 | gpu_info['Port'] = Port |
| 67 | |
| 68 | # 获取内存信息 |
| 69 | mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle) |
| 70 | # free_mem = mem_info.free |
| 71 | # total_mem = mem_info.total |
| 72 | # used_mem = total_mem - free_mem |
| 73 | gpu_info['memory_total'] = mem_info.total |
| 74 | gpu_info['memory_free'] = mem_info.free |
| 75 | gpu_info['memory_used'] = mem_info.used |
| 76 | |
| 77 | # 获取GPU利用率 |
| 78 | util = pynvml.nvmlDeviceGetUtilizationRates(handle) |
| 79 | # gpu_util = util.gpu |
| 80 | gpu_info['utilization'] = util.gpu |
| 81 | |
| 82 | gpu_info['Job_num'] = 0 |
| 83 | gpu_info['GPU_num'] = dev_cnt |
| 84 | |
| 85 | gpu_info['HandlerIp'] = config['DatasetHandlerConfig']['handlerIp_'] |
| 86 | gpu_info['HandlerPort'] = config['DatasetHandlerConfig']['handlerPort_'] |
| 87 | gpuInfo.append(gpu_info) |
| 88 | |
| 89 | # 获取GPU的properties |
| 90 | gpu_properties = get_device_properties_bytes(i) |
| 91 | compressed_data = lz4.frame.compress(gpu_properties) |
| 92 | compressed_datas.append(compressed_data) |
| 93 | |
| 94 | #插入redis, key = '192.168.0.208:0' |
| 95 | redis_key = f"{IP_addr}:{i}" |
| 96 | # r.set(redis_key, json.dumps(gpu_info))#key = '192.168.0.208:0' |
| 97 | r.hset(redis_key, mapping={'gpu_info': json.dumps(gpu_info), 'gpu_properties': compressed_data}) |
| 98 | |
| 99 | logging.debug(gpuInfo) |
| 100 | |
| 101 | # 清理，释放NVML资源 |
| 102 | pynvml.nvmlShutdown() |
| 103 | |
| 104 | #更新redis中GPU信息 |
| 105 | def update_gpu_info(): |
| 106 | # global r |
no test coverage detected