From 54cc77eaf048ac207b80033e057851e55dedef4f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Nov 2018 09:31:53 -0800
Subject: [PATCH] Add configuration rule for remote execution.
PiperOrigin-RevId: 221105195
---
WORKSPACE | 28 ++++++
third_party/toolchains/BUILD | 6 +-
.../toolchains/preconfig/generate/BUILD | 35 ++++++++
.../preconfig/generate/containers.bzl | 4 +
.../preconfig/generate/generate.bzl | 46 ++++++++++
.../toolchains/preconfig/generate/generate.sh | 85 +++++++++++++++++++
.../preconfig/generate/workspace.bzl | 25 ++++++
7 files changed, 227 insertions(+), 2 deletions(-)
create mode 100644 third_party/toolchains/preconfig/generate/BUILD
create mode 100644 third_party/toolchains/preconfig/generate/containers.bzl
create mode 100644 third_party/toolchains/preconfig/generate/generate.bzl
create mode 100755 third_party/toolchains/preconfig/generate/generate.sh
create mode 100644 third_party/toolchains/preconfig/generate/workspace.bzl
diff --git a/WORKSPACE b/WORKSPACE
index 17961829a60..0c7bc085b51 100644
--- a/WORKSPACE
+++ b/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",
],
)
+
diff --git a/third_party/toolchains/BUILD b/third_party/toolchains/BUILD
index 3fa5e261d58..a7b4687c020 100644
--- a/third_party/toolchains/BUILD
+++ b/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"],
)
diff --git a/third_party/toolchains/preconfig/generate/BUILD b/third_party/toolchains/preconfig/generate/BUILD
new file mode 100644
index 00000000000..7e3e93d6004
--- /dev/null
+++ b/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",
+)
diff --git a/third_party/toolchains/preconfig/generate/containers.bzl b/third_party/toolchains/preconfig/generate/containers.bzl
new file mode 100644
index 00000000000..1f9e29d4402
--- /dev/null
+++ b/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",
+}
diff --git a/third_party/toolchains/preconfig/generate/generate.bzl b/third_party/toolchains/preconfig/generate/generate.bzl
new file mode 100644
index 00000000000..2fb3a94cdca
--- /dev/null
+++ b/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
diff --git a/third_party/toolchains/preconfig/generate/generate.sh b/third_party/toolchains/preconfig/generate/generate.sh
new file mode 100755
index 00000000000..37c5211278a
--- /dev/null
+++ b/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}"
+
diff --git a/third_party/toolchains/preconfig/generate/workspace.bzl b/third_party/toolchains/preconfig/generate/workspace.bzl
new file mode 100644
index 00000000000..f30c2f1ae63
--- /dev/null
+++ b/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