From 660ecd394910847d64e18e1f6f647a80e41aecf6 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Fri, 6 Nov 2026 20:28:58 +0000 Subject: [PATCH 1/3] Reuse cached _TypedMetricName in MetricsContainer cell lookups --- sdks/python/apache_beam/metrics/execution.py | 24 ++++++++++++++++---- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/metrics/execution.py b/sdks/python/apache_beam/metrics/execution.py index 124031898064..92c595600f64 100644 --- a/sdks/python/apache_beam/metrics/execution.py +++ b/sdks/python/apache_beam/metrics/execution.py @@ -257,32 +257,46 @@ def __init__(self, step_name): self.lock = threading.Lock() self.metrics = {} # type: Dict[_TypedMetricName, MetricCell] + def _get_typed_metric_name(self, cell_type, metric_name): + if not isinstance(metric_name, str): + typed_name = metric_name._typed_metric_names.get(cell_type) + if typed_name is None: + typed_name = _TypedMetricName(cell_type, metric_name) + metric_name._typed_metric_names[cell_type] = typed_name + return typed_name + return _TypedMetricName(cell_type, metric_name) + def get_counter(self, metric_name): # type: (MetricName) -> CounterCell return cast( CounterCell, - self.get_metric_cell(_TypedMetricName(CounterCell, metric_name))) + self.get_metric_cell( + self._get_typed_metric_name(CounterCell, metric_name))) def get_distribution(self, metric_name): # type: (MetricName) -> DistributionCell return cast( DistributionCell, - self.get_metric_cell(_TypedMetricName(DistributionCell, metric_name))) + self.get_metric_cell( + self._get_typed_metric_name(DistributionCell, metric_name))) def get_gauge(self, metric_name): # type: (MetricName) -> GaugeCell return cast( GaugeCell, - self.get_metric_cell(_TypedMetricName(GaugeCell, metric_name))) + self.get_metric_cell( + self._get_typed_metric_name(GaugeCell, metric_name))) def get_string_set(self, metric_name): # type: (MetricName) -> StringSetCell return cast( StringSetCell, - self.get_metric_cell(_TypedMetricName(StringSetCell, metric_name))) + self.get_metric_cell( + self._get_typed_metric_name(StringSetCell, metric_name))) def get_bounded_trie(self, metric_name): - return self.get_metric_cell(_TypedMetricName(BoundedTrieCell, metric_name)) + return self.get_metric_cell( + self._get_typed_metric_name(BoundedTrieCell, metric_name)) def get_metric_cell(self, typed_metric_name): # type: (_TypedMetricName) -> MetricCell From 77728c860a346096f088449f848a9099747ecf79 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Wed, 11 Nov 2026 20:28:32 +0000 Subject: [PATCH 2/3] Make _typed_metric_names cache initialization self-contained --- sdks/python/apache_beam/metrics/execution.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/sdks/python/apache_beam/metrics/execution.py b/sdks/python/apache_beam/metrics/execution.py index 92c595600f64..625b6cb17db6 100644 --- a/sdks/python/apache_beam/metrics/execution.py +++ b/sdks/python/apache_beam/metrics/execution.py @@ -259,10 +259,14 @@ def __init__(self, step_name): def _get_typed_metric_name(self, cell_type, metric_name): if not isinstance(metric_name, str): - typed_name = metric_name._typed_metric_names.get(cell_type) + try: + cache = metric_name._typed_metric_names + except AttributeError: + cache = metric_name._typed_metric_names = {} + typed_name = cache.get(cell_type) if typed_name is None: typed_name = _TypedMetricName(cell_type, metric_name) - metric_name._typed_metric_names[cell_type] = typed_name + cache[cell_type] = typed_name return typed_name return _TypedMetricName(cell_type, metric_name) From d29f41424edabe310ae038c089baf73cd71f603c Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Fri, 13 Nov 2026 07:59:16 +0000 Subject: [PATCH 3/3] Include fast_name in _TypedMetricName.__reduce__ to support cyclic pickling and deepcopy --- sdks/python/apache_beam/metrics/execution.py | 9 ++++++--- .../apache_beam/metrics/execution_test.py | 20 +++++++++++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/sdks/python/apache_beam/metrics/execution.py b/sdks/python/apache_beam/metrics/execution.py index 625b6cb17db6..6ce6b0814389 100644 --- a/sdks/python/apache_beam/metrics/execution.py +++ b/sdks/python/apache_beam/metrics/execution.py @@ -175,12 +175,15 @@ class _TypedMetricName(object): def __init__( self, cell_type, # type: Union[Type[MetricCell], MetricCellFactory] - metric_name # type: Union[str, MetricName] + metric_name, # type: Union[str, MetricName] + fast_name=None # type: Optional[str] ): # type: (...) -> None self.cell_type = cell_type self.metric_name = metric_name - if isinstance(metric_name, str): + if fast_name is not None: + self.fast_name = fast_name + elif isinstance(metric_name, str): self.fast_name = metric_name else: self.fast_name = metric_name.fast_name() @@ -198,7 +201,7 @@ def __str__(self): return '%s %s' % (self.cell_type, self.metric_name) def __reduce__(self): - return _TypedMetricName, (self.cell_type, self.metric_name) + return _TypedMetricName, (self.cell_type, self.metric_name, self.fast_name) _DEFAULT = None # type: Any diff --git a/sdks/python/apache_beam/metrics/execution_test.py b/sdks/python/apache_beam/metrics/execution_test.py index 38e27f1f3d0c..ebcd682cc91d 100644 --- a/sdks/python/apache_beam/metrics/execution_test.py +++ b/sdks/python/apache_beam/metrics/execution_test.py @@ -80,6 +80,26 @@ def test_add_to_counter(self): counter = mc.get_counter(MetricName('namespace', 'name')) self.assertEqual(counter.value, 1) + def test_pickle_and_deepcopy_metric_name_after_caching(self): + import copy + import pickle + from apache_beam.internal import pickler + mc = MetricsContainer('astep') + name = MetricName('namespace', 'name') + mc.get_counter(name).inc() + mc.get_distribution(name).update(1) + mc.get_gauge(name).set(2) + mc.get_string_set(name).add('a') + + self.assertEqual(pickle.loads(pickle.dumps(name)), name) + self.assertEqual(pickler.loads(pickler.dumps(name)), name) + self.assertEqual(copy.deepcopy(name), name) + + key = MetricKey('astep', name) + self.assertEqual(pickle.loads(pickle.dumps(key)), key) + self.assertEqual(pickler.loads(pickler.dumps(key)), key) + self.assertEqual(copy.deepcopy(key), key) + def test_get_cumulative_or_updates(self): mc = MetricsContainer('astep')