Coverage for muutils/sysinfo.py: 82%

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())