Allow /dev/shm as sandboxed build base (#4460)

Author: yeounoh

build_torch_xla_libs.sh

@@ -35,8 +35,12 @@ if [[ "$XLA_BAZEL_VERBOSE" == "1" ]]; then
 fi
 
 BUILD_STRATEGY="standalone"
+SANDBOX_BASE="${XLA_SANDBOX_BASE}"
+if [ -z "$XLA_SANDBOX_BASE" ]; then
+  SANDBOX_BASE="/tmp"
+fi
 if [[ "$XLA_SANDBOX_BUILD" == "1" ]]; then
-  BUILD_STRATEGY="sandboxed --sandbox_base=/dev/shm"
+  BUILD_STRATEGY="sandboxed --sandbox_base=${SANDBOX_BASE}"
 else
   # We can remove this after https://github.com/bazelbuild/bazel/issues/15359 is resolved
   unset CC
@@ -84,7 +88,7 @@ else
   cp -r -u -p $THIRD_PARTY_DIR/xla_client $THIRD_PARTY_DIR/tensorflow/tensorflow/compiler/xla/
 
   pushd $THIRD_PARTY_DIR/tensorflow
-  # TensorFlow and its dependencies may introduce warning flags from newer compilers 
+  # TensorFlow and its dependencies may introduce warning flags from newer compilers
   # that PyTorch and PyTorch/XLA's default compilers don't recognize. They become error
   # while '-Werror' is used. Therefore, surpress the warnings.
   TF_EXTRA_FLAGS="--copt=-Wno-unknown-warning-option"

docker/Dockerfile

@@ -39,8 +39,9 @@ ENV BUNDLE_LIBTPU "${tpuvm}"
 # Maximum number of jobs to use for bazel build
 ENV BAZEL_JOBS "${bazel_jobs}"
 
-# This makes the bazel build behave more consistently, but it may slower.
+# This makes the bazel build behave more consistently, but runs slower.
 ENV XLA_SANDBOX_BUILD "1"
+ENV XLA_SANDBOX_BASE "/dev/shm"
 
 # To get around issue of Cloud Build with recursive submodule update
 # clone recursively from pytorch/xla if building docker image with

docker/cloudbuild.yaml

@@ -5,6 +5,7 @@ steps:
 - name: 'gcr.io/cloud-builders/docker'
   args: [
           'build',
+          '--shm-size', '16gb',
           '--build-arg', 'base_image=${_DOCKER_BASE_IMAGE}',
           '--build-arg', 'cuda=${_CUDA}',
           '--build-arg', 'python_version=${_PYTHON_VERSION}',

docker/experimental/Dockerfile

@@ -77,6 +77,7 @@ COPY build_torch_xla_libs.sh .
 
 # TODO: Remove this when it's not required anymore
 ENV XLA_SANDBOX_BUILD=1
+ENV XLA_SANDBOX_BASE "/dev/shm"
 
 ARG tpuvm
 ARG tf_cuda_compute_capabilities