Add tfio.IOTensor.from_prometheus support

In the past couple of months I have been trying to come
up with a good example of using Prometheus with TensorFlow
for infrastructure/compute usage prediction and alerting
in case of anomaly.

My plan was to use LSTM to combine with prometheus for that.
Though there is a PrometheusDataset support, the format
of dataset is not very intuitive and hard to do it in
a smooth way.

The biggest challenge is that, normally for time series data
you have a look back window to train, and you normalize
the data (with total). Both are not easily available.
Further more, to process data in normal TF operations
such as tf.roll, it requires reading the whole data into
one Tenor. But this is not very straightforward as well with
iterable dataset.

This PR adds tfio.IOTensor.from_prometheus which
allows easily convertion the prometheus observation
data into a tuple of Tensors, and do additional feature
engineering.

Signed-off-by: Yong Tang <[email protected]>

Author: yongtang

tensorflow_io/core/python/ops/io_tensor.py

@@ -22,6 +22,7 @@
 from tensorflow_io.core.python.ops import audio_io_tensor_ops
 from tensorflow_io.core.python.ops import json_io_tensor_ops
 from tensorflow_io.core.python.ops import kafka_io_tensor_ops
+from tensorflow_io.core.python.ops import prometheus_io_tensor_ops
 
 class IOTensor(io_tensor_ops._BaseIOTensor):  # pylint: disable=protected-access
   """IOTensor
@@ -264,3 +265,27 @@ def from_kafka(cls,
           servers=kwargs.get("servers", None),
           configuration=kwargs.get("configuration", None),
           internal=True)
+
+  @classmethod
+  def from_prometheus(cls,
+                      query,
+                      **kwargs):
+    """Creates an `IOTensor` from a prometheus query.
+
+    Args:
+      query: A string, the query string for prometheus.
+      endpoint: A string, the server address of prometheus, by default
+       `http://localhost:9090`.
+      name: A name prefix for the IOTensor (optional).
+
+    Returns:
+      A (`IOTensor`, `IOTensor`) tuple that represents `timestamp`
+        and `value`.
+
+    """
+    with tf.name_scope(kwargs.get("name", "IOFromPrometheus")):
+      return prometheus_io_tensor_ops.PrometheusTimestampIOTensor(
+          query, endpoint=kwargs.get("endpoint", None), internal=True,
+      ), prometheus_io_tensor_ops.PrometheusValueIOTensor(
+          query, endpoint=kwargs.get("endpoint", None), internal=True,
+      )

tensorflow_io/core/python/ops/prometheus_io_tensor_ops.py

@@ -0,0 +1,66 @@
+# 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.
+# ==============================================================================
+"""PrometheusTimestampIOTensor"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import uuid
+
+import tensorflow as tf
+from tensorflow_io.core.python.ops import io_tensor_ops
+from tensorflow_io.core.python.ops import core_ops
+
+class PrometheusTimestampIOTensor(io_tensor_ops._ColumnIOTensor): # pylint: disable=protected-access
+  """PrometheusTimestampIOTensor"""
+
+  #=============================================================================
+  # Constructor (private)
+  #=============================================================================
+  def __init__(self,
+               query,
+               endpoint=None,
+               internal=False):
+    with tf.name_scope("PrometheusTimestampIOTensor") as scope:
+      metadata = ["column: timestamp"]
+      if endpoint is not None:
+        metadata.append(["endpoint: %s" % endpoint])
+      resource, dtypes, shapes, _ = core_ops.prometheus_indexable_init(
+          query, metadata=metadata,
+          container=scope, shared_name="%s/%s" % (query, uuid.uuid4().hex))
+      super(PrometheusTimestampIOTensor, self).__init__(
+          shapes, dtypes, resource, core_ops.prometheus_indexable_get_item,
+          internal=internal)
+
+class PrometheusValueIOTensor(io_tensor_ops._ColumnIOTensor): # pylint: disable=protected-access
+  """PrometheusValueIOTensor"""
+
+  #=============================================================================
+  # Constructor (private)
+  #=============================================================================
+  def __init__(self,
+               query,
+               endpoint=None,
+               internal=False):
+    with tf.name_scope("PrometheusTimestampIOTensor") as scope:
+      metadata = ["column: value"]
+      if endpoint is not None:
+        metadata.append(["endpoint: %s" % endpoint])
+      resource, dtypes, shapes, _ = core_ops.prometheus_indexable_init(
+          query, metadata=metadata,
+          container=scope, shared_name="%s/%s" % (query, uuid.uuid4().hex))
+      super(PrometheusValueIOTensor, self).__init__(
+          shapes, dtypes, resource, core_ops.prometheus_indexable_get_item,
+          internal=internal)

tensorflow_io/json/kernels/json_kernels.cc

@@ -347,7 +347,6 @@ class JSONIndexable : public IOIndexableInterface {
   }
 
   Status GetItem(const int64 start, const int64 stop, const int64 step, std::vector<Tensor>& tensors) override {
-    Tensor& output_tensor = tensors[0];
     if (step != 1) {
       return errors::InvalidArgument("step ", step, " is not supported");
     }

tensorflow_io/prometheus/kernels/prometheus_kernels.cc

@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow_io/core/kernels/io_interface.h"
 #include "tensorflow_io/core/prometheus_go.h"
 
 namespace tensorflow {
@@ -74,5 +75,138 @@ REGISTER_KERNEL_BUILDER(Name("ReadPrometheus").Device(DEVICE_CPU),
 
 
 }  // namespace
+
+
+class PrometheusIndexable : public IOIndexableInterface {
+ public:
+  PrometheusIndexable(Env* env)
+  : env_(env) {}
+
+  ~PrometheusIndexable() {}
+  Status Init(const std::vector<string>& input, const std::vector<string>& metadata, const void* memory_data, const int64 memory_size) override {
+    if (input.size() > 1) {
+      return errors::InvalidArgument("more than 1 query is not supported");
+    }
+    const string& query = input[0];
+
+    string endpoint = "http://localhost:9090";
+    for (size_t i = 0; i < metadata.size(); i++) {
+      if (metadata[i].find_first_of("endpoint: ") == 0) {
+        endpoint = metadata[i].substr(8);
+      }
+    }
+
+    int64 ts = time(NULL);
+
+    GoString endpoint_go = {endpoint.c_str(), static_cast<int64>(endpoint.size())};
+    GoString query_go = {query.c_str(), static_cast<int64>(query.size())};
+
+    GoSlice timestamp_go = {0, 0, 0};
+    GoSlice value_go = {0, 0, 0};
+
+    GoInt returned = Query(endpoint_go, query_go, ts, timestamp_go, value_go);
+    if (returned < 0) {
+      return errors::InvalidArgument("unable to query prometheus");
+    }
+
+    timestamp_.resize(returned);
+    value_.resize(returned);
+
+    if (returned > 0) {
+      timestamp_go.data = &timestamp_[0];
+      timestamp_go.len = returned;
+      timestamp_go.cap = returned;
+      value_go.data = &value_[0];
+      value_go.len = returned;
+      value_go.cap = returned;
+
+      returned = Query(endpoint_go, query_go, ts, timestamp_go, value_go);
+      if (returned < 0) {
+        return errors::InvalidArgument("unable to query prometheus to get the value");
+      }
+    }
+
+    for (size_t i = 0; i < metadata.size(); i++) {
+      if (metadata[i].find_first_of("column: ") == 0) {
+        columns_.emplace_back(metadata[i].substr(8));
+      }
+    }
+    if (columns_.size() == 0) {
+      columns_.emplace_back("timestamp");
+      columns_.emplace_back("value");
+    }
+
+    for (size_t i = 0; i < columns_.size(); i++) {
+      if (columns_[i] == "timestamp") {
+        dtypes_.emplace_back(DT_INT64);
+        shapes_.emplace_back(TensorShape({static_cast<int64>(returned)}));
+      } else if (columns_[i] == "value") {
+        dtypes_.emplace_back(DT_DOUBLE);
+        shapes_.emplace_back(TensorShape({static_cast<int64>(returned)}));
+      } else {
+        return errors::InvalidArgument("column name other than `timestamp` or `value` is not supported: ", columns_[i]);
+      }
+    }
+
+    return Status::OK();
+  }
+  Status Spec(std::vector<DataType>& dtypes, std::vector<PartialTensorShape>& shapes) override {
+    dtypes.clear();
+    for (size_t i = 0; i < dtypes_.size(); i++) {
+      dtypes.push_back(dtypes_[i]);
+    }
+    shapes.clear();
+    for (size_t i = 0; i < shapes_.size(); i++) {
+      shapes.push_back(shapes_[i]);
+    }
+    return Status::OK();
+  }
+
+  Status Extra(std::vector<Tensor>* extra) override {
+    // Expose columns
+    Tensor columns(DT_STRING, TensorShape({static_cast<int64>(columns_.size())}));
+    for (size_t i = 0; i < columns_.size(); i++) {
+      columns.flat<string>()(i) = columns_[i];
+    }
+    extra->push_back(columns);
+    return Status::OK();
+  }
+
+  Status GetItem(const int64 start, const int64 stop, const int64 step, std::vector<Tensor>& tensors) override {
+    if (step != 1) {
+      return errors::InvalidArgument("step ", step, " is not supported");
+    }
+    for (size_t i = 0; i < columns_.size(); i++) {
+      if (columns_[i] == "timestamp") {
+        memcpy(&tensors[i].flat<int64>().data()[start], &timestamp_[0], sizeof(int64) * (stop - start));
+      } else {
+        memcpy(&tensors[i].flat<double>().data()[start], &value_[0], sizeof(double) * (stop - start));
+      }
+    }
+
+    return Status::OK();
+  }
+
+  string DebugString() const override {
+    mutex_lock l(mu_);
+    return strings::StrCat("PrometheusIndexable");
+  }
+ private:
+  mutable mutex mu_;
+  Env* env_ GUARDED_BY(mu_);
+
+  std::vector<DataType> dtypes_;
+  std::vector<TensorShape> shapes_;
+  std::vector<string> columns_;
+
+  std::vector<int64> timestamp_;
+  std::vector<double> value_;
+};
+
+REGISTER_KERNEL_BUILDER(Name("PrometheusIndexableInit").Device(DEVICE_CPU),
+                        IOInterfaceInitOp<PrometheusIndexable>);
+REGISTER_KERNEL_BUILDER(Name("PrometheusIndexableGetItem").Device(DEVICE_CPU),
+                        IOIndexableGetItemOp<PrometheusIndexable>);
+
 }  // namespace data
 }  // namespace tensorflow

tensorflow_io/prometheus/ops/prometheus_ops.cc

@@ -19,6 +19,47 @@ limitations under the License.
 
 namespace tensorflow {
 
+REGISTER_OP("PrometheusIndexableInit")
+  .Input("input: string")
+  .Input("metadata: string")
+  .Output("output: resource")
+  .Output("dtypes: int64")
+  .Output("shapes: int64")
+  .Output("columns: string")
+  .Attr("container: string = ''")
+  .Attr("shared_name: string = ''")
+  .SetIsStateful()
+  .SetShapeFn([](shape_inference::InferenceContext* c) {
+    c->set_output(0, c->Scalar());
+    c->set_output(1, c->MakeShape({c->UnknownDim()}));
+    c->set_output(2, c->MakeShape({c->UnknownDim(), c->UnknownDim()}));
+    c->set_output(3, c->MakeShape({c->UnknownDim()}));
+    return Status::OK();
+   });
+
+REGISTER_OP("PrometheusIndexableGetItem")
+  .Input("input: resource")
+  .Input("start: int64")
+  .Input("stop: int64")
+  .Input("step: int64")
+  .Output("output: dtype")
+  .Attr("dtype: list(type) >= 1")
+  .Attr("shape: list(shape) >= 1")
+  .SetShapeFn([](shape_inference::InferenceContext* c) {
+    std::vector<PartialTensorShape> shape;
+    TF_RETURN_IF_ERROR(c->GetAttr("shape", &shape));
+    if (shape.size() != c->num_outputs()) {
+      return errors::InvalidArgument("`shape` must be the same length as `types` (", shape.size(), " vs. ", c->num_outputs());
+    }
+    for (size_t i = 0; i < shape.size(); ++i) {
+      shape_inference::ShapeHandle entry;
+      TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(shape[i], &entry));
+      c->set_output(static_cast<int64>(i), entry);
+    }
+    return Status::OK();
+   });
+
+
 REGISTER_OP("ReadPrometheus")
     .Input("endpoint: string")
     .Input("query: string")

tensorflow_io/prometheus/python/ops/prometheus_ops.py

@@ -17,10 +17,19 @@
 from __future__ import division
 from __future__ import print_function
 
+import warnings
+
 import tensorflow as tf
 from tensorflow_io.core.python.ops import data_ops
 from tensorflow_io.core.python.ops import core_ops
 
+warnings.warn(
+    "The tensorflow_io.prometheus.PrometheusDataset is "
+    "deprecated. Please look for tfio.IOTensor.from_prometheus "
+    "for reading prometheus observations into tensorflow.",
+    DeprecationWarning)
+
+
 def read_prometheus(endpoint, query):
   """read_prometheus"""
   return core_ops.read_prometheus(endpoint, query)

tests/test_prometheus_eager.py

@@ -26,7 +26,7 @@
 import tensorflow as tf
 if not (hasattr(tf, "version") and tf.version.VERSION.startswith("2.")):
   tf.compat.v1.enable_eager_execution()
-import tensorflow_io.prometheus as prometheus_io # pylint: disable=wrong-import-position
+import tensorflow_io as tfio # pylint: disable=wrong-import-position
 
 if sys.platform == "darwin":
   pytest.skip(
@@ -38,25 +38,14 @@ def test_prometheus():
     subprocess.call(["dig", "@localhost", "-p", "1053", "www.google.com"])
     time.sleep(1)
   time.sleep(2)
-  prometheus_dataset = prometheus_io.PrometheusDataset(
-      "http://localhost:9090",
-      "coredns_dns_request_count_total[5s]").apply(
-          tf.data.experimental.unbatch()).batch(2)
-
-  i = 0
-  for k, v in prometheus_dataset:
-    print("K, V: ", k.numpy(), v.numpy())
-    if i == 4:
-      # Last entry guaranteed 6.0
-      assert v.numpy() == 6.0
-    i += 2
-  assert i == 6
-
-  timestamp, value = prometheus_io.read_prometheus(
-      "http://localhost:9090",
+  timestamp, value = tfio.IOTensor.from_prometheus(
       "coredns_dns_request_count_total[5s]")
   assert timestamp.shape == [5]
+  assert timestamp.dtype == tf.int64
   assert value.shape == [5]
+  assert value.dtype == tf.float64
+  # last value should be 6.0
+  assert value.to_tensor().numpy()[4] == 6.0
 
 if __name__ == "__main__":
   test.main()