1 package(default_visibility = ["//tensorflow:internal"]) 2 3 licenses(["notice"]) # Apache 2.0 4 5 exports_files(["LICENSE"]) 6 7 load( 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 17 py_library( 18 name = "kafka", 19 srcs = ["__init__.py"], 20 srcs_version = "PY2AND3", 21 deps = [ 22 ":dataset_ops", 23 ], 24 ) 25 26 tf_custom_op_library( 27 name = "_dataset_ops.so", 28 srcs = ["ops/dataset_ops.cc"], 29 deps = [":dataset_kernels"], 30 ) 31 32 tf_gen_op_libs( 33 op_lib_names = ["dataset_ops"], 34 ) 35 36 cc_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 48 py_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 63 tf_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 69 tf_kernel_library( 70 name = "dataset_ops_kernels", 71 deps = [ 72 ":dataset_kernels", 73 "//tensorflow/core:framework", 74 ], 75 alwayslink = 1, 76 ) 77 78 tf_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 105 tf_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