Skip to content

[Python] Reuse cached _TypedMetricName in MetricsContainer cell lookups #41275

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
damccorm wants to merge 3 commits into
base: master
Choose a base branch
from
+49 −8
Draft

[Python] Reuse cached _TypedMetricName in MetricsContainer cell lookups #41275

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
660ecd3
Reuse cached _TypedMetricName in MetricsContainer cell lookups
damccorm Nov 6, 2026
77728c8
Make _typed_metric_names cache initialization self-contained
damccorm Nov 11, 2026
d29f414
Include fast_name in _TypedMetricName.__reduce__ to support cyclic pi…
damccorm Nov 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 29 additions & 8 deletions sdks/python/apache_beam/metrics/execution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
Expand Down Expand Up @@ -257,32 +260,50 @@ 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):
try:
cache = metric_name._typed_metric_names
except AttributeError:
cache = metric_name._typed_metric_names = {}

@github-actions github-actions Bot Nov 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 High: Caching _TypedMetricName(cell_type, metric_name) in metric_name._typed_metric_names introduces a reference cycle (metric_name -> _typed_metric_names -> _TypedMetricName -> metric_name) that causes pickle.loads(pickle.dumps(metric_name)), dill.loads(dill.dumps(metric_name)), and copy.deepcopy(metric_name) (as well as MetricKey / MetricResult containing metric_name) to fail with AttributeError: 'MetricName' object has no attribute 'name' once metric_name has been passed to MetricsContainer.get_*.

During unpickling or deepcopy of metric_name, Python allocates an uninitialized MetricName.__new__(MetricName) and memoizes it before reconstructing its __dict__. Reconstructing _typed_metric_names invokes _TypedMetricName.__reduce__() -> _TypedMetricName(cell_type, metric_name), whose __init__ calls metric_name.fast_name() before metric_name.__dict__ has been populated.

To fix this, define __reduce__ on MetricName in sdks/python/apache_beam/metrics/metricbase.py so the transient cache is not serialized:

def __reduce__(self):
  return MetricName, (self.namespace, self.name, self.urn, self.labels)

@damccorm damccorm Nov 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great catch! Fixed in d29f414: updated _TypedMetricName.__reduce__ and __init__ to pass fast_name directly so unpickling / pickler / deepcopy of a MetricName cycle does not invoke metric_name.fast_name() before metric_name.__dict__ is restored, and added TestMetricsContainer.test_pickle_and_deepcopy_metric_name_after_caching.

typed_name = cache.get(cell_type)
if typed_name is None:
typed_name = _TypedMetricName(cell_type, metric_name)
cache[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
Expand Down
20 changes: 20 additions & 0 deletions sdks/python/apache_beam/metrics/execution_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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')

Expand Down
Loading