def _get_beam_metric(metric_filter: beam.metrics.MetricsFilter,
result: beam.runners.runner.PipelineResult,
metric_type: str = 'counters') -> Optional[int]:
"""Queries a beam pipeline result for a specificed metric.
Args:
metric_filter: an instance of apache_beam.metrics.MetricsFilter()
metric_type: A metric type (counters, distributions, etc.)
Returns:
Counter value or None
"""
query_result = result.metrics().query(metric_filter)
result_val = None
if query_result[metric_type]:
result_val = query_result[metric_type][0].result
return result_val
def ValidateMetrics(test: unittest.TestCase,
pipeline_result: beam.runners.runner.PipelineResult,
telemetry_descriptors: List[Text], logical_format: Text,
physical_format: Text):
all_metrics = pipeline_result.metrics()
maintained_metrics = all_metrics.query(
beam.metrics.metric.MetricsFilter().with_namespace(
"tfx." + ".".join(telemetry_descriptors + ["io"])))
test.assertIsNot(maintained_metrics, None)
counters = maintained_metrics[beam.metrics.metric.MetricResults.COUNTERS]
test.assertTrue(counters)
distributions = maintained_metrics[
beam.metrics.metric.MetricResults.DISTRIBUTIONS]
test.assertTrue(distributions)
for m in counters + distributions:
test.assertTrue(
m.key.metric.name.startswith(
"LogicalFormat[%s]-PhysicalFormat[%s]-" %
(logical_format, physical_format)))