[Python] Honor disableCounterMetrics, disableStringSetMetrics and disableBoundedTrieMetrics experiments

CHANGES.md

@@ -69,6 +69,7 @@
 ## New Features / Improvements
 
 * (Python) Expanded the SDK worker heap dump (`--experiments=enable_heap_dump`) with process RSS, CPython allocator/GC stats, and glibc `mallinfo2` native-heap/fragmentation stats to help distinguish native-heap from Python-object memory growth ([#39244](https://github.com/apache/beam/issues/39244)).
+* The `disableCounterMetrics`, `disableStringSetMetrics` and `disableBoundedTrieMetrics` experiments are now honored by the Python SDK, as they already were in Java (Python) ([#38746](https://github.com/apache/beam/issues/38746)).
 
 ## Breaking Changes
 

sdks/python/apache_beam/metrics/execution.pxd

@@ -31,6 +31,7 @@ cdef class _TypedMetricName(object):
 
 
 cdef object _DEFAULT
+cdef set _DISABLED_CELL_TYPES
 
 
 cdef class MetricUpdater(object):

sdks/python/apache_beam/metrics/execution.py

@@ -39,10 +39,12 @@
 from typing import Dict
 from typing import FrozenSet
 from typing import Optional
+from typing import Set
 from typing import Type
 from typing import Union
 from typing import cast
 
+from apache_beam.metrics import metrics_flag
 from apache_beam.metrics import monitoring_infos
 from apache_beam.metrics.cells import BoundedTrieCell
 from apache_beam.metrics.cells import CounterCell
@@ -201,6 +203,11 @@ def __reduce__(self):
 
 _DEFAULT = None  # type: Any
 
+# Names of the metric cell types whose updates are dropped process-wide, owned
+# by apache_beam.metrics.metrics_flag.MetricsFlag from the disable*Metrics
+# experiments. Empty (the default) means every update is delivered.
+_DISABLED_CELL_TYPES = metrics_flag.DISABLED_CELL_TYPES  # type: Set[str]
+
 
 class MetricUpdater(object):
   """A callable that updates the metric as quickly as possible."""
@@ -216,6 +223,10 @@ def __init__(
 
   def __call__(self, value=_DEFAULT):
     # type: (Any) -> None
+    if _DISABLED_CELL_TYPES and (getattr(self.typed_metric_name.cell_type,
+                                         '__name__',
+                                         None) in _DISABLED_CELL_TYPES):
+      return
     if value is _DEFAULT:
       if self.default_value is _DEFAULT:
         raise ValueError(

sdks/python/apache_beam/metrics/metrics_flag.py

@@ -0,0 +1,100 @@
+#
+# 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.
+#
+
+"""Process-wide switches that stop kinds of user metrics from being reported.
+
+This module is imported by ``apache_beam.metrics.execution`` and by the
+pipeline and worker start-up code, so it deliberately keeps its imports to a
+minimum to avoid circular imports.
+"""
+
+# pytype: skip-file
+
+import logging
+from typing import TYPE_CHECKING
+from typing import Set
+
+if TYPE_CHECKING:
+  from apache_beam.options.pipeline_options import PipelineOptions
+
+__all__ = ['MetricsFlag']
+
+_LOGGER = logging.getLogger(__name__)
+
+# Names of the metric cell types whose updates are dropped process-wide. Empty
+# (the default) means every update is delivered. apache_beam.metrics.execution
+# holds a reference to this set, so it is only ever mutated in place.
+DISABLED_CELL_TYPES = set()  # type: Set[str]
+
+
+class MetricsFlag(object):
+  """Process-wide switches that stop kinds of user metrics from being reported.
+
+  High throughput jobs may want to turn off metrics that put pressure on the
+  metrics backend. Mirroring the Java SDK, the ``disableCounterMetrics``,
+  ``disableStringSetMetrics`` and ``disableBoundedTrieMetrics`` experiments make
+  the corresponding ``Metrics.counter``, ``Metrics.string_set`` and
+  ``Metrics.bounded_trie`` updates no-ops. The metric objects themselves are
+  unchanged, so code that holds on to them keeps working.
+  """
+  _EXPERIMENTS = (
+      ('disableCounterMetrics', 'CounterCell', 'Counter'),
+      ('disableStringSetMetrics', 'StringSetCell', 'StringSet'),
+      ('disableBoundedTrieMetrics', 'BoundedTrieCell', 'BoundedTrie'),
+  )
+  _initialized = False
+
+  @classmethod
+  def set_default_pipeline_options(cls, options: 'PipelineOptions') -> None:
+    """Initializes the flags from ``options`` if not already done so.
+
+    Called when a ``Pipeline`` is constructed and at SDK worker harness
+    start-up. As in the Java SDK, the first call wins so that user code running
+    on a worker cannot change the flags the harness was started with.
+    """
+    if cls._initialized:
+      return
+    # Imported here rather than at module level, as a precaution against
+    # circular imports between apache_beam.metrics and apache_beam.options.
+    from apache_beam.options.pipeline_options import DebugOptions
+    debug_options = options.view_as(DebugOptions)
+    disabled = set()
+    for experiment, cell_type_name, kind in cls._EXPERIMENTS:
+      if debug_options.lookup_experiment(experiment):
+        disabled.add(cell_type_name)
+        _LOGGER.info('%s metrics are disabled.', kind)
+    DISABLED_CELL_TYPES.clear()
+    DISABLED_CELL_TYPES.update(disabled)
+    cls._initialized = True
+
+  @classmethod
+  def counter_disabled(cls) -> bool:
+    return 'CounterCell' in DISABLED_CELL_TYPES
+
+  @classmethod
+  def string_set_disabled(cls) -> bool:
+    return 'StringSetCell' in DISABLED_CELL_TYPES
+
+  @classmethod
+  def bounded_trie_disabled(cls) -> bool:
+    return 'BoundedTrieCell' in DISABLED_CELL_TYPES
+
+  @classmethod
+  def reset(cls) -> None:
+    """Clears the flags so the next ``set_default_pipeline_options`` applies."""
+    DISABLED_CELL_TYPES.clear()
+    cls._initialized = False

sdks/python/apache_beam/metrics/metrics_flag_test.py

@@ -0,0 +1,225 @@
+#
+# 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.
+#
+
+# pytype: skip-file
+
+import pickle
+import unittest
+
+import apache_beam as beam
+from apache_beam.metrics.cells import DistributionData
+from apache_beam.metrics.execution import MetricKey
+from apache_beam.metrics.execution import MetricsContainer
+from apache_beam.metrics.execution import MetricsEnvironment
+from apache_beam.metrics.execution import MetricUpdater
+from apache_beam.metrics.metric import Metrics
+from apache_beam.metrics.metric import MetricsFilter
+from apache_beam.metrics.metricbase import MetricName
+from apache_beam.metrics.metrics_flag import MetricsFlag
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.runners.worker import statesampler
+from apache_beam.testing.test_pipeline import TestPipeline
+from apache_beam.testing.util import assert_that
+from apache_beam.testing.util import equal_to
+from apache_beam.utils import counters
+
+
+class MetricsFlagTest(unittest.TestCase):
+  """Covers the disable*Metrics experiments.
+
+  See https://github.com/apache/beam/issues/38746.
+  """
+  def setUp(self):
+    MetricsFlag.reset()
+    self.sampler = statesampler.StateSampler('', counters.CounterFactory())
+    statesampler.set_current_tracker(self.sampler)
+    self.state = self.sampler.scoped_state(
+        'mystep', 'myState', metrics_container=MetricsContainer('mystep'))
+    self.sampler.start()
+
+  def tearDown(self):
+    self.sampler.stop()
+    MetricsFlag.reset()
+
+  @staticmethod
+  def _set_experiments(*experiments):
+    MetricsFlag.set_default_pipeline_options(
+        PipelineOptions(['--experiments=%s' % exp for exp in experiments]))
+
+  def test_flags_follow_experiments(self):
+    self.assertFalse(MetricsFlag.counter_disabled())
+    self.assertFalse(MetricsFlag.string_set_disabled())
+    self.assertFalse(MetricsFlag.bounded_trie_disabled())
+
+    for experiment, expected in [
+        ('disableCounterMetrics', (True, False, False)),
+        ('disableStringSetMetrics', (False, True, False)),
+        ('disableBoundedTrieMetrics', (False, False, True)),
+    ]:
+      MetricsFlag.reset()
+      self._set_experiments(experiment)
+      self.assertEqual((
+          MetricsFlag.counter_disabled(),
+          MetricsFlag.string_set_disabled(),
+          MetricsFlag.bounded_trie_disabled()),
+                       expected,
+                       experiment)
+
+    MetricsFlag.reset()
+    self._set_experiments(
+        'disableCounterMetrics',
+        'disableStringSetMetrics',
+        'disableBoundedTrieMetrics')
+    self.assertTrue(MetricsFlag.counter_disabled())
+    self.assertTrue(MetricsFlag.string_set_disabled())
+    self.assertTrue(MetricsFlag.bounded_trie_disabled())
+
+  def test_first_call_wins(self):
+    self._set_experiments('disableCounterMetrics')
+    # Later options, e.g. from user code constructing a Pipeline on a worker,
+    # do not change the flags the harness was started with.
+    self._set_experiments('disableStringSetMetrics')
+    self.assertTrue(MetricsFlag.counter_disabled())
+    self.assertFalse(MetricsFlag.string_set_disabled())
+
+  def test_update_call_shapes_keep_working(self):
+    # The first attempt at this feature (#38749) was reverted because it
+    # changed the signature of DelegatingCounter.inc; the metric objects must
+    # stay MetricUpdater callables that accept the value as a keyword too.
+    with self.state:
+      counter = Metrics.counter('ns', 'counter')
+      self.assertIsInstance(counter.inc, MetricUpdater)
+      counter.inc()
+      counter.inc(4)
+      counter.inc(value=5)
+      counter.dec()
+      counter.dec(2)
+      string_set = Metrics.string_set('ns', 'set')
+      string_set.add('a')
+      string_set.add(value='b')
+      container = MetricsEnvironment.current_container()
+      self.assertEqual(
+          container.get_counter(MetricName('ns', 'counter')).get_cumulative(),
+          7)
+      self.assertEqual(
+          container.get_string_set(MetricName(
+              'ns', 'set')).get_cumulative().string_set, {'a', 'b'})
+
+  def test_disabled_counter_is_noop(self):
+    with self.state:
+      container = MetricsEnvironment.current_container()
+      Metrics.counter('ns', 'before').inc()
+      self.assertEqual(len(container.metrics), 1)
+
+      self._set_experiments('disableCounterMetrics')
+      created_before = Metrics.counter('ns', 'before')
+      created_before.inc()
+      created_before.inc(value=5)
+      created_before.dec()
+      Metrics.counter('ns', 'after').inc(3)
+      self.assertEqual(len(container.metrics), 1)
+      self.assertEqual(
+          container.get_counter(MetricName('ns', 'before')).get_cumulative(), 1)
+
+      # Other kinds keep reporting.
+      Metrics.distribution('ns', 'dist').update(3)
+      Metrics.gauge('ns', 'gauge').set(2)
+      Metrics.string_set('ns', 'set').add('x')
+      Metrics.bounded_trie('ns', 'trie').add(('x', ))
+      self.assertEqual(len(container.metrics), 5)
+
+  def test_disabled_string_set_is_noop(self):
+    with self.state:
+      container = MetricsEnvironment.current_container()
+      Metrics.string_set('ns', 'before').add('seed')
+      self.assertEqual(len(container.metrics), 1)
+
+      self._set_experiments('disableStringSetMetrics')
+      Metrics.string_set('ns', 'before').add('more')
+      Metrics.string_set('ns', 'after').add('value')
+      self.assertEqual(len(container.metrics), 1)
+      self.assertEqual(
+          container.get_string_set(MetricName(
+              'ns', 'before')).get_cumulative().string_set, {'seed'})
+      Metrics.counter('ns', 'counter').inc()
+      self.assertEqual(len(container.metrics), 2)
+
+  def test_disabled_bounded_trie_is_noop(self):
+    with self.state:
+      container = MetricsEnvironment.current_container()
+      Metrics.bounded_trie('ns', 'before').add(('a', ))
+      self.assertEqual(len(container.metrics), 1)
+
+      self._set_experiments('disableBoundedTrieMetrics')
+      Metrics.bounded_trie('ns', 'before').add(('a', 'b'))
+      Metrics.bounded_trie('ns', 'after').add(('c', ))
+      self.assertEqual(len(container.metrics), 1)
+      self.assertEqual(
+          list(
+              container.get_bounded_trie(MetricName(
+                  'ns', 'before')).get_cumulative().flattened()),
+          [('a', False)])
+
+  def test_disabled_process_wide_counter_is_noop(self):
+    self._set_experiments('disableCounterMetrics')
+    name = MetricName('ns', 'process_wide')
+    counter = Metrics.DelegatingCounter(name, process_wide=True)
+    counter.inc()
+    # The process-wide container is shared with every other test in this
+    # process, so only check that this counter never reached it.
+    self.assertNotIn(
+        MetricKey(None, name),
+        MetricsEnvironment.process_wide_container().get_cumulative().counters)
+
+  def test_disabled_flag_applies_to_unpickled_metrics(self):
+    # DoFns holding metric objects are pickled at submission time and
+    # unpickled on the worker, where the harness sets the flags.
+    counter = Metrics.counter('ns', 'pickled')
+    counter = pickle.loads(pickle.dumps(counter))
+    self.assertIsInstance(counter.inc, MetricUpdater)
+    self._set_experiments('disableCounterMetrics')
+    with self.state:
+      counter.inc()
+      self.assertEqual(len(MetricsEnvironment.current_container().metrics), 0)
+
+  def test_disabled_counters_in_pipeline(self):
+    class SomeDoFn(beam.DoFn):
+      def process(self, element):
+        Metrics.counter(self.__class__, 'elements').inc()
+        Metrics.distribution(self.__class__, 'element_dist').update(element)
+        yield element
+
+    MetricsFlag.reset()
+    pipeline = TestPipeline(
+        options=PipelineOptions(['--experiments=disableCounterMetrics']))
+    results = pipeline | beam.Create([1, 2, 3]) | beam.ParDo(SomeDoFn())
+    assert_that(results, equal_to([1, 2, 3]))
+    res = pipeline.run()
+    res.wait_until_finish()
+
+    self.assertEqual(
+        res.metrics().query(MetricsFilter().with_name('elements'))['counters'],
+        [])
+    distributions = res.metrics().query(
+        MetricsFilter().with_name('element_dist'))['distributions']
+    self.assertEqual(len(distributions), 1)
+    self.assertEqual(
+        distributions[0].committed.data, DistributionData(6, 3, 1, 3))
+
+
+if __name__ == '__main__':
+  unittest.main()

sdks/python/apache_beam/pipeline.py

@@ -73,6 +73,7 @@
 from apache_beam.coders import typecoders
 from apache_beam.internal import pickler
 from apache_beam.io.filesystems import FileSystems
+from apache_beam.metrics.metrics_flag import MetricsFlag
 from apache_beam.options.pipeline_options import CrossLanguageOptions
 from apache_beam.options.pipeline_options import DebugOptions
 from apache_beam.options.pipeline_options import PipelineOptions
@@ -192,6 +193,7 @@ def __init__(
       self._options = PipelineOptions([])
 
     FileSystems.set_options(self._options)
+    MetricsFlag.set_default_pipeline_options(self._options)
 
     if runner is None:
       runner = self._options.view_as(StandardOptions).runner