STT-tensorflow/tensorflow/python/platform/sysconfig.py
Austin Anderson 0a2db3d354 Provide NVIDIA CUDA build data in metadata and API
This change is a second attempt at , which was rolled back because it was fragile.

First, cuda_configure.bzl templates a file with data it already pulled from get_cuda_config. gen_build_info loads that file to provide package
build information within TensorFlow:

    from tensorflow.python.platform import build_info
    print(build_info.build_info)
    {'cuda_version': '10.2', 'cudnn_version': '7', ... }

Also is exposed through tf.sysconfig.get_build_info(), a public API change.

setup.py pulls build_info into package metadata. The wheel's
long description ends with:

    TensorFlow 2.2.0 for NVIDIA GPUs was built with these platform
    and library versions:

      - NVIDIA CUDA 10.2
      - NVIDIA cuDNN 7
      - NVIDIA CUDA Compute Capabilities compute_30, compute_70 (etc.)

I set one of the new CUDA Classifiers, and add metadata to the "platform" tag:

    >>> import pkginfo
    >>> a = pkginfo.Wheel('./tf_nightly_gpu-2.1.0-cp36-cp36m-linux_x86_64.whl')
    >>> a.platforms
    ['cuda_version:10.2', 'cudnn_version:7', ...]

I'm not 100% confident this is the best way to accomplish this. It
still seems odd to import like this setup.py, even though it works, even in
an environment with TensorFlow installed. This method is much better than the old method as it uses data that was already gathered. It could be extended to gather tensorrt, nccl, etc. from other .bzl files, but I wanted to get feedback (and ensure this lands in 2.3) before designing something like that.

Currently tested only on Linux GPU (Remote Build) for Python 3.6. I'd
like to see more tests before merging.

The API is the same as the earlier change.

Resolves https://github.com/tensorflow/tensorflow/issues/38351.

PiperOrigin-RevId: 315018663
Change-Id: Idf68a8fe4d1585164d22b5870894c879537c280d
2020-06-05 16:44:24 -07:00

113 lines
3.4 KiB
Python

# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""System configuration library."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os.path as _os_path
import platform as _platform
from tensorflow.python.framework.versions import CXX11_ABI_FLAG as _CXX11_ABI_FLAG
from tensorflow.python.framework.versions import MONOLITHIC_BUILD as _MONOLITHIC_BUILD
from tensorflow.python.framework.versions import VERSION as _VERSION
from tensorflow.python.platform import build_info
from tensorflow.python.util.tf_export import tf_export
# pylint: disable=g-import-not-at-top
@tf_export('sysconfig.get_include')
def get_include():
"""Get the directory containing the TensorFlow C++ header files.
Returns:
The directory as string.
"""
# Import inside the function.
# sysconfig is imported from the tensorflow module, so having this
# import at the top would cause a circular import, resulting in
# the tensorflow module missing symbols that come after sysconfig.
import tensorflow as tf
return _os_path.join(_os_path.dirname(tf.__file__), 'include')
@tf_export('sysconfig.get_lib')
def get_lib():
"""Get the directory containing the TensorFlow framework library.
Returns:
The directory as string.
"""
import tensorflow as tf
return _os_path.join(_os_path.dirname(tf.__file__))
@tf_export('sysconfig.get_compile_flags')
def get_compile_flags():
"""Get the compilation flags for custom operators.
Returns:
The compilation flags.
"""
flags = []
flags.append('-I%s' % get_include())
flags.append('-D_GLIBCXX_USE_CXX11_ABI=%d' % _CXX11_ABI_FLAG)
return flags
@tf_export('sysconfig.get_link_flags')
def get_link_flags():
"""Get the link flags for custom operators.
Returns:
The link flags.
"""
is_mac = _platform.system() == 'Darwin'
ver = _VERSION.split('.')[0]
flags = []
if not _MONOLITHIC_BUILD:
flags.append('-L%s' % get_lib())
if is_mac:
flags.append('-ltensorflow_framework.%s' % ver)
else:
flags.append('-l:libtensorflow_framework.so.%s' % ver)
return flags
@tf_export('sysconfig.get_build_info')
def get_build_info():
"""Get a dictionary describing TensorFlow's build environment.
Values are generated when TensorFlow is compiled, and are static for each
TensorFlow package. The return value is a dictionary with string keys such as:
- cuda_version
- cudnn_version
- is_cuda_build
- is_rocm_build
- msvcp_dll_names
- nvcuda_dll_name
- cudart_dll_name
- cudnn_dll_name
Note that the actual keys and values returned by this function is subject to
change across different versions of TensorFlow or across platforms.
Returns:
A Dictionary describing TensorFlow's build environment.
"""
return build_info.build_info