Add prometheus dataset support

.gitignore

@@ -30,3 +30,7 @@ wheelhouse
 
 # Lint
 .pylint
+
+# Prometheus
+.coredns
+.prometheus

.kokorun/io_cpu.sh

@@ -50,7 +50,7 @@ bash -x -e tests/test_ignite/start_ignite.sh
 bash -x -e tests/test_kafka/kafka_test.sh start kafka
 bash -x -e tests/test_kinesis/kinesis_test.sh start kinesis
 bash -x -e tests/test_pubsub/pubsub_test.sh start pubsub
-
+bash -x -e tests/test_prometheus/prometheus_test.sh start
 
 export TENSORFLOW_INSTALL="$(python setup.py --package-version)"
 PYTHON_VERSION=$(python -c 'import sys; print(str(sys.version_info[0]))')
@@ -59,6 +59,11 @@ if [[ $PYTHON_VERSION == "2" ]]; then
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:14.04 bash -x -e .travis/python.release.sh "${TENSORFLOW_INSTALL}" python
 
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:16.04 bash -x -e .travis/wheel.test.sh python
+
+    ## Stop then restart prometheus
+    bash -x -e tests/test_prometheus/prometheus_test.sh stop
+    bash -x -e tests/test_prometheus/prometheus_test.sh start
+
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:18.04 bash -x -e .travis/wheel.test.sh python
 
     ## R
@@ -69,21 +74,48 @@ if [[ $PYTHON_VERSION == "2" ]]; then
     ## TF 2.0
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:14.04 bash -x -e .travis/python.release.sh "tensorflow==2.0.0b1" --preview ${KOKORO_BUILD_NUMBER} python
 
+    ## Stop then restart prometheus
+    bash -x -e tests/test_prometheus/prometheus_test.sh stop
+    bash -x -e tests/test_prometheus/prometheus_test.sh start
+
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:16.04 bash -x -e .travis/wheel.test.sh python
+
+    ## Stop then restart prometheus
+    bash -x -e tests/test_prometheus/prometheus_test.sh stop
+    bash -x -e tests/test_prometheus/prometheus_test.sh start
+
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:18.04 bash -x -e .travis/wheel.test.sh python
 else
     ## Python 2
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:14.04 bash -x -e .travis/python.release.sh "${TENSORFLOW_INSTALL}" python3.5 python3.6
 
+    ## Stop then restart prometheus
+    bash -x -e tests/test_prometheus/prometheus_test.sh stop
+    bash -x -e tests/test_prometheus/prometheus_test.sh start
+
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:16.04 bash -x -e .travis/wheel.test.sh python3.5
+
+    ## Stop then restart prometheus
+    bash -x -e tests/test_prometheus/prometheus_test.sh stop
+    bash -x -e tests/test_prometheus/prometheus_test.sh start
+
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:18.04 bash -x -e .travis/wheel.test.sh python3.6
 
     sudo rm -rf dist wheelhouse
 
     ## TF 2.0
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:14.04 bash -x -e .travis/python.release.sh "tensorflow==2.0.0b1" --preview ${KOKORO_BUILD_NUMBER} python3.5 python3.6
 
+    ## Stop then restart prometheus
+    bash -x -e tests/test_prometheus/prometheus_test.sh stop
+    bash -x -e tests/test_prometheus/prometheus_test.sh start
+
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:16.04 bash -x -e .travis/wheel.test.sh python3.5
+
+    ## Stop then restart prometheus
+    bash -x -e tests/test_prometheus/prometheus_test.sh stop
+    bash -x -e tests/test_prometheus/prometheus_test.sh start
+
     docker run -i --rm -v $PWD:/v -w /v --net=host buildpack-deps:18.04 bash -x -e .travis/wheel.test.sh python3.6
 fi
 

.travis.yml

@@ -16,6 +16,7 @@ install:
   bash -x -e tests/test_kafka/kafka_test.sh start kafka
   bash -x -e tests/test_kinesis/kinesis_test.sh start kinesis
   bash -x -e tests/test_pubsub/pubsub_test.sh start pubsub
+  bash -x -e tests/test_prometheus/prometheus_test.sh start
 - sudo -H python -m pip install -q -U twine --ignore-installed six
 - twine --version
 - export TENSORFLOW_INSTALL="$(python setup.py --package-version)"

.travis/bazel.configure.sh

@@ -13,6 +13,8 @@ python get-pip.py -q
 python -m pip --version
 if [[ $(uname) == "Darwin" ]]; then
   python -m pip install -q -U matplotlib numpy --ignore-installed six
+else
+  apt-get -y -qq install git
 fi
 python -m pip install -q --ignore-installed six "${TENSORFLOW_INSTALL}"
 python -c 'import tensorflow as tf; print(tf.version.VERSION)'

.travis/python3.7+.release.sh

@@ -21,7 +21,7 @@ set -e -x
 
 apt-get -y -qq update && apt-get -y -qq install \
     software-properties-common \
-    gcc g++ make patch \
+    gcc g++ make patch git \
     unzip curl patchelf
 
 add-apt-repository -y ppa:deadsnakes/ppa

.travis/wheel.test.sh

@@ -17,7 +17,7 @@ if [[ "$#" -gt 0 ]]; then
 fi
 if [[ $(uname) == "Linux" ]]; then
   apt-get -y -qq update
-  apt-get -y -qq install $PYTHON_VERSION ffmpeg
+  apt-get -y -qq install $PYTHON_VERSION ffmpeg  dnsutils
   curl -sSOL https://bootstrap.pypa.io/get-pip.py
   $PYTHON_VERSION get-pip.py -q
 fi

WORKSPACE

@@ -436,3 +436,49 @@ http_archive(
         "https://download.savannah.gnu.org/releases/freetype/freetype-2.10.0.tar.gz",
     ],
 )
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+# Note: patch is needed as we need to resolve multiple zlib dependencies.
+# Patch was created with:
+# diff -Naur a b > rules_go.patch
+http_archive(
+    name = "io_bazel_rules_go",
+    patch_args = ["-p1"],
+    patches = [
+        "//third_party:rules_go.patch",
+    ],
+    sha256 = "f04d2373bcaf8aa09bccb08a98a57e721306c8f6043a2a0ee610fd6853dcde3d",
+    urls = [
+        "https://storage.googleapis.com/bazel-mirror/github.com/bazelbuild/rules_go/releases/download/0.18.6/rules_go-0.18.6.tar.gz",
+        "https://github.com/bazelbuild/rules_go/releases/download/0.18.6/rules_go-0.18.6.tar.gz",
+    ],
+)
+
+http_archive(
+    name = "bazel_gazelle",
+    sha256 = "3c681998538231a2d24d0c07ed5a7658cb72bfb5fd4bf9911157c0e9ac6a2687",
+    urls = ["https://github.com/bazelbuild/bazel-gazelle/releases/download/0.17.0/bazel-gazelle-0.17.0.tar.gz"],
+)
+
+load("@io_bazel_rules_go//go:deps.bzl", "go_register_toolchains", "go_rules_dependencies")
+
+go_rules_dependencies()
+
+go_register_toolchains()
+
+load("@bazel_gazelle//:deps.bzl", "gazelle_dependencies", "go_repository")
+
+gazelle_dependencies()
+
+go_repository(
+    name = "com_github_prometheus_common",
+    importpath = "github.com/prometheus/common",
+    tag = "v0.4.1",
+)
+
+go_repository(
+    name = "com_github_prometheus_client_golang",
+    importpath = "github.com/prometheus/client_golang",
+    tag = "v0.9.3",
+)

tensorflow_io/core/BUILD

@@ -63,6 +63,7 @@ cc_binary(
         "//tensorflow_io/hdf5:hdf5_ops",
         "//tensorflow_io/lmdb:lmdb_ops",
         "//tensorflow_io/mnist:mnist_ops",
+        "//tensorflow_io/prometheus:prometheus_ops",
         "//tensorflow_io/text:text_ops",
         "@libarchive",
         "@local_config_tf//:libtensorflow_framework",

tensorflow_io/prometheus/BUILD

@@ -0,0 +1,26 @@
+licenses(["notice"])  # Apache 2.0
+
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "prometheus_ops",
+    srcs = [
+        #"//tensorflow_io/prometheus/go:prometheus.a",
+        #"//tensorflow_io/prometheus/go:prometheus.h",
+        "kernels/prometheus_input.cc",
+        "ops/prometheus_ops.cc",
+    ],
+    copts = [
+        "-pthread",
+        "-std=c++11",
+        "-DNDEBUG",
+    ],
+    includes = [
+        ".",
+    ],
+    linkstatic = True,
+    deps = [
+        "//tensorflow_io/core:dataset_ops",
+        "//tensorflow_io/prometheus/go:prometheus.cc",
+    ],
+)

tensorflow_io/prometheus/__init__.py

@@ -0,0 +1,32 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""PrometheusInput
+
+@@PrometheusDataset
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow_io.prometheus.python.ops.prometheus_ops import PrometheusDataset
+
+from tensorflow.python.util.all_util import remove_undocumented
+
+_allowed_symbols = [
+    "PrometheusDataset",
+]
+
+remove_undocumented(__name__, allowed_exception_list=_allowed_symbols)

tensorflow_io/prometheus/go/BUILD

@@ -0,0 +1,14 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_binary")
+
+go_binary(
+    name = "prometheus",
+    srcs = ["prometheus.go"],
+    cgo = True,
+    linkmode = "c-archive",
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_github_prometheus_client_golang//api:go_default_library",
+        "@com_github_prometheus_client_golang//api/prometheus/v1:go_default_library",
+        "@com_github_prometheus_common//model:go_default_library",
+    ],
+)

tensorflow_io/prometheus/go/prometheus.go

@@ -0,0 +1,50 @@
+package main
+
+import "C"
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/prometheus/client_golang/api"
+	"github.com/prometheus/client_golang/api/prometheus/v1"
+	"github.com/prometheus/common/model"
+)
+
+//export Query
+func Query(endpoint string, query string, sec int64, offset int64, key []int64, val []float64) int64 {
+	client, err := api.NewClient(api.Config{
+		Address: endpoint,
+	})
+	if err != nil {
+		return -1
+	}
+	value, err := v1.NewAPI(client).Query(context.Background(), query, time.Unix(sec, 0))
+	if err != nil {
+		return -1
+	}
+	if m, ok := value.(model.Matrix); ok && m.Len() > 0 {
+		index := int64(0)
+		for index < int64(len(key)) && offset+index < int64(len(m[0].Values)) {
+			v := m[0].Values[offset+index]
+			key[index] = v.Timestamp.Unix()
+			val[index] = float64(v.Value)
+			index++
+		}
+		return index
+	}
+	return 0
+}
+
+func main() {
+	key := make([]int64, 20, 20)
+	val := make([]float64, 20, 20)
+	sec := time.Now().Unix()
+	fmt.Println(sec)
+	returned := Query("http://localhost:9090", "coredns_dns_request_count_total[5m]", sec, 0, key, val)
+	fmt.Println(returned)
+	for i := range key {
+		fmt.Printf("%d, %q, %v\n", i, model.TimeFromUnix(key[i]).Time(), val[i])
+	}
+}

tensorflow_io/prometheus/kernels/prometheus_input.cc

@@ -0,0 +1,83 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "kernels/dataset_ops.h"
+#include "tensorflow/core/lib/io/buffered_inputstream.h"
+
+#include "go/prometheus.h"
+
+namespace tensorflow {
+namespace data {
+
+class PrometheusState {
+public:
+  PrometheusState() : time_(0), offset_(0) {}
+
+  int64 time_;
+  int64 offset_;
+};
+
+class PrometheusInput: public StreamInput<PrometheusState> {
+ public:
+  Status ReadRecord(IteratorContext* ctx, std::unique_ptr<PrometheusState>& state, int64 record_to_read, int64* record_read, std::vector<Tensor>* out_tensors) const override {
+    if (state.get() == nullptr) {
+      state.reset(new PrometheusState());
+      state.get()->time_ = time(NULL);
+    }
+    Tensor key_tensor(ctx->allocator({}), DT_INT64, {record_to_read});
+    Tensor val_tensor(ctx->allocator({}), DT_DOUBLE, {record_to_read});
+    GoSlice key_go = {key_tensor.flat<int64>().data(), record_to_read, record_to_read};
+    GoSlice val_go = {val_tensor.flat<double>().data(), record_to_read, record_to_read};
+    GoString endpoint_go = {endpoint().c_str(), static_cast<int64>(endpoint().size())};
+    GoString query_go = {schema().c_str(), static_cast<int64>(schema().size())};
+
+    GoInt returned = Query(endpoint_go, query_go, state.get()->time_, state.get()->offset_, key_go, val_go);
+    if (returned < 0) {
+        return errors::InvalidArgument("prometheus server error: ", returned);
+    }
+    if (returned > 0) {
+      state.get()->offset_ += returned;
+      *record_read = returned;
+      if (*record_read < record_to_read) {
+        Tensor key_tensor_final = key_tensor.Slice(0, *record_read);
+        Tensor val_tensor_final = val_tensor.Slice(0, *record_read);
+        out_tensors->emplace_back(std::move(key_tensor_final));
+        out_tensors->emplace_back(std::move(val_tensor_final));
+      } else {
+        out_tensors->emplace_back(std::move(key_tensor));
+        out_tensors->emplace_back(std::move(val_tensor));
+      }
+    }
+    return Status::OK();
+  }
+  Status FromEndpoint(const string& endpoint) override {
+    return Status::OK();
+  }
+  void EncodeAttributes(VariantTensorData* data) const override {
+  }
+  bool DecodeAttributes(const VariantTensorData& data) override {
+    return true;
+  }
+ protected:
+};
+
+REGISTER_UNARY_VARIANT_DECODE_FUNCTION(PrometheusInput, "tensorflow::data::PrometheusInput");
+
+REGISTER_KERNEL_BUILDER(Name("PrometheusInput").Device(DEVICE_CPU),
+                        StreamInputOp<PrometheusInput>);
+REGISTER_KERNEL_BUILDER(Name("PrometheusDataset").Device(DEVICE_CPU),
+                        StreamInputDatasetOp<PrometheusInput, PrometheusState>);
+}  // namespace data
+}  // namespace tensorflow