codekingpro/portable-devtools

1# -------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation.  All rights reserved.
3# Licensed under the MIT License.
4# --------------------------------------------------------------------------
5
6# It is used to dump machine information for Notebooks
7
8import argparse
9import importlib.metadata
10import json
11import logging
12import platform
13from os import environ
14
15import cpuinfo
16import psutil
17from py3nvml.py3nvml import (
18    NVMLError,
19    nvmlDeviceGetCount,
20    nvmlDeviceGetHandleByIndex,
21    nvmlDeviceGetMemoryInfo,
22    nvmlDeviceGetName,
23    nvmlInit,
24    nvmlShutdown,
25    nvmlSystemGetDriverVersion,
26)
27
28
29class MachineInfo:
30    """Class encapsulating Machine Info logic."""
31
32    def __init__(self, silent=False, logger=None):
33        self.silent = silent
34
35        if logger is None:
36            logging.basicConfig(
37                format="%(asctime)s - %(name)s - %(levelname)s: %(message)s",
38                level=logging.INFO,
39            )
40            self.logger = logging.getLogger(__name__)
41        else:
42            self.logger = logger
43
44        self.machine_info = None
45        try:
46            self.machine_info = self.get_machine_info()
47        except Exception:
48            self.logger.exception("Exception in getting machine info.")
49            self.machine_info = None
50
51    def get_machine_info(self):
52        """Get machine info in metric format"""
53        gpu_info = self.get_gpu_info_by_nvml()
54        cpu_info = cpuinfo.get_cpu_info()
55
56        machine_info = {
57            "gpu": gpu_info,
58            "cpu": self.get_cpu_info(),
59            "memory": self.get_memory_info(),
60            "os": platform.platform(),
61            "python": self._try_get(cpu_info, ["python_version"]),
62            "packages": self.get_related_packages(),
63            "onnxruntime": self.get_onnxruntime_info(),
64            "pytorch": self.get_pytorch_info(),
65            "tensorflow": self.get_tensorflow_info(),
66        }
67        return machine_info
68
69    def get_memory_info(self) -> dict:
70        """Get memory info"""
71        mem = psutil.virtual_memory()
72        return {"total": mem.total, "available": mem.available}
73
74    def _try_get(self, cpu_info: dict, names: list) -> str:
75        for name in names:
76            if name in cpu_info:
77                value = cpu_info[name]
78                if isinstance(value, (list, tuple)):
79                    return ",".join([str(i) for i in value])
80                return value
81        return ""
82
83    def get_cpu_info(self) -> dict:
84        """Get CPU info"""
85        cpu_info = cpuinfo.get_cpu_info()
86
87        return {
88            "brand": self._try_get(cpu_info, ["brand", "brand_raw"]),
89            "cores": psutil.cpu_count(logical=False),
90            "logical_cores": psutil.cpu_count(logical=True),
91            "hz": self._try_get(cpu_info, ["hz_actual"]),
92            "l2_cache": self._try_get(cpu_info, ["l2_cache_size"]),
93            "flags": self._try_get(cpu_info, ["flags"]),
94            "processor": platform.uname().processor,
95        }
96
97    def get_gpu_info_by_nvml(self) -> dict:
98        """Get GPU info using nvml"""
99        gpu_info_list = []
100        driver_version = None
101        try:
102            nvmlInit()
103            driver_version = nvmlSystemGetDriverVersion()
104            deviceCount = nvmlDeviceGetCount()  # noqa: N806
105            for i in range(deviceCount):
106                handle = nvmlDeviceGetHandleByIndex(i)
107                info = nvmlDeviceGetMemoryInfo(handle)
108                gpu_info = {}
109                gpu_info["memory_total"] = info.total
110                gpu_info["memory_available"] = info.free
111                gpu_info["name"] = nvmlDeviceGetName(handle)
112                gpu_info_list.append(gpu_info)
113            nvmlShutdown()
114        except NVMLError as error:
115            if not self.silent:
116                self.logger.error("Error fetching GPU information using nvml: %s", error)
117            return None
118
119        result = {"driver_version": driver_version, "devices": gpu_info_list}
120
121        if "CUDA_VISIBLE_DEVICES" in environ:
122            result["cuda_visible"] = environ["CUDA_VISIBLE_DEVICES"]
123        return result
124
125    def get_related_packages(self) -> list[str]:
126        related_packages = {
127            "onnxruntime-gpu",
128            "onnxruntime",
129            "onnx",
130            "transformers",
131            "protobuf",
132            "sympy",
133            "torch",
134            "tensorflow",
135            "flatbuffers",
136            "numpy",
137            "onnxconverter-common",
138        }
139        related_packages_list = {}
140        for dist in importlib.metadata.distributions():
141            if dist.metadata["Name"].lower() in related_packages:
142                related_packages_list[dist.metadata["Name"].lower()] = dist.version
143
144        return related_packages_list
145
146    def get_onnxruntime_info(self) -> dict:
147        try:
148            import onnxruntime  # noqa: PLC0415
149
150            return {
151                "version": onnxruntime.__version__,
152                "support_gpu": "CUDAExecutionProvider" in onnxruntime.get_available_providers(),
153            }
154        except ImportError as error:
155            if not self.silent:
156                self.logger.exception(error)
157            return None
158        except Exception as exception:
159            if not self.silent:
160                self.logger.exception(exception, False)
161            return None
162
163    def get_pytorch_info(self) -> dict:
164        try:
165            import torch  # noqa: PLC0415
166
167            return {
168                "version": torch.__version__,
169                "support_gpu": torch.cuda.is_available(),
170                "cuda": torch.version.cuda,
171            }
172        except ImportError as error:
173            if not self.silent:
174                self.logger.exception(error)
175            return None
176        except Exception as exception:
177            if not self.silent:
178                self.logger.exception(exception, False)
179            return None
180
181    def get_tensorflow_info(self) -> dict:
182        try:
183            import tensorflow as tf  # noqa: PLC0415
184
185            return {
186                "version": tf.version.VERSION,
187                "git_version": tf.version.GIT_VERSION,
188                "support_gpu": tf.test.is_built_with_cuda(),
189            }
190        except ImportError as error:
191            if not self.silent:
192                self.logger.exception(error)
193            return None
194        except ModuleNotFoundError as error:
195            if not self.silent:
196                self.logger.exception(error)
197            return None
198
199
200def parse_arguments():
201    parser = argparse.ArgumentParser()
202
203    parser.add_argument(
204        "--silent",
205        required=False,
206        action="store_true",
207        help="Do not print error message",
208    )
209    parser.set_defaults(silent=False)
210
211    args = parser.parse_args()
212    return args
213
214
215def get_machine_info(silent=True) -> str:
216    machine = MachineInfo(silent)
217    return json.dumps(machine.machine_info, indent=2)
218
219
220def get_device_info(silent=True) -> str:
221    machine = MachineInfo(silent)
222    info = machine.machine_info
223    if info:
224        info = {key: value for key, value in info.items() if key in ["gpu", "cpu", "memory"]}
225    return json.dumps(info, indent=2)
226
227
228if __name__ == "__main__":
229    args = parse_arguments()
230    print(get_machine_info(args.silent))
231