chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,256 @@
|
||||
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import paddle
|
||||
from paddle.base import core
|
||||
|
||||
|
||||
class Info:
|
||||
def __repr__(self):
|
||||
return str(self.__dict__)
|
||||
|
||||
def json(self):
|
||||
return json.dumps(self.__dict__)
|
||||
|
||||
def dict(self):
|
||||
return self.__dict__
|
||||
|
||||
def str(self, keys=None):
|
||||
if keys is None:
|
||||
keys = self.__dict__.keys()
|
||||
|
||||
if isinstance(keys, str):
|
||||
keys = keys.split(',')
|
||||
|
||||
values = [str(self.__dict__.get(k, '')) for k in keys]
|
||||
return ",".join(values)
|
||||
|
||||
|
||||
def query_smi(query=None, query_type="gpu", index=None, dtype=None):
|
||||
"""
|
||||
query_type: gpu/compute
|
||||
"""
|
||||
|
||||
if not has_nvidia_smi():
|
||||
return []
|
||||
|
||||
cmd = ["nvidia-smi", "--format=csv,noheader,nounits"]
|
||||
if isinstance(query, list) and query_type == "gpu":
|
||||
cmd.extend(["--query-gpu={}".format(",".join(query))])
|
||||
elif isinstance(query, list) and query_type.startswith("compute"):
|
||||
cmd.extend(["--query-compute-apps={}".format(",".join(query))])
|
||||
else:
|
||||
return
|
||||
|
||||
if isinstance(index, list) and len(index) > 0:
|
||||
cmd.extend(["--id={}".format(",".join(index))])
|
||||
if not isinstance(dtype, list) or len(dtype) != len(query):
|
||||
dtype = [str] * len(query)
|
||||
|
||||
output = subprocess.check_output(cmd, timeout=3)
|
||||
lines = output.decode("utf-8").split(os.linesep)
|
||||
ret = []
|
||||
for line in lines:
|
||||
if not line:
|
||||
continue
|
||||
info = Info()
|
||||
for k, v, d in zip(query, line.split(", "), dtype):
|
||||
setattr(info, k.replace(".", "_"), d(v))
|
||||
ret.append(info)
|
||||
return ret
|
||||
|
||||
|
||||
def query_rocm_smi(query=None, index=None, dtype=None, mem=32150):
|
||||
if not has_rocm_smi():
|
||||
return []
|
||||
|
||||
cmd = ["rocm-smi"]
|
||||
|
||||
if not isinstance(dtype, list) or len(dtype) != len(query):
|
||||
dtype = [str] * len(query)
|
||||
|
||||
output = subprocess.check_output(cmd, timeout=3)
|
||||
lines = output.decode("utf-8").split(os.linesep)
|
||||
ret = []
|
||||
for line in lines:
|
||||
if not line:
|
||||
continue
|
||||
if len(line.split()) != 8 or "DCU" in line.split():
|
||||
continue
|
||||
info = Info()
|
||||
line = line.split()
|
||||
line = [
|
||||
line[0],
|
||||
line[7][: len(line[7]) - 1],
|
||||
mem,
|
||||
mem * float(line[6][: len(line[6]) - 1]) / 100,
|
||||
mem - mem * float(line[6][: len(line[6]) - 1]) / 100,
|
||||
time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
|
||||
]
|
||||
for k, v, d in zip(query, line, dtype):
|
||||
setattr(info, k.replace(".", "_"), d(v))
|
||||
ret.append(info)
|
||||
return ret
|
||||
|
||||
|
||||
def query_npu_smi(query=None, index=None, dtype=None):
|
||||
if not has_npu_smi():
|
||||
return []
|
||||
|
||||
cmd = ["npu-smi", "info"]
|
||||
|
||||
if not isinstance(dtype, list) or len(dtype) != len(query):
|
||||
dtype = [str] * len(query)
|
||||
|
||||
output = subprocess.check_output(cmd, timeout=3)
|
||||
lines = output.decode("utf-8").split(os.linesep)
|
||||
ret = []
|
||||
i = 0
|
||||
|
||||
for line in lines:
|
||||
if not line:
|
||||
continue
|
||||
result = re.split(r',|/|\s+|\|', line)
|
||||
# result = [item for item in result if item]
|
||||
length = len(result)
|
||||
if length not in [18, 19] or "NPU" in result:
|
||||
continue
|
||||
result = [item for item in result if item]
|
||||
info = Info()
|
||||
result = [
|
||||
i,
|
||||
result[2],
|
||||
result[6],
|
||||
float(result[5]),
|
||||
(float(result[6]) - float(result[5])),
|
||||
time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
|
||||
]
|
||||
i += 1
|
||||
for k, v, d in zip(query, result, dtype):
|
||||
setattr(info, k.replace(".", "_"), d(v))
|
||||
ret.append(info)
|
||||
return ret
|
||||
|
||||
|
||||
def query_xpu_smi(query=None, index=None, dtype=None):
|
||||
if (
|
||||
not hasattr(core, "get_xpu_device_count")
|
||||
or core.get_xpu_device_count() == 0
|
||||
):
|
||||
return []
|
||||
if not isinstance(dtype, list) or len(dtype) != len(query):
|
||||
dtype = [str] * len(query)
|
||||
if not isinstance(index, list) or len(index) == 0:
|
||||
index = list(range(core.get_xpu_device_count()))
|
||||
ret = []
|
||||
for dev_id in index:
|
||||
dev_id = int(dev_id)
|
||||
utilization_xpu = core.get_xpu_device_utilization_rate(dev_id)
|
||||
mem_total = (
|
||||
core.get_xpu_device_total_memory(dev_id) / 1024 / 1024
|
||||
) # with MB
|
||||
mem_used = (
|
||||
core.get_xpu_device_used_memory(dev_id) / 1024 / 1024
|
||||
) # with MB
|
||||
result = [
|
||||
dev_id,
|
||||
utilization_xpu,
|
||||
mem_total,
|
||||
mem_used,
|
||||
(mem_total - mem_used),
|
||||
time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
|
||||
]
|
||||
info = Info()
|
||||
for k, v, d in zip(query, result, dtype):
|
||||
setattr(info, k.replace(".", "_"), d(v))
|
||||
ret.append(info)
|
||||
return ret
|
||||
|
||||
|
||||
def get_gpu_info(index=None):
|
||||
q = "index,uuid,driver_version,name,gpu_serial,display_active,display_mode".split(
|
||||
","
|
||||
)
|
||||
d = [int, str, str, str, str, str, str]
|
||||
index = (
|
||||
index
|
||||
if index is None or isinstance(index, list)
|
||||
else str(index).split(",")
|
||||
)
|
||||
|
||||
return query_smi(q, index=index, dtype=d)
|
||||
|
||||
|
||||
def get_gpu_util(index=None):
|
||||
q = "index,utilization.gpu,memory.total,memory.used,memory.free,timestamp".split(
|
||||
","
|
||||
)
|
||||
d = [int, int, int, int, int, str]
|
||||
index = (
|
||||
index
|
||||
if index is None or isinstance(index, list)
|
||||
else str(index).split(",")
|
||||
)
|
||||
if paddle.device.is_compiled_with_rocm():
|
||||
return query_rocm_smi(q, index=index, dtype=d)
|
||||
elif paddle.device.is_compiled_with_custom_device('npu'):
|
||||
return query_npu_smi(q, index=index, dtype=d)
|
||||
elif paddle.is_compiled_with_xpu():
|
||||
return query_xpu_smi(q, index=index, dtype=d)
|
||||
return query_smi(q, index=index, dtype=d)
|
||||
|
||||
|
||||
def get_gpu_process(index=None):
|
||||
q = "pid,process_name,gpu_uuid,gpu_name,used_memory".split(",")
|
||||
d = [int, str, str, str, int]
|
||||
index = (
|
||||
index
|
||||
if index is None or isinstance(index, list)
|
||||
else str(index).split(",")
|
||||
)
|
||||
|
||||
return query_smi(q, index=index, query_type="compute", dtype=d)
|
||||
|
||||
|
||||
def has_nvidia_smi():
|
||||
return shutil.which("nvidia-smi")
|
||||
|
||||
|
||||
def has_rocm_smi():
|
||||
return shutil.which("rocm-smi")
|
||||
|
||||
|
||||
def has_npu_smi():
|
||||
return shutil.which("npu-smi")
|
||||
|
||||
|
||||
def has_xpu_smi():
|
||||
return shutil.which("xpu-smi")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
print(get_gpu_info(0))
|
||||
print(get_gpu_util(0))
|
||||
print(get_gpu_process(0))
|
||||
|
||||
u = get_gpu_util()
|
||||
for i in u:
|
||||
print(i.str())
|
||||
Reference in New Issue
Block a user