1use super::MetricsSummary;
7use std::collections::HashMap;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum ExportFormat {
12 Prometheus,
14 Json,
16 Csv,
18}
19
20pub struct MetricsExporter {
22 prefix: String,
24 labels: HashMap<String, String>,
26}
27
28impl MetricsExporter {
29 pub fn new(prefix: impl Into<String>) -> Self {
31 Self { prefix: prefix.into(), labels: HashMap::new() }
32 }
33
34 pub fn with_label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
36 self.labels.insert(key.into(), value.into());
37 self
38 }
39
40 pub fn to_prometheus(&self, summary: &MetricsSummary) -> String {
42 let mut output = String::new();
43
44 for (metric, stats) in summary {
45 let name = format!("{}_{}", self.prefix, metric.as_str());
46 let labels = self.format_labels();
47
48 output.push_str(&format!("# HELP {} Training metric: {}\n", name, metric.as_str()));
50 output.push_str(&format!("# TYPE {name} gauge\n"));
51
52 output.push_str(&format!("{}{} {}\n", name, labels, stats.mean));
54
55 output.push_str(&format!("{}_min{} {}\n", name, labels, stats.min));
57 output.push_str(&format!("{}_max{} {}\n", name, labels, stats.max));
58 output.push_str(&format!("{}_std{} {}\n", name, labels, stats.std));
59 output.push_str(&format!("{}_count{} {}\n", name, labels, stats.count));
60 output.push('\n');
61 }
62
63 output
64 }
65
66 pub fn to_json(&self, summary: &MetricsSummary) -> Result<String, serde_json::Error> {
68 let mut export: HashMap<String, serde_json::Value> = HashMap::new();
69
70 export.insert("prefix".to_string(), self.prefix.clone().into());
71 export.insert("labels".to_string(), serde_json::to_value(&self.labels)?);
72
73 let metrics: HashMap<String, serde_json::Value> = summary
74 .iter()
75 .map(|(k, v)| {
76 (
77 k.as_str().to_string(),
78 serde_json::json!({
79 "mean": v.mean,
80 "std": v.std,
81 "min": v.min,
82 "max": v.max,
83 "count": v.count,
84 "sum": v.sum,
85 "has_nan": v.has_nan,
86 "has_inf": v.has_inf,
87 }),
88 )
89 })
90 .collect();
91
92 export.insert("metrics".to_string(), serde_json::to_value(metrics)?);
93
94 serde_json::to_string_pretty(&export)
95 }
96
97 pub fn to_csv(&self, summary: &MetricsSummary) -> String {
99 let mut output = String::from("metric,mean,std,min,max,count,sum\n");
100
101 for (metric, stats) in summary {
102 output.push_str(&format!(
103 "{},{},{},{},{},{},{}\n",
104 metric.as_str(),
105 stats.mean,
106 stats.std,
107 stats.min,
108 stats.max,
109 stats.count,
110 stats.sum
111 ));
112 }
113
114 output
115 }
116
117 fn format_labels(&self) -> String {
119 if self.labels.is_empty() {
120 String::new()
121 } else {
122 let pairs: Vec<String> =
123 self.labels.iter().map(|(k, v)| format!("{k}=\"{v}\"")).collect();
124 format!("{{{}}}", pairs.join(","))
125 }
126 }
127}
128
129impl Default for MetricsExporter {
130 fn default() -> Self {
131 Self::new("entrenar")
132 }
133}
134
135#[cfg(test)]
136mod tests {
137 use super::*;
138 use crate::monitor::{Metric, MetricStats};
139
140 fn sample_summary() -> MetricsSummary {
141 let mut summary = HashMap::new();
142 summary.insert(
143 Metric::Loss,
144 MetricStats {
145 count: 100,
146 mean: 0.25,
147 std: 0.1,
148 min: 0.1,
149 max: 0.5,
150 sum: 25.0,
151 has_nan: false,
152 has_inf: false,
153 },
154 );
155 summary.insert(
156 Metric::Accuracy,
157 MetricStats {
158 count: 100,
159 mean: 0.85,
160 std: 0.05,
161 min: 0.7,
162 max: 0.95,
163 sum: 85.0,
164 has_nan: false,
165 has_inf: false,
166 },
167 );
168 summary
169 }
170
171 #[test]
172 fn test_exporter_new() {
173 let exporter = MetricsExporter::new("test");
174 assert_eq!(exporter.prefix, "test");
175 }
176
177 #[test]
178 fn test_exporter_with_labels() {
179 let exporter =
180 MetricsExporter::new("test").with_label("model", "v1").with_label("env", "prod");
181 assert_eq!(exporter.labels.len(), 2);
182 }
183
184 #[test]
185 fn test_to_prometheus() {
186 let exporter = MetricsExporter::new("training");
187 let summary = sample_summary();
188 let prom = exporter.to_prometheus(&summary);
189
190 assert!(prom.contains("# HELP training_loss"));
191 assert!(prom.contains("# TYPE training_loss gauge"));
192 assert!(prom.contains("training_loss 0.25"));
193 assert!(prom.contains("training_loss_min 0.1"));
194 }
195
196 #[test]
197 fn test_to_prometheus_with_labels() {
198 let exporter = MetricsExporter::new("training").with_label("model", "v1");
199 let summary = sample_summary();
200 let prom = exporter.to_prometheus(&summary);
201
202 assert!(prom.contains("model=\"v1\""));
203 }
204
205 #[test]
206 fn test_to_json() {
207 let exporter = MetricsExporter::new("test");
208 let summary = sample_summary();
209 let json = exporter.to_json(&summary).expect("operation should succeed");
210
211 assert!(json.contains("\"prefix\": \"test\""));
212 assert!(json.contains("\"loss\""));
213 assert!(json.contains("\"mean\": 0.25"));
214 }
215
216 #[test]
217 fn test_to_csv() {
218 let exporter = MetricsExporter::new("test");
219 let summary = sample_summary();
220 let csv = exporter.to_csv(&summary);
221
222 assert!(csv.contains("metric,mean,std"));
223 assert!(csv.contains("loss,0.25,0.1"));
224 }
225}