muutils.sysinfo

  1"utilities for getting information about the system, see `SysInfo` class"
  2
  3from __future__ import annotations
  4
  5import subprocess
  6import sys
  7import typing
  8from importlib.metadata import distributions
  9
 10
 11def _popen(cmd: list[str], split_out: bool = False) -> dict[str, typing.Any]:
 12    p: subprocess.Popen = subprocess.Popen(
 13        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
 14    )
 15
 16    stdout, stderr = p.communicate()
 17
 18    p_out: typing.Union[str, list[str], None]
 19    if stdout:
 20        p_out = stdout.decode("utf-8")
 21        if split_out:
 22            assert isinstance(p_out, str)
 23            p_out = p_out.strip().split("\n")
 24    else:
 25        p_out = None
 26
 27    return {
 28        "stdout": p_out,
 29        "stderr": stderr.decode("utf-8") if stderr else None,
 30        "returncode": p.returncode if p.returncode is None else int(p.returncode),
 31    }
 32
 33
 34class SysInfo:
 35    """getters for various information about the system"""
 36
 37    @staticmethod
 38    def python() -> dict:
 39        """details about python version"""
 40        ver_tup = sys.version_info
 41        return {
 42            "version": sys.version,
 43            "version_info": ver_tup,
 44            "major": ver_tup[0],
 45            "minor": ver_tup[1],
 46            "micro": ver_tup[2],
 47            "releaselevel": ver_tup[3],
 48            "serial": ver_tup[4],
 49        }
 50
 51    @staticmethod
 52    def pip() -> dict:
 53        """installed packages info"""
 54        # in python <= 3.9  `Distribution` has no attribute `name`
 55        pckgs: list[tuple[str, str]] = [
 56            (
 57                (
 58                    x.metadata.get("Name", "<unknown>")  # type: ignore[attr-defined]
 59                    if sys.version_info < (3, 10)
 60                    else x.name  # type: ignore[attr-defined]
 61                ),
 62                x.version,
 63            )
 64            for x in distributions()
 65        ]
 66        return {
 67            "n_packages": len(pckgs),
 68            "packages": pckgs,
 69        }
 70
 71    @staticmethod
 72    def pytorch() -> dict:
 73        """pytorch and cuda information"""
 74        try:
 75            import torch
 76            import torch.version
 77        except Exception as e:
 78            return {
 79                "importable": False,
 80                "error": str(e),
 81            }
 82
 83        output: dict = {"importable": True}
 84
 85        output["torch.__version__"] = torch.__version__
 86        output["torch.version.cuda"] = torch.version.cuda
 87        output["torch.version.debug"] = torch.version.debug
 88        output["torch.version.git_version"] = torch.version.git_version
 89        output["torch.version.hip"] = torch.version.hip
 90        output["torch.cuda.is_available()"] = torch.cuda.is_available()
 91        output["torch.cuda.device_count()"] = torch.cuda.device_count()
 92        output["torch.cuda.is_initialized()"] = torch.cuda.is_initialized()
 93
 94        if torch.cuda.is_available():
 95            import os
 96
 97            cuda_version_nvcc: str = os.popen("nvcc --version").read()
 98            output["nvcc --version"] = cuda_version_nvcc.split("\n")
 99
100            if torch.cuda.device_count() > 0:
101                n_devices: int = torch.cuda.device_count()
102                output["torch.cuda.current_device()"] = torch.cuda.current_device()
103                output["torch devices"] = []
104                for current_device in range(n_devices):
105                    try:
106                        # print(f'checking current device {current_device} of {torch.cuda.device_count()} devices')
107                        # print(f'\tdevice {current_device}')
108                        # dev_prop = torch.cuda.get_device_properties(torch.device(0))
109                        # print(f'\t    name:                   {dev_prop.name}')
110                        # print(f'\t    version:                {dev_prop.major}.{dev_prop.minor}')
111                        # print(f'\t    total_memory:           {dev_prop.total_memory}')
112                        # print(f'\t    multi_processor_count:  {dev_prop.multi_processor_count}')
113                        # print(f'\t')
114                        dev_prop = torch.cuda.get_device_properties(current_device)
115                        output["torch devices"].append(
116                            {
117                                "device": current_device,
118                                "name": dev_prop.name,
119                                "version": {
120                                    "major": dev_prop.major,
121                                    "minor": dev_prop.minor,
122                                },
123                                "total_memory": dev_prop.total_memory,
124                                "multi_processor_count": dev_prop.multi_processor_count,
125                            }
126                        )
127                    except Exception as e:
128                        output["torch devices"].append(
129                            {
130                                "device": current_device,
131                                "error": str(e),
132                            }
133                        )
134        return output
135
136    @staticmethod
137    def platform() -> dict:
138        import platform
139
140        items = [
141            "platform",
142            "machine",
143            "processor",
144            "system",
145            "version",
146            "architecture",
147            "uname",
148            "node",
149            "python_branch",
150            "python_build",
151            "python_compiler",
152            "python_implementation",
153        ]
154
155        return {x: getattr(platform, x)() for x in items}
156
157    @staticmethod
158    def git_info(with_log: bool = False) -> dict:
159        git_version: dict = _popen(["git", "version"])
160        git_status: dict = _popen(["git", "status"])
161        if not git_status["stderr"] or git_status["stderr"].startswith(
162            "fatal: not a git repository"
163        ):
164            return {
165                "git version": git_version["stdout"],
166                "git status": git_status,
167            }
168        else:
169            output: dict = {
170                "git version": git_version["stdout"],
171                "git status": git_status,
172                "git branch": _popen(["git", "branch"], split_out=True),
173                "git remote -v": _popen(["git", "remote", "-v"], split_out=True),
174            }
175            if with_log:
176                output["git log"] = _popen(["git", "log"], split_out=False)
177
178            return output
179
180    @classmethod
181    def get_all(
182        cls,
183        include: typing.Optional[tuple[str, ...]] = None,
184        exclude: tuple[str, ...] = tuple(),
185    ) -> dict:
186        include_meta: tuple[str, ...]
187        if include is None:
188            include_meta = tuple(cls.__dict__.keys())
189        else:
190            include_meta = include
191
192        return {
193            x: getattr(cls, x)()
194            for x in include_meta
195            if all(
196                [
197                    not x.startswith("_"),
198                    x not in exclude,
199                    callable(getattr(cls, x)),
200                    x != "get_all",
201                    x in include if include is not None else True,
202                ]
203            )
204        }
205
206
207if __name__ == "__main__":
208    import pprint
209
210    pprint.pprint(SysInfo.get_all())