Home | History | Annotate | Download | only in kafka
      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