github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/direct/direct_metrics.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """
    19  DirectRunner implementation of MetricResults. It is in charge not only of
    20  responding to queries of current metrics, but also of keeping the common
    21  state consistent.
    22  """
    23  # pytype: skip-file
    24  
    25  import threading
    26  from collections import defaultdict
    27  
    28  from apache_beam.metrics.cells import CounterAggregator
    29  from apache_beam.metrics.cells import DistributionAggregator
    30  from apache_beam.metrics.cells import GaugeAggregator
    31  from apache_beam.metrics.execution import MetricKey
    32  from apache_beam.metrics.execution import MetricResult
    33  from apache_beam.metrics.metric import MetricResults
    34  
    35  
    36  class DirectMetrics(MetricResults):
    37    def __init__(self):
    38      self._counters = defaultdict(lambda: DirectMetric(CounterAggregator()))
    39      self._distributions = defaultdict(
    40          lambda: DirectMetric(DistributionAggregator()))
    41      self._gauges = defaultdict(lambda: DirectMetric(GaugeAggregator()))
    42  
    43    def _apply_operation(self, bundle, updates, op):
    44      for k, v in updates.counters.items():
    45        op(self._counters[k], bundle, v)
    46  
    47      for k, v in updates.distributions.items():
    48        op(self._distributions[k], bundle, v)
    49  
    50      for k, v in updates.gauges.items():
    51        op(self._gauges[k], bundle, v)
    52  
    53    def commit_logical(self, bundle, updates):
    54      op = lambda obj, bundle, update: obj.commit_logical(bundle, update)
    55      self._apply_operation(bundle, updates, op)
    56  
    57    def commit_physical(self, bundle, updates):
    58      op = lambda obj, bundle, update: obj.commit_physical(bundle, update)
    59      self._apply_operation(bundle, updates, op)
    60  
    61    def update_physical(self, bundle, updates):
    62      op = lambda obj, bundle, update: obj.update_physical(bundle, update)
    63      self._apply_operation(bundle, updates, op)
    64  
    65    def query(self, filter=None):
    66      counters = [
    67          MetricResult(
    68              MetricKey(k.step, k.metric),
    69              v.extract_committed(),
    70              v.extract_latest_attempted()) for k,
    71          v in self._counters.items() if self.matches(filter, k)
    72      ]
    73      distributions = [
    74          MetricResult(
    75              MetricKey(k.step, k.metric),
    76              v.extract_committed(),
    77              v.extract_latest_attempted()) for k,
    78          v in self._distributions.items() if self.matches(filter, k)
    79      ]
    80      gauges = [
    81          MetricResult(
    82              MetricKey(k.step, k.metric),
    83              v.extract_committed(),
    84              v.extract_latest_attempted()) for k,
    85          v in self._gauges.items() if self.matches(filter, k)
    86      ]
    87  
    88      return {
    89          self.COUNTERS: counters,
    90          self.DISTRIBUTIONS: distributions,
    91          self.GAUGES: gauges
    92      }
    93  
    94  
    95  class DirectMetric(object):
    96    """ Keeps a consistent state for a single metric.
    97  
    98    It keeps track of the metric's physical and logical updates.
    99    It's thread safe.
   100    """
   101    def __init__(self, aggregator):
   102      self.aggregator = aggregator
   103      self._attempted_lock = threading.Lock()
   104      self.finished_attempted = aggregator.identity_element()
   105      self.inflight_attempted = {}
   106      self._committed_lock = threading.Lock()
   107      self.finished_committed = aggregator.identity_element()
   108  
   109    def commit_logical(self, bundle, update):
   110      with self._committed_lock:
   111        self.finished_committed = self.aggregator.combine(
   112            update, self.finished_committed)
   113  
   114    def commit_physical(self, bundle, update):
   115      with self._attempted_lock:
   116        self.inflight_attempted[bundle] = update
   117        self.finished_attempted = self.aggregator.combine(
   118            update, self.finished_attempted)
   119        del self.inflight_attempted[bundle]
   120  
   121    def update_physical(self, bundle, update):
   122      self.inflight_attempted[bundle] = update
   123  
   124    def extract_committed(self):
   125      return self.aggregator.result(self.finished_committed)
   126  
   127    def extract_latest_attempted(self):
   128      res = self.finished_attempted
   129      for _, u in self.inflight_attempted.items():
   130        res = self.aggregator.combine(res, u)
   131  
   132      return self.aggregator.result(res)