Add configuration rule for remote execution.

PiperOrigin-RevId: 221105195

WORKSPACE

@@ -14,6 +14,33 @@ load("@io_bazel_rules_closure//closure:defs.bzl", "closure_repositories")
 
 closure_repositories()
 
+http_archive(
+    name = "base_images_docker",
+    sha256 = "e2b1b7254270bb7605e814a9dbf6d1e4ae04a11136ff1714fbfdabe3f87f7cf9",
+    strip_prefix = "base-images-docker-12801524f867e657fbb5d1a74f31618aff181ac6",
+    urls = ["https://github.com/GoogleCloudPlatform/base-images-docker/archive/12801524f867e657fbb5d1a74f31618aff181ac6.tar.gz"],
+)
+
+http_archive(
+    name = "bazel_toolchains",
+    sha256 = "15b5858b1b5541ec44df31b94c3b8672815b31d71215a98398761ea9f4c4eedb",
+    strip_prefix = "bazel-toolchains-6200b238c9c2d137c0d9a7262c80cc71d98e692b",
+    urls = [
+        "https://github.com/bazelbuild/bazel-toolchains/archive/6200b238c9c2d137c0d9a7262c80cc71d98e692b.tar.gz",
+    ],
+)
+
+http_archive(
+    name = "io_bazel_rules_docker",
+    sha256 = "29d109605e0d6f9c892584f07275b8c9260803bf0c6fcb7de2623b2bedc910bd",
+    strip_prefix = "rules_docker-0.5.1",
+    urls = ["https://github.com/bazelbuild/rules_docker/archive/v0.5.1.tar.gz"],
+)
+
+load("//third_party/toolchains/preconfig/generate:workspace.bzl", "remote_config_workspace")
+
+remote_config_workspace()
+
 # We must check the bazel version before trying to parse any other BUILD
 # files, in case the parsing of those build files depends on the bazel
 # version we require here.
@@ -79,3 +106,4 @@ new_http_archive(
         "http://download.tensorflow.org/models/speech_commands_v0.01.zip",
     ],
 )
+

third_party/toolchains/BUILD

@@ -2,6 +2,8 @@ licenses(["restricted"])
 
 package(default_visibility = ["//visibility:public"])
 
+load("//third_party/toolchains/preconfig/generate:containers.bzl", "container_digests")
+
 # Platform for use with remote execution with
 # custom container based off RBE Ubuntu16_04
 # http://gcr.io/cloud-marketplace/google/rbe-ubuntu16-04
@@ -30,6 +32,6 @@ platform(
     remote_execution_properties = """
         properties: {
             name: "container-image"
-            value:"docker://gcr.io/asci-toolchain/nosla-cuda9.0-cudnn7-ubuntu14.04@sha256:c26138f4c38c754da2bad44a8a068523abf7fbd71d58a57ce92e5342c5431bf5"
-        }""",
+            value:"docker://gcr.io/asci-toolchain/nosla-cuda9.0-cudnn7-ubuntu14.04@%s"
+        }""" % container_digests["cuda9.0-cudnn7-ubuntu14.04"],
 )

third_party/toolchains/preconfig/generate/BUILD

@@ -0,0 +1,35 @@
+licenses(["restricted"])
+
+load(":generate.bzl", "tensorflow_rbe_config")
+
+tensorflow_rbe_config(
+    name = "ubuntu14.04-py3-gcc-cuda9.0-cudnn7-nccl2",
+    compiler = "gcc",
+    cuda_version = "9.0",
+    cudnn_version = "7",
+    python_version = "3",
+)
+
+tensorflow_rbe_config(
+    name = "ubuntu14.04-py3-clang-cuda9.0-cudnn7-nccl2",
+    compiler = "clang",
+    cuda_version = "9.0",
+    cudnn_version = "7",
+    python_version = "3",
+)
+
+tensorflow_rbe_config(
+    name = "ubuntu14.04-py3-gcc-cuda10.0-cudnn7-nccl2",
+    compiler = "gcc",
+    cuda_version = "10.0",
+    cudnn_version = "7",
+    python_version = "3",
+)
+
+tensorflow_rbe_config(
+    name = "ubuntu14.04-py3-clang-cuda10.0-cudnn7-nccl2",
+    compiler = "clang",
+    cuda_version = "10.0",
+    cudnn_version = "7",
+    python_version = "3",
+)

third_party/toolchains/preconfig/generate/containers.bzl

@@ -0,0 +1,4 @@
+container_digests = {
+    "cuda9.0-cudnn7-ubuntu14.04": "sha256:c26138f4c38c754da2bad44a8a068523abf7fbd71d58a57ce92e5342c5431bf5",
+    "cuda10.0-cudnn7-ubuntu14.04": "sha256:34c4a55e2376b300cdc2b903775fc32e62352f6e33f927df5653743324378bfc",
+}

third_party/toolchains/preconfig/generate/generate.bzl

@@ -0,0 +1,46 @@
+load(
+    "@bazel_toolchains//rules:docker_config.bzl",
+    "docker_toolchain_autoconfig",
+)
+
+def _tensorflow_rbe_config(name, cuda_version, cudnn_version, python_version, compiler):
+    docker_toolchain_autoconfig(
+        name = name,
+        base = "@cuda%s-cudnn%s-ubuntu14.04//image" % (cuda_version, cudnn_version),
+        bazel_version = "0.16.1",
+        config_repos = [
+            "local_config_cuda",
+            "local_config_python",
+            "local_config_nccl",
+        ],
+        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,
+            "PYTHON_BIN_PATH": "/usr/bin/python%s" % python_version,
+            "TF_NEED_CUDA": "1",
+            "TF_CUDA_CLANG": "1" if compiler == "clang" else "0",
+            "CLEAR_CACHE": "1",
+            "TF_CUDA_COMPUTE_CAPABILITIES": "3.0",
+            "TF_ENABLE_XLA": "1",
+            "TF_CUDNN_VERSION": cudnn_version,
+            "TF_CUDA_VERSION": cuda_version,
+            "NCCL_INSTALL_PATH": "/usr/lib",
+            "NCCL_HDR_PATH": "/usr/include",
+            "TF_NCCL_VERSION": "2",
+            "CUDNN_INSTALL_PATH": "/usr/lib/x86_64-linux-gnu",
+        },
+        # TODO(klimek): We should use the sources that we currently work on, not
+        # just the latest snapshot of tensorflow that is checked in.
+        git_repo = "https://github.com/tensorflow/tensorflow",
+        tags = ["manual"],
+        incompatible_changes_off = True,
+    )
+
+tensorflow_rbe_config = _tensorflow_rbe_config

third_party/toolchains/preconfig/generate/generate.sh

@@ -0,0 +1,85 @@
+#!/bin/bash
+#
+# Copyright 2018 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.
+# ==============================================================================
+
+TARGET="$1"
+OUTPUT="$2"
+
+if [[ -z "${TARGET}" || -z "${OUTPUT}" ]]; then
+  echo "Usage:"
+  echo "$0 <target> <output>"
+  exit 1
+fi
+
+TEMPDIR="$(mktemp -d)"
+ROOT="${PWD}"
+PKG="third_party/toolchains/preconfig"
+IFS='-' read -ra PLATFORM <<< "${TARGET}"
+OS="${PLATFORM[0]}"
+PY_VERSION="${PLATFORM[1]}"
+COMPILER="${PLATFORM[2]}"
+CUDA_VERSION="${PLATFORM[3]}"
+CUDNN_VERSION="${PLATFORM[4]}"
+NCCL_VERSION="${PLATFORM[5]}"
+
+if [[ "${COMPILER}" == "gcc" ]]; then
+  COMPILER="gcc-nvcc-${CUDA_VERSION}"
+fi
+
+echo "OS: ${OS}"
+echo "Python: ${PY_VERSION}"
+echo "Compiler: ${COMPILER}"
+echo "CUDA: ${CUDA_VERSION}"
+echo "CUDNN: ${CUDNN_VERSION}"
+echo "NCCL: ${NCCL_VERSION}"
+
+bazel build "${PKG}/generate:${TARGET}"
+cd "${TEMPDIR}"
+tar xvf "${ROOT}/bazel-bin/${PKG}/generate/${TARGET}_outputs.tar"
+
+# Delete all empty files: configurations leave empty files around when they are
+# unnecessary.
+find . -empty -delete
+
+# We build up the following directory structure with preconfigured packages:
+# <OS>/
+#   <CUDA>-<CUDNN>/
+#   <COMPILER>/
+#   <NCCL>/
+#   <PYTHON>/
+
+# Create our toplevel output directory for the OS.
+mkdir "${OS}"
+
+# Python:
+mv local_config_python "${OS}/${PY_VERSION}"
+
+# NCCL:
+mv local_config_nccl "${OS}/${NCCL_VERSION}"
+
+# Compiler:
+mv local_config_cuda/crosstool "${OS}/${COMPILER}"
+
+# CUDA:
+mv local_config_cuda "${OS}/${CUDA_VERSION}-${CUDNN_VERSION}"
+
+# Cleanup for copybara.
+find "${OS}" -name 'BUILD' -o -name '*.bzl' |xargs buildifier
+find "${OS}" -name 'BUILD' -o -name '*.bzl' |xargs -I {} mv {} {}.oss
+
+# Tar it up:
+tar cvf "${OUTPUT}" "${OS}"
+

third_party/toolchains/preconfig/generate/workspace.bzl

@@ -0,0 +1,25 @@
+load(
+    "@io_bazel_rules_docker//container:container.bzl",
+    "container_pull",
+    container_repositories = "repositories",
+)
+load(":containers.bzl", "container_digests")
+
+def _remote_config_workspace():
+    container_repositories()
+
+    container_pull(
+        name = "cuda9.0-cudnn7-ubuntu14.04",
+        registry = "gcr.io",
+        repository = "asci-toolchain/nosla-cuda9.0-cudnn7-ubuntu14.04",
+        digest = container_digests["cuda9.0-cudnn7-ubuntu14.04"],
+    )
+
+    container_pull(
+        name = "cuda10.0-cudnn7-ubuntu14.04",
+        registry = "gcr.io",
+        repository = "asci-toolchain/nosla-cuda10.0-cudnn7-ubuntu14.04",
+        digest = container_digests["cuda10.0-cudnn7-ubuntu14.04"],
+    )
+
+remote_config_workspace = _remote_config_workspace