Add Hermetic C++ Toolchains for Linux x86_64 builds.

Hermetic toolchains provide builds that are isolated from the host system, cutting down on unexpected dependencies and side effects.

By default, JAX now builds for Linux x86_64 architectures (both CPU and CUDA-enabled GPU) using hermetic C++ toolchains. For non-hermetic builds, add the flag --config=clang_local. For remote builds with a non-hermetic toolchain, simply append _clang_local to the existing RBE flag. For example, if hermetic RBE build runs with --config=rbe_linux_cpu, the non-hermetic version would be --config=rbe_linux_cpu_clang_local.

    Example 1: Run CPU tests for Linux x86_64

    For hermetic tests run the command for Linux x86_64 build without env variables CC, CXX, etc.:

	bazel test //tests:cpu_tests \
		--config=avx_posix \
		--config=clang \
		--linkopt=-lrt \
		--host_linkopt=-lrt \
		--repo_env=HERMETIC_PYTHON_VERSION="3.13" \
		--test_tag_filters=-multiaccelerator

    For non-hermetic tests use commands with the flag "--config=clang_local" and env variables CC, CXX, etc.:

	bazel test //tests:cpu_tests \
		--config=clang_local \
		--config=avx_posix \
		--config=clang \
		--linkopt=-lrt \
		--host_linkopt=-lrt \
		--repo_env=HERMETIC_PYTHON_VERSION="3.13" \
		--test_tag_filters=-multiaccelerator \
		--action_env=CLANG_COMPILER_PATH=/usr/lib/llvm-18/bin/clang \
		--repo_env=CC=/usr/lib/llvm-18/bin/clang \
		--repo_env=CXX=/usr/lib/llvm-18/bin/clang++ \
		--repo_env=BAZEL_COMPILER=/usr/lib/llvm-18/bin/clang \
		--action_env=CCC_OVERRIDE_OPTIONS="^--gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/13"

    Example 2: Run CUDA GPU tests for Linux x86_64

    For hermetic tests, run the command without env variables CC, CXX, etc.:
	bazel test //tests:gpu_tests \
		--config=cuda \
		--config=clang \
		--config=build_cuda_with_nvcc \
		--config=avx_posix \
		--linkopt=-lrt \
		--host_linkopt=-lrt \
		--test_tag_filters=-multiaccelerator \
		--repo_env=HERMETIC_PYTHON_VERSION=3.13

    For non-hermetic tests use commands with the flag "--config=clang_local" and env variables CC, CXX, etc.:

	bazel test //tests:gpu_tests \
		--config=clang_local \
		--config=avx_posix \
		--config=build_cuda_with_nvcc \
		--config=clang \
		--config=cuda \
		--linkopt=-lrt \
		--host_linkopt=-lrt \
		--action_env=CCC_OVERRIDE_OPTIONS="^--gcc-install-dir=/usr/lib/gcc/x86_64-linux-gnu/9" \
		--test_tag_filters=-multiaccelerator \
		--action_env=CLANG_COMPILER_PATH=/usr/lib/llvm-18/bin/clang \
		--repo_env=CC=/usr/lib/llvm-18/bin/clang \
		--repo_env=CXX=/usr/lib/llvm-18/bin/clang++ \
		--repo_env=BAZEL_COMPILER=/usr/lib/llvm-18/bin/clang \
		--repo_env=HERMETIC_PYTHON_VERSION=3.13

    Example 3: Run RBE build for Linux x86_64

    For hermetic build, run the same command as before modification:
    bazel build \
        --config=ci_linux_x86_64 \
        --repo_env=HERMETIC_PYTHON_VERSION=3.13 \
        //jaxlib/tools:jaxlib_wheel

    For non-hermetic builds, add "_clang_local" suffix to RBE configuration:
    bazel build \
        --config=ci_linux_x86_64_clang_local \
        --repo_env=HERMETIC_PYTHON_VERSION=3.13 \
        //jaxlib/tools:jaxlib_wheel

    We don't support other "rbe_*_clang_local" configuration flags, and you could construct your own combinations.

PiperOrigin-RevId: 774963837

Author: Google-ML-Automation

.bazelrc

@@ -4,9 +4,6 @@
 # TODO: Enable Bzlmod
 common --noenable_bzlmod
 
-# TODO: Migrate for https://github.com/bazelbuild/bazel/issues/7260
-common --noincompatible_enable_cc_toolchain_resolution
-
 # Make Bazel print out all options from rc files.
 common --announce_rc
 
@@ -19,6 +16,13 @@ build --enable_platform_specific_config
 
 common --experimental_cc_shared_library
 
+build --incompatible_enable_cc_toolchain_resolution
+build --repo_env USE_HERMETIC_CC_TOOLCHAIN=1
+
+# TODO: Migrate for https://github.com/bazelbuild/bazel/issues/7260
+build:clang_local --noincompatible_enable_cc_toolchain_resolution
+build:clang_local --repo_env USE_HERMETIC_CC_TOOLCHAIN=0
+
 # Do not use C-Ares when building gRPC.
 build --define=grpc_no_ares=true
 
@@ -61,10 +65,13 @@ build:macos --host_linkopt=-Wl,-undefined,dynamic_lookup
 
 # Use cc toolchains from apple_support for Apple builds.
 # https://github.com/bazelbuild/apple_support/tree/master?tab=readme-ov-file#bazel-6-setup
+build:macos --config=clang_local
 build:macos --apple_crosstool_top=@local_config_apple_cc//:toolchain
 build:macos --crosstool_top=@local_config_apple_cc//:toolchain
 build:macos --host_crosstool_top=@local_config_apple_cc//:toolchain
 
+build:windows --config=clang_local
+
 # Windows has a relatively short command line limit, which JAX has begun to hit.
 # See https://docs.bazel.build/versions/main/windows.html
 build:windows --features=compiler_param_file
@@ -159,7 +166,6 @@ build:cuda --repo_env TF_NCCL_USE_STUB=1
 # "sm" means we emit only cubin, which is forward compatible within a GPU generation.
 # "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations.
 build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,sm_90,sm_100,compute_120"
-build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
 build:cuda --@local_config_cuda//:enable_cuda
 
 # Default hermetic CUDA, CUDNN and NVSHMEM versions.
@@ -194,10 +200,12 @@ build:build_cuda_with_nvcc --action_env=TF_NVCC_CLANG="1"
 build:build_cuda_with_nvcc --@local_config_cuda//:cuda_compiler=nvcc
 
 # Requires MSVC and LLVM to be installed
+build:win_clang --config=clang_local
 build:win_clang --extra_toolchains=@local_config_cc//:cc-toolchain-x64_windows-clang-cl
 build:win_clang --extra_execution_platforms=//jax/tools/toolchains:x64_windows-clang-cl
 build:win_clang --compiler=clang-cl
 
+build:rocm_base --config=clang_local
 build:rocm_base --crosstool_top=@local_config_rocm//crosstool:toolchain
 build:rocm_base --define=using_rocm=true --define=using_rocm_hipcc=true
 build:rocm_base --repo_env TF_NEED_ROCM=1
@@ -241,15 +249,18 @@ build:ci_linux_x86_64 --config=mkl_open_source_only
 build:ci_linux_x86_64 --config=clang --verbose_failures=true
 build:ci_linux_x86_64 --color=yes
 
+# Deprecated CI config with non-hermetic toolchains.
 # TODO(b/356695103): We do not have a CPU only toolchain so we use the CUDA
 # toolchain for both CPU and GPU builds.
-build:ci_linux_x86_64 --host_crosstool_top="@local_config_cuda//crosstool:toolchain"
-build:ci_linux_x86_64 --crosstool_top="@local_config_cuda//crosstool:toolchain"
-build:ci_linux_x86_64 --extra_toolchains="@local_config_cuda//crosstool:toolchain-linux-x86_64"
-build:ci_linux_x86_64 --repo_env=TF_SYSROOT="/dt9"
+build:ci_linux_x86_64_clang_local --config=ci_linux_x86_64
+build:ci_linux_x86_64_clang_local --config=clang_local
+build:ci_linux_x86_64_clang_local --host_crosstool_top="@local_config_cuda//crosstool:toolchain"
+build:ci_linux_x86_64_clang_local --crosstool_top="@local_config_cuda//crosstool:toolchain"
+build:ci_linux_x86_64_clang_local --extra_toolchains="@local_config_cuda//crosstool:toolchain-linux-x86_64"
+build:ci_linux_x86_64_clang_local --repo_env=TF_SYSROOT="/dt9"
 
 # Clang path needs to be set for remote toolchain to be configured correctly.
-build:ci_linux_x86_64 --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-18/bin/clang"
+build:ci_linux_x86_64_clang_local --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-18/bin/clang"
 
 # The toolchain in `--config=cuda` needs to be read before the toolchain in
 # `--config=ci_linux_x86_64`. Otherwise, we run into issues with manylinux
@@ -258,6 +269,7 @@ build:ci_linux_x86_64_cuda --config=cuda --config=build_cuda_with_nvcc
 build:ci_linux_x86_64_cuda --config=ci_linux_x86_64
 
 # Linux Aarch64 CI configs
+build:ci_linux_aarch64_base --config=clang_local
 build:ci_linux_aarch64_base --config=clang --verbose_failures=true
 build:ci_linux_aarch64_base --action_env=TF_SYSROOT="/dt10"
 build:ci_linux_aarch64_base --color=yes
@@ -279,12 +291,14 @@ build:ci_linux_aarch64_cuda --config=cuda --config=build_cuda_with_nvcc
 build:ci_linux_aarch64_cuda --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-18/bin/clang"
 
 # Mac Arm64 CI configs
+build:ci_darwin_arm64 --config=clang_local
 build:ci_darwin_arm64 --macos_minimum_os=11.0
 build:ci_darwin_arm64 --config=macos_cache_push
 build:ci_darwin_arm64 --verbose_failures=true
 build:ci_darwin_arm64 --color=yes
 
 # Windows x86 CI configs
+build:ci_windows_amd64 --config=clang_local
 build:ci_windows_amd64 --config=avx_windows
 build:ci_windows_amd64 --compiler=clang-cl --config=clang --verbose_failures=true
 build:ci_windows_amd64 --crosstool_top="@xla//tools/toolchains/win2022/20241118:toolchain"
@@ -356,6 +370,7 @@ build:rbe_linux_x86_64_cuda --repo_env=USE_NVSHMEM_TAR_ARCHIVE_FILES=1
 # Set the remote worker pool
 common:rbe_windows_amd64 --remote_instance_name=projects/tensorflow-testing/instances/windows
 
+build:rbe_windows_amd64 --config=clang_local
 build:rbe_windows_amd64 --config=rbe
 
 # Set the host, execution, and target platform
@@ -382,6 +397,7 @@ build:rbe_windows_amd64 --config=ci_windows_amd64
 # flags seem to be actually used to specify the execution platform details. It
 # seems it is this way because these flags are old and predate the distinction
 # between host and execution platform.
+build:cross_compile_base --config=clang_local
 build:cross_compile_base --host_cpu=k8
 build:cross_compile_base --host_crosstool_top=@xla//tools/toolchains/cross_compile/cc:cross_compile_toolchain_suite
 build:cross_compile_base --extra_execution_platforms=@xla//tools/toolchains/cross_compile/config:linux_x86_64

WORKSPACE

@@ -54,17 +54,6 @@ load("@pypi//:requirements.bzl", "install_deps")
 
 install_deps()
 
-# Optional, to facilitate testing against newest versions of Python
-load("@xla//third_party/py:python_repo.bzl", "custom_python_interpreter")
-
-custom_python_interpreter(
-    name = "python_dev",
-    strip_prefix = "Python-{version_variant}",
-    urls = ["https://www.python.org/ftp/python/{version}/Python-{version_variant}.tgz"],
-    version = "3.13.0",
-    version_variant = "3.13.0rc2",
-)
-
 load("@xla//:workspace2.bzl", "xla_workspace2")
 
 xla_workspace2()
@@ -99,7 +88,18 @@ python_wheel_version_suffix_repository(
 )
 
 load(
-    "@xla//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
+    "@rules_ml_toolchain//cc_toolchain/deps:cc_toolchain_deps.bzl",
+    "cc_toolchain_deps",
+)
+
+cc_toolchain_deps()
+
+register_toolchains("@rules_ml_toolchain//cc_toolchain:lx64_lx64")
+
+register_toolchains("@rules_ml_toolchain//cc_toolchain:lx64_lx64_cuda")
+
+load(
+    "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
     "cuda_json_init_repository",
 )
 
@@ -111,7 +111,7 @@ load(
     "CUDNN_REDISTRIBUTIONS",
 )
 load(
-    "@xla//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
+    "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
     "cuda_redist_init_repositories",
     "cudnn_redist_init_repository",
 )
@@ -125,28 +125,28 @@ cudnn_redist_init_repository(
 )
 
 load(
-    "@xla//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
+    "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
     "cuda_configure",
 )
 
 cuda_configure(name = "local_config_cuda")
 
 load(
-    "@xla//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
+    "@rules_ml_toolchain//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
     "nccl_redist_init_repository",
 )
 
 nccl_redist_init_repository()
 
 load(
-    "@xla//third_party/nccl/hermetic:nccl_configure.bzl",
+    "@rules_ml_toolchain//third_party/nccl/hermetic:nccl_configure.bzl",
     "nccl_configure",
 )
 
 nccl_configure(name = "local_config_nccl")
 
 load(
-    "@xla//third_party/nvshmem/hermetic:nvshmem_json_init_repository.bzl",
+    "@rules_ml_toolchain//third_party/nvshmem/hermetic:nvshmem_json_init_repository.bzl",
     "nvshmem_json_init_repository",
 )
 
@@ -157,7 +157,7 @@ load(
     "NVSHMEM_REDISTRIBUTIONS",
 )
 load(
-    "@xla//third_party/nvshmem/hermetic:nvshmem_redist_init_repository.bzl",
+    "@rules_ml_toolchain//third_party/nvshmem/hermetic:nvshmem_redist_init_repository.bzl",
     "nvshmem_redist_init_repository",
 )
 
@@ -166,7 +166,7 @@ nvshmem_redist_init_repository(
 )
 
 load(
-    "@xla//third_party/nvshmem/hermetic:nvshmem_configure.bzl",
+    "@rules_ml_toolchain//third_party/nvshmem/hermetic:nvshmem_configure.bzl",
     "nvshmem_configure",
 )
 

ci/utilities/run_auditwheel.sh

@@ -37,14 +37,16 @@ for wheel in $WHEELS; do
     wheel_name=$(basename $wheel)
     OUTPUT=${OUTPUT_FULL//${wheel_name}/}
 
-    # If a wheel is manylinux2014 compliant, `auditwheel show` will return the
-    # platform tag as manylinux_2_17. manylinux2014 is an alias for
-    # manylinux_2_17.
-    if echo "$OUTPUT" | grep -q "manylinux_2_17"; then
+    # If a wheel is manylinux_2_27 or manylinux2014 compliant, `auditwheel show`
+    # will return platform tag as manylinux_2_27 or manylinux_2_17 respectively.
+    # manylinux2014 is an alias for manylinux_2_17.
+    if echo "$OUTPUT" | grep -q "manylinux_2_27"; then
+        printf "\n$wheel_name is manylinux_2_27 compliant.\n"
+    elif echo "$OUTPUT" | grep -q "manylinux_2_17"; then
         printf "\n$wheel_name is manylinux2014 compliant.\n"
     else
         echo "$OUTPUT_FULL"
-        printf "\n$wheel_name is NOT manylinux2014 compliant.\n"
+        printf "\n$wheel_name is NOT manylinux_2_27 or manylinux2014 compliant.\n"
         exit 1
     fi
 done

docs/developer.md

@@ -299,7 +299,7 @@ and re-run the requirements updater command for a selected version of Python.
 For example:
 
 ```
-echo -e "\n$(realpath jaxlib-0.4.27.dev20240416-cp312-cp312-manylinux2014_x86_64.whl)" >> build/requirements.in
+echo -e "\n$(realpath jaxlib-0.4.27.dev20240416-cp312-cp312-manylinux_2_27_x86_64.whl)" >> build/requirements.in
 python build/build.py requirements_update --python_version=3.12
 ```
 
@@ -542,7 +542,7 @@ python build/build.py requirements_update
 Alternatively, to install `jaxlib` from a local wheel (assuming Python 3.12):
 
 ```
-echo -e "\n$(realpath jaxlib-0.4.26-cp312-cp312-manylinux2014_x86_64.whl)" >> build/requirements.in
+echo -e "\n$(realpath jaxlib-0.4.26-cp312-cp312-manylinux_2_27_x86_64.whl)" >> build/requirements.in
 python build/build.py requirements_update --python_version=3.12
 ```
 

jaxlib/jax.bzl

@@ -59,7 +59,7 @@ jax_test_util_visibility = []
 loops_visibility = []
 
 PLATFORM_TAGS_DICT = {
-    ("Linux", "x86_64"): ("manylinux2014", "x86_64"),
+    ("Linux", "x86_64"): ("manylinux_2_27", "x86_64"),
     ("Linux", "aarch64"): ("manylinux2014", "aarch64"),
     ("Linux", "ppc64le"): ("manylinux2014", "ppc64le"),
     ("Darwin", "x86_64"): ("macosx_11_0", "x86_64"),