Merge branch 'master' into TFS_Load_TfLite_BE

Author: yasiribmcon

.bazelrc

@@ -11,20 +11,20 @@ build:cuda --@local_config_cuda//:enable_cuda
 build:cuda_clang --config=cuda
 build:cuda_clang --copt=-Wno-gnu-offsetof-extensions
 build:cuda_clang --copt=-Wno-error=unused-command-line-argument
+build --linkopt=-Wno-unused-command-line-argument
+build --host_linkopt=-Wno-unused-command-line-argument
 build:cuda_clang --host_copt=-Wno-error=unused-command-line-argument
 build:cuda_clang --repo_env TF_NEED_TENSORRT=0
 build:cuda_clang --action_env=TF_CUDA_CLANG="1"
 build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
-build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_60,sm_70,sm_80,compute_90"
+build:cuda_clang --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_60,sm_70,sm_80,sm_89,compute_90"
 
-build:cuda_clang --repo_env=HERMETIC_CUDA_VERSION="12.2.0"
-build:cuda_clang --repo_env=HERMETIC_CUDNN_VERSION="8.9.4.25"
+build:cuda_clang --repo_env=HERMETIC_CUDA_VERSION="12.5.1"
+build:cuda_clang --repo_env=HERMETIC_CUDNN_VERSION="9.3.0"
+build:cuda_clang --repo_env=HERMETIC_NVSHMEM_VERSION="3.2.5"
 build:cuda_clang --action_env=GCC_HOST_COMPILER_PATH="/usr/bin/gcc-10"
 build:cuda_clang --action_env=CLANG_CUDA_COMPILER_PATH="/usr/lib/llvm-17/bin/clang"
 
-build:kokoro --repo_env=CC="/usr/bin/gcc-10"
-build:kokoro --action_env=CC="/usr/bin/gcc-10"
-
 # Options used to build with TPU support.
 build:tpu --define=with_tpu_support=true --define=framework_shared_object=false
 build:tpu --copt=-DLIBTPU_ON_GCE
@@ -110,4 +110,16 @@ build --copt=-Wno-unused-but-set-variable # due to `ifdefs` in ml_dtypes
 # Revert to the legacy WORKSPACE dependency resolution system
 # and disable Bzlmod.
 common --enable_bzlmod=false
-common --noincompatible_enable_cc_toolchain_resolution
+common --incompatible_enable_cc_toolchain_resolution
+common --repo_env USE_HERMETIC_CC_TOOLCHAIN=1
+common --repo_env=HERMETIC_PYTHON_VERSION="3.10"
+build --features=-layering_check
+build --define=framework_shared_object=false
+build --dynamic_mode=off
+
+
+common:clang_local --noincompatible_enable_cc_toolchain_resolution
+common:clang_local --noincompatible_enable_android_toolchain_resolution
+common:clang_local --@rules_ml_toolchain//cc_toolchain/config:enable_hermetic_cc=False
+common:clang_local --repo_env USE_HERMETIC_CC_TOOLCHAIN=0
+

WORKSPACE

@@ -24,9 +24,17 @@ local_repository(
 load("//tensorflow_serving:repo.bzl", "tensorflow_http_archive")
 tensorflow_http_archive(
     name = "org_tensorflow",
-    sha256 = "e2dea5a1873f900a6af48ae16eccc20b275b5f3edd5b1a3b47da3feadfb513d9",
-    git_commit = "5fb3b1fefda9320202da184752a3366fbeddfeac",
+    sha256 = "86150d55ce57b2298d8ed42caa7b91c466ad33d9f7f347117c2257cc576d3413",
+    git_commit = "72fbba3d20f4616d7312b5e2b7f79daf6e82f2fa",
     patch = "//third_party/tensorflow:tensorflow.patch",
+    patch_cmds = [
+        "sed -i '/cc_library = _cc_library/d' tensorflow/core/platform/rules_cc.bzl",
+        "echo -e \"\\ndef cc_library_oss(deps=[], **kwargs):\\n    if kwargs.get(\\\"name\\\") == \\\"lib_internal_impl\\\" or \\\"protobuf\\\" in kwargs.get(\\\"name\\\", \\\"\\\"):\\n        _cc_library(deps = deps, **kwargs)\\n        return\\n    if type(deps) == \\\"list\\\":\\n        if \\\"@com_google_protobuf//:protobuf\\\" not in deps:\\n            deps = deps + [\\\"@com_google_protobuf//:protobuf\\\"]\\n    else:\\n        deps = deps + [\\\"@com_google_protobuf//:protobuf\\\"]\\n    _cc_library(deps = deps, **kwargs)\\ncc_library = cc_library_oss\" >> tensorflow/core/platform/rules_cc.bzl",
+        "sed -i 's#deps = \\[op_gen\\] + deps#deps = [op_gen] + deps + [clean_dep(\"//tensorflow/core/framework:kernel_shape_util\"), clean_dep(\"//tensorflow/core/framework:full_type_util\")]#' tensorflow/tensorflow.bzl",
+        "sed -i '/name = \"kernel_shape_util\",/a \\    visibility = [\"//visibility:public\"],' tensorflow/core/framework/BUILD",
+        "echo -e '\\nalias(name = \"tensorflow_libtensorflow_framework\", actual = \"//tensorflow/core:tensorflow\", visibility = [\"//visibility:public\"])' >> BUILD",
+        "echo -e '\\nalias(name = \"tensorflow_tf_header_lib\", actual = \"//tensorflow/core:tensorflow\", visibility = [\"//visibility:public\"])' >> BUILD",
+    ],
 )
 
 # Import all of TensorFlow Serving's external dependencies.
@@ -44,10 +52,10 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
 http_archive(
     name = "bazel_skylib",
-    sha256 = "74d544d96f4a5bb630d465ca8bbcfe231e3594e5aae57e1edbf17a6eb3ca2506",
+    sha256 = "97e709db2e97b646263b5c5e83e3b00de48c1ae55b9e421e3b5e3f9467d02a3a",
     urls = [
-        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
-        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.8.0/bazel-skylib-1.8.0.tar.gz",
+        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.8.0/bazel-skylib-1.8.0.tar.gz",
     ],
 )
 
@@ -58,34 +66,54 @@ http_archive(
     url = "https://github.com/bazelbuild/rules_cc/releases/download/0.1.5/rules_cc-0.1.5.tar.gz",
 )
 
-http_archive(
-    name = "rules_python",
-    sha256 = "84aec9e21cc56fbc7f1335035a71c850d1b9b5cc6ff497306f84cced9a769841",
-    strip_prefix = "rules_python-0.23.1",
-    url = "https://github.com/bazelbuild/rules_python/releases/download/0.23.1/rules_python-0.23.1.tar.gz",
+# Initialize hermetic Python
+load("@org_tensorflow//third_party/py:python_init_rules.bzl", "python_init_rules")
+python_init_rules()
+
+load("@org_tensorflow//third_party/py:python_init_repositories.bzl", "python_init_repositories")
+python_init_repositories(
+    default_python_version = "system",
+    requirements = {
+        "3.10": "@org_tensorflow//:requirements_lock_3_10.txt",
+        "3.11": "@org_tensorflow//:requirements_lock_3_11.txt",
+        "3.12": "@org_tensorflow//:requirements_lock_3_12.txt",
+        "3.13": "@org_tensorflow//:requirements_lock_3_13.txt",
+        "3.14": "@org_tensorflow//:requirements_lock_3_14.txt",
+    },
 )
 
-load("@rules_python//python:repositories.bzl", "python_register_toolchains")
+load("@org_tensorflow//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
+python_init_toolchains()
 
-load(
-    "@rules_python//python:repositories.bzl",
-    "py_repositories",
-    "python_register_toolchains",
+load("@org_tensorflow//third_party/py:python_init_pip.bzl", "python_init_pip")
+python_init_pip()
+
+load("@pypi//:requirements.bzl", "install_deps")
+install_deps()
+
+# Toolchains for ML projects hermetic builds.
+# Details: https://github.com/google-ml-infra/rules_ml_toolchain
+http_archive(
+    name = "rules_ml_toolchain",
+    sha256 = "de3b14418657eeacd8afc2aa89608be6ec8d66cd6a5de81c4f693e77bc41bee1",
+    strip_prefix = "rules_ml_toolchain-5653e5a0ca87c1272069b4b24864e55ce7f129a1",
+    urls = [
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/google-ml-infra/rules_ml_toolchain/archive/5653e5a0ca87c1272069b4b24864e55ce7f129a1.tar.gz",
+        "https://github.com/google-ml-infra/rules_ml_toolchain/archive/5653e5a0ca87c1272069b4b24864e55ce7f129a1.tar.gz",
+    ],
 )
-py_repositories()
 
 load(
-    "@org_tensorflow//tensorflow/tools/toolchains/python:python_repo.bzl",
-    "python_repository",
+    "@rules_ml_toolchain//cc_toolchain/deps:cc_toolchain_deps.bzl",
+    "cc_toolchain_deps",
 )
-python_repository(name = "python_version_repo")
-load("@python_version_repo//:py_version.bzl", "HERMETIC_PYTHON_VERSION")
 
-python_register_toolchains(
-    name = "python",
-    ignore_root_user_error = True,
-    python_version = HERMETIC_PYTHON_VERSION,
-)
+cc_toolchain_deps()
+
+register_toolchains("@rules_ml_toolchain//cc_toolchain:lx64_lx64")
+register_toolchains("@rules_ml_toolchain//cc_toolchain:lx64_lx64_cuda")
+# register_toolchains("@rules_ml_toolchain//cc_toolchain:la64_la64")
+# register_toolchains("@rules_ml_toolchain//cc_toolchain:la64_la64_cuda")
 
 # Initialize TensorFlow's external dependencies.
 load("@org_tensorflow//tensorflow:workspace3.bzl", "tf_workspace3")
@@ -101,10 +129,20 @@ tf_workspace0()
 load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
 rules_pkg_dependencies()
 
-load("@rules_proto//proto:repositories.bzl", "rules_proto_dependencies", "rules_proto_toolchains")
+load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
+
+http_archive(
+    name = "rules_shell",
+    sha256 = "0d0c56d01c3c40420bf7bf14d73113f8a92fbd9f5cd13205a3b89f72078f0321",
+    strip_prefix = "rules_shell-0.1.1",
+    urls = [
+        "https://github.com/bazelbuild/rules_shell/releases/download/v0.1.1/rules_shell-v0.1.1.tar.gz",
+    ],
+)
+
+load("@rules_proto//proto:repositories.bzl", "rules_proto_dependencies")
 
 rules_proto_dependencies()
-rules_proto_toolchains()
 
 load(
     "@local_xla//third_party/py:python_wheel.bzl",
@@ -114,7 +152,7 @@ load(
 python_wheel_version_suffix_repository(name = "tf_wheel_version_suffix")
 
 load(
-    "@local_xla//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
+    "@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
     "cuda_json_init_repository",
 )
 
@@ -126,7 +164,7 @@ load(
     "CUDNN_REDISTRIBUTIONS",
 )
 load(
-    "@local_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",
 )
@@ -140,23 +178,50 @@ cudnn_redist_init_repository(
 )
 
 load(
-    "@local_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(
-    "@local_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(
-    "@local_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(
+    "@rules_ml_toolchain//third_party/nvshmem/hermetic:nvshmem_json_init_repository.bzl",
+    "nvshmem_json_init_repository",
+)
+
+nvshmem_json_init_repository()
+
+load(
+    "@nvshmem_redist_json//:distributions.bzl",
+    "NVSHMEM_REDISTRIBUTIONS",
+)
+load(
+    "@rules_ml_toolchain//third_party/nvshmem/hermetic:nvshmem_redist_init_repository.bzl",
+    "nvshmem_redist_init_repository",
+)
+
+nvshmem_redist_init_repository(
+    nvshmem_redistributions = NVSHMEM_REDISTRIBUTIONS,
+)
+
+load(
+    "@rules_ml_toolchain//third_party/nvshmem/hermetic:nvshmem_configure.bzl",
+    "nvshmem_configure",
+)
+
+nvshmem_configure(name = "local_config_nvshmem")
+

tensorflow_serving/apis/predict.proto

@@ -72,6 +72,11 @@ message PredictRequest {
     // but will not return them in the response.
     repeated int64 return_stoptokens = 4;
 
+    // Only supported in disaggregated serving. When set, returns all stop
+    // tokens in response if the model stops at them. Overrides
+    // return_stoptokens.
+    optional bool return_all_stoptokens = 9;
+
     // When set and supported by servable, the model server will cache the
     // prefix of request up to this length.
     optional int64 max_cache_length = 6;
@@ -92,7 +97,6 @@ message PredictStreamedOptions {
   // CANCEL is used for early termination of a stream. It can be sent anytime to
   // cancel the stream.
   //
-  //
   // SPLIT is used when multiple streamed requests are used to generate a
   // logical request. END_SPLIT should be called for the last split of the
   // multi-turn request to start the processing of the current turn. NONE can

tensorflow_serving/batching/BUILD

@@ -5,7 +5,7 @@ load("@rules_cc//cc:cc_test.bzl", "cc_test")
 
 package(
     default_visibility = ["//tensorflow_serving:internal"],
-    features = ["-layering_check"],
+    features = ["layering_check"],
 )
 
 licenses(["notice"])
@@ -44,6 +44,7 @@ cc_test(
         ":streaming_batch_scheduler",
         "//tensorflow_serving/core/test_util:test_main",
         "@com_google_absl//absl/synchronization",
+        "@com_google_googletest//:gtest",
         "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core:test",
@@ -123,6 +124,7 @@ cc_test(
         "//tensorflow_serving/core/test_util:test_main",
         "//tensorflow_serving/servables/tensorflow:serving_session",
         "//tensorflow_serving/test_util",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/synchronization",
         "@com_google_googletest//:gtest",
         "@org_tensorflow//tensorflow/cc/saved_model:loader",
@@ -155,6 +157,7 @@ cc_test(
         ":batch_scheduler_retrier",
         "//tensorflow_serving/core/test_util:test_main",
         "@com_google_absl//absl/synchronization",
+        "@com_google_googletest//:gtest",
         "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core:test",
@@ -182,9 +185,9 @@ cc_test(
     deps = [
         ":batching_util",
         "//tensorflow_serving/core/test_util:test_main",
+        "@com_google_googletest//:gtest",
         "@org_tensorflow//tensorflow/core:framework",
         "@org_tensorflow//tensorflow/core:lib",
-        "@org_tensorflow//tensorflow/core:protos_all_cc",
     ],
 )
 
@@ -198,6 +201,7 @@ cc_test(
         "@com_google_absl//absl/time",
         "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:test",
+        "@org_tensorflow//tensorflow/core/platform",
     ],
 )
 
@@ -241,6 +245,7 @@ cc_test(
         "//tensorflow_serving/core/test_util:test_main",
         "//tensorflow_serving/servables/tensorflow/test_util:mock_tfrt_saved_model",
         "@com_google_absl//absl/functional:bind_front",
+        "@com_google_absl//absl/log:check",
         "@com_google_googletest//:gtest",
         "@org_tensorflow//tensorflow/core:lib",
         "@org_tensorflow//tensorflow/core:test",