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 import torch.version 

67 except Exception as e: 

68 return { 

69 "importable": False, 

70 "error": str(e), 

71 } 

72 

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

74 

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

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

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

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

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

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

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

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

83 

84 if torch.cuda.is_available(): 

85 import os 

86 

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

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

89 

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

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

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

93 output["torch devices"] = [] 

94 for current_device in range(n_devices): 

95 try: 

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

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

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

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

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

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

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

103 # print(f'\t') 

104 dev_prop = torch.cuda.get_device_properties(current_device) 

105 output["torch devices"].append( 

106 { 

107 "device": current_device, 

108 "name": dev_prop.name, 

109 "version": { 

110 "major": dev_prop.major, 

111 "minor": dev_prop.minor, 

112 }, 

113 "total_memory": dev_prop.total_memory, 

114 "multi_processor_count": dev_prop.multi_processor_count, 

115 } 

116 ) 

117 except Exception as e: 

118 output["torch devices"].append( 

119 { 

120 "device": current_device, 

121 "error": str(e), 

122 } 

123 ) 

124 return output 

125 

126 @staticmethod 

127 def platform() -> dict: 

128 import platform 

129 

130 items = [ 

131 "platform", 

132 "machine", 

133 "processor", 

134 "system", 

135 "version", 

136 "architecture", 

137 "uname", 

138 "node", 

139 "python_branch", 

140 "python_build", 

141 "python_compiler", 

142 "python_implementation", 

143 ] 

144 

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

146 

147 @staticmethod 

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

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

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

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

152 "fatal: not a git repository" 

153 ): 

154 return { 

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

156 "git status": git_status, 

157 } 

158 else: 

159 output: dict = { 

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

161 "git status": git_status, 

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

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

164 } 

165 if with_log: 

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

167 

168 return output 

169 

170 @classmethod 

171 def get_all( 

172 cls, 

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

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

175 ) -> dict: 

176 include_meta: tuple[str, ...] 

177 if include is None: 

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

179 else: 

180 include_meta = include 

181 

182 return { 

183 x: getattr(cls, x)() 

184 for x in include_meta 

185 if all( 

186 [ 

187 not x.startswith("_"), 

188 x not in exclude, 

189 callable(getattr(cls, x)), 

190 x != "get_all", 

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

192 ] 

193 ) 

194 } 

195 

196 

197if __name__ == "__main__": 

198 import pprint 

199 

200 pprint.pprint(SysInfo.get_all())