experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

nccl_configure: introduce environment variable TF_NCCL_CONFIG_REPO

Browse Source 
TF_NCCL_CONFIG_REPO follows the same pattern as used in the other *_configure rules. If set TF_NCCL_CONFIG_REPO should point to a package with pregenerated configuration files.

PiperOrigin-RevId: 295804343
Change-Id: Ie1a69732fc3a538ccc3ed158c8ae79bda280514a
This commit is contained in:
Jakob Buchgraber 2020-02-18 13:10:12 -08:00 committed by TensorFlower Gardener
parent fbdfc9db01
commit ecb8befb32
1 changed files with 27 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
third_party/nccl/nccl_configure.bzl vendored
View File

@ -63,14 +63,7 @@ alias(
def _label(file):
return Label("//third_party/nccl:{}".format(file))
def _nccl_configure_impl(repository_ctx):
"""Implementation of the nccl_configure repository rule."""
if (not enable_cuda(repository_ctx) or
get_cpu_value(repository_ctx) not in ("Linux", "FreeBSD")):
# Add a dummy build file to make bazel query happy.
repository_ctx.file("BUILD", _NCCL_DUMMY_BUILD_CONTENT)
return
def _create_local_nccl_repository(repository_ctx):
# Resolve all labels before doing any real work. Resolving causes the
# function to be restarted with all previous state being lost. This
# can easily lead to a O(n^2) runtime in the number of labels.
@ -120,8 +113,33 @@ def _nccl_configure_impl(repository_ctx):
}
repository_ctx.template("BUILD", _label("system.BUILD.tpl"), config_wrap)
def _create_remote_nccl_repository(repository_ctx, remote_config_repo):
repository_ctx.template(
"BUILD",
Label(remote_config_repo + ":BUILD"),
{},
)
nccl_version = get_host_environ(repository_ctx, _TF_NCCL_VERSION, "")
if nccl_version == "":
repository_ctx.template(
"build_defs.bzl",
Label(remote_config_repo + ":build_defs.bzl"),
{},
)
def _nccl_autoconf_impl(repository_ctx):
if (not enable_cuda(repository_ctx) or
get_cpu_value(repository_ctx) not in ("Linux", "FreeBSD")):
# Add a dummy build file to make bazel query happy.
repository_ctx.file("BUILD", _NCCL_DUMMY_BUILD_CONTENT)
elif get_host_environ(repository_ctx, "TF_NCCL_CONFIG_REPO") != None:
_create_remote_nccl_repository(repository_ctx, get_host_environ(repository_ctx, "TF_NCCL_CONFIG_REPO"))
else:
_create_local_nccl_repository(repository_ctx)
nccl_configure = repository_rule(
implementation = _nccl_configure_impl,
implementation = _nccl_autoconf_impl,
environ = [
_CUDA_TOOLKIT_PATH,
_NCCL_HDR_PATH,