Migrate the remote config toolchains from tensorflow to the toolchain…

…s repo. PiperOrigin-RevId: 345158948 Change-Id: Id09fabbee85baaf2cfec6124457ddc1a13259f2e
ggkgg · Dec 2, 2020 · 0445859 · 0445859
1 parent ffe615e
commit 0445859
Show file tree

Hide file tree

Showing 7 changed files with 352 additions and 11 deletions.
diff --git a/WORKSPACE b/WORKSPACE
@@ -12,16 +12,6 @@ http_archive(
     ],
 )
 
-http_archive(
-    name = "tf_toolchains",
-    sha256 = "d60f9637c64829e92dac3f4477a2c45cdddb9946c5da0dd46db97765eb9de08e",
-    strip_prefix = "toolchains-1.1.5",
-    urls = [
-        "http://mirror.tensorflow.org/github.com/tensorflow/toolchains/archive/v1.1.5.tar.gz",
-        "https://github.com/tensorflow/toolchains/archive/v1.1.5.tar.gz",
-    ],
-)
-
 # Load tf_repositories() before loading dependencies for other repository so
 # that dependencies like com_google_protobuf won't be overridden.
 load("//tensorflow:workspace.bzl", "tf_repositories")
@@ -131,3 +121,4 @@ grpc_extra_deps()
 load("//third_party/googleapis:repository_rules.bzl", "config_googleapis")
 
 config_googleapis()
+
diff --git a/tensorflow/opensource_only.files b/tensorflow/opensource_only.files
@@ -272,6 +272,10 @@ tensorflow/third_party/toolchains/remote/BUILD
 tensorflow/third_party/toolchains/remote/BUILD.tpl
 tensorflow/third_party/toolchains/remote/configure.bzl
 tensorflow/third_party/toolchains/remote/execution.bzl.tpl
+tensorflow/third_party/toolchains/remote_config/BUILD
+tensorflow/third_party/toolchains/remote_config/configs.bzl
+tensorflow/third_party/toolchains/remote_config/containers.bzl
+tensorflow/third_party/toolchains/remote_config/rbe_config.bzl
 tensorflow/third_party/typing_extensions.BUILD
 tensorflow/third_party/wrapt.BUILD
 tensorflow/third_party/zlib.BUILD

diff --git a/tensorflow/workspace.bzl b/tensorflow/workspace.bzl
@@ -39,7 +39,7 @@ load("//third_party/psimd:workspace.bzl", psimd = "repo")
 load("//third_party/ruy:workspace.bzl", ruy = "repo")
 load("//third_party/sobol_data:workspace.bzl", sobol_data = "repo")
 load("//third_party/vulkan_headers:workspace.bzl", vulkan_headers = "repo")
-load("@tf_toolchains//toolchains/remote_config:configs.bzl", "initialize_rbe_configs")
+load("//third_party/toolchains/remote_config:configs.bzl", "initialize_rbe_configs")
 
 def initialize_third_party():
     """ Load third party repositories.  See above load() statements. """
@@ -1138,6 +1138,17 @@ def tf_repositories(path_prefix = "", tf_repo_name = ""):
         ],
     )
 
+    tf_http_archive(
+        name = "tf_toolchains",
+        sha256 = "eb175afa73e5a33d2b5d2aabcfde6c8c3395fd7001eb5ba765a5cd98cce714ba",
+        strip_prefix = "toolchains-0.0.2",
+        build_file = clean_dep("//third_party:tf_toolchains.BUILD"),
+        urls = [
+            "http://mirror.tensorflow.org/github.com/tensorflow/toolchains/archive/v0.0.2.tar.gz",
+            "https://github.com/tensorflow/toolchains/archive/v0.0.2.tar.gz",
+        ],
+    )
+
 def tf_bind():
     """Bind targets for some external repositories"""
     ##############################################################################

diff --git a/third_party/toolchains/remote_config/BUILD b/third_party/toolchains/remote_config/BUILD
diff --git a/third_party/toolchains/remote_config/configs.bzl b/third_party/toolchains/remote_config/configs.bzl
@@ -0,0 +1,109 @@
+"""Configurations of RBE builds used with remote config."""
+
+load("//third_party/toolchains/remote_config:rbe_config.bzl", "tensorflow_local_config", "tensorflow_rbe_config", "tensorflow_rbe_win_config")
+
+def initialize_rbe_configs():
+    tensorflow_local_config(
+        name = "local_execution",
+    )
+
+    tensorflow_rbe_config(
+        name = "ubuntu16.04-manylinux2010-py3",
+        os = "ubuntu16.04-manylinux2010",
+        python_versions = ["3"],
+        compiler = "",
+    )
+
+    tensorflow_rbe_config(
+        name = "ubuntu16.04-py3-gcc7_manylinux2010-cuda10.0-cudnn7-tensorrt5.1",
+        compiler = "/dt7/usr/bin/gcc",
+        compiler_prefix = "/usr/bin",
+        cuda_version = "10.0",
+        cudnn_version = "7",
+        os = "ubuntu16.04-manylinux2010",
+        python_versions = ["3"],
+        tensorrt_install_path = "/usr",
+        tensorrt_version = "5.1",
+    )
+
+    tensorflow_rbe_config(
+        name = "ubuntu16.04-gcc7_manylinux2010-cuda10.1-cudnn7-tensorrt6.0",
+        compiler = "/dt7/usr/bin/gcc",
+        compiler_prefix = "/usr/bin",
+        cuda_version = "10.1",
+        cudnn_version = "7",
+        os = "ubuntu16.04-manylinux2010-multipython",
+        python_versions = ["2.7", "3.5", "3.6", "3.7", "3.8"],
+        tensorrt_install_path = "/usr",
+        tensorrt_version = "6.0",
+        python_install_path = "/usr/local",
+    )
+
+    tensorflow_rbe_config(
+        name = "ubuntu18.04-gcc7_manylinux2010-cuda10.1-cudnn7-tensorrt6.0",
+        compiler = "/dt7/usr/bin/gcc",
+        compiler_prefix = "/usr/bin",
+        cuda_version = "10.1",
+        cudnn_version = "7",
+        os = "ubuntu18.04-manylinux2010-multipython",
+        python_versions = ["2.7", "3.5", "3.6", "3.7", "3.8"],
+        tensorrt_install_path = "/usr",
+        tensorrt_version = "6.0",
+        python_install_path = "/usr/local",
+    )
+
+    tensorflow_rbe_config(
+        name = "ubuntu18.04-gcc7_manylinux2010-cuda11.0-cudnn8-tensorrt7.1",
+        compiler = "/dt7/usr/bin/gcc",
+        compiler_prefix = "/usr/bin",
+        cuda_version = "11.0",
+        cudnn_version = "8",
+        os = "ubuntu18.04-manylinux2010-multipython",
+        python_versions = ["2.7", "3.5", "3.6", "3.7", "3.8"],
+        tensorrt_install_path = "/usr",
+        tensorrt_version = "7.1",
+        python_install_path = "/usr/local",
+    )
+
+    # TODO(klimek): Delete this once all users are migrated to a python-version
+    # independent configuration. In the future, use
+    # "ubuntu16.04-gcc7_manylinux2010-cuda10.1-cudnn7-tensorrt6.0" instead.
+    tensorflow_rbe_config(
+        name = "ubuntu16.04-py3-gcc7_manylinux2010-cuda10.1-cudnn7-tensorrt6.0",
+        compiler = "/dt7/usr/bin/gcc",
+        compiler_prefix = "/usr/bin",
+        cuda_version = "10.1",
+        cudnn_version = "7",
+        os = "ubuntu16.04-manylinux2010",
+        python_versions = ["3"],
+        tensorrt_install_path = "/usr",
+        tensorrt_version = "6.0",
+    )
+
+    tensorflow_rbe_config(
+        name = "ubuntu16.04-clang_manylinux2010-cuda10.1-cudnn7-tensorrt6.0",
+        compiler = "/clang_r42cab985fd95ba4f3f290e7bb26b93805edb447d/bin/clang",
+        cuda_version = "10.1",
+        cudnn_version = "7",
+        os = "ubuntu16.04-manylinux2010-multipython",
+        python_versions = ["2.7", "3.5", "3.6", "3.7", "3.8"],
+        tensorrt_install_path = "/usr",
+        tensorrt_version = "6.0",
+        sysroot = "/dt7",
+        python_install_path = "/usr/local",
+    )
+
+    tensorflow_rbe_config(
+        name = "ubuntu18.04-gcc7_manylinux2010-rocm",
+        compiler = "/dt7/usr/bin/gcc",
+        compiler_prefix = "/usr/bin",
+        rocm_version = "3.5",  # Any version will do.
+        os = "ubuntu18.04-manylinux2010-multipython",
+        python_versions = ["2.7", "3.5", "3.6", "3.7", "3.8"],
+        python_install_path = "/usr/local",
+    )
+
+    tensorflow_rbe_win_config(
+        name = "windows_py37",
+        python_bin_path = "C:/Python37/python.exe",
+    )
diff --git a/third_party/toolchains/remote_config/containers.bzl b/third_party/toolchains/remote_config/containers.bzl
@@ -0,0 +1,61 @@
+"""Docker images used with remote config and RBE."""
+
+load("//third_party/toolchains/preconfig/generate:containers.bzl", "container_digests")
+
+containers = {
+    # Built with //tensorflow/tools/ci_build/Dockerfile.rbe.ubuntu16.04-manylinux2010.
+    "ubuntu16.04-manylinux2010": {
+        "registry": "gcr.io",
+        "repository": "tensorflow-testing/nosla-ubuntu16.04-manylinux2010",
+        "digest": container_digests["ubuntu16.04-manylinux2010"],
+    },
+
+    # Built with //tensorflow/tools/ci_build/Dockerfile.rbe.cuda10.0-cudnn7-ubuntu16.04-manylinux2010.
+    "cuda10.0-cudnn7-ubuntu16.04-manylinux2010": {
+        "registry": "gcr.io",
+        "repository": "tensorflow-testing/nosla-cuda10.0-cudnn7-ubuntu16.04-manylinux2010",
+        "digest": container_digests["cuda10.0-cudnn7-ubuntu16.04-manylinux2010"],
+    },
+
+    # Built with //tensorflow/tools/ci_build/Dockerfile.rbe.cuda10.1-cudnn7-ubuntu16.04-manylinux2010.
+    "cuda10.1-cudnn7-ubuntu16.04-manylinux2010": {
+        "registry": "gcr.io",
+        "repository": "tensorflow-testing/nosla-cuda10.1-cudnn7-ubuntu16.04-manylinux2010",
+        "digest": container_digests["cuda10.1-cudnn7-ubuntu16.04-manylinux2010"],
+    },
+
+    # Built with //tensorflow/tools/ci_build/Dockerfile.rbe.cuda10.1-cudnn7-ubuntu16.04-manylinux2010-multipython.
+    "cuda10.1-cudnn7-ubuntu16.04-manylinux2010-multipython": {
+        "registry": "gcr.io",
+        "repository": "tensorflow-testing/nosla-cuda10.1-cudnn7-ubuntu16.04-manylinux2010-multipython",
+        "digest": container_digests["cuda10.1-cudnn7-ubuntu16.04-manylinux2010-multipython"],
+    },
+
+    # Built with //tensorflow/tools/ci_build/Dockerfile.rbe.cuda10.1-cudnn7-ubuntu18.04-manylinux2010-multipython.
+    "cuda10.1-cudnn7-ubuntu18.04-manylinux2010-multipython": {
+        "registry": "gcr.io",
+        "repository": "tensorflow-testing/nosla-cuda10.1-cudnn7-ubuntu18.04-manylinux2010-multipython",
+        "digest": container_digests["cuda10.1-cudnn7-ubuntu18.04-manylinux2010-multipython"],
+    },
+
+    # Built with //tensorflow/tools/ci_build/Dockerfile.rbe.cuda11.0-cudnn8-ubuntu18.04-manylinux2010-multipython.
+    "cuda11.0-cudnn8-ubuntu18.04-manylinux2010-multipython": {
+        "registry": "gcr.io",
+        "repository": "tensorflow-testing/nosla-cuda11.0-cudnn8-ubuntu18.04-manylinux2010-multipython",
+        "digest": container_digests["cuda11.0-cudnn8-ubuntu18.04-manylinux2010-multipython"],
+    },
+
+    # Built with //tensorflow/tools/ci_build/Dockerfile.rbe.rocm-ubuntu18.04-manylinux2010-multipython.
+    "rocm-ubuntu18.04-manylinux2010-multipython": {
+        "registry": "gcr.io",
+        "repository": "tensorflow-testing/nosla-rocm-ubuntu18.04-manylinux2010-multipython",
+        "digest": container_digests["rocm-ubuntu18.04-manylinux2010-multipython"],
+    },
+
+    # Built by gunan@ from a private Dockerfile.
+    "windows-1803": {
+        "registry": "gcr.io",
+        "repository": "tensorflow-testing/tf-win-rbe",
+        "digest": container_digests["windows-1803"],
+    },
+}
diff --git a/third_party/toolchains/remote_config/rbe_config.bzl b/third_party/toolchains/remote_config/rbe_config.bzl
@@ -0,0 +1,165 @@
+"""Macro that creates external repositories for remote config."""
+
+load("//third_party/py:python_configure.bzl", "local_python_configure", "remote_python_configure")
+load("//third_party/gpus:cuda_configure.bzl", "remote_cuda_configure")
+load("//third_party/nccl:nccl_configure.bzl", "remote_nccl_configure")
+load("//third_party/gpus:rocm_configure.bzl", "remote_rocm_configure")
+load("//third_party/tensorrt:tensorrt_configure.bzl", "remote_tensorrt_configure")
+load("//third_party/toolchains/remote_config:containers.bzl", "containers")
+load("//third_party/remote_config:remote_platform_configure.bzl", "remote_platform_configure")
+
+def _container_image_uri(container_name):
+    container = containers[container_name]
+    return "docker://%s/%s@%s" % (container["registry"], container["repository"], container["digest"])
+
+def _tensorflow_rbe_config(name, compiler, python_versions, os, rocm_version = None, cuda_version = None, cudnn_version = None, tensorrt_version = None, tensorrt_install_path = None, cudnn_install_path = None, compiler_prefix = None, sysroot = None, python_install_path = "/usr"):
+    if cuda_version != None and rocm_version != None:
+        fail("Specifying both cuda_version and rocm_version is not supported.")
+
+    env = {
+        "ABI_VERSION": "gcc",
+        "ABI_LIBC_VERSION": "glibc_2.19",
+        "BAZEL_COMPILER": compiler,
+        "BAZEL_HOST_SYSTEM": "i686-unknown-linux-gnu",
+        "BAZEL_TARGET_LIBC": "glibc_2.19",
+        "BAZEL_TARGET_CPU": "k8",
+        "BAZEL_TARGET_SYSTEM": "x86_64-unknown-linux-gnu",
+        "CC_TOOLCHAIN_NAME": "linux_gnu_x86",
+        "CC": compiler,
+        "CLEAR_CACHE": "1",
+        "HOST_CXX_COMPILER": compiler,
+        "HOST_C_COMPILER": compiler,
+    }
+
+    if cuda_version != None:
+        # The cuda toolchain currently contains its own C++ toolchain definition,
+        # so we do not fetch local_config_cc.
+        env.update({
+            "TF_NEED_CUDA": "1",
+            "TF_CUDA_CLANG": "1" if compiler.endswith("clang") else "0",
+            "TF_CUDA_COMPUTE_CAPABILITIES": "3.5,6.0",
+            "TF_ENABLE_XLA": "1",
+            "TF_CUDNN_VERSION": cudnn_version,
+            "TF_CUDA_VERSION": cuda_version,
+            "CUDNN_INSTALL_PATH": cudnn_install_path if cudnn_install_path != None else "/usr/lib/x86_64-linux-gnu",
+            "TF_NEED_TENSORRT": "1",
+            "TF_TENSORRT_VERSION": tensorrt_version,
+            "TENSORRT_INSTALL_PATH": tensorrt_install_path if tensorrt_install_path != None else "/usr/lib/x86_64-linux-gnu",
+            "GCC_HOST_COMPILER_PATH": compiler if not compiler.endswith("clang") else "",
+            "GCC_HOST_COMPILER_PREFIX": compiler_prefix if compiler_prefix != None else "/usr/bin",
+            "CLANG_CUDA_COMPILER_PATH": compiler if compiler.endswith("clang") else "",
+            "TF_SYSROOT": sysroot if sysroot else "",
+        })
+
+        container_name = "cuda%s-cudnn%s-%s" % (cuda_version, cudnn_version, os)
+        container_image = _container_image_uri(container_name)
+        exec_properties = {
+            "container-image": container_image,
+            "Pool": "default",
+        }
+
+        remote_cuda_configure(
+            name = "%s_config_cuda" % name,
+            environ = env,
+            exec_properties = exec_properties,
+        )
+
+        remote_nccl_configure(
+            name = "%s_config_nccl" % name,
+            environ = env,
+            exec_properties = exec_properties,
+        )
+
+        remote_tensorrt_configure(
+            name = "%s_config_tensorrt" % name,
+            environ = env,
+            exec_properties = exec_properties,
+        )
+    elif rocm_version != None:
+        # The rocm toolchain currently contains its own C++ toolchain definition,
+        # so we do not fetch local_config_cc.
+        env.update({
+            "TF_NEED_ROCM": "1",
+            "TF_ENABLE_XLA": "0",
+        })
+
+        container_name = "rocm-%s" % (os)
+        container_image = _container_image_uri(container_name)
+        exec_properties = {
+            "container-image": container_image,
+            "Pool": "default",
+        }
+
+        remote_rocm_configure(
+            name = "%s_config_rocm" % name,
+            environ = env,
+            exec_properties = exec_properties,
+        )
+    elif python_versions != None:
+        container_image = _container_image_uri(os)
+        exec_properties = {
+            "container-image": container_image,
+            "Pool": "default",
+        }
+
+    else:
+        fail("Neither cuda_version, rocm_version nor python_version specified.")
+
+    remote_platform_configure(
+        name = "%s_config_platform" % name,
+        platform = "linux",
+        platform_exec_properties = exec_properties,
+    )
+    for python_version in python_versions:
+        env.update({
+            "PYTHON_BIN_PATH": "%s/bin/python%s" % (python_install_path, python_version),
+        })
+
+        # For backwards compatibility do not add the python version to the name
+        # if we only create a single python configuration.
+        version = python_version if len(python_versions) > 1 else ""
+        remote_python_configure(
+            name = "%s_config_python%s" % (name, version),
+            environ = env,
+            exec_properties = exec_properties,
+            platform_constraint = "@%s_config_platform//:platform_constraint" % name,
+        )
+
+def _tensorflow_rbe_win_config(name, python_bin_path, container_name = "windows-1803"):
+    container_image = _container_image_uri(container_name)
+    exec_properties = {
+        "container-image": container_image,
+        "OSFamily": "Windows",
+    }
+
+    env = {
+        "PYTHON_BIN_PATH": python_bin_path,
+    }
+
+    remote_platform_configure(
+        name = "%s_config_platform" % name,
+        platform = "windows",
+        platform_exec_properties = exec_properties,
+    )
+
+    remote_python_configure(
+        name = "%s_config_python" % name,
+        environ = env,
+        exec_properties = exec_properties,
+        platform_constraint = "@%s_config_platform//:platform_constraint" % name,
+    )
+
+def _tensorflow_local_config(name):
+    remote_platform_configure(
+        name = "%s_config_platform" % name,
+        platform = "local",
+        platform_exec_properties = {},
+    )
+    local_python_configure(
+        name = "%s_config_python" % name,
+        platform_constraint = "@%s_config_platform//:platform_constraint" % name,
+    )
+
+tensorflow_rbe_config = _tensorflow_rbe_config
+tensorflow_rbe_win_config = _tensorflow_rbe_win_config
+tensorflow_local_config = _tensorflow_local_config