[Feat] Add tf-serving config shell script

Author: rhdong

README.md

@@ -278,41 +278,56 @@ de = tfra.dynamic_embedding.get_variable("VariableOnGpu",
 sess_config.gpu_options.allow_growth = True
 ```
 
-## Inference with TensorFlow Serving
+## Inference 
+
+### With TensorFlow Serving
 
 #### Compatibility Matrix
-| TFRA  | TensorFlow | Serving | Compiler  | CUDA | CUDNN | Compute Capability |
-|:------|:-----------|:---- |:---------| :------------ | :---- | :------------ |
-| 0.6.0 | 2.8.3      | 2.5.2  | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
-| 0.5.1 | 2.8.3      | 2.5.2  | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
-| 0.5.0 | 2.8.3      | 2.5.2  | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
-| 0.4.0 | 2.5.1      | 2.5.2  | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
-| 0.3.1 | 2.5.1      | 2.5.2  | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
-| 0.2.0 | 2.4.1      | 2.4.0  | GCC 7.3.1 | 11.0 | 8.0 | 6.0, 6.1, 7.0, 7.5, 8.0 |
-| 0.2.0 | 1.15.2     | 1.15.0  | GCC 7.3.1 | 10.0 | 7.6 | 6.0, 6.1, 7.0, 7.5 |
-| 0.1.0 | 2.4.1      | 2.4.0  | GCC 7.3.1 | - | - | - |
-
-**NOTICE**：Reference documents: https://www.tensorflow.org/tfx/serving/custom_op
-
-#### CPU or GPU Serving TensorFlow models with custom ops
-When compiling, set the environment variable:
+| TFRA  | TensorFlow | Serving branch | Compiler  | CUDA | CUDNN | Compute Capability |
+|:------|:-----------|:---------------|:---------| :------------ | :---- | :------------ |
+| 0.6.0 | 2.8.3      | r2.8           | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
+| 0.5.1 | 2.8.3      | r2.8           | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
+| 0.5.0 | 2.8.3      | r2.8           | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
+| 0.4.0 | 2.5.1      | r2.5           | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
+| 0.3.1 | 2.5.1      | r2.5           | GCC 7.3.1 | 11.2| 8.1 | 6.0, 6.1, 7.0, 7.5, 8.0, 8.6 |
+| 0.2.0 | 2.4.1      | r2.4           | GCC 7.3.1 | 11.0 | 8.0 | 6.0, 6.1, 7.0, 7.5, 8.0 |
+| 0.2.0 | 1.15.2     | r1.15          | GCC 7.3.1 | 10.0 | 7.6 | 6.0, 6.1, 7.0, 7.5 |
+| 0.1.0 | 2.4.1      | r2.4           | GCC 7.3.1 | - | - | - |
+
+Serving TFRA-enable models by custom ops in TensorFlow Serving. 
+
 ```sh
-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",
-    "//tensorflow_recommenders_addons/dynamic_embedding/core:_math_ops.so",
-]
+## If enable GPU OPs
+export SERVING_WITH_GPU=1 
+
+## Specifiy the branch of TFRA
+export TFRA_BRANCH="master" # The `master` and `r0.6` are available.
+
+## Create workspace, modify the directory as you prefer to.
+export TFRA_SERVING_WORKSPACE=~/tfra_serving_workspace/
+mkdir -p $TFRA_SERVING_WORKSPACE && cd $TFRA_SERVING_WORKSPACE
+
+## Clone the release branches of serving and TFRA according to `Compatibility Matrix`.
+git clone -b r2.8 https://github.com/tensorflow/serving.git
+git clone -b $TFRA_BRANCH https://github.com/tensorflow/recommenders-addons.git
+
+## Run config shell script
+cd $TFRA_SERVING_WORKSPACE/recommenders-addons/tools
+bash config_tfserving.sh $TFRA_BRANCH $TFRA_SERVING_WORKSPACE/serving $SERVING_WITH_GPU
+
+## Build serving with TFRA OPs.
+cd $TFRA_SERVING_WORKSPACE/serving
+./tools/run_in_docker.sh bazel build tensorflow_serving/model_servers:tensorflow_model_server
+
 ```
+
+For more detail, please refer to the shell script `./tools/config_tfserving.sh`.
+
 **NOTICE**
 - Distributed inference is only supported when using Redis as Key-Value storage. 
+- Reference documents: https://www.tensorflow.org/tfx/serving/custom_op
 
+### With Triton(W.I.P)
 
 ## Community
 

tensorflow_recommenders_addons/dynamic_embedding/core/ops/redis_table_ops.cc

@@ -44,11 +44,11 @@ Status ScalarAndTwoElementVectorInputsAndScalarOutputs(InferenceContext *c) {
 
 }  // namespace
 
-Status ValidateTableResourceHandle(InferenceContext *c, ShapeHandle keys,
-                                   const string &key_dtype_attr,
-                                   const string &value_dtype_attr,
-                                   bool is_lookup,
-                                   ShapeAndType *output_shape_and_type) {
+Status ValidateTableResourceHandleRedis(InferenceContext *c, ShapeHandle keys,
+                                        const string &key_dtype_attr,
+                                        const string &value_dtype_attr,
+                                        bool is_lookup,
+                                        ShapeAndType *output_shape_and_type) {
   auto *handle_data = c->input_handle_shapes_and_types(0);
   if (handle_data == nullptr || handle_data->size() != 2) {
     output_shape_and_type->shape = c->UnknownShape();
@@ -124,7 +124,7 @@ REGISTER_OP(PREFIX_OP_NAME(RedisTableFind))
       TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
 
       ShapeAndType value_shape_and_type;
-      TF_RETURN_IF_ERROR(ValidateTableResourceHandle(
+      TF_RETURN_IF_ERROR(ValidateTableResourceHandleRedis(
           c,
           /*keys=*/c->input(1),
           /*key_dtype_attr=*/"Tin",
@@ -149,7 +149,7 @@ REGISTER_OP(PREFIX_OP_NAME(RedisTableFindWithExists))
 
       ShapeHandle keys = c->UnknownShapeOfRank(1);
       ShapeAndType value_shape_and_type;
-      TF_RETURN_IF_ERROR(ValidateTableResourceHandle(
+      TF_RETURN_IF_ERROR(ValidateTableResourceHandleRedis(
           c,
           /*keys=*/c->input(1),
           /*key_dtype_attr=*/"Tin",
@@ -224,7 +224,7 @@ REGISTER_OP(PREFIX_OP_NAME(RedisTableExport))
       TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
       ShapeHandle keys = c->UnknownShapeOfRank(1);
       ShapeAndType value_shape_and_type;
-      TF_RETURN_IF_ERROR(ValidateTableResourceHandle(
+      TF_RETURN_IF_ERROR(ValidateTableResourceHandleRedis(
           c,
           /*keys=*/keys,
           /*key_dtype_attr=*/"Tkeys",

tools/config_tfserving.sh

@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+
+# Config the serving for build serving with TFRA OPs.
+# Usage:
+#   Param $1 is the branch name of TFRA
+#   Param $2 is the serving root directory.
+#   Param $3 is the flag indicating if enable CUDA for GPU.
+#   ```shell
+#     cd /recommenders_addons/tools
+#     bash ./config_tfserving.sh "r0.6" /root/of/serving/ "1"
+#     cd /root/of/serving/
+#     ./tools/run_in_docker.sh bazel build tensorflow_serving/model_servers:tensorflow_model_server
+#   ```
+
+tfra_branch=$1
+tfserving_root=$2
+is_gpu=$3
+
+# Pleas modify this mapping when update TFRA.
+declare -A serving_version
+serving_version["master"]="2.8.3"
+
+# 1. copy directory to tfserving root
+cp -r ../tensorflow_recommenders_addons ${tfserving_root}/
+cp -r ../build_deps ${tfserving_root}/
+
+# 2. Padding the WORKSPACE
+cat "../WORKSPACE"| tail -n +2 >> ${tfserving_root}/WORKSPACE
+
+# 3. Padding the tools/run_in_docker.sh
+file="${tfserving_root}/tools/run_in_docker.sh"
+original_docker="tensorflow\/serving:nightly-devel"
+replacement_docker="tfra\/serving:${serving_version[$tfra_branch]}-devel"
+if [[ "$is_gpu" == "1" ]]; then
+  replacement_docker="tfra\/serving:${serving_version[$tfra_branch]}-devel-gpu"
+fi
+
+sed -i "s/$original_docker/$replacement_docker/g" $file
+
+# 4. Padding .bazelrc
+file=${tfserving_root}/.bazelrc
+cat "serving_padding/.bazelrc_padding" >> $file
+
+if [[ "$is_gpu" == "1" ]]; then
+  cat "serving_padding/.bazelrc_gpu_padding" >> $file
+fi
+
+# 5. Padding tensorflow_serving/model_servers/BUILD
+file=${tfserving_root}/tensorflow_serving/model_servers/BUILD
+target_string="org_tensorflow_text\/\/tensorflow_text:ops_lib"
+
+sed -i "/$target_string/a \\
+    \"\/\/tensorflow_recommenders_addons\/dynamic_embedding\/core:_cuckoo_hashtable_ops.so\", \\
+    \"\/\/tensorflow_recommenders_addons\/dynamic_embedding\/core:_math_ops.so\", \\
+    \"\/\/tensorflow_recommenders_addons\/dynamic_embedding\/core:_redis_table_ops.so\",
+    " $file

tools/docker/serving-devel-gpu.Dockerfile

@@ -0,0 +1,10 @@
+# Dockerfile to build a serving with TFRA OPs.
+#
+# To push a new version, run:
+# $ docker build -f serving-devel.Dockerfile . --tag "tfra/serving:2.8.3-devel-gpu"
+# $ docker push tfra/serving:2.8.3-devel-gpu
+
+FROM tensorflow/serving:2.6.3-devel-gpu
+
+RUN pip install --upgrade pip && pip install tensorflow==2.8.3
+

tools/docker/serving-devel.Dockerfile

@@ -0,0 +1,10 @@
+# Dockerfile to build a serving with TFRA OPs.
+#
+# To push a new version, run:
+# $ docker build -f serving-devel.Dockerfile . --tag "tfra/serving:2.8.3-devel"
+# $ docker push tfra/serving:2.8.3-devel
+
+FROM tensorflow/serving:2.6.3-devel
+
+RUN pip install --upgrade pip && pip install tensorflow==2.8.3
+

tools/serving_padding/.bazelrc_gpu_padding

@@ -0,0 +1,9 @@
+build --action_env TF_NEED_CUDA="1"
+build --action_env CUDA_TOOLKIT_PATH="/usr/local/cuda"
+build --action_env CUDNN_INSTALL_PATH="/usr/lib/x86_64-linux-gnu"
+build --action_env TF_CUDA_VERSION="11.2"
+build --action_env TF_CUDNN_VERSION="8.1"
+test --config=cuda
+build --config=cuda
+build:cuda --define=using_cuda=true --define=using_cuda_nvcc=true
+build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain

tools/serving_padding/.bazelrc_padding

@@ -0,0 +1,11 @@
+build --action_env TF_HEADER_DIR="/usr/local/lib/python3.7/dist-packages/tensorflow/include"
+build --action_env TF_SHARED_LIBRARY_DIR="/usr/local/lib/python3.7/dist-packages/tensorflow"
+build --action_env TF_SHARED_LIBRARY_NAME="libtensorflow_framework.so.2"
+build --action_env TF_CXX11_ABI_FLAG="0"
+build --action_env TF_CXX_STANDARD="c++14"
+build --action_env TF_VERSION_INTEGER="2083"
+build --action_env FOR_TF_SERVING="1"
+build --spawn_strategy=standalone
+build --strategy=Genrule=standalone
+build -c opt
+build --copt=-mavx