Major tf-keras update:

1) Update Bazel version to 7.4.1
2) Introduce hermetic C++ and CUDA toolchains
3) Update protobuf to 5.28.3
4) Introduce the new wheel building process.

The instructions on how to build and test the wheels using Bazel commands are provided in `wheel/README.md`.

Note that there are the changes in the wheel content:
1) The majority of the files in `tf_keras/layers/preprocessing/benchmarks` are not added in `tf_keras/src` because they are actually unit tests.
2) Same about `tf_keras/src/tests/keras_doctest.py`.
3) `tf_keras/src/saving/legacy/saved_model/create_test_saved_model.py` is absent because `create_test_saved_model` target has the tag `"no_oss"`.

PiperOrigin-RevId: 800518252

Author: tensorflower-gardener

.bazelrc

@@ -28,6 +28,9 @@
 #     nohdfs:       Disable hadoop hdfs support.
 #     nonccl:       Disable nccl support.
 
+# No MODULE.bzl file.
+common --noenable_bzlmod
+
 # Sets the default Apple platform to macOS.
 build --apple_platform_type=macos
 
@@ -39,11 +42,15 @@ test --define open_source_build=true
 build --define=use_fast_cpp_protos=false
 test --define=use_fast_cpp_protos=false
 
+build --incompatible_enable_cc_toolchain_resolution
+build --repo_env USE_HERMETIC_CC_TOOLCHAIN=1
+
 # This config refers to building with CUDA available. It does not necessarily
 # mean that we build CUDA op kernels.
 build:using_cuda --define=using_cuda=true
 build:using_cuda --action_env TF_NEED_CUDA=1
 build:using_cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
+build:using_cuda --@rules_ml_toolchain//common:enable_cuda
 
 # Enable the mlir generated GPU kernels only for cuda builds.
 build --define=tensorflow_enable_mlir_generated_gpu_kernels=0
@@ -53,6 +60,9 @@ build:using_cuda --define=tensorflow_enable_mlir_generated_gpu_kernels=1
 # This config refers to building CUDA op kernels with nvcc.
 build:cuda --config=using_cuda
 build:cuda --define=using_cuda_nvcc=true
+# Build CUDA with NVCC and other C++ targets with Clang
+build:cuda --action_env=TF_NVCC_CLANG="1"
+build:cuda --@local_config_cuda//:cuda_compiler=nvcc
 
 # dbg config, as a shorthand for '--config=opt -c dbg'
 build:dbg --config=opt -c dbg

.bazelversion

@@ -1 +1 @@
-5.4.0
+7.4.1

BUILD

@@ -1,118 +1,56 @@
-py_library(
-    name = "expect_absl_installed",
-    # This is a dummy rule used as a absl dependency in open-source.
-    # We expect absl to already be installed on the system, e.g. via
-    # `pip install absl`
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_h5py_installed",
-    # This is a dummy rule used as a h5 dependency in open-source.
-    # We expect h5py to already be installed on the system, e.g. via
-    # `pip install h5py'
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_numpy_installed",
-    # This is a dummy rule used as a numpy dependency in open-source.
-    # We expect numpy to already be installed on the system, e.g. via
-    # `pip install numpy`
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_pandas_installed",
-    # This is a dummy rule used as a pandas dependency in open-source.
-    # We expect pandas to already be installed on the system, e.g. via
-    # `pip install pandas'
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_pillow_installed",
-    # This is a dummy rule used as a pillow dependency in open-source.
-    # We expect pillow to already be installed on the system, e.g. via
-    # `pip install Pillow'
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-# Note that this dependency is for testing only.
-py_library(
-    name = "expect_portpicker_installed",
-    # This is a dummy rule used as a pandas dependency in open-source.
-    # We expect portpicker to already be installed on the system, e.g. via
-    # `pip install portpicker'
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_pydot_installed",
-    # This is a dummy rule used as a pydot dependency in open-source.
-    # We expect pydot to already be installed on the system, e.g. via
-    # `pip install pydot'
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_scipy_installed",
-    # This is a dummy rule used as a scipy dependency in open-source.
-    # We expect scipy to already be installed on the system, e.g. via
-    # `pip install scipy'
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_six_installed",
-    # This is a dummy rule used as a six dependency in open-source.
-    # We expect six to already be installed on the system, e.g. via
-    # `pip install six`
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_tensorboard_installed",
-    # This is a dummy rule used as a tensorboard dependency in open-source.
-    # We expect tensorboard to already be installed on the system, e.g. via
-    # `pip install tensorflow`
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_tensorflow_installed",
-    # This is a dummy rule used as a tensorflow dependency in open-source.
-    # We expect tensorflow to already be installed on the system, e.g. via
-    # `pip install tensorflow`
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-py_library(
-    name = "expect_yaml_installed",
-    # This is a dummy rule used as a yaml dependency in open-source.
-    # We expect yaml to already be installed on the system, e.g. via
-    # `pip install yaml`
-    visibility = ["//visibility:public"],
-    deps = [],
-)
-
-# Note that this dependency is for testing only.
-py_library(
-    name = "expect_tensorflow_io_installed",
-    # This is a dummy rule used as a tensorflow_io dependency in open-source.
-    # We expect tensorflow_io to already be installed on the system, e.g. via
-    # `pip install tensorflow-io`
-    visibility = ["//visibility:public"],
-    deps = [],
+load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
+load("@python//:defs.bzl", "compile_pip_requirements")
+load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+load(
+    "@xla//third_party/py:py_import.bzl",
+    "py_import",
+)
+
+exports_files(["oss_setup.py"])
+
+compile_pip_requirements(
+    name = "requirements",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--rebuild",
+    ],
+    generate_hashes = True,
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+)
+
+py_import(
+    name = "tf_nightly_with_deps",
+    wheel = "@pypi//tf_nightly:whl",
+    wheel_deps = if_cuda([
+        "@pypi_nvidia_cublas_cu12//:pkg",
+        "@pypi_nvidia_cuda_cupti_cu12//:pkg",
+        "@pypi_nvidia_cuda_nvcc_cu12//:pkg",
+        "@pypi_nvidia_cuda_nvrtc_cu12//:pkg",
+        "@pypi_nvidia_cuda_runtime_cu12//:pkg",
+        "@pypi_nvidia_cudnn_cu12//:pkg",
+        "@pypi_nvidia_cufft_cu12//:pkg",
+        "@pypi_nvidia_curand_cu12//:pkg",
+        "@pypi_nvidia_cusolver_cu12//:pkg",
+        "@pypi_nvidia_cusparse_cu12//:pkg",
+        "@pypi_nvidia_nccl_cu12//:pkg",
+        "@pypi_nvidia_nvjitlink_cu12//:pkg",
+    ]),
+    deps = [
+        "@pypi_absl_py//:pkg",
+        "@pypi_astunparse//:pkg",
+        "@pypi_flatbuffers//:pkg",
+        "@pypi_gast//:pkg",
+        "@pypi_ml_dtypes//:pkg",
+        "@pypi_numpy//:pkg",
+        "@pypi_opt_einsum//:pkg",
+        "@pypi_packaging//:pkg",
+        "@pypi_protobuf//:pkg",
+        "@pypi_requests//:pkg",
+        "@pypi_tensorboard//:pkg",
+        "@pypi_termcolor//:pkg",
+        "@pypi_typing_extensions//:pkg",
+        "@pypi_wrapt//:pkg",
+    ],
 )

WORKSPACE

@@ -2,8 +2,126 @@ workspace(name = "org_keras")
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
+# Toolchains for ML projects hermetic builds.
+# Details: https://github.com/google-ml-infra/rules_ml_toolchain
+http_archive(
+    name = "rules_ml_toolchain",
+    sha256 = "59d7eb36a02cbe3c2e2fa67fda5e8f1ab7e274bc4773bbd207c51fe199e11c19",
+    strip_prefix = "rules_ml_toolchain-ffd9e3d7b84e43c2686c803cb08ce790ffd58baa",
+    urls = [
+        "https://github.com/google-ml-infra/rules_ml_toolchain/archive/ffd9e3d7b84e43c2686c803cb08ce790ffd58baa.tar.gz",
+    ],
+)
+
+load(
+    "@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl",
+    "cc_toolchain_deps",
+)
+
+cc_toolchain_deps()
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_cuda")
+
+http_archive(
+    name = "xla",
+    sha256 = "5b20a5fd981a23a5dea092f13f9279b5c2e8005b509ac141fb65f22a42ad35a9",
+    strip_prefix = "xla-32b7537aee5a1e9145a6e11fc258347c41d6f5f8",
+    urls = [
+      "https://storage.googleapis.com/mirror.tensorflow.org/github.com/openxla/xla/archive/32b7537aee5a1e9145a6e11fc258347c41d6f5f8.tar.gz",
+      "https://github.com/openxla/xla/archive/32b7537aee5a1e9145a6e11fc258347c41d6f5f8.tar.gz",
+      ],
+)
+
+load("@xla//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
+
+# Initialize hermetic Python
+load("@xla//third_party/py:python_init_rules.bzl", "python_init_rules")
+
+python_init_rules()
+
+load("@xla//third_party/py:python_init_repositories.bzl", "python_init_repositories")
+
+python_init_repositories(
+    requirements = {
+        "3.9": "//:requirements_lock_3_9.txt",
+    },
+    local_wheel_workspaces = ["//:WORKSPACE"],
+    local_wheel_dist_folder = "dist",
+    local_wheel_inclusion_list = [
+        "tf_keras*",
+        "tf_nightly*",
+        "tf-keras*",
+        "tf-nightly*",
+        "tensorflow*",
+        "keras*",
+    ],
+)
+
+load("@xla//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
+
+python_init_toolchains()
+
+load("@xla//third_party/py:python_init_pip.bzl", "python_init_pip")
+
+python_init_pip()
+
+load("@pypi//:requirements.bzl", "install_deps")
+
+install_deps()
+
+load("//:tf_keras_python_wheel.bzl", "tf_keras_python_wheel_repository")
+
+tf_keras_python_wheel_repository(
+    name = "tf_keras_wheel",
+    version_key = "__version__",
+    version_source = "//tf_keras:__init__.py",
+)
+
+load(
+    "@xla//third_party/py:python_wheel.bzl",
+    "python_wheel_version_suffix_repository",
+)
+
+python_wheel_version_suffix_repository(
+    name = "tf_keras_wheel_version_suffix",
+)
+
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_json_init_repository.bzl",
+    "cuda_json_init_repository",
+)
+
+cuda_json_init_repository()
+
+load(
+    "@cuda_redist_json//:distributions.bzl",
+    "CUDA_REDISTRIBUTIONS",
+    "CUDNN_REDISTRIBUTIONS",
+)
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_redist_init_repositories.bzl",
+    "cuda_redist_init_repositories",
+    "cudnn_redist_init_repository",
+)
+
+cuda_redist_init_repositories(
+    cuda_redistributions = CUDA_REDISTRIBUTIONS,
+)
+
+cudnn_redist_init_repository(
+    cudnn_redistributions = CUDNN_REDISTRIBUTIONS,
+)
+
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_configure.bzl",
+    "cuda_configure",
+)
+
+cuda_configure(name = "local_config_cuda")
+
 # Needed by protobuf
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 http_archive(
     name = "bazel_skylib",
     urls = [
@@ -29,11 +147,12 @@ bind(
     actual = "@six_archive//:six",
 )
 
-http_archive(
+tf_http_archive(
     name = "com_google_protobuf",
-    sha256 = "f66073dee0bc159157b0bd7f502d7d1ee0bc76b3c1eac9836927511bdc4b3fc1",
-    strip_prefix = "protobuf-3.21.9",
-    urls = ["https://github.com/protocolbuffers/protobuf/archive/v3.21.9.zip"],
+    patch_file = ["@xla//third_party/protobuf:protobuf.patch"],
+    sha256 = "f645e6e42745ce922ca5388b1883ca583bafe4366cc74cf35c3c9299005136e2",
+    strip_prefix = "protobuf-5.28.3",
+    urls = tf_mirror_urls("https://github.com/protocolbuffers/protobuf/archive/refs/tags/v5.28.3.zip"),
 )
 
 # ZLIB. Need by com_google_protobuf.

pip_build.py

@@ -1,5 +1,7 @@
 """Build the TF-Keras pip package.
 
+DEPRECATED: Use the Bazel commands in `wheel/README.md` instead.
+
 The steps are as follows:
 
 0. Run bazel build in TF-Keras root directory to obtain protobuf Python files.