Addons: Create Bazel rules for linking custom ops to TensorFlow Serving

[guillaumekln]

TensorFlow Serving can statically link to custom ops. We should define Bazel rules that produce compatible static libraries.

Assuming tf.contrib is a good example to follow, here's what would be needed:

1. For each custom op:
   a. Define a kernel library (see tf_kernel_library and example in contrib)
   b. Define an op library (see tf_gen_op_libs and example in contrib)
2. In tensorflow_addons/custom_ops/BUILD, define 2 global Bazel rules that TensorFlow Serving can depend on:
   a. addons_kernels that depends on all kernel libraries (see example in contrib)
   b. addons_ops_op_lib that depends on all op libraries (see example in contrib)

[seanpmorgan]

Hmmm I'm wondering if this issue is not better handled in https://github.com/tensorflow/custom-op and then we can just follow the upstream template.

I imagine other custom-op's would want to have an easy path to link with Serving, but I'm not sure that following the contrib kernel_library is the correct path.

cc @yifeif @unclepeddy to see if you have any thoughts on how this should be handled

[helinwang]

I tried to link addons/text/parse_time with TF serving, the experience lead me to think that we can improve tensorflow/addons's build infra so that tensorflow/serving could easily link operators in addons (either official, or users manually patch tensorflow/serving).
EDIT: Sorry @guillaumekln , I did not fully read the issue you posted before posting this, I think we are trying to do the same thing :)

The main thing that we can improve on are (tensorflow/text just made similar changes, their changes are a good example to follow, see the "Reference" below):

• [ ] Add cc_library rule to build operators. Currently, we only have cc_binary rule: https://github.com/tensorflow/addons/blob/master/tensorflow_addons/tensorflow_addons.bzl#L45 Reference: https://github.com/tensorflow/text/commit/64c678c7ef682057491f7d7118c395fe3c5682ce#diff-044f4228ee13caa326f04499e79907e9, https://github.com/tensorflow/text/commit/830b2e465427906840395cf750909a70e2b9f47e#diff-044f4228ee13caa326f04499e79907e9
• [ ] Add tutorial for how to modify tensorflow/serving to link operators from tensorflow/addons. Reference: https://github.com/tensorflow/serving/commit/142d0adb5e2975689d80d8fc608c9684e96de078

Once the above points has been done, it would be nice if we can ask tensorflow/serving to link whitelisted ops from addons. Even if tensorflow/serving does not agree, users would have a clear tutorial to follow to build themselves.

Here is what I did to link addons/text/parse_time (I followed https://github.com/tensorflow/serving/commit/142d0adb5e2975689d80d8fc608c9684e96de078):

diff --git a/tensorflow_serving/model_servers/BUILD b/tensorflow_serving/model_servers/BUILD
index 5d032965..88c00a08 100644
--- a/tensorflow_serving/model_servers/BUILD
+++ b/tensorflow_serving/model_servers/BUILD
@@ -302,6 +302,8 @@ SUPPORTED_TENSORFLOW_OPS = if_v2([]) + if_not_v2([
     "@org_tensorflow_text//tensorflow_text:unicode_script_tokenizer_cc",
     "@org_tensorflow_text//tensorflow_text:whitespace_tokenizer_cc",
     "@org_tensorflow_text//tensorflow_text:wordpiece_tokenizer_cc",
+] + [
+    "@org_tensorflow_addons//tensorflow_addons/custom_ops/text:_parse_time_op.so",
 ]

 TENSORFLOW_DEPS = [
diff --git a/tensorflow_serving/workspace.bzl b/tensorflow_serving/workspace.bzl
index ba8725bd..a45ece44 100644
--- a/tensorflow_serving/workspace.bzl
+++ b/tensorflow_serving/workspace.bzl
@@ -70,3 +70,15 @@ def tf_serving_workspace():
         patches = ["@//third_party/tf_text:tftext.patch"],
         patch_args = ["-p1"],
     )
+
+    # ===== TF.Addons dependencies
+    http_archive(
+        name = "org_tensorflow_addons",
+        sha256 = "b0c91389fef50bdd50648ae04e8ae6011027964d99dd6501268aae82142b0e49",
+        strip_prefix = "addons-dfaa5c0f0fd441ab7d14f31a80e3de5bddc3ef9e",
+        urls = [
+            "https://github.com/tensorflow/addons/archive/dfaa5c0f0fd441ab7d14f31a80e3de5bddc3ef9e.zip",
+        ],
+        patches = ["@//third_party/tf_addons:tfaddons.patch"],
+        patch_args = ["-p1"],
+    )
diff --git a/third_party/tf_addons/BUILD b/third_party/tf_addons/BUILD
new file mode 100644
index 00000000..2d59128b
--- /dev/null
+++ b/third_party/tf_addons/BUILD
@@ -0,0 +1 @@
+# Empty BUILD so this is treated like a package.
diff --git a/third_party/tf_addons/tfaddons.patch b/third_party/tf_addons/tfaddons.patch
new file mode 100644
index 00000000..8325a9e6
--- /dev/null
+++ b/third_party/tf_addons/tfaddons.patch
@@ -0,0 +1,56 @@
+--- a/tensorflow_addons/tensorflow_addons.bzl
++++ b/tensorflow_addons/tensorflow_addons.bzl
+@@ -1,6 +1,3 @@
+-load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI")
+-load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda")
+-
+ def custom_op_library(
+         name,
+         srcs = [],
+@@ -10,43 +7,18 @@ 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",
+-            "-nvcc_options=relaxed-constexpr",
+-            "-nvcc_options=ftz=true",
+-        ])
+-        cuda_deps = deps + if_cuda_is_configured(cuda_deps) + if_cuda_is_configured([
+-            "@local_config_cuda//cuda:cuda_headers",
+-            "@local_config_cuda//cuda:cudart_static",
+-        ])
+-        basename = name.split(".")[0]
+-        native.cc_library(
+-            name = basename + "_gpu",
+-            srcs = cuda_srcs,
+-            deps = cuda_deps,
+-            copts = cuda_copts,
+-            alwayslink = 1,
+-            **kwargs
+-        )
+-        deps = deps + if_cuda_is_configured([":" + basename + "_gpu"])
+-
+     copts = copts + [
+         "-pthread",
+-        "-std=c++11",
+-        D_GLIBCXX_USE_CXX11_ABI,
+     ]
+ 
+-    native.cc_binary(
++    native.cc_library(
+         name = name,
+         srcs = srcs,
+         copts = copts,
+-        linkshared = 1,
++        alwayslink = 1,
+         deps = deps,
+         **kwargs
+     )

[seanpmorgan]

@helinwang thanks for the summary of your experience. This pattern looks good, and we would be happy to review a PR (though we might want to upstream it to tensorflow/custom-op). @guillaumekln does this strategy look okay to you?

[helinwang]

Thanks, @seanpmorgan . Sorry that currently I don't have bandwidth for a PR. Would love to contribute if no one worked on it by the time I have bandwidth though.

Some more context, I workaround the following problems in my experiment:
Removed the following lines in https://github.com/tensorflow/addons/blob/master/tensorflow_addons/tensorflow_addons.bzl:

load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI")
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda")

As it's causing conflicts when building from tensorflow/serving.
Removed the cuda block:

if cuda_srcs:
  // ...

As I was not sure how to pass the cuda related flag from tensorflow/serving to addons.

[guillaumekln]

Thanks @helinwang, that's some useful information! Would you by any chance have some time to address this issue in the near future? If not, I would do my best to move forward on this as it is a blocker issue for some of my users.

Once the above points has been done, it would be nice if we can ask tensorflow/serving to link whitelisted ops from addons.

According to https://github.com/tensorflow/serving/issues/1381, this will not happen anytime soon.

[helinwang]

Sorry, @guillaumekln I won't have capacity in the coming few weeks. Please feel free to move this forward. Thanks in advance!