Source code for apache_beam.utils.counters

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

# cython: profile=False
# cython: overflowcheck=True

"""Counters collect the progress of the Worker for reporting to the service.

For internal use only; no backwards-compatibility guarantees.
"""

import threading
from apache_beam.transforms import cy_combiners


[docs]class Counter(object): """A counter aggregates a series of values. The aggregation kind of the Counter is specified when the Counter is created. The values aggregated must be of an appropriate for the aggregation used. Aggregations supported are listed in the code. (The aggregated value will be reported to the Dataflow service.) Do not create directly; call CounterFactory.get_counter instead. Attributes: name: the name of the counter, a string combine_fn: the CombineFn to use for aggregation accumulator: the accumulator created for the combine_fn """ # Handy references to common counters. SUM = cy_combiners.SumInt64Fn() MEAN = cy_combiners.MeanInt64Fn() def __init__(self, name, combine_fn): """Creates a Counter object. Args: name: the name of this counter. Typically has three parts: "step-output-counter". combine_fn: the CombineFn to use for aggregation """ self.name = name self.combine_fn = combine_fn self.accumulator = combine_fn.create_accumulator() self._add_input = self.combine_fn.add_input
[docs] def update(self, value): self.accumulator = self._add_input(self.accumulator, value)
[docs] def value(self): return self.combine_fn.extract_output(self.accumulator)
def __str__(self): return '<%s>' % self._str_internal() def __repr__(self): return '<%s at %s>' % (self._str_internal(), hex(id(self))) def _str_internal(self): return '%s %s %s' % (self.name, self.combine_fn.__class__.__name__, self.value())
[docs]class AccumulatorCombineFnCounter(Counter): """Counter optimized for a mutating accumulator that holds all the logic.""" def __init__(self, name, combine_fn): assert isinstance(combine_fn, cy_combiners.AccumulatorCombineFn) super(AccumulatorCombineFnCounter, self).__init__(name, combine_fn) self._fast_add_input = self.accumulator.add_input
[docs] def update(self, value): self._fast_add_input(value)
# Counters that represent Accumulators have names starting with this USER_COUNTER_PREFIX = 'user-'
[docs]class CounterFactory(object): """Keeps track of unique counters.""" def __init__(self): self.counters = {} # Lock to be acquired when accessing the counters map. self._lock = threading.Lock()
[docs] def get_counter(self, name, combine_fn): """Returns a counter with the requested name. Passing in the same name will return the same counter; the combine_fn must agree. Args: name: the name of this counter. Typically has three parts: "step-output-counter". combine_fn: the CombineFn to use for aggregation Returns: A new or existing counter with the requested name. """ with self._lock: counter = self.counters.get(name, None) if counter: assert counter.combine_fn == combine_fn else: if isinstance(combine_fn, cy_combiners.AccumulatorCombineFn): counter = AccumulatorCombineFnCounter(name, combine_fn) else: counter = Counter(name, combine_fn) self.counters[name] = counter return counter
[docs] def get_aggregator_counter(self, step_name, aggregator): """Returns an AggregationCounter for this step's aggregator. Passing in the same values will return the same counter. Args: step_name: the name of this step. aggregator: an Aggregator object. Returns: A new or existing counter. """ return self.get_counter( '%s%s-%s' % (USER_COUNTER_PREFIX, step_name, aggregator.name), aggregator.combine_fn)
[docs] def get_counters(self): """Returns the current set of counters. Returns: An iterable that contains the current set of counters. To make sure that multiple threads can iterate over the set of counters, we return a new iterable here. Note that the actual set of counters may get modified after this method returns hence the returned iterable may be stale. """ with self._lock: return self.counters.values()
[docs] def get_aggregator_values(self, aggregator_or_name): """Returns dict of step names to values of the aggregator.""" with self._lock: return get_aggregator_values( aggregator_or_name, self.counters, lambda counter: counter.value())
[docs]def get_aggregator_values(aggregator_or_name, counter_dict, value_extractor=None): """Extracts the named aggregator value from a set of counters. Args: aggregator_or_name: an Aggregator object or the name of one. counter_dict: a dict object of {name: value_wrapper} value_extractor: a function to convert the value_wrapper into a value. If None, no extraction is done and the value is return unchanged. Returns: dict of step names to values of the aggregator. """ name = aggregator_or_name if value_extractor is None: value_extractor = lambda x: x if not isinstance(aggregator_or_name, basestring): name = aggregator_or_name.name return {n: value_extractor(c) for n, c in counter_dict.iteritems() if n.startswith(USER_COUNTER_PREFIX) and n.endswith('-%s' % name)}