xref: /aosp_15_r20/external/pytorch/.ci/docker/common/install_cache.sh (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1#!/bin/bash
2
3set -ex
4
5install_ubuntu() {
6  echo "Preparing to build sccache from source"
7  apt-get update
8  # libssl-dev will not work as it is upgraded to libssl3 in Ubuntu-22.04.
9  # Instead use lib and headers from OpenSSL1.1 installed in `install_openssl.sh``
10  apt-get install -y cargo
11  echo "Checking out sccache repo"
12  git clone https://github.com/pytorch/sccache
13  cd sccache
14  echo "Building sccache"
15  cargo build --release
16  cp target/release/sccache /opt/cache/bin
17  echo "Cleaning up"
18  cd ..
19  rm -rf sccache
20  apt-get remove -y cargo rustc
21  apt-get autoclean && apt-get clean
22}
23
24install_binary() {
25  echo "Downloading sccache binary from S3 repo"
26  curl --retry 3 https://s3.amazonaws.com/ossci-linux/sccache -o /opt/cache/bin/sccache
27}
28
29mkdir -p /opt/cache/bin
30mkdir -p /opt/cache/lib
31sed -e 's|PATH="\(.*\)"|PATH="/opt/cache/bin:\1"|g' -i /etc/environment
32export PATH="/opt/cache/bin:$PATH"
33
34# Setup compiler cache
35if [ -n "$ROCM_VERSION" ]; then
36  curl --retry 3 http://repo.radeon.com/misc/.sccache_amd/sccache -o /opt/cache/bin/sccache
37else
38  ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
39  # TODO: Install the pre-built binary from S3 as building from source
40  # https://github.com/pytorch/sccache has started failing mysteriously
41  # in which sccache server couldn't start with the following error:
42  #   sccache: error: Invalid argument (os error 22)
43  install_binary
44fi
45chmod a+x /opt/cache/bin/sccache
46
47function write_sccache_stub() {
48  # Unset LD_PRELOAD for ps because of asan + ps issues
49  # https://gcc.gnu.org/bugzilla/show_bug.cgi?id=90589
50  printf "#!/bin/sh\nif [ \$(env -u LD_PRELOAD ps -p \$PPID -o comm=) != sccache ]; then\n  exec sccache $(which $1) \"\$@\"\nelse\n  exec $(which $1) \"\$@\"\nfi" > "/opt/cache/bin/$1"
51  chmod a+x "/opt/cache/bin/$1"
52}
53
54write_sccache_stub cc
55write_sccache_stub c++
56write_sccache_stub gcc
57write_sccache_stub g++
58
59# NOTE: See specific ROCM_VERSION case below.
60if [ "x$ROCM_VERSION" = x ]; then
61  write_sccache_stub clang
62  write_sccache_stub clang++
63fi
64
65if [ -n "$CUDA_VERSION" ]; then
66  # TODO: This is a workaround for the fact that PyTorch's FindCUDA
67  # implementation cannot find nvcc if it is setup this way, because it
68  # appears to search for the nvcc in PATH, and use its path to infer
69  # where CUDA is installed.  Instead, we install an nvcc symlink outside
70  # of the PATH, and set CUDA_NVCC_EXECUTABLE so that we make use of it.
71
72  write_sccache_stub nvcc
73  mv /opt/cache/bin/nvcc /opt/cache/lib/
74fi
75
76if [ -n "$ROCM_VERSION" ]; then
77  # ROCm compiler is hcc or clang. However, it is commonly invoked via hipcc wrapper.
78  # hipcc will call either hcc or clang using an absolute path starting with /opt/rocm,
79  # causing the /opt/cache/bin to be skipped. We must create the sccache wrappers
80  # directly under /opt/rocm while also preserving the original compiler names.
81  # Note symlinks will chain as follows: [hcc or clang++] -> clang -> clang-??
82  # Final link in symlink chain must point back to original directory.
83
84  # Original compiler is moved one directory deeper. Wrapper replaces it.
85  function write_sccache_stub_rocm() {
86    OLDCOMP=$1
87    COMPNAME=$(basename $OLDCOMP)
88    TOPDIR=$(dirname $OLDCOMP)
89    WRAPPED="$TOPDIR/original/$COMPNAME"
90    mv "$OLDCOMP" "$WRAPPED"
91    printf "#!/bin/sh\nexec sccache $WRAPPED \"\$@\"" > "$OLDCOMP"
92    chmod a+x "$OLDCOMP"
93  }
94
95  if [[ -e "/opt/rocm/hcc/bin/hcc" ]]; then
96    # ROCm 3.3 or earlier.
97    mkdir /opt/rocm/hcc/bin/original
98    write_sccache_stub_rocm /opt/rocm/hcc/bin/hcc
99    write_sccache_stub_rocm /opt/rocm/hcc/bin/clang
100    write_sccache_stub_rocm /opt/rocm/hcc/bin/clang++
101    # Fix last link in symlink chain, clang points to versioned clang in prior dir
102    pushd /opt/rocm/hcc/bin/original
103    ln -s ../$(readlink clang)
104    popd
105  elif [[ -e "/opt/rocm/llvm/bin/clang" ]]; then
106    # ROCm 3.5 and beyond.
107    mkdir /opt/rocm/llvm/bin/original
108    write_sccache_stub_rocm /opt/rocm/llvm/bin/clang
109    write_sccache_stub_rocm /opt/rocm/llvm/bin/clang++
110    # Fix last link in symlink chain, clang points to versioned clang in prior dir
111    pushd /opt/rocm/llvm/bin/original
112    ln -s ../$(readlink clang)
113    popd
114  else
115    echo "Cannot find ROCm compiler."
116    exit 1
117  fi
118fi
119