1#!/usr/bin/env python3
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()
329