Hermetic CUDA repeatedly download cuda redistribution.

[yliu120]

Hey,

Currently when we use the hermetic cuda when compiling xla, the code in the cuda redistribution now repeated download cuda redistribution at almost every build. This is extremely annoying and because the bazel file can check whether the extracted file already exists... It doesn't need to download this again and again.

Can you basically fix that?

[ybaturina]

Hi @yliu120 the files shouldn't be downloaded again if the repositories are already created and there are no flag/env var changes that affect the repository rule.
Would you provide command examples please, so we could troubleshoot the issue?

[yliu120]

I think if you simply use the external, non-google, environment and build any XLA target. Then you modify the build file a bit, and rebuild that target, you will see. This is just very easily reproducible.

[ybaturina]

This is the command I use (the same flags and options are used in the XLA GPU presubmit job):

bazel build \
  --build_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only,-oneapi-only \
  --test_tag_filters=-no_oss,requires-gpu-nvidia,gpu,-rocm-only,-oneapi-only,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-amd,-requires-gpu-intel \
  --config=warnings \
  --config=rbe_linux_cuda_nvcc \
  --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 \
  --run_under=//build_tools/ci:parallel_gpu_execute \
  --//xla/tsl:ci_build \
  --color=yes \
  --test_output=errors \
  --verbose_failures \
  --keep_going --flaky_test_attempts=3 \
  --jobs=150 \
  --bes_upload_mode=fully_async \
  -- //xla/tests:reduce_hlo_test_nvgpu_any

I've executed it on a clean cache, and confirmed that CUDA archives were downloaded. Then I've removed one of the dependencies of //xla/tests:reduce_hlo_test_nvgpu_any target. On the second execution no CUDA archives were downloaded, but bazel build failed as expected.

I'm unable to reproduce the issue, this is why I need a full command that you are using to troubleshoot the problem.

[yliu120]

Weird. I also cannot reproduce this from the upstream HEAD.

But I ran a little old XLA version. Let me try a bit.

[yliu120]

OK. I tried a bunch of things internally and I think it links to our weird issues we see in our internal NFS system.

I tested a lot of setups and commits and branches. This only happens in some folder in the NFS file system but not the local file system. I think the theory would be bazel trying to read from the cache but somehow the NFS system tells it the cache doesn't exist because of FS issues.

https://cs.opensource.google/bazel/bazel/+/master:src/main/java/com/google/devtools/build/lib/bazel/repository/downloader/DownloadManager.java;l=68;drc=906b91b824092e1e5ec156afb5794e4a8a5119df;bpv=1;bpt=1

I read all the code here and i think if we can print out logs from this bazel download manager, we can have some explanation from bazel.

[ybaturina]

In this case I would expect that the cache is invalidated for all the repositories, not only CUDA ones. CUDA/CUDNN archives are the heaviest though, and thus they are the most noticeable in the logs.

While this issue affects your builds, I recommend using --config=use_tar_archive_files to save some time for the recreation of the repositories.

Overall, it seems that this is not a problem with the rules_ml_toolchain repository rules implementation, but rather Bazel itself. So I recommend creating an issue in Bazel repository.

[yliu120]

Python 3.11 will be used.
To select Python version, either set HERMETIC_PYTHON_VERSION env variable in
your shell:
  export HERMETIC_PYTHON_VERSION=3.12
OR pass it as an argument to bazel command directly or inside your .bazelrc
file:
  --repo_env=HERMETIC_PYTHON_VERSION=3.12
DEBUG: /root/xla-private/third_party/py/python_repo.bzl:87:10: 
=============================
Hermetic Python configuration:
Version: "3.11"
Kind: ""
Interpreter: "default" (provided by rules_python)
Requirements_lock label: "@//:requirements_lock_3_11.txt"
=====================================
DEBUG: /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/rules_ml_toolchain/third_party/gpus/nvidia_common_rules.bzl:296:10: Downloading and extracting https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-9.8.0.87_cuda12-archive.tar.xz
DEBUG: /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/rules_ml_toolchain/third_party/gpus/nvidia_common_rules.bzl:296:10: Downloading and extracting https://developer.download.nvidia.com/compute/cuda/redist/libcublas/linux-x86_64/libcublas-linux-x86_64-12.8.3.14-archive.tar.xz
Computing main repo mapping: 
    Fetching repository @@cuda_cublas; starting 11s
    Fetching repository @@cuda_cudnn; starting 11s
    Fetching /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/cuda_cudnn; Extracting cudnn-linux-x86_64-9.8.0.87_cuda12-archive.tar.xz 10s
    Fetching /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/cuda_cublas; Extracting libcublas-linux-x86_64-12.8.3.14-archive.tar.xz 10s

Yes I think this probably some bazel behavior. After I changed to another new directory, the repository cache still exists so in principle, it should not download again...

[yliu120]

Yes this is not specific to cudnn, this also happens for llvm.

[ybaturina]

Would you agree that this bug can be closed? According to the discussion above, the issue is not caused by rules_ml_toolchain.

[yliu120]

Yes I think it is not related to this package. I now confirm it is some bazel behavior.

I actually compiled bazel from source and add some logs in the DownloadManager.java file. Then I use the bazel-dev version to compile and I found that bazel is not able to find any tar.xz file in the repository cache so it goes ahead to download the file again.

Checksum is present for https://github.com/bazelbuild/rules_cc/releases/download/0.0.13/rules_cc-0.0.13.tar.gz
Checking in SHA-256 cache for /root/.cache/bazel/_bazel_root/3f64c2b1168bb75626afeec28663b347/external/rules_cc/temp1782993122967135130/rules_cc-0.0.13.tar.gz
java.io.FileNotFoundException: /root/.cache/bazel/_bazel_root/3f64c2b1168bb75626afeec28663b347/external/rules_cc/temp1782993122967135130/rules_cc-0.0.13.tar.gz (No such file or directory)
        at java.base/java.io.FileInputStream.open0(Native Method)
        at java.base/java.io.FileInputStream.open(FileInputStream.java:213)
        at java.base/java.io.FileInputStream.<init>(FileInputStream.java:152)
        at com.google.devtools.build.lib.unix.UnixFileSystem.createFileInputStream(UnixFileSystem.java:489)
        at com.google.devtools.build.lib.vfs.AbstractFileSystem.createMaybeProfiledInputStream(AbstractFileSystem.java:96)
        at com.google.devtools.build.lib.vfs.AbstractFileSystem.getInputStream(AbstractFileSystem.java:59)
        at com.google.devtools.build.lib.vfs.Path.getInputStream(Path.java:763)
        at com.google.devtools.build.lib.bazel.repository.cache.RepositoryCache.getChecksum(RepositoryCache.java:387)
        at com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager.downloadInExecutor(DownloadManager.java:213)
        at com.google.devtools.build.lib.bazel.repository.downloader.DownloadManager.lambda$startDownload$0(DownloadManager.java:127)
        at java.base/java.util.concurrent.FutureTask.run(FutureTask.java:317)
        at java.base/java.lang.VirtualThread.run(VirtualThread.java:329)

Then I commented out the repository_ctx.delete() call from this package. I found all the tar.xz files from cuda are kepted in the repository cache.

root@memx-aar-55-sr1:~/bazel# ls /root/.cache/bazel/_bazel_root/3f64c2b1168bb75626afeec28663b347/external/cuda_nvcc/
BUILD  LICENSE  REPO.bazel  WORKSPACE  bin  cuda_nvcc-linux-x86_64-12.8.61-archive.tar.xz  include  lib  nvvm  version.bzl

Then I modify a BUILD file (a random BUILD file in XLA) and rerun the command, I stopped before the hermetic cuda kicks in but I found bazel removed all the tar.xz files automatically from the cache,

root@memx-aar-55-sr1:~/bazel# ls /root/.cache/bazel/_bazel_root/3f64c2b1168bb75626afeec28663b347/external/cuda_nvcc/
BUILD  REPO.bazel  WORKSPACE  version.bzl
root@memx-aar-55-sr1:~/bazel# ls /root/.cache/bazel/_bazel_root/3f64c2b1168bb75626afeec28663b347/external/cuda_cublas/
BUILD  REPO.bazel  WORKSPACE  version.bzl

Then the DownloadManager won't be able to see them even I intended to save all the tar.xz files.

I found the bazel repo cache behavior is very very weird and hard to understand. How come it can just automatically delete the archive files? I think it is not only removing all the archive files, it is basically just removed all extracted files from the repo cache.

And this doesn't relate to only cuda packages. For all the llvm-raw, it redownloads again.

[ybaturina]

There is indeed a code that deletes .tar.xz files from the hermetic CUDA repositories: link

However, the repository shouldn't be recreated (e.g. the archives shouldn't be downloaded and unpacked again) unless there are changes in the env var values that the repository rule uses (such as HERMETIC_CUDA_VERSION, TF_NEED_CUDA etc.) or in the rule implementation code.

Could it be that the bazel cache dir has root-only permissions?

I genuinely can't reproduce the error on my Linux machine.

What might be helpful: would you send me the Bazel log for the case when the archives are downloaded again though they shouldn't be?

[yliu120]

At the very beginning, when you clone the repo and build. And then modifying something in BUILD file (for instance, adding some empty lines), then rebuild, it won't do this. So in this way, we cannot repro.

But after we built some GPU targets, doing some real things. This problem shows up.

I fully understand the logic of bazel on repository download now. I think first it checks the output_base and then repo cache, if both don't have files, it will start to download from the webserver.

However, we do see that we hit repo cache. but the cost of Extracting and copying is still huge for these CUDA libraries.

I then did an experiment to use chattr +i chattr +i /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/cuda_cublas/lib/ to make the external repo inside output_base non-removable. Then I rerun bazel build,

DEBUG: /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/rules_ml_toolchain/third_party/gpus/nvidia_common_rules.bzl:296:10: Downloading and extracting https://developer.download.nvidia.com/compute/cuda/redist/cuda_cupti/linux-x86_64/cuda_cupti-linux-x86_64-12.8.57-archive.tar.xz
DEBUG: /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/rules_ml_toolchain/third_party/gpus/nvidia_common_rules.bzl:296:10: Downloading and extracting https://developer.download.nvidia.com/compute/cuda/redist/cuda_cccl/linux-x86_64/cuda_cccl-linux-x86_64-12.8.55-archive.tar.xz
DEBUG: /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/rules_ml_toolchain/third_party/gpus/nvidia_common_rules.bzl:296:10: Downloading and extracting https://developer.download.nvidia.com/compute/cuda/redist/libnvjitlink/linux-x86_64/libnvjitlink-linux-x86_64-12.8.61-archive.tar.xz
DEBUG: /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/rules_ml_toolchain/third_party/gpus/nvidia_common_rules.bzl:296:10: Downloading and extracting https://developer.download.nvidia.com/compute/cuda/redist/libcusparse/linux-x86_64/libcusparse-linux-x86_64-12.5.7.53-archive.tar.xz
DEBUG: /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/rules_ml_toolchain/third_party/gpus/nvidia_common_rules.bzl:296:10: Downloading and extracting https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64-12.8.61-archive.tar.xz
DEBUG: /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/rules_ml_toolchain/third_party/gpus/nvidia_common_rules.bzl:296:10: Downloading and extracting https://developer.download.nvidia.com/compute/cuda/redist/libcurand/linux-x86_64/libcurand-linux-x86_64-10.3.9.55-archive.tar.xz
DEBUG: /local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/rules_ml_toolchain/third_party/gpus/nvidia_common_rules.bzl:296:10: Downloading and extracting https://developer.download.nvidia.com/compute/cuda/redist/libcufft/linux-x86_64/libcufft-linux-x86_64-11.3.3.41-archive.tar.xz
ERROR: /root/xla-private/WORKSPACE:74:30: fetching _redist_repo rule //external:cuda_cublas: com.google.devtools.build.lib.unix.FilePermissionException: fchmod (/local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/cuda_cublas/lib) (Operation not permitted)
ERROR: Error computing the main repository mapping: no such package '@@cuda_cublas//': fchmod (/local/bazel_user_root/3f64c2b1168bb75626afeec28663b347/external/cuda_cublas/lib) (Operation not permitted)
Computing main repo mapping: 
    Fetching repository @@cuda_cccl; starting
    Fetching repository @@cuda_cudnn; starting
    Fetching repository @@cuda_cupti; starting
    Fetching repository @@cuda_cufft; starting
    Fetching repository @@cuda_cusparse; starting
    Fetching repository @@cuda_cusolver; starting
    Fetching repository @@cuda_curand; starting
    Fetching repository @@cuda_nvcc; starting ... (20 fetches)

Then I got the above error.

[yliu120]

What I meant is that if bazel doesn't remove files in output_base when it recomputes the repo mapping. It won't re-fetch the external repo.

[yliu120]

And now I am pretty sure that only the cuda libraries are re-fetched...
let's say if I set other external repo as non-removable, then the command runs fine. It means bazel doesn't try to remove them.