1#!/usr/bin/env python 

2# This file is part of Xpra. 

3# Copyright (C) 2013-2020 Antoine Martin <antoine@xpra.org> 

4# Xpra is released under the terms of the GNU GPL v2, or, at your option, any 

5# later version. See the file COPYING for details. 

6 

7import sys 

8import os 

9 

10from xpra.util import pver, print_nested_dict, engs, envbool, csv 

11from xpra.os_util import bytestostr, strtobytes, POSIX 

12from xpra.log import Logger 

13 

14log = Logger("encoder", "util") 

15 

16MIN_VERSION = 375 

17 

18nvml_init_warned = False 

19def wrap_nvml_init(nvmlInit) -> bool: 

20 try: 

21 nvmlInit() 

22 return True 

23 except Exception as e: 

24 log("get_nvml_driver_version() pynvml error", exc_info=True) 

25 global nvml_init_warned 

26 if not nvml_init_warned: 

27 log.warn("Warning: failed to initialize NVML:") 

28 log.warn(" %s", e) 

29 nvml_init_warned = True 

30 return False 

31 

32def get_nvml_driver_version(): 

33 try: 

34 from pynvml import nvmlInit, nvmlShutdown, nvmlSystemGetDriverVersion 

35 except ImportError as e: 

36 log("cannot use nvml to query the kernel module version:") 

37 log(" %s", e) 

38 else: 

39 try: 

40 if wrap_nvml_init(nvmlInit): 

41 try: 

42 v = nvmlSystemGetDriverVersion() 

43 finally: 

44 nvmlShutdown() 

45 log("nvmlSystemGetDriverVersion=%s", bytestostr(v)) 

46 return v.split(b".") 

47 except Exception as e: 

48 log("get_nvml_driver_version() pynvml error", exc_info=True) 

49 log.warn("Warning: failed to query the NVidia kernel module version using NVML:") 

50 log.warn(" %s", e) 

51 return () 

52 

53 

54def get_proc_driver_version(): 

55 if not POSIX: 

56 return () 

57 from xpra.os_util import load_binary_file 

58 proc_file = "/proc/driver/nvidia/version" 

59 v = load_binary_file(proc_file) 

60 if not v: 

61 log.warn("Warning: NVidia kernel module not installed?") 

62 log.warn(" cannot open '%s'", proc_file) 

63 return () 

64 KSTR = b"Kernel Module" 

65 p = v.find(KSTR) 

66 if not p: 

67 log.warn("unknown NVidia kernel module version") 

68 return "" 

69 v = v[p+len(KSTR):].strip().split(b" ")[0] 

70 v = v.split(b".") 

71 return v 

72 

73 

74def identify_nvidia_module_version(): 

75 v = get_nvml_driver_version() or get_proc_driver_version() 

76 #only keep numeric values: 

77 numver = [] 

78 try: 

79 for x in v: 

80 try: 

81 numver.append(int(x)) 

82 except ValueError: 

83 if not numver: 

84 raise 

85 if numver: 

86 log.info("NVidia driver version %s", pver(numver)) 

87 return tuple(numver) 

88 except Exception as e: 

89 log.warn("failed to parse Nvidia driver version '%s': %s", v, e) 

90 return () 

91 

92nvidia_module_version = None 

93def get_nvidia_module_version(probe=True): 

94 global nvidia_module_version 

95 if nvidia_module_version is None and probe: 

96 nvidia_module_version = identify_nvidia_module_version() 

97 return nvidia_module_version 

98 

99 

100def identify_cards(): 

101 devices = {} 

102 try: 

103 import pynvml 

104 from pynvml import nvmlInit, nvmlShutdown, nvmlDeviceGetCount, nvmlDeviceGetHandleByIndex 

105 deviceCount = None 

106 try: 

107 if not wrap_nvml_init(nvmlInit): 

108 return devices 

109 deviceCount = nvmlDeviceGetCount() 

110 log("identify_cards() will probe %i cards", deviceCount) 

111 for i in range(deviceCount): 

112 handle = nvmlDeviceGetHandleByIndex(i) 

113 log("identify_cards() handle(%i)=%s", i, handle) 

114 props = {} 

115 def meminfo(memory): 

116 return { 

117 "total" : int(memory.total), 

118 "free" : int(memory.free), 

119 "used" : int(memory.used), 

120 } 

121 def pciinfo(pci): 

122 i = {} 

123 for nvname, pubname in { 

124 "domain" : "domain", 

125 "bus" : "bus", 

126 "device" : "device", 

127 "pciDeviceId" : "pci-device-id", 

128 "pciSubSystemId" : "pci-subsystem-id", 

129 }.items(): 

130 try: 

131 i[pubname] = int(getattr(pci, nvname)) 

132 except (ValueError, AttributeError): 

133 pass 

134 try: 

135 i["bus-id"] = bytestostr(pci.busId) 

136 except AttributeError: 

137 pass 

138 return i 

139 for prefix, prop, fn_name, args, conv in ( 

140 ("", "name", "nvmlDeviceGetName", (), strtobytes), 

141 ("", "serial", "nvmlDeviceGetSerial", (), strtobytes), 

142 ("", "uuid", "nvmlDeviceGetUUID", (), strtobytes), 

143 ("", "pci", "nvmlDeviceGetPciInfo", (), pciinfo), 

144 ("", "memory", "nvmlDeviceGetMemoryInfo", (), meminfo), 

145 ("pcie-link", "generation-max", "nvmlDeviceGetMaxPcieLinkGeneration", (), int), 

146 ("pcie-link", "width-max", "nvmlDeviceGetMaxPcieLinkWidth", (), int), 

147 ("pcie-link", "generation", "nvmlDeviceGetCurrPcieLinkGeneration", (), int), 

148 ("pcie-link", "width", "nvmlDeviceGetCurrPcieLinkWidth", (), int), 

149 ("clock-info", "graphics", "nvmlDeviceGetClockInfo", (0,), int), 

150 ("clock-info", "sm", "nvmlDeviceGetClockInfo", (1,), int), 

151 ("clock-info", "mem", "nvmlDeviceGetClockInfo", (2,), int), 

152 ("clock-info", "graphics-max", "nvmlDeviceGetMaxClockInfo", (0,), int), 

153 ("clock-info", "sm-max", "nvmlDeviceGetMaxClockInfo", (1,), int), 

154 ("clock-info", "mem-max", "nvmlDeviceGetMaxClockInfo", (2,), int), 

155 ("", "fan-speed", "nvmlDeviceGetFanSpeed", (), int), 

156 ("", "temperature", "nvmlDeviceGetTemperature", (0,), int), 

157 ("", "power-state", "nvmlDeviceGetPowerState", (), int), 

158 ("", "vbios-version", "nvmlDeviceGetVbiosVersion", (), strtobytes), 

159 ): 

160 try: 

161 fn = getattr(pynvml, fn_name) 

162 v = fn(handle, *args) 

163 if conv: 

164 v = conv(v) 

165 if prefix: 

166 d = props.setdefault(prefix, {}) 

167 else: 

168 d = props 

169 d[prop] = v 

170 except Exception as e: 

171 log("identify_cards() cannot query %s using %s on device %i with handle %s: %s", 

172 prop, fn, i, handle, e) 

173 continue 

174 log("identify_cards() [%i]=%s", i, props) 

175 devices[i] = props 

176 #unitCount = nvmlUnitGetCount() 

177 #log.info("unitCount=%s", unitCount) 

178 except Exception as e: 

179 log("identify_cards() pynvml error", exc_info=True) 

180 log.warn("Warning: failed to query the NVidia cards using NVML:") 

181 log.warn(" %s", e) 

182 finally: 

183 if deviceCount is not None: 

184 nvmlShutdown() 

185 except ImportError as e: 

186 log("cannot use nvml to query the kernel module version:") 

187 log(" %s", e) 

188 return devices 

189 

190 

191_cards = None 

192def get_cards(probe=True): 

193 global _cards 

194 if _cards is None and probe: 

195 _cards = identify_cards() 

196 return _cards 

197 

198 

199def is_blacklisted(): 

200 v = get_nvidia_module_version(True) 

201 try: 

202 if v[0]>MIN_VERSION: 

203 return False 

204 except Exception as e: 

205 log.warn("Warning: error checking driver version:") 

206 log.warn(" %s", e) 

207 return None #we don't know: unreleased / untested 

208 

209 

210_version_warning = False 

211def validate_driver_yuv444lossless(): 

212 #this should log the kernel module version 

213 v = get_nvidia_module_version() 

214 if not v: 

215 log.warn("Warning: unknown NVidia driver version") 

216 bl = None 

217 else: 

218 bl = is_blacklisted() 

219 if bl is True: 

220 raise Exception("NVidia driver version %s is blacklisted, it does not work with NVENC" % pver(v)) 

221 elif bl is None: 

222 global _version_warning 

223 if _version_warning: 

224 l = log 

225 else: 

226 l = log.warn 

227 _version_warning = True 

228 if v: 

229 l("Warning: NVidia driver version %s is untested with NVENC", pver(v)) 

230 l(" (this encoder has been tested with versions %s.x and later only)", MIN_VERSION) 

231 if not envbool("XPRA_NVENC_YUV444P", False): 

232 l(" disabling YUV444P and lossless mode") 

233 l(" use XPRA_NVENC_YUV444P=1 to force enable") 

234 return False 

235 l(" force enabling YUV444P and lossless mode") 

236 return True 

237 

238 

239def parse_nvfbc_hex_key(s): 

240 #ie: 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10 

241 #ie: 0102030405060708090A0B0C0D0E0F10 

242 #start by removing spaces and 0x: 

243 hexstr = s.replace("0x", "").replace(",", "").replace(" ", "") 

244 import binascii 

245 return binascii.unhexlify(hexstr) 

246 

247 

248license_keys = {} 

249def get_license_keys(version=0, basefilename="nvenc"): 

250 global license_keys 

251 filename = "%s%s.keys" % (basefilename, version or "") 

252 keys = license_keys.get(filename) 

253 if keys is not None: 

254 return keys 

255 env_name = "XPRA_%s_CLIENT_KEY" % basefilename.upper() 

256 env_keys = os.environ.get(env_name, "") 

257 if env_keys: 

258 keys = [x.strip() for x in env_keys.split(",")] 

259 log("using %s keys from environment variable %s: %s", basefilename, env_name, csv(keys)) 

260 else: 

261 #try to load the license file 

262 keys = [] 

263 try: 

264 #see read_xpra_defaults for an explanation of paths 

265 from xpra.platform.paths import get_default_conf_dirs, get_system_conf_dirs, get_user_conf_dirs 

266 dirs = get_default_conf_dirs() + get_system_conf_dirs() + get_user_conf_dirs() 

267 for d in dirs: 

268 if not d: 

269 continue 

270 keys_file = os.path.join(d, filename) 

271 keys_file = os.path.expanduser(keys_file) 

272 if not os.path.exists(keys_file): 

273 log("get_license_keys(%s, %s) '%s' does not exist", basefilename, version, keys_file) 

274 continue 

275 log("loading %s version %s keys from %s", basefilename, version, keys_file) 

276 with open(keys_file, "rb") as f: 

277 fkeys = [] 

278 for line in f: 

279 sline = line.strip().rstrip(b'\r\n').strip().decode("latin1") 

280 if not sline: 

281 log("skipping empty line") 

282 continue 

283 if sline[0] in ('!', '#'): 

284 log("skipping comments") 

285 continue 

286 fkeys.append(sline) 

287 log("added key: %s", sline) 

288 log("added %i key%s from %s", len(fkeys), engs(fkeys), keys_file) 

289 keys += fkeys 

290 except Exception: 

291 log.error("Error loading %s license keys", basefilename, exc_info=True) 

292 license_keys[filename] = keys 

293 log("get_nvenc_license_keys(%s)=%s", version, keys) 

294 return keys 

295 

296 

297def main(): 

298 if "-v" in sys.argv or "--verbose" in sys.argv: 

299 log.enable_debug() 

300 

301 from xpra.platform import program_context 

302 with program_context("Nvidia-Info", "Nvidia Info"): 

303 #this will log the version number: 

304 get_nvidia_module_version() 

305 if is_blacklisted(): 

306 log.warn("Warning: this driver version is blacklisted") 

307 log.info("NVENC license keys:") 

308 for v in (0, 8): 

309 keys = get_license_keys(v) 

310 log.info("* version %s: %s key(s)", v or "common", len(keys)) 

311 for k in keys: 

312 log.info(" %s", k) 

313 try: 

314 import pynvml 

315 assert pynvml 

316 except ImportError: 

317 log.warn("Warning: the pynvml library is missing") 

318 log.warn(" cannot identify the GPUs installed") 

319 else: 

320 cards = get_cards() 

321 if cards: 

322 log.info("") 

323 log.info("%i card%s:", len(cards), engs(cards)) 

324 print_nested_dict(cards, print_fn=log.info) 

325 

326 

327if __name__ == "__main__": 

328 main()