[tf.data service] Package dataset ops that depend on grpc separately.

PiperOrigin-RevId: 306518696
Change-Id: I118f6599d93b149ae21f3eab63595008ded2692f

tensorflow/BUILD

@@ -709,8 +709,8 @@ tf_cc_shared_object(
         "//tensorflow/c:version_script.lds",
         "//tensorflow/c/eager:c_api",
         "//tensorflow/c/eager:c_api_experimental",
+        "//tensorflow/core:distributed_tensorflow_dependencies",
         "//tensorflow/core:tensorflow",
-        "//tensorflow/core/distributed_runtime/rpc:grpc_session",
     ],
 )
 

tensorflow/core/BUILD

@@ -1140,6 +1140,15 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "distributed_tensorflow_dependencies",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core/distributed_runtime/rpc:grpc_session",
+        "//tensorflow/core/kernels:data_service_ops",
+    ],
+)
+
 cc_library(
     name = "testlib_kernels_impl",
     deps = [

tensorflow/core/kernels/BUILD

@@ -8365,6 +8365,13 @@ tf_kernel_library(
     ],
 )
 
+tf_kernel_library(
+    name = "data_service_ops",
+    deps = [
+        "//tensorflow/core/kernels/data/experimental:data_service_kernels",
+    ],
+)
+
 tf_kernel_library(
     name = "dataset_ops",
     deps = [

tensorflow/core/kernels/data/experimental/BUILD

@@ -670,6 +670,16 @@ tf_cc_test(
     ],
 )
 
+tf_kernel_library(
+    # data service kernels depend on GRPC, so we package them separately
+    # so that downstream rules can avoid depending on GRPC.
+    name = "data_service_kernels",
+    deps = [
+        ":data_service_dataset_op",
+        ":data_service_ops",
+    ],
+)
+
 tf_kernel_library(
     name = "dataset_kernels",
     deps = [

tensorflow/python/BUILD

@@ -5891,6 +5891,7 @@ pywrap_tensorflow_macro(
         "//tensorflow/core/grappler/optimizers:meta_optimizer",
         "//tensorflow/core:lib",
         "//tensorflow/core:reader_base",
+        "//tensorflow/core/kernels:data_service_ops",
         "//tensorflow/core/debug",
         "//tensorflow/core/distributed_runtime:server_lib",
         "//tensorflow/core/platform:stacktrace_handler",