Skip to content

Commit 42c6133

Browse files
saetatensorflower-gardener
saeta
authored and
tensorflower-gardener
committed
[tf.data / Bigtable] Initial tf.data Bigtable integration
This change allows TensorFlow users to stream data directly from Cloud Bigtable into the TensorFlow runtime using tf.data. PiperOrigin-RevId: 202664219
1 parent 55259c7 commit 42c6133

31 files changed

+3073
-5
lines changed

‎tensorflow/contrib/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ py_library(
2525
"//tensorflow/contrib/all_reduce",
2626
"//tensorflow/contrib/batching:batch_py",
2727
"//tensorflow/contrib/bayesflow:bayesflow_py",
28+
"//tensorflow/contrib/bigtable",
2829
"//tensorflow/contrib/boosted_trees:init_py",
2930
"//tensorflow/contrib/checkpoint/python:checkpoint",
3031
"//tensorflow/contrib/cloud:cloud_py",

‎tensorflow/contrib/bigtable/BUILD

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
# Cloud Bigtable client for TensorFlow
2+
3+
package(
4+
default_visibility = ["//tensorflow:internal"],
5+
)
6+
7+
licenses(["notice"]) # Apache 2.0
8+
9+
load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
10+
load(
11+
"//tensorflow:tensorflow.bzl",
12+
"tf_copts",
13+
"tf_custom_op_library",
14+
"tf_gen_op_libs",
15+
"tf_gen_op_wrapper_py",
16+
"tf_kernel_library",
17+
"tf_cc_test",
18+
"tf_py_test",
19+
)
20+
21+
tf_custom_op_py_library(
22+
name = "bigtable",
23+
srcs = ["__init__.py"] + glob(["python/ops/*.py"]),
24+
dso = [
25+
":python/ops/_bigtable.so",
26+
],
27+
kernels = [
28+
":bigtable_kernels",
29+
":bigtable_ops_op_lib",
30+
],
31+
srcs_version = "PY2AND3",
32+
deps = [
33+
":bigtable_ops",
34+
"//tensorflow/contrib/util:util_py",
35+
"//tensorflow/python:framework_for_generated_wrappers",
36+
"//tensorflow/python:platform",
37+
"//tensorflow/python:util",
38+
"//tensorflow/python/data",
39+
],
40+
)
41+
42+
tf_custom_op_library(
43+
name = "python/ops/_bigtable.so",
44+
srcs = [
45+
"kernels/bigtable_kernels.cc",
46+
"kernels/bigtable_lookup_dataset_op.cc",
47+
"kernels/bigtable_prefix_key_dataset_op.cc",
48+
"kernels/bigtable_range_key_dataset_op.cc",
49+
"kernels/bigtable_scan_dataset_op.cc",
50+
"ops/bigtable_ops.cc",
51+
],
52+
deps = [
53+
":bigtable_lib_cc",
54+
"@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
55+
],
56+
)
57+
58+
tf_gen_op_wrapper_py(
59+
name = "bigtable_ops",
60+
deps = [":bigtable_ops_op_lib"],
61+
)
62+
63+
tf_gen_op_libs(
64+
op_lib_names = [
65+
"bigtable_ops",
66+
"bigtable_test_ops",
67+
],
68+
)
69+
70+
tf_kernel_library(
71+
name = "bigtable_kernels",
72+
srcs = [
73+
"kernels/bigtable_kernels.cc",
74+
"kernels/bigtable_lookup_dataset_op.cc",
75+
"kernels/bigtable_prefix_key_dataset_op.cc",
76+
"kernels/bigtable_range_key_dataset_op.cc",
77+
"kernels/bigtable_scan_dataset_op.cc",
78+
],
79+
deps = [
80+
":bigtable_lib_cc",
81+
"//tensorflow/core:framework_headers_lib",
82+
"//third_party/eigen3",
83+
"@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
84+
],
85+
)
86+
87+
# A library for use in the bigtable kernels.
88+
cc_library(
89+
name = "bigtable_lib_cc",
90+
srcs = ["kernels/bigtable_lib.cc"],
91+
hdrs = ["kernels/bigtable_lib.h"],
92+
deps = [
93+
"//tensorflow/core:framework_headers_lib",
94+
"//third_party/eigen3",
95+
"@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
96+
],
97+
)
98+
99+
cc_library(
100+
name = "bigtable_test_client",
101+
srcs = ["kernels/test_kernels/bigtable_test_client.cc"],
102+
hdrs = ["kernels/test_kernels/bigtable_test_client.h"],
103+
deps = [
104+
"//tensorflow/core:framework_headers_lib",
105+
"@com_github_googleapis_googleapis//:bigtable_protos",
106+
"@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
107+
"@com_googlesource_code_re2//:re2",
108+
],
109+
)
110+
111+
tf_cc_test(
112+
name = "bigtable_test_client_test",
113+
srcs = ["kernels/test_kernels/bigtable_test_client_test.cc"],
114+
tags = ["manual"],
115+
deps = [
116+
":bigtable_test_client",
117+
"//tensorflow/core:test",
118+
"//tensorflow/core:test_main",
119+
"@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
120+
],
121+
)
122+
123+
tf_gen_op_wrapper_py(
124+
name = "bigtable_test_ops",
125+
deps = [":bigtable_test_ops_op_lib"],
126+
)
127+
128+
tf_custom_op_library(
129+
name = "python/kernel_tests/_bigtable_test.so",
130+
srcs = [
131+
"kernels/test_kernels/bigtable_test_client_op.cc",
132+
"ops/bigtable_test_ops.cc",
133+
],
134+
deps = [
135+
":bigtable_lib_cc",
136+
":bigtable_test_client",
137+
"@com_googlesource_code_re2//:re2",
138+
],
139+
)
140+
141+
# Don't use tf_kernel_library because it prevents access to strings/stringprintf.h
142+
cc_library(
143+
name = "bigtable_test_kernels",
144+
srcs = [
145+
"kernels/test_kernels/bigtable_test_client_op.cc",
146+
],
147+
copts = tf_copts(),
148+
linkstatic = 1,
149+
deps = [
150+
":bigtable_lib_cc",
151+
":bigtable_test_client",
152+
"//tensorflow/core:framework_headers_lib",
153+
"//third_party/eigen3",
154+
"@com_googlesource_code_re2//:re2",
155+
],
156+
alwayslink = 1,
157+
)
158+
159+
tf_custom_op_py_library(
160+
name = "bigtable_test_py",
161+
dso = [
162+
":python/kernel_tests/_bigtable_test.so",
163+
],
164+
kernels = [
165+
":bigtable_test_kernels",
166+
":bigtable_test_ops_op_lib",
167+
],
168+
srcs_version = "PY2AND3",
169+
deps = [
170+
":bigtable_test_ops",
171+
# "//tensorflow/contrib/util:util_py",
172+
# "//tensorflow/python:framework_for_generated_wrappers",
173+
# "//tensorflow/python:platform",
174+
# "//tensorflow/python:util",
175+
# "//tensorflow/python/data",
176+
],
177+
)
178+
179+
tf_py_test(
180+
name = "bigtable_ops_test",
181+
size = "small",
182+
srcs = ["python/kernel_tests/bigtable_ops_test.py"],
183+
additional_deps = [
184+
":bigtable",
185+
":bigtable_test_py",
186+
"//tensorflow/core:protos_all_py",
187+
"//tensorflow/contrib/util:util_py",
188+
"//tensorflow/python:array_ops",
189+
"//tensorflow/python:client_testlib",
190+
"//tensorflow/python:dtypes",
191+
"//tensorflow/python:framework_for_generated_wrappers",
192+
"//tensorflow/python:platform",
193+
"//tensorflow/python:util",
194+
],
195+
data = [":python/kernel_tests/_bigtable_test.so"],
196+
tags = ["manual"],
197+
)

‎tensorflow/contrib/bigtable/README.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
# Bigtable #
2+
3+
[Google Cloud Bigtable](https://cloud.google.com/bigtable/) is a high
4+
performance storage system that can store and serve training data. This contrib
5+
package contains an experimental integration with TensorFlow.
6+
7+
> **Status: Highly experimental.** The current implementation is very much in
8+
> flux. Please use at your own risk! :-)
9+
10+
<!-- TODO(saeta): Document usage / methods / etc. -->

‎tensorflow/contrib/bigtable/__init__.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ==============================================================================
15+
"""Cloud Bigtable Client for TensorFlow.
16+
17+
This contrib package allows TensorFlow to interface directly with Cloud Bigtable
18+
for high-speed data loading.
19+
20+
@@BigtableClient
21+
@@BigTable
22+
23+
"""
24+
25+
from __future__ import absolute_import
26+
from __future__ import division
27+
from __future__ import print_function
28+
29+
from tensorflow.contrib.bigtable.python.ops.bigtable_api import BigTable
30+
from tensorflow.contrib.bigtable.python.ops.bigtable_api import BigtableClient
31+
32+
from tensorflow.python.util.all_util import remove_undocumented
33+
34+
_allowed_symbols = [
35+
'BigTable',
36+
'BigtableClient',
37+
]
38+
39+
remove_undocumented(__name__, _allowed_symbols)

0 commit comments

Comments
 (0)