Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
jantic
GitHub Repository: jantic/deoldify
 Path: blob/master/fastai/utils/pynvml_gate.py
781 views
1
"""Get OS specific nvml wrapper. On OSX we use pynvx as drop in replacement for pynvml"""

2


3
import platform

4
from ..script import *

5


6
#

7
# BEGIN: Temporary workaround for nvml.dll load issue in Win10

8
#

9
# Remove once nicolargo/nvidia-ml-py3#2 and a new version of the module is released 

10
# (OR fbcotter/py3nvml#10 but will require extra work to rename things)

11
# Refer https://forums.fast.ai/t/nvml-dll-loading-issue-in-nvidia-ml-py3-7-352-0-py-0/39684/8

12
import threading

13
from ctypes import *

14


15
nvmlLib = None

16
libLoadLock = threading.Lock()

17


18
def _LoadNvmlLibrary():

19
    '''

20
    Load the library if it isn't loaded already

21
    '''

22


23
    global nvmlLib

24


25
    if (nvmlLib == None):

26
        libLoadLock.acquire()

27


28
        try:

29
            if (nvmlLib == None):

30
                try:

31
                    if (sys.platform[:3] == "win"):

32
                        searchPaths = [

33
                            os.path.join(os.getenv("ProgramFiles", r"C:\Program Files"), r"NVIDIA Corporation\NVSMI\nvml.dll"),

34
                            os.path.join(os.getenv("WinDir", r"C:\Windows"), r"System32\nvml.dll"),

35
                        ]

36
                        nvmlPath = next((x for x in searchPaths if os.path.isfile(x)), None)

37
                        if (nvmlPath == None):

38
                            nvmlLib = None

39
                        else:

40
                            nvmlLib = CDLL(nvmlPath)

41
                    else:

42
                        nvmlLib = None

43
                except OSError as ose:

44
                    nvmlLib = None

45
        finally:

46
            libLoadLock.release()

47
#

48
# END: Temporary workaround for nvml.dll load issue in Win10

49
#

50


51
def load_pynvml_env():

52
    import pynvml # nvidia-ml-py3

53


54
    #

55
    # BEGIN: Temporary workaround for nvml.dll load issue in Win10 (continued)

56
    _LoadNvmlLibrary()

57
    pynvml.nvmlLib = nvmlLib

58
    #

59
    # END: Temporary workaround for nvml.dll load issue in Win10

60
    #

61


62
    if platform.system() == "Darwin":

63
        try:

64
            from pynvx import pynvml

65
        except:

66
            print("please install pynvx on OSX: pip install pynvx")

67
            sys.exit(1)

68


69
        pynvml.nvmlInit()

70
        return pynvml

71


72
    pynvml.nvmlInit()

73


74
    return pynvml

75


76