Use NVSHMEM tar files in RBE CUDA builds for XLA and Tensorflow.

PiperOrigin-RevId: 769761344
This commit is contained in:
A. Unique TensorFlower 2025-06-10 13:02:57 -07:00 committed by TensorFlower Gardener
parent 461c82dbad
commit 428f3034bb
3 changed files with 6 additions and 0 deletions

View File

@ -625,6 +625,7 @@ common:rbe_linux_cpu_old --remote_instance_name=projects/tensorflow-testing/inst
build:rbe_linux_cuda --config=cuda_clang_official
build:rbe_linux_cuda --config=rbe_linux_cpu
build:rbe_linux_cuda --repo_env=USE_CUDA_TAR_ARCHIVE_FILES=1
build:rbe_linux_cuda --repo_env=USE_NVSHMEM_TAR_ARCHIVE_FILES=1
# For Remote build execution -- GPU configuration
build:rbe_linux_cuda --repo_env=REMOTE_GPU_TESTING=1

View File

@ -479,6 +479,7 @@ common:rbe_linux_cpu --remote_instance_name=projects/tensorflow-testing/instance
build:rbe_linux_cuda --config=cuda_clang_official
build:rbe_linux_cuda --config=rbe_linux_cpu
build:rbe_linux_cuda --repo_env=USE_CUDA_TAR_ARCHIVE_FILES=1
build:rbe_linux_cuda --repo_env=USE_NVSHMEM_TAR_ARCHIVE_FILES=1
# For Remote build execution -- GPU configuration
build:rbe_linux_cuda --repo_env=REMOTE_GPU_TESTING=1

View File

@ -273,6 +273,10 @@ NVSHMEM_REDIST_JSON_DICT = {
}
MIRRORED_TARS_NVSHMEM_REDIST_JSON_DICT = {
"3.2.5": [
"https://storage.googleapis.com/mirror.tensorflow.org/developer.download.nvidia.com/compute/nvshmem/redist/redistrib_3.2.5_tar.json",
"641f7ca7048e4acfb466ce8be722f4828b2fa6b8671c28f6e8c230344484fd1c",
],
}
CUDA_12_NCCL_WHEEL_DICT = {