Upload prometheus_adapter.py

aworld/metrics/prometheus/prometheus_adapter.py

@@ -0,0 +1,386 @@
+import time
+import threading
+from typing import Sequence, Optional, Dict, List
+from prometheus_client import Counter as PCounter, Gauge as PGauge, Histogram as PHistogram, CollectorRegistry
+from prometheus_client import start_http_server, REGISTRY
+from aworld.metrics.metric import(
+    MetricProvider,
+    Counter,
+    UpDownCounter,
+    MetricExporter,
+    Gauge,
+    Histogram,
+    set_metric_provider
+)
+
+
+class PrometheusMetricProvider(MetricProvider):
+    """
+    PrometheusMetricProvider is a subclass of MetricProvider, representing a metric provider for Prometheus.
+    """
+
+    def __init__(self, exporter: MetricExporter):
+        """
+        Initialize the PrometheusMetricProvider.
+        Args:
+            port: The port to use for the Prometheus server.
+        """
+        super().__init__()
+        self.exporter = exporter
+
+    def shutdown(self) -> None:
+        """
+        Shutdown the PrometheusMetricProvider.
+        """
+        self.exporter.shutdown()
+
+    def create_counter(self, name: str, description: str, unit: str,
+                       labelnames: Optional[Sequence[str]] = None) -> Counter:
+        """
+        Create a counter metric.
+        Args:
+            name: The name of the metric.
+            description: The description of the metric.
+            unit: The unit of the metric.
+        Returns:
+            The counter metric.
+        """
+        return PrometheusCounter(name, description, unit, self, labelnames)
+
+    def create_un_down_counter(self, name: str, description: str, unit: str,
+                               labelnames: Optional[Sequence[str]] = None) -> UpDownCounter:
+        """
+        Create an up-down counter metric.
+        Args:
+            name: The name of the metric.
+            description: The description of the metric.
+            unit: The unit of the metric.
+        Returns:
+            The up-down counter metric.
+        """
+        return PrometheusUpDownCounter(name, description, unit, self, labelnames)
+
+    def create_gauge(self, name: str, description: str, unit: str, labelnames: Optional[Sequence[str]] = None) -> Gauge:
+        """
+        Create a gauge metric.
+        Args:
+            name: The name of the metric.
+            description: The description of the metric.
+            unit: The unit of the metric.
+        Returns:
+            The gauge metric.
+        """
+        return PrometheusGauge(name, description, unit, self, labelnames)
+
+    def create_histogram(self,
+                         name: str,
+                         description: str,
+                         unit: str,
+                         buckets: Optional[Sequence[float]] = None,
+                         labelnames: Optional[Sequence[str]] = None) -> Histogram:
+        """
+        Create a histogram metric.
+        Args:
+            name: The name of the metric.
+            description: The description of the metric.
+            unit: The unit of the metric.
+            buckets: The buckets of the histogram.
+        Returns:
+            The histogram metric.
+        """
+        return PrometheusHistogram(name, description, unit, self, buckets, labelnames)
+
+
+class PrometheusCounter(Counter):
+    """
+    PrometheusCounter is a subclass of Counter, representing a counter metric for Prometheus.
+    """
+
+    def __init__(self,
+                 name: str,
+                 description: str,
+                 unit: str,
+                 provider: MetricProvider,
+                 labelnames: Optional[Sequence[str]] = None):
+        """
+        Initialize the PrometheusCounter.
+        Args:
+            name: The name of the metric.
+            description: The description of the metric.
+            unit: The unit of the metric.
+            provider: The provider of the metric.
+        """
+        labelnames = labelnames or []
+        super().__init__(name, description, unit, provider, labelnames)
+        self._counter = PCounter(name=name, documentation=description, labelnames=labelnames, unit=unit)
+
+    def add(self, value: int, labels: dict = None) -> None:
+        """
+        Add a value to the counter.
+        Args:
+            value: The value to add to the counter.
+            labels: The labels to associate with the value.
+        """
+        if labels:
+            self._counter.labels(**labels).inc(value)
+        else:
+            self._counter.inc(value)
+
+
+class PrometheusUpDownCounter(UpDownCounter):
+    """
+    PrometheusUpDownCounter is a subclass of UpDownCounter, representing an up-down counter metric for Prometheus.
+    """
+
+    def __init__(self,
+                 name: str,
+                 description: str,
+                 unit: str,
+                 provider: MetricProvider,
+                 labelnames: Optional[Sequence[str]] = None):
+        """
+        Initialize the PrometheusUpDownCounter.
+        Args:
+            name: The name of the metric.
+            description: The description of the metric.
+            unit: The unit of the metric.
+            provider: The provider of the metric.
+        """
+        labelnames = labelnames or []
+        super().__init__(name, description, unit, provider, labelnames)
+        self._gauge = PGauge(name=name, documentation=description, labelnames=labelnames, unit=unit)
+
+    def inc(self, value: int, labels: dict = None) -> None:
+        """
+        Add a value to the counter.
+        Args:
+            value: The value to add to the counter.
+            labels: The labels to associate with the value.
+        """
+        if labels:
+            self._gauge.labels(**labels).inc(value)
+        else:
+            self._gauge.inc(value)
+
+    def dec(self, value: int, labels: dict = None) -> None:
+        """
+        Subtract a value from the counter.
+        Args:
+            value: The value to subtract from the counter.
+            labels: The labels to associate with the value.
+        """
+        if labels:
+            self._gauge.labels(**labels).dec(value)
+        else:
+            self._gauge.dec(value)
+
+
+class PrometheusGauge(Gauge):
+    """
+    PrometheusGauge is a subclass of Gauge, representing a gauge metric for Prometheus.
+    """
+
+    def __init__(self,
+                 name: str,
+                 description: str,
+                 unit: str,
+                 provider: MetricProvider,
+                 labelnames: Optional[Sequence[str]] = None):
+        """
+        Initialize the PrometheusGauge.
+        Args:
+            name: The name of the metric.
+            description: The description of the metric.
+            unit: The unit of the metric.
+            provider: The provider of the metric.
+        """
+        labelnames = labelnames or []
+        super().__init__(name, description, unit, provider, labelnames)
+        self._gauge = PGauge(name=name, documentation=description, labelnames=labelnames, unit=unit)
+
+    def set(self, value: int, labels: dict = None) -> None:
+        """
+        Set the value of the gauge.
+        Args:
+            value: The value to set the gauge to.
+            labels: The labels to associate with the value.
+        """
+        if labels:
+            self._gauge.labels(**labels).set(value)
+        else:
+            self._gauge.set(value)
+
+    def inc(self, value: int, labels: dict = None) -> None:
+        """
+        Add a value to the gauge.
+        Args:
+            value: The value to add to the gauge.
+            labels: The labels to associate with the value.
+        """
+        if labels:
+            self._gauge.labels(**labels).inc(value)
+        else:
+            self._gauge.inc(value)
+
+    def dec(self, value: int, labels: dict = None) -> None:
+        """
+        Subtract a value from the gauge.
+        Args:
+            value: The value to subtract from the gauge.
+            labels: The labels to associate with the value.
+        """
+        if labels:
+            self._gauge.labels(**labels).dec(value)
+        else:
+            self._gauge.dec(value)
+
+
+class PrometheusHistogram(Histogram):
+    """
+    PrometheusHistogram is a subclass of Histogram, representing a histogram metric for Prometheus.
+    """
+
+    def __init__(self,
+                 name: str,
+                 description: str,
+                 unit: str,
+                 provider: MetricProvider,
+                 buckets: Sequence[float] = None,
+                 labelnames: Optional[Sequence[str]] = None):
+        """
+        Initialize the PrometheusHistogram.
+        Args:
+            name: The name of the metric.
+            description: The description of the metric.
+            unit: The unit of the metric.
+            provider: The provider of the metric.
+        """
+        labelnames = labelnames or []
+        super().__init__(name, description, unit, provider, buckets, labelnames)
+        if buckets:
+            self._histogram = PHistogram(name=name, documentation=description, labelnames=labelnames, unit=unit,
+                                         buckets=buckets)
+        else:
+            self._histogram = PHistogram(name=name, documentation=description, labelnames=labelnames, unit=unit)
+
+    def record(self, value: int, labels: dict = None) -> None:
+        """
+        Record a value in the histogram.
+        Args:
+            value: The value to record in the histogram.
+            labels: The labels to associate with the value.
+        """
+        if labels:
+            self._histogram.labels(**labels).observe(value)
+        else:
+            self._histogram.observe(value)
+
+
+class PrometheusMetricExporter(MetricExporter):
+    """
+    PrometheusMetricExporter is a class for exporting metrics to Prometheus.
+    """
+
+    def __init__(self, port: int = 8000):
+        """
+        Initialize the PrometheusMetricExporter.
+        Args:
+            port: The port to use for the Prometheus server.
+        """
+        self.port = port
+        server, server_thread = start_http_server(self.port)
+        self.server = server
+        self.server_thread = server_thread
+
+    def shutdown(self) -> None:
+        """
+        Shutdown the PrometheusMetricExporter.
+        """
+        self.server.shutdown()
+        self.server_thread.join()
+
+
+class PrometheusConsoleMetricExporter(MetricExporter):
+    """Implementation of :class:`MetricExporter` that prints metrics to the
+    console.
+
+    This class can be used for diagnostic purposes. It prints the exported
+    metrics to the console STDOUT.
+    """
+
+    def __init__(self, out_interval_secs: float = 1.0):
+        """Initialize the console exporter."""
+        self._should_shutdown = False
+        self.out_interval_secs = out_interval_secs
+        self.metrics_thread = threading.Thread(target=self._output_metrics_to_console)
+        self.metrics_thread.daemon = True
+        self.metrics_thread.start()
+
+    def generate_latest(self, registry: CollectorRegistry = REGISTRY) -> bytes:
+        """Returns the metrics from the registry in latest text format as a string."""
+
+        def sample_line(line):
+            if line.labels:
+                labelstr = '{{{0}}}'.format(','.join(
+                    ['{}="{}"'.format(
+                        k, v.replace('\\', r'\\').replace('\n', r'\n').replace('"', r'\"'))
+                        for k, v in sorted(line.labels.items())]))
+            else:
+                labelstr = ''
+            timestamp = ''
+            if line.timestamp is not None:
+                # Convert to milliseconds.
+                timestamp = f' {int(float(line.timestamp) * 1000):d}'
+            return f'{line.name}{labelstr} {line.value}{timestamp}\n'
+
+        output = []
+        for metric in registry.collect():
+            try:
+                om_samples: Dict[str, List[str]] = {}
+                for s in metric.samples:
+                    for suffix in ['_gsum', '_gcount']:
+                        if s.name == metric.name + suffix:
+                            # OpenMetrics specific sample, put in a gauge at the end.
+                            om_samples.setdefault(suffix, []).append(sample_line(s))
+                            break
+                    else:
+                        output.append(sample_line(s))
+            except Exception as exception:
+                exception.args = (exception.args or ('',)) + (metric,)
+                raise
+
+            for suffix, lines in sorted(om_samples.items()):
+                output.extend(lines)
+        return ''.join(output).encode('utf-8')
+
+    def _output_metrics_to_console(self):
+        while not self._should_shutdown:
+            metrics_text = self.generate_latest(REGISTRY)
+            print(metrics_text.decode('utf-8'))
+            time.sleep(self.out_interval_secs)
+
+    def shutdown(self) -> None:
+        """
+        Shutdown the PrometheusConsoleMetricExporter.
+        """
+        self._should_shutdown = True
+
+def configure_prometheus_provider(backend: str,
+                                 base_url: str = None,
+                                 write_token: str = None,
+                                 **kwargs
+):
+
+    """
+    Initialize the prometheus metric provider.
+    Args:
+    backend: The backend of the metric provider.
+    base_url: The base url of the metric provider.
+    write_token: The write token of the metric provider.
+    """
+    if backend == "console":
+        exporter = PrometheusConsoleMetricExporter(out_interval_secs=2)
+        set_metric_provider(PrometheusMetricProvider(exporter))
+    elif backend == "prometheus":
+        exporter = PrometheusMetricExporter()
+        set_metric_provider(PrometheusMetricProvider(exporter))