Serving: Link to tensorflow/addons kernels

[guillaumekln]

Feature Request

This issue is to track the support of tensorflow/addons ops in TensorFlow Serving.

Describe the problem the feature is intended to solve

TensorFlow Serving is currently built against contrib ops. However, this module will be replaced by tensorflow/addons in TensorFlow 2.0. TensorFlow Serving should support running addons ops out-of-the-box in the TensorFlow 2.0 ecosystem.

Describe the solution

TensorFlow Serving should replace the contrib dependency by addons by the time TensorFlow 2.0 is released.

Describe alternatives you've considered

None.

Additional context

See the addons repository for more details on the existing custom ops.

cc @facaiy @seanpmorgan

[misterpeddy]

Thanks for opening this issue. It was my impression that many of the most widely used contrib ops will be in fact moving to core, for example:
"//tensorflow/contrib/boosted_trees:boosted_trees_ops_op_lib",
"//tensorflow/contrib/tensor_forest:model_ops_op_lib",
"//tensorflow/contrib/tensor_forest:stats_ops_op_lib",
"//tensorflow/contrib/tensor_forest:tensor_forest_ops_op_lib",
"//tensorflow/contrib/seq2seq:beam_search_ops_op_lib"

The current plan is to continue linking in TF core ops (which will contain these ops) since those will be the only ops officially supported by the Tensorflow team.

Anyone who needs other custom ops (like ones in the addons repo) can easily build the modelserver
(which is now fairly painless with the run_in_docker utility we provide) with the ops linked in.

Please let me know if you have any feedback!

[guillaumekln]

As far as I know, the beam search op in particular will not be moved to core. It would be surprising to users that Serving is not supporting seq2seq models out of the box.

Some addons modules (like seq2seq) are in fact supported by TensorFlow members.

Maybe @ewilderj has additional inputs on the global strategy here.

[misterpeddy]

Ah you're absolutely correct on seq2seq not moving to core:
https://github.com/tensorflow/addons/tree/master/tensorflow_addons/custom_ops/seq2seq/cc

I think there are 2 main concerns with adding arbitrary ops to the official build:

• Robustness - these ops have the ability to crash the model server and thus we need to make sure they are well-tested, and well-supported. The TF team, AFAIU, only officially supports the ops in core.
• Binary size - each op that gets added to the model server will make the binary larger; if an op is only used by a few modelserver users, it may make more sense for those users to build their own modelserver with the op linked in (again, operating on the assumption that this is relatively painless..) so that all other users who don't need the op don't pay the cost of unnecessary ops taking space.

What do you think about that? Would this fact be less "surprising" considering the fact that Model Server is statically built and thus all the ops need to be declared at build time (thus forcing a direct tradeoff between the set of ops supported and the size of the binary)?

/cc @gunan @netfs as well in case you guys have thoughts on the global strategy front :)

[ewilderj]

This definitely merits broader discussion cc/ @martinwicke

[seanpmorgan]

I think at a minimum we need to resolve some issues with the interaction between custom-op template and tf-serving so that it can be statically linked at all:
https://github.com/tensorflow/custom-op/issues/21

[facaiy]

cc SIG/IO @yongtang @terrytangyuan for visibility.

[yongtang]

/cc @dmitrievanthony @terrytangyuan @BryanCutler

[guillaumekln]

What do you think about that?

These are very valid concerns. If a technical solution would prove difficult to implement by the time TensorFlow 2.0 is released (I read about dynamic ops loading as a possible candidate), there should be at least a clear path for users of addons and other packages like tensorflow/text to use Serving. This could be a simple documentation on how to use the devel Docker image to build the Model Server with additional ops.

[misterpeddy]

@seanpmorgan can you elaborate what you mean by custom-op template? Are you talking about the fact that the bazel macro [1] only generates .so files (no static libs) or something else.

@guillaumekln - Docs would indeed be awesome. I can think of 2 concrete use cases to cover:

• I've written an op, and used it in a tf graph - how do I serve this graph with tf serving?
• I want to use an op built and shipped by someone else (ex. tf.text) in a tf graph - how do I use w/ tfs?

The first one is pretty easy since you can pull in the op code into your TF Serving repo, and build the static library and link it into serving.

The second one is slightly more subtle. You'd have to make sure the library you'd like to link to is creating a static library (.a) in addition to the .so (which currently [1] does not create and I believe is actually a pretty tricky thing to get right for all possible cases/platforms and bazel magic); you'd have to also ensure the op and the serving binary are built using the same compiler flags, against matching ABIs, etc.

So my intuition is that while the second guide would be much more helpful for people, getting it to work for all (even most?) cases may be hard (impossible?).

What do you think?

[1] https://chromium.googlesource.com/external/github.com/tensorflow/tensorflow/+/refs/heads/master/tensorflow/tensorflow.bzl#1741
[2] https://www.tensorflow.org/guide/extend/op

[seanpmorgan]

@unclepeddy https://github.com/tensorflow/custom-op/ is the provided template for building new custom-ops with TensorFlow. I haven't tried it myself but https://github.com/tensorflow/custom-op/issues/21 would seem to show that there is an issue building this custom-op template with tf-serving. The solution may just be some documentation on how to modify the template but I think it's important for those building custom-ops who wish to run them using tf-serving.

[guillaumekln]

I want to use an op built and shipped by someone else (ex. tf.text) in a tf graph - how do I use w/ tfs?

In the context of the current issue, the "someone else" is actually TensorFlow as we are mentioning official TensorFlow packages like tensorflow_addons, tensorflow_text, etc (third party packages are out of scope). So maybe a parallel effort is to ensure that selected official packages can be used with Serving as mentioned by @seanpmorgan.

[misterpeddy]

Just a soft update that we're discussing this internally. If you have any other thoughts, please do voice them.

[misterpeddy]

Just an update that we've added a quick guide to unblock folks needing to build tf serving with custom ops.

I will keep this issue open to track the longer-term question of
1) Whether (and how) we will add dynamic linking to the project's roadmap
2) If not (or if 1 is a long way away, in the meantime), how can we serve users who use text, addons and other common libraries without asking them to build the binary.

Also @guillaumekln sorry I did not reply to your last point - that makes sense - the main issue with the non TF core packages (more than the who) is that they're not bound to the same protocols (ex. forward and backward compatibility guarantees).

[helinwang]

Thanks for the discussion.
@unclepeddy regarding "robustness" and "binary size" in your comment https://github.com/tensorflow/serving/issues/1381#issuecomment-501452606. How does the following proposal sound?

1. Make TF serving able to link with Google open sourced ops (ops from tensorflow_text, tensorflow_addon) with a compiler flag default to false.
2. Make sure CI covers the case for setting the compiler flag to true.
3. Optionally, TF serving can publish two official images, with and without the custom ops.

Then "binary size" won't become an issue, since the official release won't have the operators. "robustness" seems to be fine given only Google open sourced repos are linked.

What do you think?

[guillaumekln]

Looks like https://github.com/tensorflow/serving/commit/142d0adb5e2975689d80d8fc608c9684e96de078 added TensorFlow Text kernels in Serving. Should the same be done for Addons?

[misterpeddy]

Unfortunately the same cannot be done for Addons as currently we only support linking in ops that are owned and supported by Google-internal teams. If you have any feedback on specific pain points in building your ops statically into the server please let us know.

[372046933]

I have successfully added Beam search ops to Serving 2.2.0. Here is the patch

diff --git a/tensorflow_serving/model_servers/BUILD b/tensorflow_serving/model_servers/BUILD
index cda8cf1d..a1977057 100644
--- a/tensorflow_serving/model_servers/BUILD
+++ b/tensorflow_serving/model_servers/BUILD
@@ -298,6 +298,10 @@ SUPPORTED_TENSORFLOW_OPS = if_v2([]) + if_not_v2([
     "@org_tensorflow//tensorflow/contrib:contrib_ops_op_lib",
 ]) + [
     "@org_tensorflow_text//tensorflow_text:ops_lib",
+    "@org_tensorflow_addons//tensorflow_addons/custom_ops/seq2seq:_beam_search_ops_gpu",
 ]

 cc_library(


diff --git a/tensorflow_serving/workspace.bzl b/tensorflow_serving/workspace.bzl
index b314e41f..cbfdd3c1 100644
--- a/tensorflow_serving/workspace.bzl
+++ b/tensorflow_serving/workspace.bzl
@@ -87,6 +87,18 @@ def tf_serving_workspace():
         repo_mapping = {"@com_google_re2": "@com_googlesource_code_re2"},
     )

+    # ===== TF.Addons dependencies
+    http_archive(
+        name = "org_tensorflow_addons",
+        sha256 = "a41cc626bca1f3b2020b5d85fe93a2eec5eeda841a7399ff6c7a571f1a9e53b4",
+        strip_prefix = "addons-0.10.0",
+        urls = [
+            "https://github.com/tensorflow/addons/archive/v0.10.0.zip",
+        ],
+        patches = ["@//third_party/tf_addons:tfaddons.patch"],
+        patch_args = ["-p1"],
+    )
+
     http_archive(
         name = "com_google_sentencepiece",
         strip_prefix = "sentencepiece-1.0.0",

Then add tfaddons.patch to third_party/tf_addons/tfaddons.patch with the following content

--- a/tensorflow_addons/tensorflow_addons.bzl
+++ b/tensorflow_addons/tensorflow_addons.bzl
@@ -1,4 +1,3 @@
-load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda", "if_cuda_is_configured")

 def custom_op_library(
@@ -10,14 +9,14 @@ def custom_op_library(
         copts = [],
         **kwargs):
     deps = deps + [
-        "@local_config_tf//:libtensorflow_framework",
-        "@local_config_tf//:tf_header_lib",
+        "@org_tensorflow//tensorflow/core:tensorflow_opensource"
     ]

     if cuda_srcs:
         copts = copts + if_cuda(["-DGOOGLE_CUDA=1"])
         cuda_copts = copts + if_cuda_is_configured([
             "-x cuda",
+            "--cuda_log",
             "-nvcc_options=relaxed-constexpr",
             "-nvcc_options=ftz=true",
         ])
@@ -51,7 +50,7 @@ def custom_op_library(
             "/DNOGDI",
             "/UTF_COMPILE_LIBRARY",
         ],
-        "//conditions:default": ["-pthread", "-std=c++11", D_GLIBCXX_USE_CXX11_ABI],
+        "//conditions:default": ["-pthread", "-D_GLIBCXX_USE_CXX11_ABI=0",],
     })

     native.cc_binary(

[haimin777]

You just save me! Thanks!