upd pynvml.py

This commit is contained in:
iperov 2019-04-05 21:33:53 +04:00
parent 406ef91851
commit 41601f50b8

View file

@ -636,9 +636,16 @@ def _LoadNvmlLibrary():
if (nvmlLib == None): if (nvmlLib == None):
try: try:
if (sys.platform[:3] == "win"): if (sys.platform[:3] == "win"):
# cdecl calling convention searchPaths = [
# load nvml.dll from %ProgramFiles%/NVIDIA Corporation/NVSMI/nvml.dll os.path.join(os.getenv("ProgramFiles", r"C:\Program Files"), r"NVIDIA Corporation\NVSMI\nvml.dll"),
nvmlLib = CDLL(os.path.join(os.getenv("ProgramFiles", "C:/Program Files"), "NVIDIA Corporation/NVSMI/nvml.dll")) os.path.join(os.getenv("WinDir", r"C:\Windows"), r"System32\nvml.dll"),
]
nvmlPath = next((x for x in searchPaths if os.path.isfile(x)), None)
if (nvmlPath == None):
_nvmlCheckReturn(NVML_ERROR_LIBRARY_NOT_FOUND)
else:
# cdecl calling convention
nvmlLib = CDLL(nvmlPath)
else: else:
# assume linux # assume linux
nvmlLib = CDLL("libnvidia-ml.so.1") nvmlLib = CDLL("libnvidia-ml.so.1")
@ -1668,7 +1675,7 @@ def nvmlSystemGetTopologyGpuSet(cpuNumber):
if ret != NVML_SUCCESS: if ret != NVML_SUCCESS:
raise NVMLError(ret) raise NVMLError(ret)
print ("c_count.value") print(c_count.value)
# call again with a buffer # call again with a buffer
device_array = c_nvmlDevice_t * c_count.value device_array = c_nvmlDevice_t * c_count.value
c_devices = device_array() c_devices = device_array()
@ -1717,4 +1724,4 @@ def nvmlDeviceGetCudaComputeCapability(device):
if (ret != NVML_SUCCESS): if (ret != NVML_SUCCESS):
raise NVMLError(ret) raise NVMLError(ret)
return c_major.value, c_minor.value return c_major.value, c_minor.value