1package(default_visibility = ["//tensorflow:internal"])
2
3licenses(["notice"])  # Apache 2.0
4
5exports_files(["LICENSE"])
6
7load(
8    "//tensorflow:tensorflow.bzl",
9    "tf_gen_op_wrapper_py",
10    "tf_kernel_library",
11    "tf_custom_op_library",
12    "tf_custom_op_py_library",
13    "tf_gen_op_libs",
14    "tf_py_test",
15)
16
17py_library(
18    name = "kafka",
19    srcs = ["__init__.py"],
20    srcs_version = "PY2AND3",
21    deps = [
22        ":dataset_ops",
23    ],
24)
25
26tf_custom_op_library(
27    name = "_dataset_ops.so",
28    srcs = ["ops/dataset_ops.cc"],
29    deps = [":dataset_kernels"],
30)
31
32tf_gen_op_libs(
33    op_lib_names = ["dataset_ops"],
34)
35
36cc_library(
37    name = "dataset_kernels",
38    srcs = ["kernels/kafka_dataset_ops.cc"],
39    deps = [
40        "//tensorflow/core:framework_headers_lib",
41        "//third_party/eigen3",
42        "@kafka",
43        "@protobuf_archive//:protobuf_headers",
44    ],
45    alwayslink = 1,
46)
47
48py_library(
49    name = "dataset_ops",
50    srcs = [
51        "python/ops/kafka_dataset_ops.py",
52    ],
53    srcs_version = "PY2AND3",
54    deps = [
55        ":kafka_op_loader",
56        "//tensorflow/python:dataset_ops_gen",
57        "//tensorflow/python:util",
58        "//tensorflow/python/data/ops:dataset_ops",
59        "//tensorflow/python/data/util:nest",
60    ],
61)
62
63tf_gen_op_wrapper_py(
64    name = "gen_dataset_ops",
65    out = "python/ops/gen_dataset_ops.py",
66    deps = ["//tensorflow/contrib/kafka:dataset_ops_op_lib"],
67)
68
69tf_kernel_library(
70    name = "dataset_ops_kernels",
71    deps = [
72        ":dataset_kernels",
73        "//tensorflow/core:framework",
74    ],
75    alwayslink = 1,
76)
77
78tf_custom_op_py_library(
79    name = "kafka_op_loader",
80    srcs = ["python/ops/kafka_op_loader.py"],
81    dso = ["//tensorflow/contrib/kafka:_dataset_ops.so"],
82    kernels = [
83        ":dataset_ops_kernels",
84        "//tensorflow/contrib/kafka:dataset_ops_op_lib",
85    ],
86    srcs_version = "PY2AND3",
87    deps = [
88        ":gen_dataset_ops",
89        "//tensorflow/contrib/util:util_py",
90        "//tensorflow/python:platform",
91    ],
92)
93
94# The Kafka server has to be setup before running the test.
95# The Kafka server is setup through Docker so the Docker engine
96# has to be installed.
97#
98# Once the Docker engine is ready:
99# To setup the Kafka server:
100# $ bash tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh start kafka
101#
102# After the test is complete:
103# To team down the Kafka server:
104# $ bash tensorflow/contrib/kafka/python/kernel_tests/kafka_test.sh stop kafka
105tf_py_test(
106    name = "kafka_test",
107    srcs = ["python/kernel_tests/kafka_test.py"],
108    additional_deps = [
109        ":kafka",
110        "//third_party/py/numpy",
111        "//tensorflow/python:client_testlib",
112        "//tensorflow/python:framework",
113        "//tensorflow/python:framework_test_lib",
114        "//tensorflow/python:platform_test",
115    ],
116    tags = [
117        "manual",
118        "no_windows",
119        "notap",
120    ],
121)
122