32
32
import glob
33
33
import os
34
34
import tempfile
35
+ from typing import Collection
35
36
36
37
from prometheus_client import (
38
+ Metric ,
37
39
CollectorRegistry ,
38
40
core ,
39
41
generate_latest ,
@@ -55,28 +57,29 @@ def collect_metrics():
55
57
return generate_latest (registry )
56
58
57
59
58
- def prometheus_setup_worker (labels = None ):
60
+ def prometheus_setup_worker (labels : None | dict [ str , str ] = None ) -> None :
59
61
"""
60
62
Monkey-patch mmap_key and ValueClass to add implicit labels. This must be
61
63
done before any metrics are instantiated.
62
64
"""
63
- if labels :
65
+ if labels is not None :
64
66
from prometheus_client import values
65
67
66
- def mmap_key (metric_name , name , labelnames , labelvalues ) :
68
+ def mmap_key (metric_name : str , name : str , labelnames : list [ str ] , labelvalues : list [ str ], help_text : str ) -> str :
67
69
return mmap_dict .mmap_key (
68
70
metric_name ,
69
71
name ,
70
- tuple (labels .keys ()) + labelnames ,
71
- tuple (labels .values ()) + labelvalues ,
72
+ list (labels .keys ()) + list (labelnames ) if labels else labelnames ,
73
+ list (labels .values ()) + list (labelvalues ) if labels else labelvalues ,
74
+ help_text ,
72
75
)
73
76
74
77
values .mmap_key = mmap_key
75
78
# synthesize a new ValueClass (captures mmap_key)
76
79
values .ValueClass = values .get_value_class ()
77
80
78
81
79
- def prometheus_cleanup_worker (pid ) :
82
+ def prometheus_cleanup_worker (pid : int ) -> None :
80
83
"""
81
84
Aggregate dead worker's metrics into a single archive file, preventing
82
85
collection time from growing without bound as pointed out in
@@ -109,7 +112,7 @@ def prometheus_cleanup_worker(pid):
109
112
collect_paths = paths + archive_paths
110
113
collector = multiprocess .MultiProcessCollector (None )
111
114
112
- metrics = collector .merge (collect_paths , accumulate = False )
115
+ metrics : Collection [ Metric ] = collector .merge (collect_paths , accumulate = False )
113
116
114
117
tmp_histogram = tempfile .NamedTemporaryFile (delete = False )
115
118
tmp_counter = tempfile .NamedTemporaryFile (delete = False )
@@ -124,7 +127,7 @@ def prometheus_cleanup_worker(pid):
124
127
os .unlink (path )
125
128
126
129
127
- def write_metrics (metrics , histogram_file , counter_file ) :
130
+ def write_metrics (metrics : Collection [ Metric ] , histogram_file : str , counter_file : str ) -> None :
128
131
129
132
histograms = mmap_dict .MmapedDict (histogram_file )
130
133
counters = mmap_dict .MmapedDict (counter_file )
@@ -142,7 +145,7 @@ def write_metrics(metrics, histogram_file, counter_file):
142
145
# prometheus_client 0.4+ adds extra fields
143
146
name , labels , value = sample [:3 ]
144
147
key = mmap_dict .mmap_key (
145
- metric .name , name , tuple (labels ), tuple (labels .values ()),
148
+ metric .name , name , list (labels . keys ()), list (labels .values ()), metric . documentation ,
146
149
)
147
150
sink .write_value (key , value )
148
151
finally :
0 commit comments