Bazel build of JAX fails with missing dependency declarations

[yantaow]

I am trying build cuda-enabled JAX from source on a cluster with CentOS version7. In the jax home directory, I run:

python build/build.py --enable_cuda --cuda_path=$CUDA_HOME --cudnn_path=$CUDNN_HOME

Here are my specs:

• cuda version: 11.6
• cudnn version 7.5
• gcc version: 11.2.0
• Bazel binary path: ./bazel-5.1.1-linux-x86_64
• Bazel version: 5.1.1
• Python binary path: ~/.conda/envs/JAX/bin/python
• Python version: 3.9
• NumPy version: 1.21.5
• MKL-DNN enabled: yes
• Target CPU: x86_64
• Target CPU features: release
• CUDA enabled: yes

Here are my environment variables:

CC=/usr/local/app/compiler/gcc/11.2.0/bin/gcc CXX=/usr/local/app/compiler/gcc/11.2.0/bin/g++
INCLUDE=/usr/local/app/compiler/gcc/11.2.0/include LIBRARY_PATH=/usr/local/app/compiler/gcc/11.2.0/lib64:/usr/local/app/lib/nvidia/cuda/11.6.1/lib64:/usr/local/app/lib/nvidia/driver/510.47/lib64/nvidia LD_LIBRARY_PATH=/usr/local/app/compiler/gcc/11.2.0/lib64:/usr/local/app/lib/nvidia/cuda/11.6.1/lib64:/usr/local/app/lib/nvidia/driver/510.47/lib64/nvidia

The Bazel build throws this error:

ERROR: /lustre1/home/rice_cake/.cache/bazel/_bazel_rice_cake/8366d0a62cbb3b115627233e356374ab/external/zlib/BUILD.bazel:5:11: Compiling uncompr.c failed: undeclared inclusion(s) in rule '@zlib//:zlib': this rule is missing dependency declarations for the following files included by 'uncompr.c': '/usr/local/app/compiler/gcc/11.2.0/lib/gcc/x86_64-redhat-linux/11.2.0/include/stddef.h' '/usr/local/app/compiler/gcc/11.2.0/lib/gcc/x86_64-redhat-linux/11.2.0/include-fixed/limits.h' '/usr/local/app/compiler/gcc/11.2.0/lib/gcc/x86_64-redhat-linux/11.2.0/include-fixed/syslimits.h' '/usr/local/app/compiler/gcc/11.2.0/lib/gcc/x86_64-redhat-linux/11.2.0/include/stdarg.h'

It seems that Bazel is complaining that I do not declare dependency for the C++ library files. Is there a way to solve this?

The trouble is that these Bazel build files are automatically generated by jax/build/build.py, and I don't know how to fix them manually.

P.S. Building JAX from source is successful on my local machine. I went into the Bazel cache files on the cluster and checked the BUILD file, and it's the same as the on my local machine, so I'm a bit stuck.

Thanks very much for any help,
Rice