Enable setting remote configuration for cuda and python as an env variable

PiperOrigin-RevId: 159176334

2 files changed, 19 insertions, 10 deletions

diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl
index 77dc602fd9..83a377dde5 100644
--- a/third_party/gpus/cuda_configure.bzl
+++ b/third_party/gpus/cuda_configure.bzl
@@ -26,6 +26,7 @@ _TF_CUDA_VERSION = "TF_CUDA_VERSION"
 _TF_CUDNN_VERSION = "TF_CUDNN_VERSION"
 _CUDNN_INSTALL_PATH = "CUDNN_INSTALL_PATH"
 _TF_CUDA_COMPUTE_CAPABILITIES = "TF_CUDA_COMPUTE_CAPABILITIES"
+_TF_CUDA_CONFIG_REPO = "TF_CUDA_CONFIG_REPO"
 
 _DEFAULT_CUDA_VERSION = ""
 _DEFAULT_CUDNN_VERSION = ""
@@ -1001,8 +1002,7 @@ def _create_local_cuda_repository(repository_ctx):
                "%{cuda_toolkit_path}": cuda_config.cuda_toolkit_path,
        })
 
-
-def _create_remote_cuda_repository(repository_ctx):
+def _create_remote_cuda_repository(repository_ctx, remote_config_repo):
   """Creates pointers to a remotely configured repo set up to build with CUDA."""
   _tpl(repository_ctx, "cuda:build_defs.bzl",
        {
@@ -1013,10 +1013,10 @@ def _create_remote_cuda_repository(repository_ctx):
        })
   _tpl(repository_ctx, "cuda:remote.BUILD",
        {
-           "%{remote_cuda_repo}": repository_ctx.attr.remote_config_repo,
+           "%{remote_cuda_repo}": remote_config_repo,
        }, "cuda/BUILD")
   _tpl(repository_ctx, "crosstool:remote.BUILD", {
-           "%{remote_cuda_repo}": repository_ctx.attr.remote_config_repo,
+           "%{remote_cuda_repo}": remote_config_repo,
        }, "crosstool/BUILD")
 
 def _cuda_autoconf_impl(repository_ctx):
@@ -1024,8 +1024,12 @@ def _cuda_autoconf_impl(repository_ctx):
   if not _enable_cuda(repository_ctx):
     _create_dummy_repository(repository_ctx)
   else:
-    if repository_ctx.attr.remote_config_repo != "":
-      _create_remote_cuda_repository(repository_ctx)
+    if _TF_CUDA_CONFIG_REPO in repository_ctx.os.environ:
+      _create_remote_cuda_repository(repository_ctx,
+          repository_ctx.os.environ[_TF_CUDA_CONFIG_REPO])
+    elif repository_ctx.attr.remote_config_repo != "":
+      _create_remote_cuda_repository(repository_ctx,
+          repository_ctx.attr.remote_config_repo)
     else:
       _create_local_cuda_repository(repository_ctx)
 
@@ -1043,6 +1047,7 @@ cuda_configure = repository_rule(
         _TF_CUDA_VERSION,
         _TF_CUDNN_VERSION,
         _TF_CUDA_COMPUTE_CAPABILITIES,
+        _TF_CUDA_CONFIG_REPO,
     ],
 )
 
diff --git a/third_party/py/python_configure.bzl b/third_party/py/python_configure.bzl
index b4a98af7b6..19a6f1e749 100644
--- a/third_party/py/python_configure.bzl
+++ b/third_party/py/python_configure.bzl
@@ -13,6 +13,7 @@ _NUMPY_INCLUDE_PATH = "NUMPY_INCLUDE_PATH"
 _PYTHON_BIN_PATH = "PYTHON_BIN_PATH"
 _PYTHON_INCLUDE_PATH = "PYTHON_INCLUDE_PATH"
 _PYTHON_LIB_PATH = "PYTHON_LIB_PATH"
+_TF_PYTHON_CONFIG_REPO = "TF_PYTHON_CONFIG_REPO"
 
 
 def _tpl(repository_ctx, tpl, substitutions={}, out=None):
@@ -278,18 +279,20 @@ def _create_local_python_repository(repository_ctx):
     })
 
 
-def _create_remote_python_repository(repository_ctx):
+def _create_remote_python_repository(repository_ctx, remote_config_repo):
   """Creates pointers to a remotely configured repo set up to build with Python.
   """
   _tpl(repository_ctx, "remote.BUILD", {
-      "%{REMOTE_PYTHON_REPO}": repository_ctx.attr.remote_config_repo,
+      "%{REMOTE_PYTHON_REPO}": remote_config_repo,
   }, "BUILD")
 
 
 def _python_autoconf_impl(repository_ctx):
   """Implementation of the python_autoconf repository rule."""
-  if repository_ctx.attr.remote_config_repo != "":
-    _create_remote_python_repository(repository_ctx)
+  remote_config_repo = _get_env_var(repository_ctx, _TF_PYTHON_CONFIG_REPO,
+      repository_ctx.attr.remote_config_repo, False)
+  if remote_config_repo != "":
+    _create_remote_python_repository(repository_ctx, remote_config_repo)
   else:
     _create_local_python_repository(repository_ctx)
 
@@ -307,6 +310,7 @@ python_configure = repository_rule(
         _PYTHON_INCLUDE_PATH,
         _PYTHON_LIB_PATH,
         _NUMPY_INCLUDE_PATH,
+        _TF_PYTHON_CONFIG_REPO,
     ],
 )
 """Detects and configures the local Python.