Skip to content

Commit

Permalink
Support TF Serving compilation
Browse files Browse the repository at this point in the history
  • Loading branch information
@Mr-Nineteen @rhdong
Mr-Nineteen authored and rhdong committed Jun 1, 2021
1 parent b0e4e9f commit 2eadd95
Show file tree
Hide file tree
Showing 6 changed files with 72 additions and 73 deletions.
59 changes: 4 additions & 55 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -188,66 +188,15 @@ sess_config.gpu_options.allow_growth = True
| 0.2.0 | 2.4.1 | 2.4.0 | GCC 7.3.1 | 11.0 | 8.0 | 3.5, 5.2, 6.0, 6.1, 7.0, 7.5, 8.0 |
| 0.1.0 | 2.4.1 | 2.4.0 | GCC 7.3.1 | - | - | - |

#### CPU Serving TensorFlow models with custom ops
Reference documents: https://www.tensorflow.org/tfx/serving/custom_op
**NOTICE**：Reference documents: https://www.tensorflow.org/tfx/serving/custom_op

TFRA modification([`tensorflow_recommenders_addons.bzl`](tensorflow_recommenders_addons/tensorflow_recommenders_addons.bzl)):
#### CPU or GPU Serving TensorFlow models with custom ops
When compiling, set the environment variable:
```
deps = deps + [
# "@local_config_tf//:libtensorflow_framework",
"@local_config_tf//:tf_header_lib",
]
native.cc_library(
name = name,
srcs = srcs,
copts = copts,
alwayslink = 1,
features = select({
"//tensorflow_recommenders_addons:windows": ["windows_export_all_symbols"],
"//conditions:default": [],
}),
deps = deps,
**kwargs
)
export FOR_TF_SERVING = "1"
```
Tensorflow Serving modification(**model_servers/BUILD**):
```
SUPPORTED_TENSORFLOW_OPS = if_v2([]) + if_not_v2([
"@org_tensorflow//tensorflow/contrib:contrib_kernels",
"@org_tensorflow//tensorflow/contrib:contrib_ops_op_lib",
]) + [
"@org_tensorflow_text//tensorflow_text:ops_lib",
"//tensorflow_recommenders_addons/dynamic_embedding/core:_cuckoo_hashtable_ops.so",
]
```
#### GPU Serving TensorFlow models with custom ops
Including CPU Serving modification, TFRA modification([`tensorflow_recommenders_addons.bzl`](tensorflow_recommenders_addons/tensorflow_recommenders_addons.bzl)) again:
```
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(["//tensorflow_recommenders_addons/dynamic_embedding/core/lib/nvhash:nvhashtable"]) + 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"])
```
Tensorflow Serving modification(**model_servers/BUILD**) again:
```
SUPPORTED_TENSORFLOW_OPS = if_v2([]) + if_not_v2([
"@org_tensorflow//tensorflow/contrib:contrib_kernels",
"@org_tensorflow//tensorflow/contrib:contrib_ops_op_lib",
Expand Down
2 changes: 2 additions & 0 deletions build_deps/tf_dependency/build_defs.bzl.tpl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,5 @@
D_GLIBCXX_USE_CXX11_ABI = "%{tf_cx11_abi}"

DTF_VERSION = "%{tf_version}"

FOR_TF_SERVING = "%{for_tf_serving}"
5 changes: 5 additions & 0 deletions build_deps/tf_dependency/tf_configure.bzl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ _TF_SHARED_LIBRARY_NAME = "TF_SHARED_LIBRARY_NAME"

_TF_CXX11_ABI_FLAG = "TF_CXX11_ABI_FLAG"

_FOR_TF_SERVING = "FOR_TF_SERVING"

TF_VERSION = "TF_VERSION"

def _tpl(repository_ctx, tpl, substitutions = {}, out = None):
Expand Down Expand Up @@ -207,6 +209,7 @@ def _tf_pip_impl(repository_ctx):
tf_shared_library_path = "%s/%s" % (tf_shared_library_dir, tf_shared_library_name)
tf_cx11_abi = "-D_GLIBCXX_USE_CXX11_ABI=%s" % (repository_ctx.os.environ[_TF_CXX11_ABI_FLAG])
tf_version = "-DTF_VERSION=%s" % (repository_ctx.os.environ[TF_VERSION])
for_tf_serving = repository_ctx.os.environ[_FOR_TF_SERVING]

tf_shared_library_rule = _symlink_genrule_for_dir(
repository_ctx,
Expand All @@ -229,6 +232,7 @@ def _tf_pip_impl(repository_ctx):
{
"%{tf_cx11_abi}": tf_cx11_abi,
"%{tf_version}": tf_version,
"%{for_tf_serving}": for_tf_serving,
},
)

Expand All @@ -238,6 +242,7 @@ tf_configure = repository_rule(
_TF_SHARED_LIBRARY_DIR,
_TF_SHARED_LIBRARY_NAME,
_TF_CXX11_ABI_FLAG,
_FOR_TF_SERVING,
],
implementation = _tf_pip_impl,
)
2 changes: 2 additions & 0 deletions configure.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,8 @@ def create_build_configuration():
# TODO(Lifann) write them to enviroment variables.
write_action_env("TF_VERSION", tf_version)

write_action_env("FOR_TF_SERVING", os.getenv("FOR_TF_SERVING", "0"))

write("build --spawn_strategy=standalone")
write("build --strategy=Genrule=standalone")
write("build -c opt")
Expand Down
9 changes: 7 additions & 2 deletions tensorflow_recommenders_addons/dynamic_embedding/core/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,9 @@ load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda", "if_cuda_is_configure

package(default_visibility = ["//visibility:public"])

load("//tensorflow_recommenders_addons:tensorflow_recommenders_addons.bzl", "custom_op_library")
load("//tensorflow_recommenders_addons:tensorflow_recommenders_addons.bzl", "custom_op_library", "if_cuda_for_tf_serving")
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
load("@local_config_tf//:build_defs.bzl", "FOR_TF_SERVING")

custom_op_library(
name = "_cuckoo_hashtable_ops.so",
Expand All @@ -14,7 +15,11 @@ custom_op_library(
"kernels/cuckoo_hashtable_op.cc",
"ops/cuckoo_hashtable_ops.cc",
] + glob(["kernels/lookup_impl/lookup_table_op_cpu*"]),
cuda_deps = if_cuda(["//tensorflow_recommenders_addons/dynamic_embedding/core/lib/nvhash:nvhashtable"]),
cuda_deps = if_cuda_for_tf_serving(
["//tensorflow_recommenders_addons/dynamic_embedding/core/lib/nvhash:nvhashtable"],
[],
FOR_TF_SERVING,
),
cuda_srcs = if_cuda([
"kernels/cuckoo_hashtable_op_gpu.h",
"kernels/cuckoo_hashtable_op_gpu.cu.cc",
Expand Down
68 changes: 52 additions & 16 deletions tensorflow_recommenders_addons/tensorflow_recommenders_addons.bzl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ load(
"@local_config_tf//:build_defs.bzl",
"DTF_VERSION",
"D_GLIBCXX_USE_CXX11_ABI",
"FOR_TF_SERVING",
)
load(
"@local_config_cuda//cuda:build_defs.bzl",
Expand All @@ -17,10 +18,15 @@ def custom_op_library(
cuda_deps = [],
copts = [],
**kwargs):
deps = deps + [
"@local_config_tf//:libtensorflow_framework",
"@local_config_tf//:tf_header_lib",
]
if FOR_TF_SERVING == "1":
deps = deps + [
"@local_config_tf//:tf_header_lib",
]
else:
deps = deps + [
"@local_config_tf//:libtensorflow_framework",
"@local_config_tf//:tf_header_lib",
]

if cuda_srcs:
copts = copts + if_cuda(["-DGOOGLE_CUDA=1"])
Expand Down Expand Up @@ -67,15 +73,45 @@ def custom_op_library(
],
})

native.cc_binary(
name = name,
srcs = srcs,
copts = copts,
linkshared = 1,
features = select({
"//tensorflow_recommenders_addons:windows": ["windows_export_all_symbols"],
"//conditions:default": [],
}),
deps = deps,
**kwargs
)
if FOR_TF_SERVING == "1":
native.cc_library(
name = name,
srcs = srcs,
copts = copts,
alwayslink = 1,
features = select({
"//tensorflow_recommenders_addons:windows": ["windows_export_all_symbols"],
"//conditions:default": [],
}),
deps = deps,
**kwargs
)
else:
native.cc_binary(
name = name,
srcs = srcs,
copts = copts,
linkshared = 1,
features = select({
"//tensorflow_recommenders_addons:windows": ["windows_export_all_symbols"],
"//conditions:default": [],
}),
deps = deps,
**kwargs
)

def if_cuda_for_tf_serving(if_true, if_false = [], for_tf_serving = "0"):
"""Shorthand for select()'ing on whether we're building with CUDA.
Returns a select statement which evaluates to if_true if we're building
with CUDA enabled. Otherwise, the select statement evaluates to if_false.
"""
if for_tf_serving == "1":
return if_true

return select({
"@local_config_cuda//cuda:using_nvcc": if_true,
"@local_config_cuda//cuda:using_clang": if_true,
"//conditions:default": if_false,
})

0 comments on commit 2eadd95

Please sign in to comment.