1"utilities for getting information about the system, see `SysInfo` class"

2

3 from __future__ import annotations

4

5 import subprocess

6 import sys

7 import typing

8 from importlib.metadata import distributions

9

10

11 def _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

34 class 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 # for some reason, python 3.8 thinks `Distribution` has no attribute `name`?

55 pckgs: list[tuple[str, str]] = [(x.name, x.version) for x in distributions()] # type: ignore[attr-defined]

56 return {

57 "n_packages": len(pckgs),

58 "packages": pckgs,

59 }

60

61 @staticmethod

62 def pytorch() -> dict:

63 """pytorch and cuda information"""

64 try:

65 import torch

66 except Exception as e:

67 return {

68 "importable": False,

69 "error": str(e),

70 }

71

72 output: dict = {"importable": True}

73

74 output["torch.__version__"] = torch.__version__

75 output["torch.version.cuda"] = torch.version.cuda

76 output["torch.version.debug"] = torch.version.debug

77 output["torch.version.git_version"] = torch.version.git_version

78 output["torch.version.hip"] = torch.version.hip

79 output["torch.cuda.is_available()"] = torch.cuda.is_available()

80 output["torch.cuda.device_count()"] = torch.cuda.device_count()

81 output["torch.cuda.is_initialized()"] = torch.cuda.is_initialized()

82

83 if torch.cuda.is_available():

84 import os

85

86 cuda_version_nvcc: str = os.popen("nvcc --version").read()

87 output["nvcc --version"] = cuda_version_nvcc.split("\n")

88

89 if torch.cuda.device_count() > 0:

90 n_devices: int = torch.cuda.device_count()

91 output["torch.cuda.current_device()"] = torch.cuda.current_device()

92 output["torch devices"] = []

93 for current_device in range(n_devices):

94 try:

95 # print(f'checking current device {current_device} of {torch.cuda.device_count()} devices')

96 # print(f'\tdevice {current_device}')

97 # dev_prop = torch.cuda.get_device_properties(torch.device(0))

98 # print(f'\t name: {dev_prop.name}')

99 # print(f'\t version: {dev_prop.major}.{dev_prop.minor}')

100 # print(f'\t total_memory: {dev_prop.total_memory}')

101 # print(f'\t multi_processor_count: {dev_prop.multi_processor_count}')

102 # print(f'\t')

103 dev_prop = torch.cuda.get_device_properties(current_device)

104 output["torch devices"].append(

105 {

106 "device": current_device,

107 "name": dev_prop.name,

108 "version": {

109 "major": dev_prop.major,

110 "minor": dev_prop.minor,

111 },

112 "total_memory": dev_prop.total_memory,

113 "multi_processor_count": dev_prop.multi_processor_count,

114 }

115 )

116 except Exception as e:

117 output["torch devices"].append(

118 {

119 "device": current_device,

120 "error": str(e),

121 }

122 )

123 return output

124

125 @staticmethod

126 def platform() -> dict:

127 import platform

128

129 items = [

130 "platform",

131 "machine",

132 "processor",

133 "system",

134 "version",

135 "architecture",

136 "uname",

137 "node",

138 "python_branch",

139 "python_build",

140 "python_compiler",

141 "python_implementation",

142 ]

143

144 return {x: getattr(platform, x)() for x in items}

145

146 @staticmethod

147 def git_info(with_log: bool = False) -> dict:

148 git_version: dict = _popen(["git", "version"])

149 git_status: dict = _popen(["git", "status"])

150 if not git_status["stderr"] or git_status["stderr"].startswith(

151 "fatal: not a git repository"

152 ):

153 return {

154 "git version": git_version["stdout"],

155 "git status": git_status,

156 }

157 else:

158 output: dict = {

159 "git version": git_version["stdout"],

160 "git status": git_status,

161 "git branch": _popen(["git", "branch"], split_out=True),

162 "git remote -v": _popen(["git", "remote", "-v"], split_out=True),

163 }

164 if with_log:

165 output["git log"] = _popen(["git", "log"], split_out=False)

166

167 return output

168

169 @classmethod

170 def get_all(

171 cls,

172 include: typing.Optional[tuple[str, ...]] = None,

173 exclude: tuple[str, ...] = tuple(),

174 ) -> dict:

175 include_meta: tuple[str, ...]

176 if include is None:

177 include_meta = tuple(cls.__dict__.keys())

178 else:

179 include_meta = include

180

181 return {

182 x: getattr(cls, x)()

183 for x in include_meta

184 if all(

185 [

186 not x.startswith("_"),

187 x not in exclude,

188 callable(getattr(cls, x)),

189 x != "get_all",

190 x in include if include is not None else True,

191 ]

192 )

193 }

194

195

196 if __name__ == "__main__":

197 import pprint

198

199 pprint.pprint(SysInfo.get_all())