diff --git a/digits/device_query.py b/digits/device_query.py index 560729022..07bb0076c 100755 --- a/digits/device_query.py +++ b/digits/device_query.py @@ -73,9 +73,11 @@ class c_cudaDeviceProp(ctypes.Structure): ('managedMemSupported', ctypes.c_int), ('isMultiGpuBoard', ctypes.c_int), ('multiGpuBoardGroupID', ctypes.c_int), + # Extra space for new fields in future toolkits + ('__future_buffer', ctypes.c_int * 128), # added later with cudart.cudaDeviceGetPCIBusId # (needed by NVML) - ('pciBusID_str', ctypes.c_char * 13), + ('pciBusID_str', ctypes.c_char * 16), ] class struct_c_nvmlDevice_t(ctypes.Structure): @@ -93,6 +95,8 @@ class c_nvmlMemory_t(ctypes.Structure): ('total', ctypes.c_ulonglong), ('free', ctypes.c_ulonglong), ('used', ctypes.c_ulonglong), + # Extra space for new fields in future toolkits + ('__future_buffer', ctypes.c_ulonglong * 8), ] class c_nvmlUtilization_t(ctypes.Structure): @@ -102,6 +106,8 @@ class c_nvmlUtilization_t(ctypes.Structure): _fields_ = [ ('gpu', ctypes.c_uint), ('memory', ctypes.c_uint), + # Extra space for new fields in future toolkits + ('__future_buffer', ctypes.c_uint * 8), ] def get_library(name): @@ -128,14 +134,12 @@ def get_cudart(): if cudart is not None: return cudart else: - for name in ( - 'libcudart.so.7.0', - 'libcudart.so.7.5', - 'libcudart.so.8.0', - 'libcudart.so'): - cudart = get_library(name) - if cudart is not None: - return cudart + for major in xrange(9,5,-1): + for minor in (5,0): + cudart = get_library('libcudart.so.%d.%d' % (major, minor)) + if cudart is not None: + return cudart + return get_library('libcudart.so') return None def get_nvml(): @@ -196,9 +200,9 @@ def get_devices(force_reload=False): properties = c_cudaDeviceProp() rc = cudart.cudaGetDeviceProperties(ctypes.byref(properties), x) if rc == 0: - pciBusID_str = ' ' * 13 + pciBusID_str = ' ' * 16 # also save the string representation of the PCI bus ID - rc = cudart.cudaDeviceGetPCIBusId(ctypes.c_char_p(pciBusID_str), 13, x) + rc = cudart.cudaDeviceGetPCIBusId(ctypes.c_char_p(pciBusID_str), 16, x) if rc == 0: properties.pciBusID_str = pciBusID_str devices.append(properties) @@ -281,6 +285,8 @@ def get_nvml_info(device_id): print 'Device #%d:' % i print '>>> CUDA attributes:' for name, t in device._fields_: + if name in ['__future_buffer']: + continue if not args.verbose and name not in [ 'name', 'totalGlobalMem', 'clockRate', 'major', 'minor',]: continue