0a2db3d354
This change is a second attempt at #38964, 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
39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
# Copyright 2019 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.
|
|
# ==============================================================================
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
from tensorflow.python.platform import googletest
|
|
from tensorflow.python.platform import sysconfig
|
|
from tensorflow.python.platform import test
|
|
|
|
|
|
class SysconfigTest(googletest.TestCase):
|
|
|
|
def test_get_build_info_works(self):
|
|
build_info = sysconfig.get_build_info()
|
|
self.assertIsInstance(build_info, dict)
|
|
|
|
def test_rocm_cuda_info_matches(self):
|
|
build_info = sysconfig.get_build_info()
|
|
self.assertEqual(build_info["is_rocm_build"], test.is_built_with_rocm())
|
|
self.assertEqual(build_info["is_cuda_build"], test.is_built_with_cuda())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
googletest.main()
|