Run find_cuda_config.py and return cuda_toolkit_path, or None.
(environ_cp)
| 1314 | |
| 1315 | |
| 1316 | def validate_cuda_config(environ_cp): |
| 1317 | """Run find_cuda_config.py and return cuda_toolkit_path, or None.""" |
| 1318 | |
| 1319 | def maybe_encode_env(env): |
| 1320 | """Encodes unicode in env to str on Windows python 2.x.""" |
| 1321 | if not is_windows() or sys.version_info[0] != 2: |
| 1322 | return env |
| 1323 | for k, v in env.items(): |
| 1324 | if isinstance(k, unicode): |
| 1325 | k = k.encode('ascii') |
| 1326 | if isinstance(v, unicode): |
| 1327 | v = v.encode('ascii') |
| 1328 | env[k] = v |
| 1329 | return env |
| 1330 | |
| 1331 | cuda_libraries = ['cuda', 'cudnn'] |
| 1332 | if is_linux(): |
| 1333 | if int(environ_cp.get('TF_NEED_TENSORRT', False)): |
| 1334 | cuda_libraries.append('tensorrt') |
| 1335 | if environ_cp.get('TF_NCCL_VERSION', None): |
| 1336 | cuda_libraries.append('nccl') |
| 1337 | |
| 1338 | proc = subprocess.Popen( |
| 1339 | [environ_cp['PYTHON_BIN_PATH'], 'third_party/gpus/find_cuda_config.py'] + |
| 1340 | cuda_libraries, |
| 1341 | stdout=subprocess.PIPE, |
| 1342 | env=maybe_encode_env(environ_cp)) |
| 1343 | |
| 1344 | if proc.wait(): |
| 1345 | # Errors from find_cuda_config.py were sent to stderr. |
| 1346 | print('Asking for detailed CUDA configuration...\n') |
| 1347 | return False |
| 1348 | |
| 1349 | config = dict( |
| 1350 | tuple(line.decode('ascii').rstrip().split(': ')) for line in proc.stdout) |
| 1351 | |
| 1352 | print('Found CUDA %s in:' % config['cuda_version']) |
| 1353 | print(' %s' % config['cuda_library_dir']) |
| 1354 | print(' %s' % config['cuda_include_dir']) |
| 1355 | |
| 1356 | print('Found cuDNN %s in:' % config['cudnn_version']) |
| 1357 | print(' %s' % config['cudnn_library_dir']) |
| 1358 | print(' %s' % config['cudnn_include_dir']) |
| 1359 | |
| 1360 | if 'tensorrt_version' in config: |
| 1361 | print('Found TensorRT %s in:' % config['tensorrt_version']) |
| 1362 | print(' %s' % config['tensorrt_library_dir']) |
| 1363 | print(' %s' % config['tensorrt_include_dir']) |
| 1364 | |
| 1365 | if config.get('nccl_version', None): |
| 1366 | print('Found NCCL %s in:' % config['nccl_version']) |
| 1367 | print(' %s' % config['nccl_library_dir']) |
| 1368 | print(' %s' % config['nccl_include_dir']) |
| 1369 | |
| 1370 | print('\n') |
| 1371 | |
| 1372 | environ_cp['CUDA_TOOLKIT_PATH'] = config['cuda_toolkit_path'] |
| 1373 | return True |