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 # 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 

196if __name__ == "__main__": 

197 import pprint 

198 

199 pprint.pprint(SysInfo.get_all())