Skip to content

Commit

Permalink
fix(deps): prometheus_client 0.16 compat
Browse files Browse the repository at this point in the history
Pass through extra arguments to mmap_key after prometheus/client_python#804
  • Loading branch information
jvansanten committed Feb 6, 2023
1 parent 5ce819a commit 3129ec7
Showing 1 changed file with 12 additions and 9 deletions.
21 changes: 12 additions & 9 deletions ampel/metrics/prometheus.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,8 +32,10 @@
import glob
import os
import tempfile
from typing import Collection

from prometheus_client import (
Metric,
CollectorRegistry,
core,
generate_latest,
Expand All @@ -55,28 +57,29 @@ def collect_metrics():
return generate_latest(registry)


def prometheus_setup_worker(labels=None):
def prometheus_setup_worker(labels: None | dict[str,str] = None) -> None:
"""
Monkey-patch mmap_key and ValueClass to add implicit labels. This must be
done before any metrics are instantiated.
"""
if labels:
if labels is not None:
from prometheus_client import values

def mmap_key(metric_name, name, labelnames, labelvalues):
def mmap_key(metric_name: str, name: str, labelnames: list[str], labelvalues: list[str], help_text: str) -> str:
return mmap_dict.mmap_key(
metric_name,
name,
tuple(labels.keys()) + labelnames,
tuple(labels.values()) + labelvalues,
list(labels.keys()) + list(labelnames) if labels else labelnames,
list(labels.values()) + list(labelvalues) if labels else labelvalues,
help_text,
)

values.mmap_key = mmap_key
# synthesize a new ValueClass (captures mmap_key)
values.ValueClass = values.get_value_class()


def prometheus_cleanup_worker(pid):
def prometheus_cleanup_worker(pid: int) -> None:
"""
Aggregate dead worker's metrics into a single archive file, preventing
collection time from growing without bound as pointed out in
Expand Down Expand Up @@ -109,7 +112,7 @@ def prometheus_cleanup_worker(pid):
collect_paths = paths + archive_paths
collector = multiprocess.MultiProcessCollector(None)

metrics = collector.merge(collect_paths, accumulate=False)
metrics: Collection[Metric] = collector.merge(collect_paths, accumulate=False)

tmp_histogram = tempfile.NamedTemporaryFile(delete=False)
tmp_counter = tempfile.NamedTemporaryFile(delete=False)
Expand All @@ -124,7 +127,7 @@ def prometheus_cleanup_worker(pid):
os.unlink(path)


def write_metrics(metrics, histogram_file, counter_file):
def write_metrics(metrics: Collection[Metric], histogram_file: str, counter_file: str) -> None:

histograms = mmap_dict.MmapedDict(histogram_file)
counters = mmap_dict.MmapedDict(counter_file)
Expand All @@ -142,7 +145,7 @@ def write_metrics(metrics, histogram_file, counter_file):
# prometheus_client 0.4+ adds extra fields
name, labels, value = sample[:3]
key = mmap_dict.mmap_key(
metric.name, name, tuple(labels), tuple(labels.values()),
metric.name, name, list(labels.keys()), list(labels.values()), metric.documentation,
)
sink.write_value(key, value)
finally:
Expand Down

0 comments on commit 3129ec7

Please sign in to comment.