MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / _gather_gpu_devices_cudart

Function _gather_gpu_devices_cudart

tensorflow/tools/test/gpu_info_lib.py:118–159  ·  view source on GitHub ↗

Try to gather NVidia GPU device information via libcudart.

()

Source from the content-addressed store, hash-verified

116
117
118def _gather_gpu_devices_cudart():
119 """Try to gather NVidia GPU device information via libcudart."""
120 dev_info = []
121
122 system = platform.system()
123 if system == "Linux":
124 libcudart = ct.cdll.LoadLibrary("libcudart.so")
125 elif system == "Darwin":
126 libcudart = ct.cdll.LoadLibrary("libcudart.dylib")
127 elif system == "Windows":
128 libcudart = ct.windll.LoadLibrary("libcudart.dll")
129 else:
130 raise NotImplementedError("Cannot identify system.")
131
132 version = ct.c_int()
133 rc = libcudart.cudaRuntimeGetVersion(ct.byref(version))
134 if rc != 0:
135 raise ValueError("Could not get version")
136 if version.value < 6050:
137 raise NotImplementedError("CUDA version must be between >= 6.5")
138
139 device_count = ct.c_int()
140 libcudart.cudaGetDeviceCount(ct.byref(device_count))
141
142 for i in range(device_count.value):
143 properties = CUDADeviceProperties()
144 rc = libcudart.cudaGetDeviceProperties(ct.byref(properties), i)
145 if rc != 0:
146 raise ValueError("Could not get device properties")
147 pci_bus_id = " " * 13
148 rc = libcudart.cudaDeviceGetPCIBusId(ct.c_char_p(pci_bus_id), 13, i)
149 if rc != 0:
150 raise ValueError("Could not get device PCI bus id")
151
152 info = test_log_pb2.GPUInfo() # No UUID available
153 info.model = properties.name
154 info.bus_id = pci_bus_id
155 dev_info.append(info)
156
157 del properties
158
159 return dev_info
160
161
162def gather_gpu_devices():

Callers 1

gather_gpu_devicesFunction · 0.85

Calls 4

rangeFunction · 0.50
LoadLibraryMethod · 0.45
appendMethod · 0.45

Tested by

no test coverage detected