Implement histogram user metric for python SDK (#36335)

Author: hjtran

sdks/python/apache_beam/internal/metrics/cells.py

@@ -30,6 +30,7 @@
 
 from apache_beam.metrics.cells import MetricCell
 from apache_beam.metrics.cells import MetricCellFactory
+from apache_beam.portability.api import metrics_pb2
 from apache_beam.utils.histogram import Histogram
 
 if TYPE_CHECKING:
@@ -66,10 +67,12 @@ def get_cumulative(self) -> 'HistogramData':
     return self.data.get_cumulative()
 
   def to_runner_api_monitoring_info(self, name, transform_id):
-    # Histogram metric is currently worker-local and internal
-    # use only. This method should be implemented when runners
-    # support Histogram metric reporting.
-    return None
+    from apache_beam.metrics import monitoring_infos
+    return monitoring_infos.user_histogram(
+        name.namespace,
+        name.name,
+        self.get_cumulative(),
+        ptransform=transform_id)
 
 
 class HistogramCellFactory(MetricCellFactory):
@@ -150,3 +153,13 @@ def combine(self, other: Optional['HistogramData']) -> 'HistogramData':
   @staticmethod
   def identity_element(bucket_type) -> 'HistogramData':
     return HistogramData(Histogram(bucket_type))
+
+  def to_proto(self) -> metrics_pb2.HistogramValue:
+    return self.histogram.to_runner_api()
+
+  @classmethod
+  def from_proto(cls, proto: metrics_pb2.HistogramValue):
+    return cls(Histogram.from_runner_api(proto))
+
+  def get_result(self):
+    return self.histogram

sdks/python/apache_beam/internal/metrics/metric_test.py

@@ -17,17 +17,21 @@
 
 # pytype: skip-file
 
+import re
 import unittest
 
 from mock import patch
 
+import apache_beam as beam
 from apache_beam.internal.metrics.cells import HistogramCellFactory
 from apache_beam.internal.metrics.metric import Metrics as InternalMetrics
 from apache_beam.internal.metrics.metric import MetricLogger
 from apache_beam.metrics.execution import MetricsContainer
 from apache_beam.metrics.execution import MetricsEnvironment
 from apache_beam.metrics.metric import Metrics
+from apache_beam.metrics.metric import MetricsFilter
 from apache_beam.metrics.metricbase import MetricName
+from apache_beam.runners.direct.direct_runner import BundleBasedDirectRunner
 from apache_beam.runners.worker import statesampler
 from apache_beam.utils import counters
 from apache_beam.utils.histogram import LinearBucket
@@ -87,5 +91,40 @@ def test_create_process_wide(self):
       sampler.stop()
 
 
+class HistogramTest(unittest.TestCase):
+  def test_histogram(self):
+    class WordExtractingDoFn(beam.DoFn):
+      def __init__(self):
+        super().__init__()
+        self.word_lengths_dist = InternalMetrics.histogram(
+            self.__class__,
+            'latency_histogram_ms',
+            LinearBucket(0, 1, num_buckets=10))
+
+      def process(self, element):
+        text_line = element.strip()
+        words = re.findall(r'[\w\']+', text_line, re.UNICODE)
+        for w in words:
+          self.word_lengths_dist.update(len(w))
+        return words
+
+    with beam.Pipeline(runner=BundleBasedDirectRunner()) as p:
+      lines = p | 'read' >> beam.Create(["x x x yyyyyy yyyyyy yyyyyy"])
+      _ = (
+          lines
+          | 'split' >>
+          (beam.ParDo(WordExtractingDoFn()).with_output_types(str)))
+
+    result = p.result
+
+    filter = MetricsFilter().with_name('latency_histogram_ms')
+    query_result = result.metrics().query(filter)
+    histogram = query_result['histograms'][0].committed
+    assert histogram._buckets == {1: 3, 6: 3}
+    assert histogram.total_count() == 6
+    assert 1 < histogram.get_linear_interpolation(0.50) < 3
+    assert histogram.get_linear_interpolation(0.99) > 3
+
+
 if __name__ == '__main__':
   unittest.main()

sdks/python/apache_beam/metrics/execution.py

@@ -42,6 +42,8 @@
 from typing import Union
 from typing import cast
 
+from apache_beam.internal.metrics.cells import HistogramCellFactory
+from apache_beam.internal.metrics.cells import HistogramData
 from apache_beam.metrics import monitoring_infos
 from apache_beam.metrics.cells import BoundedTrieCell
 from apache_beam.metrics.cells import CounterCell
@@ -310,8 +312,14 @@ def get_cumulative(self):
         for k, v in self.metrics.items() if k.cell_type == BoundedTrieCell
     }
 
+    histograms = {
+        MetricKey(self.step_name, k.metric_name): v.get_cumulative()
+        for k, v in self.metrics.items()
+        if isinstance(k.cell_type, HistogramCellFactory)
+    }
+
     return MetricUpdates(
-        counters, distributions, gauges, string_sets, bounded_tries)
+        counters, distributions, gauges, string_sets, bounded_tries, histograms)
 
   def to_runner_api(self):
     return [
@@ -365,6 +373,7 @@ def __init__(
       gauges=None,  # type: Optional[Dict[MetricKey, GaugeData]]
       string_sets=None,  # type: Optional[Dict[MetricKey, StringSetData]]
       bounded_tries=None,  # type: Optional[Dict[MetricKey, BoundedTrieData]]
+      histograms=None,  # type: Optional[Dict[MetricKey, HistogramData]]
   ):
     # type: (...) -> None
 
@@ -382,3 +391,4 @@ def __init__(
     self.gauges = gauges or {}
     self.string_sets = string_sets or {}
     self.bounded_tries = bounded_tries or {}
+    self.histograms = histograms or {}

sdks/python/apache_beam/metrics/metric.py

@@ -195,6 +195,7 @@ class MetricResults(object):
   GAUGES = "gauges"
   STRINGSETS = "string_sets"
   BOUNDED_TRIES = "bounded_tries"
+  HISTOGRAMS = "histograms"
 
   @staticmethod
   def _matches_name(filter: 'MetricsFilter', metric_key: 'MetricKey') -> bool:

sdks/python/apache_beam/metrics/monitoring_infos.py

@@ -27,6 +27,7 @@
 
 from apache_beam.coders import coder_impl
 from apache_beam.coders import coders
+from apache_beam.internal.metrics.cells import HistogramData
 from apache_beam.metrics.cells import BoundedTrieData
 from apache_beam.metrics.cells import DistributionData
 from apache_beam.metrics.cells import DistributionResult
@@ -47,6 +48,7 @@
     common_urns.monitoring_info_specs.FINISH_BUNDLE_MSECS.spec.urn)
 TOTAL_MSECS_URN = common_urns.monitoring_info_specs.TOTAL_MSECS.spec.urn
 USER_COUNTER_URN = common_urns.monitoring_info_specs.USER_SUM_INT64.spec.urn
+USER_HISTOGRAM_URN = common_urns.monitoring_info_specs.USER_HISTOGRAM.spec.urn
 USER_DISTRIBUTION_URN = (
     common_urns.monitoring_info_specs.USER_DISTRIBUTION_INT64.spec.urn)
 USER_GAUGE_URN = common_urns.monitoring_info_specs.USER_LATEST_INT64.spec.urn
@@ -59,6 +61,7 @@
     USER_GAUGE_URN,
     USER_STRING_SET_URN,
     USER_BOUNDED_TRIE_URN,
+    USER_HISTOGRAM_URN
 ])
 WORK_REMAINING_URN = common_urns.monitoring_info_specs.WORK_REMAINING.spec.urn
 WORK_COMPLETED_URN = common_urns.monitoring_info_specs.WORK_COMPLETED.spec.urn
@@ -77,12 +80,14 @@
 PROGRESS_TYPE = common_urns.monitoring_info_types.PROGRESS_TYPE.urn
 STRING_SET_TYPE = common_urns.monitoring_info_types.SET_STRING_TYPE.urn
 BOUNDED_TRIE_TYPE = common_urns.monitoring_info_types.BOUNDED_TRIE_TYPE.urn
+HISTOGRAM_TYPE = common_urns.monitoring_info_types.HISTOGRAM.urn
 
 COUNTER_TYPES = set([SUM_INT64_TYPE])
 DISTRIBUTION_TYPES = set([DISTRIBUTION_INT64_TYPE])
 GAUGE_TYPES = set([LATEST_INT64_TYPE])
 STRING_SET_TYPES = set([STRING_SET_TYPE])
 BOUNDED_TRIE_TYPES = set([BOUNDED_TRIE_TYPE])
+HISTOGRAM_TYPES = set([HISTOGRAM_TYPE])
 
 # TODO(migryz) extract values from beam_fn_api.proto::MonitoringInfoLabels
 PCOLLECTION_LABEL = (
@@ -177,6 +182,14 @@ def extract_bounded_trie_value(monitoring_info_proto):
       metrics_pb2.BoundedTrie.FromString(monitoring_info_proto.payload))
 
 
+def extract_histogram_value(monitoring_info_proto):
+  if not is_histogram(monitoring_info_proto):
+    raise ValueError('Unsupported type %s' % monitoring_info_proto.type)
+
+  return HistogramData.from_proto(
+      metrics_pb2.HistogramValue.FromString(monitoring_info_proto.payload))
+
+
 def create_labels(ptransform=None, namespace=None, name=None, pcollection=None):
   """Create the label dictionary based on the provided values.
 
@@ -334,6 +347,25 @@ def user_set_string(namespace, name, metric, ptransform=None):
       USER_STRING_SET_URN, STRING_SET_TYPE, metric, labels)
 
 
+def user_histogram(namespace, name, metric: HistogramData, ptransform=None):
+  """Return the histogram monitoring info for the URN, metric and labels.
+
+  Args:
+    namespace: User-defined namespace of Histogram.
+    name: Name of Histogram.
+    metric: The Histogram representing the metrics.
+    ptransform: The ptransform id used as a label.
+  """
+  labels = create_labels(ptransform=ptransform, namespace=namespace, name=name)
+  metric_proto = metric.to_proto()
+
+  return create_monitoring_info(
+      USER_HISTOGRAM_URN,
+      HISTOGRAM_TYPE,
+      metric_proto.SerializeToString(),
+      labels)
+
+
 def user_bounded_trie(namespace, name, metric, ptransform=None):
   """Return the string set monitoring info for the URN, metric and labels.
 
@@ -353,7 +385,7 @@ def user_bounded_trie(namespace, name, metric, ptransform=None):
 
 def create_monitoring_info(
     urn, type_urn, payload, labels=None) -> metrics_pb2.MonitoringInfo:
-  """Return the gauge monitoring info for the URN, type, metric and labels.
+  """Return the monitoring info for the URN, type, metric and labels.
 
   Args:
     urn: The URN of the monitoring info/metric.
@@ -386,6 +418,11 @@ def is_distribution(monitoring_info_proto):
   return monitoring_info_proto.type in DISTRIBUTION_TYPES
 
 
+def is_histogram(monitoring_info_proto):
+  """Returns true if the monitoring info is a distrbution metric."""
+  return monitoring_info_proto.type in HISTOGRAM_TYPES
+
+
 def is_string_set(monitoring_info_proto):
   """Returns true if the monitoring info is a StringSet metric."""
   return monitoring_info_proto.type in STRING_SET_TYPES

sdks/python/apache_beam/metrics/monitoring_infos_test.py

@@ -18,10 +18,14 @@
 
 import unittest
 
+from apache_beam.internal.metrics.cells import HistogramCell
+from apache_beam.internal.metrics.cells import HistogramData
 from apache_beam.metrics import monitoring_infos
 from apache_beam.metrics.cells import CounterCell
 from apache_beam.metrics.cells import GaugeCell
 from apache_beam.metrics.cells import StringSetCell
+from apache_beam.utils.histogram import Histogram
+from apache_beam.utils.histogram import LinearBucket
 
 
 class MonitoringInfosTest(unittest.TestCase):
@@ -76,6 +80,17 @@ def test_parse_namespace_and_name_for_user_string_set_metric(self):
     self.assertEqual(namespace, "stringsetnamespace")
     self.assertEqual(name, "stringsetname")
 
+  def test_parse_namespace_and_name_for_user_histogram_metric(self):
+    urn = monitoring_infos.USER_HISTOGRAM_URN
+    labels = {}
+    labels[monitoring_infos.NAMESPACE_LABEL] = "histogramnamespace"
+    labels[monitoring_infos.NAME_LABEL] = "histogramname"
+    input = monitoring_infos.create_monitoring_info(
+        urn, "typeurn", None, labels)
+    namespace, name = monitoring_infos.parse_namespace_and_name(input)
+    self.assertEqual(name, "histogramname")
+    self.assertEqual(namespace, "histogramnamespace")
+
   def test_int64_user_gauge(self):
     metric = GaugeCell().get_cumulative()
     result = monitoring_infos.int64_user_gauge(
@@ -130,6 +145,26 @@ def test_user_set_string(self):
     self.assertEqual(set(), string_set_value)
     self.assertEqual(result.labels, expected_labels)
 
+  def test_user_histogram(self):
+    datapoints = [5, 50, 90]
+    expected_labels = {}
+    expected_labels[monitoring_infos.NAMESPACE_LABEL] = "histogramnamespace"
+    expected_labels[monitoring_infos.NAME_LABEL] = "histogramname"
+
+    cell = HistogramCell(LinearBucket(0, 1, 100))
+    for point in datapoints:
+      cell.update(point)
+    metric = cell.get_cumulative()
+    result = monitoring_infos.user_histogram(
+        'histogramnamespace', 'histogramname', metric)
+    histogramvalue = monitoring_infos.extract_histogram_value(result)
+
+    self.assertEqual(result.labels, expected_labels)
+    exp_histogram = Histogram(LinearBucket(0, 1, 100))
+    for point in datapoints:
+      exp_histogram.record(point)
+    self.assertEqual(HistogramData(exp_histogram), histogramvalue)
+
 
 if __name__ == '__main__':
   unittest.main()

sdks/python/apache_beam/runners/direct/direct_metrics.py

@@ -80,6 +80,27 @@ def result(self, x):
     return int(x)
 
 
+_IDENTITY_HISTOGRAM = object()
+
+
+class HistogramAggregator(MetricAggregator):
+  @staticmethod
+  def identity_element():
+    return _IDENTITY_HISTOGRAM
+
+  def combine(self, x, y):
+    if x is _IDENTITY_HISTOGRAM:
+      return y
+    if y is _IDENTITY_HISTOGRAM:
+      return x
+    return x.combine(y)
+
+  def result(self, x):
+    if x is _IDENTITY_HISTOGRAM:
+      raise TypeError
+    return x.get_result()
+
+
 class GenericAggregator(MetricAggregator):
   def __init__(self, data_class):
     self._data_class = data_class
@@ -105,6 +126,7 @@ def __init__(self):
         lambda: DirectMetric(GenericAggregator(StringSetData)))
     self._bounded_tries = defaultdict(
         lambda: DirectMetric(GenericAggregator(BoundedTrieData)))
+    self._histograms = defaultdict(lambda: DirectMetric(HistogramAggregator()))
 
   def _apply_operation(self, bundle, updates, op):
     for k, v in updates.counters.items():
@@ -122,6 +144,9 @@ def _apply_operation(self, bundle, updates, op):
     for k, v in updates.bounded_tries.items():
       op(self._bounded_tries[k], bundle, v)
 
+    for k, v in updates.histograms.items():
+      op(self._histograms[k], bundle, v)
+
   def commit_logical(self, bundle, updates):
     op = lambda obj, bundle, update: obj.commit_logical(bundle, update)
     self._apply_operation(bundle, updates, op)
@@ -170,13 +195,21 @@ def query(self, filter=None):
             v.extract_latest_attempted())
         for k, v in self._bounded_tries.items() if self.matches(filter, k)
     ]
+    histograms = [
+        MetricResult(
+            MetricKey(k.step, k.metric),
+            v.extract_committed(),
+            v.extract_latest_attempted()) for k, v in self._histograms.items()
+        if self.matches(filter, k)
+    ]
 
     return {
         self.COUNTERS: counters,
         self.DISTRIBUTIONS: distributions,
         self.GAUGES: gauges,
         self.STRINGSETS: string_sets,
         self.BOUNDED_TRIES: bounded_tries,
+        self.HISTOGRAMS: histograms,
     }